Add DynamicUpdateSlice.

internal/gobackend/capabilities.go

@@ -157,9 +157,9 @@ var Capabilities = compute.Capabilities{
 		// TODO: not implemented yet:
 		compute.OpTypeSelectAndScatterMax: true,
 		compute.OpTypeSelectAndScatterMin: true,
+		compute.OpTypeDynamicUpdateSlice:  true,
 		// compute.OpTypeSelectAndScatterSum: true,
 		// compute.OpTypeDynamicSlice: true,
-		// compute.OpTypeDynamicUpdateSlice: true,
 
 		// Lower priority ops:
 		// compute.OpTypeBatchNormForInference: true,

internal/gobackend/executable.go

@@ -227,6 +227,12 @@ func (fe *FunctionExecutable) Execute(backend *Backend, inputs []*Buffer, donate
 	// Set up parameters from inputs using idx directly
 	for i, inputNode := range funcParams {
 		inputIdx := inputNode.Index
+		if inputIdx >= fe.NumNodesToProcess || fe.NumUses[inputIdx] == 0 {
+			if donate[i] {
+				backend.PutBuffer(inputs[i])
+			}
+			continue
+		}
 		execBuf.results[inputIdx] = inputs[i]
 		execBuf.owned[inputIdx] = donate[i]
 	}
@@ -235,6 +241,12 @@ func (fe *FunctionExecutable) Execute(backend *Backend, inputs []*Buffer, donate
 	// If donateCaptures[i] is true, the closure takes ownership of the buffer.
 	for i, captureNode := range fe.Function.CapturedLocalNodes {
 		captureIdx := captureNode.Index
+		if captureIdx >= fe.NumNodesToProcess || fe.NumUses[captureIdx] == 0 {
+			if donateCaptures[i] {
+				backend.PutBuffer(capturedInputs[i])
+			}
+			continue
+		}
 		execBuf.results[captureIdx] = capturedInputs[i]
 		execBuf.owned[captureIdx] = donateCaptures[i]
 	}

internal/gobackend/ops/dynamicupdateslice.go

@@ -0,0 +1,189 @@
+// Copyright 2023-2026 The GoMLX Authors. SPDX-License-Identifier: Apache-2.0
+
+package ops
+
+import (
+	"github.com/gomlx/compute"
+	"github.com/gomlx/compute/dtypes"
+	"github.com/gomlx/compute/dtypes/bfloat16"
+	"github.com/gomlx/compute/dtypes/float16"
+	"github.com/gomlx/compute/internal/gobackend"
+	"github.com/gomlx/compute/shapeinference"
+	"github.com/pkg/errors"
+)
+
+func init() {
+	gobackend.RegisterDynamicUpdateSlice.Register(DynamicUpdateSlice, gobackend.PriorityGeneric)
+	gobackend.SetNodeExecutor(compute.OpTypeDynamicUpdateSlice, gobackend.PriorityGeneric, execDynamicUpdateSlice)
+}
+
+// DynamicUpdateSlice implements compute.Function.
+func DynamicUpdateSlice(f *gobackend.Function, operandOp, updateOp compute.Value, startIndicesOps []compute.Value) (compute.Value, error) {
+	allOps := make([]compute.Value, 0, 2+len(startIndicesOps))
+	allOps = append(allOps, operandOp, updateOp)
+	allOps = append(allOps, startIndicesOps...)
+	inputs, err := f.VerifyAndCastValues("DynamicUpdateSlice", allOps...)
+	if err != nil {
+		return nil, err
+	}
+	operand, update := inputs[0], inputs[1]
+	startIndices := inputs[2:]
+	rank := operand.Shape.Rank()
+	if len(startIndices) == 1 && !startIndices[0].Shape.IsScalar() && startIndices[0].Shape.Rank() == 1 {
+		if startIndices[0].Shape.Dimensions[0] != rank {
+			return nil, errors.Errorf("DynamicUpdateSlice: 1D startIndices has length %d, but operand rank is %d",
+				startIndices[0].Shape.Dimensions[0], rank)
+		}
+	} else if len(startIndices) != rank {
+		return nil, errors.Errorf("DynamicUpdateSlice: len(startIndices) (%d) must match operand rank (%d)",
+			len(startIndices), rank)
+	}
+	outputShape, err := shapeinference.DynamicUpdateSlice(operand.Shape, update.Shape)
+	if err != nil {
+		return nil, err
+	}
+	node, _ := f.GetOrCreateNode(compute.OpTypeDynamicUpdateSlice, outputShape, inputs, nil)
+	return node, nil
+}
+
+func scalarToIntAt(buf *gobackend.Buffer, idx int) int {
+	switch buf.RawShape.DType {
+	case dtypes.Int8:
+		return int(buf.Flat.([]int8)[idx])
+	case dtypes.Int16:
+		return int(buf.Flat.([]int16)[idx])
+	case dtypes.Int32:
+		return int(buf.Flat.([]int32)[idx])
+	case dtypes.Int64:
+		return int(buf.Flat.([]int64)[idx])
+	case dtypes.Uint8:
+		return int(buf.Flat.([]uint8)[idx])
+	case dtypes.Uint16:
+		return int(buf.Flat.([]uint16)[idx])
+	case dtypes.Uint32:
+		return int(buf.Flat.([]uint32)[idx])
+	case dtypes.Uint64:
+		return int(buf.Flat.([]uint64)[idx])
+	case dtypes.Float32:
+		return int(buf.Flat.([]float32)[idx])
+	case dtypes.Float64:
+		return int(buf.Flat.([]float64)[idx])
+	default:
+		panic(errors.Errorf("unsupported dtype %s for index", buf.RawShape.DType))
+	}
+}
+
+func copyUpdateSlice[T any](outputSlice, updateSlice []T, starts, updateDims, outStrides []int) {
+	rank := len(updateDims)
+	if rank == 0 {
+		outputSlice[0] = updateSlice[0]
+		return
+	}
+	lastDim := updateDims[rank-1]
+	if lastDim == 0 {
+		return
+	}
+	if rank == 1 {
+		copy(outputSlice[starts[0]:starts[0]+lastDim], updateSlice[:lastDim])
+		return
+	}
+	numSlices := len(updateSlice) / lastDim
+	coord := make([]int, rank-1)
+	updateIdx := 0
+	for sliceIdx := 0; sliceIdx < numSlices; sliceIdx++ {
+		outIdx := starts[rank-1]
+		for a := 0; a < rank-1; a++ {
+			outIdx += (starts[a] + coord[a]) * outStrides[a]
+		}
+		copy(outputSlice[outIdx:outIdx+lastDim], updateSlice[updateIdx:updateIdx+lastDim])
+		updateIdx += lastDim
+
+		for a := rank - 2; a >= 0; a-- {
+			coord[a]++
+			if coord[a] < updateDims[a] {
+				break
+			}
+			coord[a] = 0
+		}
+	}
+}
+
+func execDynamicUpdateSlice(backend *gobackend.Backend, node *gobackend.Node, inputs []*gobackend.Buffer, inputsOwned []bool) (*gobackend.Buffer, error) {
+	operand := inputs[0]
+	update := inputs[1]
+	startIndices := inputs[2:]
+
+	var output *gobackend.Buffer
+	var err error
+	if inputsOwned[0] {
+		output = operand
+		inputs[0] = nil
+	} else {
+		output, err = backend.CloneBuffer(operand)
+		if err != nil {
+			return nil, err
+		}
+	}
+
+	if backend.NoOps || update.RawShape.Size() == 0 {
+		return output, nil
+	}
+
+	rank := operand.RawShape.Rank()
+	starts := make([]int, rank)
+	if len(startIndices) == 1 && !startIndices[0].RawShape.IsScalar() && startIndices[0].RawShape.Rank() == 1 {
+		for i := 0; i < rank; i++ {
+			val := scalarToIntAt(startIndices[0], i)
+			maxStart := operand.RawShape.Dimensions[i] - update.RawShape.Dimensions[i]
+			if maxStart < 0 {
+				maxStart = 0
+			}
+			starts[i] = min(max(val, 0), maxStart)
+		}
+	} else {
+		for i := 0; i < rank; i++ {
+			val := scalarToIntAt(startIndices[i], 0)
+			maxStart := operand.RawShape.Dimensions[i] - update.RawShape.Dimensions[i]
+			if maxStart < 0 {
+				maxStart = 0
+			}
+			starts[i] = min(max(val, 0), maxStart)
+		}
+	}
+
+	updateDims := update.RawShape.Dimensions
+	outStrides := output.RawShape.Strides()
+
+	switch output.RawShape.DType {
+	case dtypes.Float32:
+		copyUpdateSlice(output.Flat.([]float32), update.Flat.([]float32), starts, updateDims, outStrides)
+	case dtypes.Float64:
+		copyUpdateSlice(output.Flat.([]float64), update.Flat.([]float64), starts, updateDims, outStrides)
+	case dtypes.BFloat16:
+		copyUpdateSlice(output.Flat.([]bfloat16.BFloat16), update.Flat.([]bfloat16.BFloat16), starts, updateDims, outStrides)
+	case dtypes.Float16:
+		copyUpdateSlice(output.Flat.([]float16.Float16), update.Flat.([]float16.Float16), starts, updateDims, outStrides)
+	case dtypes.Int32:
+		copyUpdateSlice(output.Flat.([]int32), update.Flat.([]int32), starts, updateDims, outStrides)
+	case dtypes.Int64:
+		copyUpdateSlice(output.Flat.([]int64), update.Flat.([]int64), starts, updateDims, outStrides)
+	case dtypes.Int16:
+		copyUpdateSlice(output.Flat.([]int16), update.Flat.([]int16), starts, updateDims, outStrides)
+	case dtypes.Int8:
+		copyUpdateSlice(output.Flat.([]int8), update.Flat.([]int8), starts, updateDims, outStrides)
+	case dtypes.Uint32:
+		copyUpdateSlice(output.Flat.([]uint32), update.Flat.([]uint32), starts, updateDims, outStrides)
+	case dtypes.Uint64:
+		copyUpdateSlice(output.Flat.([]uint64), update.Flat.([]uint64), starts, updateDims, outStrides)
+	case dtypes.Uint16:
+		copyUpdateSlice(output.Flat.([]uint16), update.Flat.([]uint16), starts, updateDims, outStrides)
+	case dtypes.Uint8:
+		copyUpdateSlice(output.Flat.([]uint8), update.Flat.([]uint8), starts, updateDims, outStrides)
+	case dtypes.Bool:
+		copyUpdateSlice(output.Flat.([]bool), update.Flat.([]bool), starts, updateDims, outStrides)
+	default:
+		return nil, errors.Errorf("DynamicUpdateSlice: unsupported dtype %s", output.RawShape.DType)
+	}
+
+	return output, nil
+}

internal/gobackend/ops/reducewindow.go

@@ -5,6 +5,7 @@ import (
 	"sync"
 
 	"github.com/gomlx/compute"
+	"github.com/gomlx/compute/dtypes"
 	"github.com/gomlx/compute/dtypes/gotype"
 	"github.com/gomlx/compute/internal/gobackend"
 	"github.com/gomlx/compute/shapeinference"
@@ -16,6 +17,8 @@ import (
 func init() {
 	gobackend.RegisterReduceWindow.Register(ReduceWindow, gobackend.PriorityGeneric)
 	gobackend.SetNodeExecutor(compute.OpTypeReduceWindow, gobackend.PriorityGeneric, execReduceWindow)
+	reduceWindowMinDTypeMap.Register(dtypes.Bool, gobackend.PriorityGeneric, reduceWindowMinBuildUpdateFnBool)
+	reduceWindowMaxDTypeMap.Register(dtypes.Bool, gobackend.PriorityGeneric, reduceWindowMaxBuildUpdateFnBool)
 }
 
 type reduceWindowNode struct {
@@ -351,3 +354,19 @@ func reduceWindowProductBuildUpdateFnHalf[T gotype.HalfPrecision[T], P gotype.Ha
 			outputFlat[outputFlatIdx].Float32() * operandFlat[operandFlatIdx].Float32())
 	}
 }
+
+func reduceWindowMinBuildUpdateFnBool(operand, output *gobackend.Buffer) reduceWindowUpdateFn {
+	operandFlat := operand.Flat.([]bool)
+	outputFlat := output.Flat.([]bool)
+	return func(operandFlatIdx, outputFlatIdx int) {
+		outputFlat[outputFlatIdx] = outputFlat[outputFlatIdx] && operandFlat[operandFlatIdx]
+	}
+}
+
+func reduceWindowMaxBuildUpdateFnBool(operand, output *gobackend.Buffer) reduceWindowUpdateFn {
+	operandFlat := operand.Flat.([]bool)
+	outputFlat := output.Flat.([]bool)
+	return func(operandFlatIdx, outputFlatIdx int) {
+		outputFlat[outputFlatIdx] = outputFlat[outputFlatIdx] || operandFlat[operandFlatIdx]
+	}
+}

internal/gobackend/ops/slice.go

@@ -98,21 +98,15 @@ func Slice(f *gobackend.Function, operandOp compute.Value, starts, limits, strid
 
 		// Start
 		start := starts[axis]
-		if dimSize != shapes.DynamicDim {
-			if start < 0 {
-				start = dimSize + start
-			}
-			start = min(max(start, 0), dimSize)
+		if dimSize != shapes.DynamicDim && start < 0 {
+			start = dimSize + start
 		}
 		data.starts[axis] = start
 
 		// Limit
 		limit := limits[axis]
-		if dimSize != shapes.DynamicDim {
-			if limit < 0 {
-				limit = dimSize + limit
-			}
-			limit = min(max(limit, 0), dimSize)
+		if dimSize != shapes.DynamicDim && limit < 0 {
+			limit = dimSize + limit
 		}
 		data.limits[axis] = limit
 

internal/gobackend/workerspool/workerspool_test.go

@@ -53,17 +53,35 @@ func TestPool_Saturate(t *testing.T) {
 
 	// Test Unlimited
 	pool.SetMaxParallelism(-1)
-	count.Store(0)
+	wantUnlimited := runtime.GOMAXPROCS(0)
 	var started atomic.Int32
-	pool.Saturate(func() {
-		started.Add(1)
-		runtime.Gosched()
-		count.Add(1)
-	})
+	doneUnlimited := xsync.NewLatch()
+	doneTestUnlimited := xsync.NewLatch()
+
+	go func() {
+		pool.Saturate(func() {
+			got := started.Add(1)
+			runtime.Gosched()
+			if int(got) == wantUnlimited {
+				doneUnlimited.Trigger()
+				return
+			}
+			doneUnlimited.Wait()
+		})
+		doneTestUnlimited.Trigger()
+	}()
+
+	select {
+	case <-doneTestUnlimited.WaitChan():
+		// Success
+	case <-time.After(100 * time.Millisecond):
+		t.Fatal("Timeout before all unlimited tasks were executed.")
+	}
+
 	if runtime.GOMAXPROCS(0) > 1 && started.Load() <= 1 {
 		t.Errorf("Expected more than 1 started task for unlimited parallelism, got %d", started.Load())
 	}
-	if count.Load() != started.Load() {
-		t.Errorf("Expected count %d to match started %d", count.Load(), started.Load())
+	if int(started.Load()) != wantUnlimited {
+		t.Errorf("Expected started %d to match wantUnlimited %d", started.Load(), wantUnlimited)
 	}
 }