diff --git a/internal/gobackend/capabilities.go b/internal/gobackend/capabilities.go index 5164b58..4765a90 100644 --- a/internal/gobackend/capabilities.go +++ b/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, diff --git a/internal/gobackend/executable.go b/internal/gobackend/executable.go index 8630af9..92c0108 100644 --- a/internal/gobackend/executable.go +++ b/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] } diff --git a/internal/gobackend/ops/dynamicupdateslice.go b/internal/gobackend/ops/dynamicupdateslice.go new file mode 100644 index 0000000..803e00b --- /dev/null +++ b/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 +} diff --git a/internal/gobackend/ops/reducewindow.go b/internal/gobackend/ops/reducewindow.go index 7e48675..d63dd23 100644 --- a/internal/gobackend/ops/reducewindow.go +++ b/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] + } +} diff --git a/internal/gobackend/ops/slice.go b/internal/gobackend/ops/slice.go index 78b4441..30e0b07 100644 --- a/internal/gobackend/ops/slice.go +++ b/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 diff --git a/internal/gobackend/workerspool/workerspool_test.go b/internal/gobackend/workerspool/workerspool_test.go index 915acc3..dcc932b 100644 --- a/internal/gobackend/workerspool/workerspool_test.go +++ b/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) } }