binary_ops.go

package main

import (
	"fmt"
	"os"
	"os/exec"
	"path"
	"text/template"

	"github.com/gomlx/stablehlo/internal/utils"
	"github.com/gomlx/stablehlo/shapeinference"
)

const (
	binaryOpsFile = "gen_binary_ops.go"
)

var (
	binaryOpsTemplate = template.Must(
		template.
			New(binaryOpsFile).
			Parse(
				`/***** File generated by ./internal/cmd/ops_generator. Don't edit it directly. *****/

package stablehlo

import (
	"github.com/gomlx/stablehlo/internal/optypes"
)

{{- range .}}
// {{.Name}} implements the corresponding standard binary operation.
func {{.Name}}(lhs, rhs *Value) (*Value, error) {
	fn := lhs.fn
	return fn.binaryOp(optypes.{{.Name}}, lhs, rhs)
}
{{- end}}
`))
)

type BinaryOp struct {
	Name string
}

func GenerateBinaryOps() {
	binaryOps := shapeinference.StandardBinaryOperations
	data := make([]BinaryOp, 0, len(binaryOps))
	for _, k := range utils.SortedKeys(binaryOps) {
		data = append(data, BinaryOp{Name: k.String()})
	}

	fileName := binaryOpsFile
	f := must1(os.Create(fileName))
	must(binaryOpsTemplate.Execute(f, data))
	must(f.Close())

	cmd := exec.Command("gofmt", "-w", fileName)
	must(cmd.Run())
	fmt.Printf("✅ Successfully generated %s\n", path.Join(must1(os.Getwd()), fileName))
}