diff --git a/go.mod b/go.mod index 064f37f..45dba11 100644 --- a/go.mod +++ b/go.mod @@ -9,7 +9,7 @@ require ( github.com/pkg/errors v0.9.1 github.com/x448/float16 v0.8.4 golang.org/x/term v0.38.0 - google.golang.org/protobuf v1.36.10 + google.golang.org/protobuf v1.36.11 k8s.io/klog/v2 v2.130.1 ) @@ -18,9 +18,10 @@ require ( github.com/go-logr/logr v1.4.3 // indirect github.com/pascaldekloe/name v1.0.0 // indirect github.com/streadway/quantile v0.0.0-20220407130108-4246515d968d // indirect + github.com/stretchr/testify v1.11.1 // indirect golang.org/x/mod v0.29.0 // indirect golang.org/x/sync v0.17.0 // indirect - golang.org/x/sys v0.39.0 // indirect + golang.org/x/sys v0.40.0 // indirect golang.org/x/tools v0.38.0 // indirect ) diff --git a/go.sum b/go.sum index 1228263..7745178 100644 --- a/go.sum +++ b/go.sum @@ -16,22 +16,22 @@ github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZb github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= github.com/streadway/quantile v0.0.0-20220407130108-4246515d968d h1:X4+kt6zM/OVO6gbJdAfJR60MGPsqCzbtXNnjoGqdfAs= github.com/streadway/quantile v0.0.0-20220407130108-4246515d968d/go.mod h1:lbP8tGiBjZ5YWIc2fzuRpTaz0b/53vT6PEs3QuAWzuU= -github.com/stretchr/testify v1.10.0 h1:Xv5erBjTwe/5IxqUQTdXv5kgmIvbHo3QQyRwhJsOfJA= -github.com/stretchr/testify v1.10.0/go.mod h1:r2ic/lqez/lEtzL7wO/rwa5dbSLXVDPFyf8C91i36aY= +github.com/stretchr/testify v1.11.1 h1:7s2iGBzp5EwR7/aIZr8ao5+dra3wiQyKjjFuvgVKu7U= +github.com/stretchr/testify v1.11.1/go.mod h1:wZwfW3scLgRK+23gO65QZefKpKQRnfz6sD981Nm4B6U= github.com/x448/float16 v0.8.4 h1:qLwI1I70+NjRFUR3zs1JPUCgaCXSh3SW62uAKT1mSBM= github.com/x448/float16 v0.8.4/go.mod h1:14CWIYCyZA/cWjXOioeEpHeN/83MdbZDRQHoFcYsOfg= golang.org/x/mod v0.29.0 h1:HV8lRxZC4l2cr3Zq1LvtOsi/ThTgWnUk/y64QSs8GwA= golang.org/x/mod v0.29.0/go.mod h1:NyhrlYXJ2H4eJiRy/WDBO6HMqZQ6q9nk4JzS3NuCK+w= golang.org/x/sync v0.17.0 h1:l60nONMj9l5drqw6jlhIELNv9I0A4OFgRsG9k2oT9Ug= golang.org/x/sync v0.17.0/go.mod h1:9KTHXmSnoGruLpwFjVSX0lNNA75CykiMECbovNTZqGI= -golang.org/x/sys v0.39.0 h1:CvCKL8MeisomCi6qNZ+wbb0DN9E5AATixKsvNtMoMFk= -golang.org/x/sys v0.39.0/go.mod h1:OgkHotnGiDImocRcuBABYBEXf8A9a87e/uXjp9XT3ks= +golang.org/x/sys v0.40.0 h1:DBZZqJ2Rkml6QMQsZywtnjnnGvHza6BTfYFWY9kjEWQ= +golang.org/x/sys v0.40.0/go.mod h1:OgkHotnGiDImocRcuBABYBEXf8A9a87e/uXjp9XT3ks= golang.org/x/term v0.38.0 h1:PQ5pkm/rLO6HnxFR7N2lJHOZX6Kez5Y1gDSJla6jo7Q= golang.org/x/term v0.38.0/go.mod h1:bSEAKrOT1W+VSu9TSCMtoGEOUcKxOKgl3LE5QEF/xVg= golang.org/x/tools v0.38.0 h1:Hx2Xv8hISq8Lm16jvBZ2VQf+RLmbd7wVUsALibYI/IQ= golang.org/x/tools v0.38.0/go.mod h1:yEsQ/d/YK8cjh0L6rZlY8tgtlKiBNTL14pGDJPJpYQs= -google.golang.org/protobuf v1.36.10 h1:AYd7cD/uASjIL6Q9LiTjz8JLcrh/88q5UObnmY3aOOE= -google.golang.org/protobuf v1.36.10/go.mod h1:HTf+CrKn2C3g5S8VImy6tdcUvCska2kB7j23XfzDpco= +google.golang.org/protobuf v1.36.11 h1:fV6ZwhNocDyBLK0dj+fg8ektcVegBBuEolpbTQyBNVE= +google.golang.org/protobuf v1.36.11/go.mod h1:HTf+CrKn2C3g5S8VImy6tdcUvCska2kB7j23XfzDpco= gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA= gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= k8s.io/klog/v2 v2.130.1 h1:n9Xl7H1Xvksem4KFG4PYbdQCQxqc/tTUyrgXaOhHSzk= diff --git a/pkg/stablehlo/ops.go b/pkg/stablehlo/ops.go index 9c2df8f..08ae4a4 100644 --- a/pkg/stablehlo/ops.go +++ b/pkg/stablehlo/ops.go @@ -1719,6 +1719,468 @@ func GetDimensionSize(operand *Value, dimension int) (*Value, error) { return stmt.Outputs[0], nil } +// DynamicReshape reshapes the operand to a shape specified by a 1D tensor. +// +// Parameters: +// - operand: the tensor to reshape. +// - outputShapeTensor: a 1D tensor of i32 or i64 values specifying the target shape dimensions. +// - bounds: dimension bounds for the output shape. Must have length equal to output rank. +// Use the actual dimension value for known dimensions, or an upper bound for dynamic ones. +// +// This is the dynamic version of Reshape where the output shape is determined at runtime. +// The total number of elements must remain the same. +func DynamicReshape(operand *Value, outputShapeTensor *Value, bounds []int) (*Value, error) { + op := optypes.DynamicReshape + fn := operand.fn + if fn.Returned { + return nil, errors.Errorf("cannot add operation %s after returning, in function %q", + op, fn.Name) + } + if outputShapeTensor.fn != fn { + return nil, errors.Errorf("cannot add operation %s to function %q, because operand and outputShapeTensor are from different functions", + op, fn.Name) + } + + // Validate outputShapeTensor is 1D tensor of integer type + if outputShapeTensor.shape.Rank() != 1 { + return nil, errors.Errorf("outputShapeTensor must be a 1D tensor, got rank %d", + outputShapeTensor.shape.Rank()) + } + if !outputShapeTensor.shape.DType.IsInt() { + return nil, errors.Errorf("outputShapeTensor must be integer type, got %s", + outputShapeTensor.shape.DType) + } + + // Get output rank from the shape tensor + outputRank := outputShapeTensor.shape.Dimensions[0] + if outputRank < 0 { + return nil, errors.Errorf("outputShapeTensor has dynamic size, cannot determine output rank") + } + + // Validate bounds length + if len(bounds) != outputRank { + return nil, errors.Errorf("bounds length (%d) must match output rank (%d)", + len(bounds), outputRank) + } + + // Create output shape with dynamic dimensions and caller-provided bounds + resultShape := operand.shape.Clone() + resultShape.Dimensions = make([]int, outputRank) + resultShape.DimensionBounds = make([]int, outputRank) + for i := range outputRank { + resultShape.Dimensions[i] = shapes.DimUnknown // Dynamic + resultShape.DimensionBounds[i] = bounds[i] + } + resultShape.EncodeBounds = true + + stmt := fn.addOp(op, resultShape, operand, outputShapeTensor) + return stmt.Outputs[0], nil +} + +// DynamicBroadcastInDim broadcasts the operand to a shape specified by a 1D tensor. +// +// Parameters: +// - operand: the tensor to broadcast. +// - outputDimensions: a 1D tensor of i32 or i64 values specifying the target shape. +// - broadcastDimensions: maps operand axes to output axes (like BroadcastInDim). +// - bounds: dimension bounds for the output shape. Must have length equal to output rank. +// +// This is the dynamic version of BroadcastInDim where the output shape is determined at runtime. +func DynamicBroadcastInDim(operand *Value, outputDimensions *Value, broadcastDimensions []int, bounds []int) (*Value, error) { + op := optypes.DynamicBroadcastInDim + fn := operand.fn + if fn.Returned { + return nil, errors.Errorf("cannot add operation %s after returning, in function %q", + op, fn.Name) + } + if outputDimensions.fn != fn { + return nil, errors.Errorf("cannot add operation %s to function %q, because operand and outputDimensions are from different functions", + op, fn.Name) + } + + // Validate outputDimensions is 1D tensor of integer type + if outputDimensions.shape.Rank() != 1 { + return nil, errors.Errorf("outputDimensions must be a 1D tensor, got rank %d", + outputDimensions.shape.Rank()) + } + if !outputDimensions.shape.DType.IsInt() { + return nil, errors.Errorf("outputDimensions must be integer type, got %s", + outputDimensions.shape.DType) + } + + // Validate broadcastDimensions length matches operand rank + if len(broadcastDimensions) != operand.shape.Rank() { + return nil, errors.Errorf("broadcastDimensions length (%d) must match operand rank (%d)", + len(broadcastDimensions), operand.shape.Rank()) + } + + // Get output rank from the dimensions tensor + outputRank := outputDimensions.shape.Dimensions[0] + if outputRank < 0 { + return nil, errors.Errorf("outputDimensions has dynamic size, cannot determine output rank") + } + + // Validate bounds length + if len(bounds) != outputRank { + return nil, errors.Errorf("bounds length (%d) must match output rank (%d)", + len(bounds), outputRank) + } + + // Create output shape with dynamic dimensions and caller-provided bounds + outputShape := operand.shape.Clone() + outputShape.Dimensions = make([]int, outputRank) + outputShape.DimensionBounds = make([]int, outputRank) + for i := range outputRank { + outputShape.Dimensions[i] = shapes.DimUnknown // Dynamic + outputShape.DimensionBounds[i] = bounds[i] + } + outputShape.EncodeBounds = true + + stmt := fn.addOp(op, outputShape, operand, outputDimensions) + stmt.Attributes = map[string]any{ + "broadcast_dimensions": intSliceToArrayI64StableHLO(broadcastDimensions), + } + return stmt.Outputs[0], nil +} + +// DynamicIota creates a tensor filled with values increasing along the specified dimension, +// where the output shape is determined at runtime. +// +// Parameters: +// - fn: the function to add this operation to. +// - dtype: the data type of the output tensor. +// - outputShape: a 1D tensor of i32 or i64 values specifying the output dimensions. +// - iotaDimension: the axis along which values increase (0, 1, 2, ...). +// - bounds: dimension bounds for the output shape. Must have length equal to output rank. +// +// This is the dynamic version of Iota where the output shape is determined at runtime. +func DynamicIota(fn *Function, dtype dtypes.DType, outputShape *Value, iotaDimension int, bounds []int) (*Value, error) { + op := optypes.DynamicIota + if fn.Returned { + return nil, errors.Errorf("cannot add operation %s after returning, in function %q", + op, fn.Name) + } + if outputShape.fn != fn { + return nil, errors.Errorf("cannot add operation %s to function %q, because outputShape is from a different function", + op, fn.Name) + } + + // Validate outputShape is 1D tensor of integer type + if outputShape.shape.Rank() != 1 { + return nil, errors.Errorf("outputShape must be a 1D tensor, got rank %d", + outputShape.shape.Rank()) + } + if !outputShape.shape.DType.IsInt() { + return nil, errors.Errorf("outputShape must be integer type, got %s", + outputShape.shape.DType) + } + + // Get output rank from the shape tensor + outputRank := outputShape.shape.Dimensions[0] + if outputRank < 0 { + return nil, errors.Errorf("outputShape has dynamic size, cannot determine output rank") + } + + // Validate iotaDimension + if iotaDimension < 0 || iotaDimension >= outputRank { + return nil, errors.Errorf("iotaDimension %d out of bounds for output rank %d", + iotaDimension, outputRank) + } + + // Validate bounds length + if len(bounds) != outputRank { + return nil, errors.Errorf("bounds length (%d) must match output rank (%d)", + len(bounds), outputRank) + } + + // Create output shape with dynamic dimensions and caller-provided bounds + resultShape := shapes.Make(dtype) + resultShape.Dimensions = make([]int, outputRank) + resultShape.DimensionBounds = make([]int, outputRank) + for i := range outputRank { + resultShape.Dimensions[i] = shapes.DimUnknown + resultShape.DimensionBounds[i] = bounds[i] + } + resultShape.EncodeBounds = true + + stmt := fn.addOp(op, resultShape, outputShape) + stmt.Attributes = map[string]any{ + "iota_dimension": int64(iotaDimension), + } + return stmt.Outputs[0], nil +} + +// DynamicGather performs a gather operation with dynamically-specified slice sizes. +// +// This is the dynamic version of Gather where slice_sizes is a 1D tensor instead of +// static values. All other parameters work the same as Gather. +// +// Parameters: +// - operand: the tensor to gather from. +// - startIndices: indices specifying where to gather from. +// - sliceSizes: a 1D tensor specifying the size of each slice (length = operand.Rank()). +// - indexVectorAxis: the axis in startIndices that contains the index vectors. +// - offsetOutputAxes: output axes corresponding to non-collapsed, non-batched operand axes. +// - collapsedSliceAxes: operand axes to collapse (must have slice size 1). +// - operandBatchingAxes: operand's batching axes. +// - startIndicesBatchingAxes: startIndices' batching axes. +// - startIndexMap: maps index vector elements to operand axes. +// - indicesAreSorted: whether indices are guaranteed sorted. +// - bounds: dimension bounds for the output shape. +func DynamicGather(operand, startIndices, sliceSizes *Value, indexVectorAxis int, + offsetOutputAxes, collapsedSliceAxes, operandBatchingAxes, + startIndicesBatchingAxes, startIndexMap []int, + indicesAreSorted bool, bounds []int) (*Value, error) { + op := optypes.DynamicGather + fn, err := innerMostFunction(operand, startIndices, sliceSizes) + if err != nil { + return nil, err + } + if fn.Returned { + return nil, errors.Errorf("cannot add operation %s after returning, in function %q", + op, fn.Name) + } + + // Validate sliceSizes is 1D tensor of integer type + if sliceSizes.shape.Rank() != 1 { + return nil, errors.Errorf("sliceSizes must be a 1D tensor, got rank %d", + sliceSizes.shape.Rank()) + } + if !sliceSizes.shape.DType.IsInt() { + return nil, errors.Errorf("sliceSizes must be integer type, got %s", + sliceSizes.shape.DType) + } + + // sliceSizes length should match operand rank + if sliceSizes.shape.Dimensions[0] != operand.shape.Rank() { + return nil, errors.Errorf("sliceSizes length (%d) must match operand rank (%d)", + sliceSizes.shape.Dimensions[0], operand.shape.Rank()) + } + + // Validate bounds length + if len(bounds) == 0 { + return nil, errors.Errorf("bounds must be provided for dynamic gather") + } + + // Calculate output rank: batch dimensions from startIndices + offset dimensions + // Output rank = len(offsetOutputAxes) + startIndices.Rank() - 1 + len(operandBatchingAxes) + startIndicesRank := startIndices.shape.Rank() + outputRank := len(offsetOutputAxes) + startIndicesRank - 1 + len(startIndicesBatchingAxes) + if len(bounds) != outputRank { + return nil, errors.Errorf("bounds length (%d) must match output rank (%d)", + len(bounds), outputRank) + } + + // Create output shape with dynamic dimensions + outputShape := shapes.Shape{DType: operand.shape.DType} + outputShape.Dimensions = make([]int, outputRank) + outputShape.DimensionBounds = make([]int, outputRank) + for i := range outputRank { + outputShape.Dimensions[i] = shapes.DimUnknown + outputShape.DimensionBounds[i] = bounds[i] + } + outputShape.EncodeBounds = true + + stmt := fn.addOp(op, outputShape, operand, startIndices, sliceSizes) + stmt.Attributes = map[string]any{ + "dimension_numbers": literalStrF( + "#stablehlo.gather<\n"+ + "\toffset_dims = %s,\n"+ + "\tcollapsed_slice_dims = %s,\n"+ + "\toperand_batching_dims = %s,\n"+ + "\tstart_indices_batching_dims = %s,\n"+ + "\tstart_index_map = %s,\n"+ + "\tindex_vector_dim = %d>", + intSliceToStableHLO(offsetOutputAxes), + intSliceToStableHLO(collapsedSliceAxes), + intSliceToStableHLO(operandBatchingAxes), + intSliceToStableHLO(startIndicesBatchingAxes), + intSliceToStableHLO(startIndexMap), + indexVectorAxis), + "indices_are_sorted": indicesAreSorted, + } + return stmt.Outputs[0], nil +} + +// DynamicPad pads the operand with dynamically-specified padding amounts. +// +// This is the dynamic version of Pad where the padding amounts are 1D tensors +// instead of static values. +// +// Parameters: +// - operand: the tensor to pad. +// - paddingValue: scalar value to use for padding (must match operand dtype). +// - edgePaddingLow: 1D tensor specifying low-end padding for each axis. +// - edgePaddingHigh: 1D tensor specifying high-end padding for each axis. +// - interiorPadding: 1D tensor specifying interior padding for each axis. +// - bounds: dimension bounds for the output shape. +func DynamicPad(operand, paddingValue, edgePaddingLow, edgePaddingHigh, interiorPadding *Value, bounds []int) (*Value, error) { + op := optypes.DynamicPad + fn, err := innerMostFunction(operand, paddingValue, edgePaddingLow, edgePaddingHigh, interiorPadding) + if err != nil { + return nil, err + } + if fn.Returned { + return nil, errors.Errorf("cannot add operation %s after returning, in function %q", + op, fn.Name) + } + + // Validate paddingValue is scalar + if !paddingValue.shape.IsScalar() { + return nil, errors.Errorf("paddingValue must be a scalar, got shape %s", paddingValue.shape) + } + if paddingValue.shape.DType != operand.shape.DType { + return nil, errors.Errorf("paddingValue dtype (%s) must match operand dtype (%s)", + paddingValue.shape.DType, operand.shape.DType) + } + + rank := operand.shape.Rank() + + // Validate padding tensors are 1D with correct length + for name, tensor := range map[string]*Value{ + "edgePaddingLow": edgePaddingLow, + "edgePaddingHigh": edgePaddingHigh, + "interiorPadding": interiorPadding, + } { + if tensor.shape.Rank() != 1 { + return nil, errors.Errorf("%s must be a 1D tensor, got rank %d", name, tensor.shape.Rank()) + } + if !tensor.shape.DType.IsInt() { + return nil, errors.Errorf("%s must be integer type, got %s", name, tensor.shape.DType) + } + if tensor.shape.Dimensions[0] != rank { + return nil, errors.Errorf("%s length (%d) must match operand rank (%d)", + name, tensor.shape.Dimensions[0], rank) + } + } + + // Validate bounds length + if len(bounds) != rank { + return nil, errors.Errorf("bounds length (%d) must match operand rank (%d)", + len(bounds), rank) + } + + // Create output shape with dynamic dimensions + outputShape := operand.shape.Clone() + outputShape.Dimensions = make([]int, rank) + outputShape.DimensionBounds = make([]int, rank) + for i := range rank { + outputShape.Dimensions[i] = shapes.DimUnknown + outputShape.DimensionBounds[i] = bounds[i] + } + outputShape.EncodeBounds = true + + stmt := fn.addOp(op, outputShape, operand, paddingValue, edgePaddingLow, edgePaddingHigh, interiorPadding) + return stmt.Outputs[0], nil +} + +// DynamicConv performs a convolution with dynamically-specified padding. +// +// This is the dynamic version of Convolution where the padding amounts are specified +// via a tensor instead of static values. All other parameters work the same as Convolution. +// +// Parameters: +// - input: the input tensor. +// - kernel: the convolution kernel/filter. +// - padding: a 2D tensor of shape [numSpatialDims, 2] specifying [low, high] padding per spatial axis. +// - strides, inputDilations, kernelDilations: same as Convolution (static). +// - axis configuration parameters: same as Convolution. +// - bounds: dimension bounds for the output shape. +func DynamicConv(input, kernel, padding *Value, + strides, inputDilations, kernelDilations []int, + inputBatchAxis, inputChannelsAxis int, inputSpatialAxes []int, + kernelInputChannelsAxis, kernelOutputChannelsAxis int, kernelSpatialAxes []int, + outputBatchAxis, outputChannelsAxis int, outputSpatialAxes []int, + channelGroupCount, batchGroupCount int, + inputPrecision, kernelPrecision types.DotGeneralPrecisionType, + bounds []int) (*Value, error) { + op := optypes.DynamicConv + fn, err := innerMostFunction(input, kernel, padding) + if err != nil { + return nil, err + } + if fn.Returned { + return nil, errors.Errorf("cannot add operation %s after returning, in function %q", + op, fn.Name) + } + + rank := input.shape.Rank() + rankSpatial := rank - 2 + + // Validate padding tensor shape: should be [numSpatialDims, 2] + if padding.shape.Rank() != 2 { + return nil, errors.Errorf("padding must be a 2D tensor, got rank %d", padding.shape.Rank()) + } + if padding.shape.Dimensions[0] != rankSpatial || padding.shape.Dimensions[1] != 2 { + return nil, errors.Errorf("padding must have shape [%d, 2], got %v", + rankSpatial, padding.shape.Dimensions) + } + + // Set default for any missing slices + windowReversal := make([]bool, rankSpatial) + for _, s := range []*[]int{&strides, &inputDilations, &kernelDilations} { + if len(*s) == 0 { + *s = slices.Repeat([]int{1}, rankSpatial) + } + } + + // Fix negative axes + for _, axisConfig := range []*int{&inputBatchAxis, &inputChannelsAxis, &kernelInputChannelsAxis, &kernelOutputChannelsAxis, &outputBatchAxis, &outputChannelsAxis} { + adjustedAxis, err := shapeinference.AdjustAxisToRank(*axisConfig, rank) + if err != nil { + return nil, errors.Errorf("invalid channel/batch axis %d for rank %d", *axisConfig, rank) + } + *axisConfig = adjustedAxis + } + for _, s := range []*[]int{&inputSpatialAxes, &kernelSpatialAxes, &outputSpatialAxes} { + *s = slices.Clone(*s) + for i, axis := range *s { + adjustedAxis, err := shapeinference.AdjustAxisToRank(axis, rank) + if err != nil { + return nil, errors.Errorf("invalid spatial axis %d for rank %d", axis, rank) + } + (*s)[i] = adjustedAxis + } + } + + // Validate bounds length + if len(bounds) != rank { + return nil, errors.Errorf("bounds length (%d) must match output rank (%d)", + len(bounds), rank) + } + + // Create output shape with dynamic dimensions + outputShape := shapes.Shape{DType: input.shape.DType} + outputShape.Dimensions = make([]int, rank) + outputShape.DimensionBounds = make([]int, rank) + for i := range rank { + outputShape.Dimensions[i] = shapes.DimUnknown + outputShape.DimensionBounds[i] = bounds[i] + } + outputShape.EncodeBounds = true + + // Build statement + stmt := fn.addOp(op, outputShape, input, kernel, padding) + + convConfig := getConvAxesConfig(inputBatchAxis, inputChannelsAxis, inputSpatialAxes, + kernelInputChannelsAxis, kernelOutputChannelsAxis, kernelSpatialAxes, + outputBatchAxis, outputChannelsAxis, outputSpatialAxes) + precisionConfig := literalStrF("[#stablehlo, #stablehlo]", + inputPrecision.ToStableHLO(), kernelPrecision.ToStableHLO()) + + stmt.Attributes = map[string]any{ + "window_strides": intSliceToArrayI64StableHLO(strides), + "lhs_dilation": intSliceToArrayI64StableHLO(inputDilations), + "rhs_dilation": intSliceToArrayI64StableHLO(kernelDilations), + "window_reversal": boolSliceToArrayI1StableHLO(windowReversal), + "dimension_numbers": convConfig, + "feature_group_count": int64(channelGroupCount), + "batch_group_count": int64(batchGroupCount), + "precision_config": precisionConfig, + } + return stmt.Outputs[0], nil +} + // While executes body repeatedly while condition returns true. // // The While operation implements a loop that continues executing the body function diff --git a/pkg/stablehlo/stablehlo_test.go b/pkg/stablehlo/stablehlo_test.go index 6913638..c1a7ff6 100644 --- a/pkg/stablehlo/stablehlo_test.go +++ b/pkg/stablehlo/stablehlo_test.go @@ -221,7 +221,7 @@ func TestDynamicOperations(t *testing.T) { } program := string(must1(b.Build())) - fmt.Printf("%s program:\n%s", t.Name(), program) + fmt.Printf("%s program:\n%s\n", t.Name(), program) // Verify the operation is present in the program if !strings.Contains(program, "stablehlo.get_dimension_size") { @@ -232,4 +232,195 @@ func TestDynamicOperations(t *testing.T) { } }) + t.Run("DynamicReshape", func(t *testing.T) { + b := New(t.Name()) + fn := b.Main() + + // Create input tensor [2, 3] = 6 elements + operand := must1(fn.NamedInput("operand", shapes.Make(dtypes.Float32, 2, 3))) + + // Create shape tensor specifying new shape [3, 2] - a 1D tensor with 2 elements + newShape := must1(fn.ConstantFromFlatAndDimensions([]int64{3, 2}, 2)) + + // Dynamic reshape with bounds + reshaped := must1(DynamicReshape(operand, newShape, []int{3, 2})) + + if err := fn.Return(reshaped); err != nil { + t.Fatalf("expected no error, got %v", err) + } + + program := string(must1(b.Build())) + fmt.Printf("%s program:\n%s\n", t.Name(), program) + + if !strings.Contains(program, "stablehlo.dynamic_reshape") { + t.Fatal("expected program to contain dynamic_reshape operation") + } + // Check for bounded dynamic shape encoding + if !strings.Contains(program, "#stablehlo.bounds") { + t.Fatal("expected program to contain bounds encoding") + } + }) + + t.Run("DynamicBroadcastInDim", func(t *testing.T) { + b := New(t.Name()) + fn := b.Main() + + // Create input tensor [3] + operand := must1(fn.NamedInput("operand", shapes.Make(dtypes.Float32, 3))) + + // Create output dimensions tensor [2, 3] - a 1D tensor with 2 elements + outputDims := must1(fn.ConstantFromFlatAndDimensions([]int64{2, 3}, 2)) + + // Broadcast with dimension mapping: input dim 0 -> output dim 1 + broadcasted := must1(DynamicBroadcastInDim(operand, outputDims, []int{1}, []int{2, 3})) + + if err := fn.Return(broadcasted); err != nil { + t.Fatalf("expected no error, got %v", err) + } + + program := string(must1(b.Build())) + fmt.Printf("%s program:\n%s\n", t.Name(), program) + + if !strings.Contains(program, "stablehlo.dynamic_broadcast_in_dim") { + t.Fatal("expected program to contain dynamic_broadcast_in_dim operation") + } + if !strings.Contains(program, "broadcast_dimensions") { + t.Fatal("expected program to contain broadcast_dimensions attribute") + } + }) + + t.Run("DynamicIota", func(t *testing.T) { + b := New(t.Name()) + fn := b.Main() + + // Create shape tensor [4, 3] - a 1D tensor with 2 elements + outputShape := must1(fn.ConstantFromFlatAndDimensions([]int64{4, 3}, 2)) + + // Create iota along dimension 0 + iota := must1(DynamicIota(fn, dtypes.Int32, outputShape, 0, []int{4, 3})) + + if err := fn.Return(iota); err != nil { + t.Fatalf("expected no error, got %v", err) + } + + program := string(must1(b.Build())) + fmt.Printf("%s program:\n%s\n", t.Name(), program) + + if !strings.Contains(program, "stablehlo.dynamic_iota") { + t.Fatal("expected program to contain dynamic_iota operation") + } + if !strings.Contains(program, "iota_dimension = 0") { + t.Fatal("expected program to contain iota_dimension attribute") + } + }) + + t.Run("DynamicGather", func(t *testing.T) { + b := New(t.Name()) + fn := b.Main() + + // Create operand [3, 4, 2] + operand := must1(fn.NamedInput("operand", shapes.Make(dtypes.Float32, 3, 4, 2))) + + // Create start indices [2, 3, 2] - 2 batches of 3 indices, each index has 2 elements + startIndices := must1(fn.NamedInput("indices", shapes.Make(dtypes.Int64, 2, 3, 2))) + + // Create dynamic slice sizes [1, 1, 2] - a 1D tensor with 3 elements + sliceSizes := must1(fn.ConstantFromFlatAndDimensions([]int64{1, 1, 2}, 3)) + + // Gather with dimension configuration + gathered := must1(DynamicGather(operand, startIndices, sliceSizes, + 2, // indexVectorAxis + []int{2}, // offsetOutputAxes + []int{0, 1}, // collapsedSliceAxes + []int{}, // operandBatchingAxes + []int{}, // startIndicesBatchingAxes + []int{0, 1}, // startIndexMap + false, // indicesAreSorted + []int{2, 3, 2})) // bounds + + if err := fn.Return(gathered); err != nil { + t.Fatalf("expected no error, got %v", err) + } + + program := string(must1(b.Build())) + fmt.Printf("%s program:\n%s\n", t.Name(), program) + + if !strings.Contains(program, "stablehlo.dynamic_gather") { + t.Fatal("expected program to contain dynamic_gather operation") + } + if !strings.Contains(program, "#stablehlo.gather") { + t.Fatal("expected program to contain gather dimension_numbers") + } + }) + + t.Run("DynamicPad", func(t *testing.T) { + b := New(t.Name()) + fn := b.Main() + + // Create operand [2, 3] + operand := must1(fn.NamedInput("operand", shapes.Make(dtypes.Float32, 2, 3))) + + // Create padding value + padValue := must1(fn.ConstantFromScalar(float32(0))) + + // Create padding tensors - 1D tensors with 2 elements each + edgeLow := must1(fn.ConstantFromFlatAndDimensions([]int64{1, 2}, 2)) + edgeHigh := must1(fn.ConstantFromFlatAndDimensions([]int64{1, 2}, 2)) + interior := must1(fn.ConstantFromFlatAndDimensions([]int64{0, 0}, 2)) + + // Dynamic pad with bounds for output shape [4, 7] + padded := must1(DynamicPad(operand, padValue, edgeLow, edgeHigh, interior, []int{4, 7})) + + if err := fn.Return(padded); err != nil { + t.Fatalf("expected no error, got %v", err) + } + + program := string(must1(b.Build())) + fmt.Printf("%s program:\n%s\n", t.Name(), program) + + if !strings.Contains(program, "stablehlo.dynamic_pad") { + t.Fatal("expected program to contain dynamic_pad operation") + } + }) + + t.Run("DynamicConv", func(t *testing.T) { + b := New(t.Name()) + fn := b.Main() + + // Create input [1, 4, 4, 1] - NHWC format + input := must1(fn.NamedInput("input", shapes.Make(dtypes.Float32, 1, 4, 4, 1))) + + // Create kernel [3, 3, 1, 1] - HWIO format + kernel := must1(fn.NamedInput("kernel", shapes.Make(dtypes.Float32, 3, 3, 1, 1))) + + // Create dynamic padding tensor [2, 2] for 2 spatial dims with [low, high] each + padding := must1(fn.ConstantFromFlatAndDimensions([]int64{1, 1, 1, 1}, 2, 2)) + + // Dynamic conv with bounds + convolved := must1(DynamicConv(input, kernel, padding, + []int{1, 1}, // strides + []int{1, 1}, // inputDilations + []int{1, 1}, // kernelDilations + 0, 3, []int{1, 2}, // input: batch=0, channels=3, spatial=[1,2] + 2, 3, []int{0, 1}, // kernel: inChannels=2, outChannels=3, spatial=[0,1] + 0, 3, []int{1, 2}, // output: batch=0, channels=3, spatial=[1,2] + 1, 1, // channelGroupCount, batchGroupCount + 0, 0, // precision (DEFAULT) + []int{1, 4, 4, 1})) // bounds + + if err := fn.Return(convolved); err != nil { + t.Fatalf("expected no error, got %v", err) + } + + program := string(must1(b.Build())) + fmt.Printf("%s program:\n%s\n", t.Name(), program) + + if !strings.Contains(program, "stablehlo.dynamic_conv") { + t.Fatal("expected program to contain dynamic_conv operation") + } + if !strings.Contains(program, "dimension_numbers") { + t.Fatal("expected program to contain dimension_numbers attribute") + } + }) + } diff --git a/pkg/types/shapes/shape.go b/pkg/types/shapes/shape.go index e5b115f..ac2ca16 100644 --- a/pkg/types/shapes/shape.go +++ b/pkg/types/shapes/shape.go @@ -87,10 +87,12 @@ const DimUnknown = -1 // // Use Make to create a new shape. See example in package shapes documentation. type Shape struct { - DType dtypes.DType - Dimensions []int - TupleShapes []Shape // Shapes of the tuple, if this is a tuple. - Quantization *Quantization // Quantization metadata for quantized types. + DType dtypes.DType + Dimensions []int + DimensionBounds []int // Upper bounds for dynamic dimensions (when Dimensions[i] == DimUnknown). nil or len(DimensionBounds)==len(Dimensions). 0 means no bound. + TupleShapes []Shape // Shapes of the tuple, if this is a tuple. + Quantization *Quantization // Quantization metadata for quantized types. + EncodeBounds bool // Whether to encode DimensionBounds in StableHLO output. Set to true for ops that require bounded dynamism (e.g., dynamic_reshape). } // Make returns a Shape structure filled with the values given. @@ -134,6 +136,20 @@ func (s Shape) IsDynamic() bool { return slices.Contains(s.Dimensions, DimUnknown) } +// HasBoundedDynamism returns true if the shape has at least one dynamic dimension +// with a known upper bound. This is used for StableHLO bounds encoding. +func (s Shape) HasBoundedDynamism() bool { + if len(s.DimensionBounds) == 0 { + return false + } + for i, dim := range s.Dimensions { + if dim == DimUnknown && i < len(s.DimensionBounds) && s.DimensionBounds[i] > 0 { + return true + } + } + return false +} + // Dim returns the dimension of the given axis. axis can take negative numbers, in which // case it counts as starting from the end -- so axis=-1 refers to the last axis. // Like with a slice indexing, it panics for an out-of-bound axis. @@ -263,6 +279,7 @@ func (s Shape) EqualDimensions(s2 Shape) bool { func (s Shape) Clone() (s2 Shape) { s2.DType = s.DType s2.Dimensions = slices.Clone(s.Dimensions) + s2.DimensionBounds = slices.Clone(s.DimensionBounds) if s.TupleSize() > 0 { s2.TupleShapes = make([]Shape, 0, len(s.TupleShapes)) for _, subShape := range s.TupleShapes { @@ -270,6 +287,7 @@ func (s Shape) Clone() (s2 Shape) { } } s2.Quantization = s.Quantization.Clone() + s2.EncodeBounds = s.EncodeBounds return } diff --git a/pkg/types/shapes/stablehlo.go b/pkg/types/shapes/stablehlo.go index 700931a..18ae06a 100644 --- a/pkg/types/shapes/stablehlo.go +++ b/pkg/types/shapes/stablehlo.go @@ -62,6 +62,28 @@ func (s Shape) WriteStableHLO(writer io.Writer) error { w("%s", s.DType.ToStableHLO()) } + // Encode bounds only if explicitly requested via EncodeBounds flag. + // XLA requires bounded dynamic dimensions for compilation. + if s.EncodeBounds && s.HasBoundedDynamism() { + w(", #stablehlo.bounds<") + for i, dim := range s.Dimensions { + if i > 0 { + w(", ") + } + if dim == DimUnknown { + // Dynamic dimension - check if we have a bound + if i < len(s.DimensionBounds) && s.DimensionBounds[i] > 0 { + w("%d", s.DimensionBounds[i]) + } else { + w("?") + } + } else { + // Static dimension - use ? (no bound needed, the dimension is already known) + w("?") + } + } + w(">") + } w(">") return err }