Implement OptimizationBarrier and SchedulingBarrier for the XLA backend

cmd/pjrt_installer/go.mod

@@ -9,7 +9,7 @@ require (
 	github.com/charmbracelet/huh/spinner v0.0.0-20260223110133-9dc45e34a40b
 	github.com/charmbracelet/lipgloss v1.1.0
 	github.com/charmbracelet/x/term v0.2.2
-	github.com/gomlx/go-xla v0.2.3-0.20260618154843-d9c3ca970fe6
+	github.com/gomlx/go-xla v0.2.3-0.20260621191029-4015225fdf7c
 	github.com/pkg/errors v0.9.1
 	k8s.io/klog/v2 v2.140.0
 )
@@ -28,7 +28,7 @@ require (
 	github.com/erikgeiser/coninput v0.0.0-20211004153227-1c3628e74d0f // indirect
 	github.com/go-logr/logr v1.4.3 // indirect
 	github.com/gofrs/flock v0.13.0 // indirect
-	github.com/gomlx/compute v0.0.0-20260615181056-6e3a924a46cb // indirect
+	github.com/gomlx/compute v0.0.0-20260621190903-a14fa9673b04 // indirect
 	github.com/lucasb-eyer/go-colorful v1.3.0 // indirect
 	github.com/mattn/go-isatty v0.0.20 // indirect
 	github.com/mattn/go-localereader v0.0.1 // indirect

cmd/pjrt_installer/go.sum

@@ -54,10 +54,10 @@ github.com/go-logr/logr v1.4.3 h1:CjnDlHq8ikf6E492q6eKboGOC0T8CDaOvkHCIg8idEI=
 github.com/go-logr/logr v1.4.3/go.mod h1:9T104GzyrTigFIr8wt5mBrctHMim0Nb2HLGrmQ40KvY=
 github.com/gofrs/flock v0.13.0 h1:95JolYOvGMqeH31+FC7D2+uULf6mG61mEZ/A8dRYMzw=
 github.com/gofrs/flock v0.13.0/go.mod h1:jxeyy9R1auM5S6JYDBhDt+E2TCo7DkratH4Pgi8P+Z0=
-github.com/gomlx/compute v0.0.0-20260615181056-6e3a924a46cb h1:ZclLuMEjLP+N5oV7By2mtrPZOSmX8cSbpAvczuH76HM=
-github.com/gomlx/compute v0.0.0-20260615181056-6e3a924a46cb/go.mod h1:MDTT683Wvq7IMSKXDGLvV3Co7HDpx/Xsgp2RbF5lhNM=
-github.com/gomlx/go-xla v0.2.3-0.20260618154843-d9c3ca970fe6 h1:srYXjBoc2miAKTfpJsihQrnDlr64YwEytwlWDJBs6Ww=
-github.com/gomlx/go-xla v0.2.3-0.20260618154843-d9c3ca970fe6/go.mod h1:zMOV0+tZBqrAulL9XG6KWIzEwtavwIVdNkCFV9zi+5U=
+github.com/gomlx/compute v0.0.0-20260621190903-a14fa9673b04 h1:FU1UPQTpAbmc+BO140UXpTiDDMbAO0Af4TFBtsbOvrw=
+github.com/gomlx/compute v0.0.0-20260621190903-a14fa9673b04/go.mod h1:MDTT683Wvq7IMSKXDGLvV3Co7HDpx/Xsgp2RbF5lhNM=
+github.com/gomlx/go-xla v0.2.3-0.20260621191029-4015225fdf7c h1:q7lj3LFkBBGH6ouMxp13wy67N3+RUGv5rMNZheKbVjs=
+github.com/gomlx/go-xla v0.2.3-0.20260621191029-4015225fdf7c/go.mod h1:hV90YGdvcfC8PKjHQrC73rUrxy15QOMguFAMjjV4FQM=
 github.com/lucasb-eyer/go-colorful v1.3.0 h1:2/yBRLdWBZKrf7gB40FoiKfAWYQ0lqNcbuQwVHXptag=
 github.com/lucasb-eyer/go-colorful v1.3.0/go.mod h1:R4dSotOR9KMtayYi1e77YzuveK+i7ruzyGqttikkLy0=
 github.com/mattn/go-isatty v0.0.20 h1:xfD0iDuEKnDkl03q4limB+vH+GxLEtL/jb4xVJSWWEY=

compute/xla/capabilities.go

@@ -120,6 +120,8 @@ var Capabilities = compute.Capabilities{
 		compute.OpTypeSlice:                 true,
 		compute.OpTypeTranspose:             true,
 		compute.OpTypeWhere:                 true,
+		compute.OpTypeOptimizationBarrier:   true,
+		compute.OpTypeSchedulingBarrier:     true,
 
 		// Collective (distributed across devices) operations:
 		compute.OpTypeAllReduce: true,

compute/xla/ops.go

@@ -5,6 +5,8 @@ package xla
 // This file contains manually implemented operations.
 
 import (
+	"reflect"
+
 	"github.com/gomlx/compute"
 	"github.com/gomlx/compute/dtypes"
 	"github.com/gomlx/compute/shapes"
@@ -795,13 +797,13 @@ func (f *Function) Call(target compute.Function, inputs ...compute.Value) ([]com
 }
 
 // OptimizationBarrier returned values are identity to the operands, but they become only available
-// in compilation time once all the inptus are calculated.
+// in compilation time once all the inputs are calculated.
 func (f *Function) OptimizationBarrier(operands ...compute.Value) ([]compute.Value, error) {
 	if len(operands) == 0 {
 		return nil, errors.New("OptimizationBarrier requires at least one operand")
 	}
-	if len(operands) == 1 {
-		// No-op: a value already depends on itself.
+	if len(operands) == 1 && f.builder.backend.plugin.IsCPU() {
+		// No-op: XLA CPU compiler has a bug with single-operand OptimizationBarrier and produces garbage/crashes.
 		return operands, nil
 	}
 	nodes, err := f.verifyAndCastValues("OptimizationBarrier", operands...)
@@ -816,3 +818,85 @@ func (f *Function) OptimizationBarrier(operands ...compute.Value) ([]compute.Val
 	outputNodes := xslices.Map(values, func(v *stablehlo.Value) compute.Value { return f.newNode(v) })
 	return outputNodes, nil
 }
+
+// SchedulingBarrier introduces a scheduling barrier.
+// Returned value is identity to the operand, but it is guaranteed to depend on all the dependencies.
+//
+// Since XLA/StableHLO don't support this operation, we "hack" it by means of a no-op: taking a scalar of each depdency,
+// multiply by 0 and add them to the operand.
+//
+// See: https://github.com/openxla/stablehlo/issues/2923.
+func (f *Function) SchedulingBarrier(operand compute.Value, dependencies ...compute.Value) (compute.Value, error) {
+	if len(dependencies) == 0 {
+		return operand, nil
+	}
+	nodes, err := f.verifyAndCastValues("SchedulingBarrier", append([]compute.Value{operand}, dependencies...)...)
+	if err != nil {
+		return nil, err
+	}
+	opNode := nodes[0]
+	depNodes := nodes[1:]
+
+	opDType := opNode.shape.DType
+
+	// Create a constant zero of the operand's dtype.
+	zeroFlat := reflect.MakeSlice(reflect.SliceOf(opDType.GoType()), 1, 1)
+	zeroConst, err := f.Constant(zeroFlat.Interface())
+	if err != nil {
+		return nil, errors.WithMessage(err, "SchedulingBarrier: failed to create constant zero")
+	}
+
+	// Add OptimizationBarrier on the constant zero so it doesn't get optimized away.
+	zeroBarriers, err := f.OptimizationBarrier(zeroConst)
+	if err != nil {
+		return nil, errors.WithMessage(err, "SchedulingBarrier: failed to apply OptimizationBarrier to zero")
+	}
+	zero := zeroBarriers[0]
+
+	accumulatedZero := zero
+	for _, dep := range depNodes {
+		var scalarDep compute.Value
+		if dep.shape.Rank() == 0 {
+			scalarDep = dep
+		} else {
+			starts := make([]int, dep.shape.Rank())
+			limits := xslices.SliceWithValue(dep.shape.Rank(), 1)
+			slice, err := f.Slice(dep, starts, limits, nil)
+			if err != nil {
+				return nil, err
+			}
+			scalarDep, err = f.Reshape(slice)
+			if err != nil {
+				return nil, err
+			}
+		}
+
+		if dep.shape.DType != opDType {
+			var err error
+			scalarDep, err = f.ConvertDType(scalarDep, opDType)
+			if err != nil {
+				return nil, err
+			}
+		}
+
+		// Multiply the scalar dependency by the constant zero.
+		term, err := f.Mul(scalarDep, zero)
+		if err != nil {
+			return nil, err
+		}
+
+		// Add it to our accumulated zero.
+		accumulatedZero, err = f.Add(accumulatedZero, term)
+		if err != nil {
+			return nil, err
+		}
+	}
+
+	// Finally, add the accumulated zero (which depends on all dependencies) to the operand.
+	operandReady, err := f.Add(operand, accumulatedZero)
+	if err != nil {
+		return nil, err
+	}
+
+	return operandReady, nil
+}

compute/xla/ops_dotgeneral.go

@@ -3,11 +3,8 @@
 package xla
 
 import (
-	"reflect"
-
 	"github.com/gomlx/compute"
 	"github.com/gomlx/compute/dtypes"
-	"github.com/gomlx/compute/support/xslices"
 	"github.com/gomlx/go-xla/stablehlo"
 	stablehlotypes "github.com/gomlx/go-xla/types"
 	"github.com/pkg/errors"
@@ -144,34 +141,14 @@ func dotGeneralAddDependency(f *Function, lhs, rhs *Node) (*Node, *Node, error)
 		lhs, rhs = rhs, lhs
 	}
 
-	// Create a fake barrier, by creating a dependency:
-	starts := make([]int, lhs.shape.Rank())
-	limits := xslices.SliceWithValue(lhs.shape.Rank(), 1)
-	dependency, err := f.Slice(lhs, starts, limits, nil)
-	if err != nil {
-		return nil, nil, errors.WithMessagef(err, "failed to create fake barrier for DotGeneral operands")
-	}
-	dependency, err = f.Reshape(dependency) // Convert to scalar.
+	// Create a scheduling barrier: rhs (the parameter) depends on lhs (the weight/non-parameter).
+	rhsVal, err := f.SchedulingBarrier(rhs, lhs)
 	if err != nil {
-		return nil, nil, errors.WithMessagef(err, "failed to create fake barrier for DotGeneral operands")
+		return nil, nil, errors.WithMessagef(err, "failed to create scheduling barrier for DotGeneral operands")
 	}
-	zeroFlatV := reflect.MakeSlice(reflect.SliceOf(lhs.shape.DType.GoType()), 1, 1)
-	zero, err := f.Constant(zeroFlatV.Interface())
-	if err != nil {
-		return nil, nil, errors.WithMessagef(err, "failed to create fake barrier for DotGeneral operands")
-	}
-	dependency, err = f.Mul(dependency, zero)
-	if err != nil {
-		return nil, nil, errors.WithMessagef(err, "failed to create fake barrier for DotGeneral operands")
-	}
-	rhsReady, err := f.Add(rhs, dependency) // Adds a dependency from the weights to the LHS.
-	if err != nil {
-		return nil, nil, errors.WithMessagef(err, "failed to create fake barrier for DotGeneral operands")
-	}
-
-	lhsReady := compute.Value(lhs)
+	rhs = rhsVal.(*Node)
 	if isSwapped {
-		lhsReady, rhsReady = rhsReady, lhsReady
+		lhs, rhs = rhs, lhs
 	}
-	return lhsReady.(*Node), rhsReady.(*Node), nil
+	return lhs, rhs, nil
 }

go.mod

@@ -4,7 +4,7 @@ go 1.26
 
 require (
 	github.com/gofrs/flock v0.13.0
-	github.com/gomlx/compute v0.0.0-20260615181056-6e3a924a46cb
+	github.com/gomlx/compute v0.0.0-20260621195442-7cf34e76eacb
 	github.com/janpfeifer/go-benchmarks v0.1.1
 	github.com/pkg/errors v0.9.1
 	github.com/stretchr/testify v1.11.1

go.sum

@@ -7,8 +7,8 @@ github.com/go-logr/logr v1.4.3 h1:CjnDlHq8ikf6E492q6eKboGOC0T8CDaOvkHCIg8idEI=
 github.com/go-logr/logr v1.4.3/go.mod h1:9T104GzyrTigFIr8wt5mBrctHMim0Nb2HLGrmQ40KvY=
 github.com/gofrs/flock v0.13.0 h1:95JolYOvGMqeH31+FC7D2+uULf6mG61mEZ/A8dRYMzw=
 github.com/gofrs/flock v0.13.0/go.mod h1:jxeyy9R1auM5S6JYDBhDt+E2TCo7DkratH4Pgi8P+Z0=
-github.com/gomlx/compute v0.0.0-20260615181056-6e3a924a46cb h1:ZclLuMEjLP+N5oV7By2mtrPZOSmX8cSbpAvczuH76HM=
-github.com/gomlx/compute v0.0.0-20260615181056-6e3a924a46cb/go.mod h1:MDTT683Wvq7IMSKXDGLvV3Co7HDpx/Xsgp2RbF5lhNM=
+github.com/gomlx/compute v0.0.0-20260621195442-7cf34e76eacb h1:kTvVmJYXNBmh1sw4PL7uA9Qood0B/JzyWVpNUWnLDlo=
+github.com/gomlx/compute v0.0.0-20260621195442-7cf34e76eacb/go.mod h1:MDTT683Wvq7IMSKXDGLvV3Co7HDpx/Xsgp2RbF5lhNM=
 github.com/google/go-cmp v0.7.0 h1:wk8382ETsv4JYUZwIsn6YpYiWiBsYLSJiTsyBybVuN8=
 github.com/google/go-cmp v0.7.0/go.mod h1:pXiqmnSA92OHEEa9HXL2W4E7lf9JzCmGVUdgjX3N/iU=
 github.com/janpfeifer/go-benchmarks v0.1.1 h1:gLLy07/JrOKSnMWeUxSnjTdhkglgmrNR2IBDnR4kRqw=