diff --git a/stablehlo/dialect/Base.cpp b/stablehlo/dialect/Base.cpp index e5cf93f97d..0545ead529 100644 --- a/stablehlo/dialect/Base.cpp +++ b/stablehlo/dialect/Base.cpp @@ -611,5 +611,29 @@ bool isSplatArray(ArrayRef arr, int64_t val) { [val](int64_t x) { return x == val; }); } +mlir::Speculation::Speculatability getShapedSpeculatability(Operation* op, + int64_t count) { + auto resultType = cast(op->getResult(0).getType()); + // The result type's shape is fully dynamic, so there cannot be a mismatch + // with the output shape operand at runtime (the type has no expectations). + if (llvm::all_of(llvm::seq(resultType.getRank()), + [&](int64_t i) { return resultType.isDynamicDim(i); })) + return mlir::Speculation::Speculatable; + + // If all inputs are static and the shape-related operands are constant + // then any relationship between the input and the shapes can be + // verified statically. Shapes must be static due to ODS constraints. + bool allInputsStatic = llvm::all_of(op->getOperandTypes(), [](Type t) { + return cast(t).hasStaticShape(); + }); + bool allShapesConstant = llvm::all_of(llvm::seq(count), [&](int64_t i) { + return matchPattern(op->getOperand(op->getNumOperands() - 1 - i), + m_Constant()); + }); + return allInputsStatic && allShapesConstant + ? mlir::Speculation::Speculatable + : mlir::Speculation::NotSpeculatable; +} + } // namespace hlo } // namespace mlir diff --git a/stablehlo/dialect/Base.h b/stablehlo/dialect/Base.h index 2d621c8f42..8ffa4df803 100644 --- a/stablehlo/dialect/Base.h +++ b/stablehlo/dialect/Base.h @@ -251,6 +251,13 @@ void writeEnumAttribute(EnumTypeAttr val, DialectBytecodeWriter &writer) { } } // namespace bytecode +// Determines the speculatability for a shaped operation `op` with `count` +// shape operands. The last `count` operands are assumed to be shape operands. +// To be speculatable, such an op must either have a fully dynamic result type +// or have only static inputs and constant shape operands. +mlir::Speculation::Speculatability getShapedSpeculatability(Operation *op, + int64_t count); + namespace OpTrait { template @@ -456,24 +463,7 @@ struct SpeculatableIfAllInputsStaticAndShapeConstantImplTrait SpeculatableIfAllInputsStaticAndShapeConstantImplTrait> { mlir::Speculation::Speculatability getSpeculatability() { auto op = this->getOperation(); - auto resultType = cast(op->getResult(0).getType()); - // The result type's shape is fully dynamic, so there cannot be a mismatch - // with the output shape operand at runtime (the type has no expectations). - if (llvm::all_of(llvm::seq(resultType.getRank()), - [&](int64_t i) { return resultType.isDynamicDim(i); })) - return mlir::Speculation::Speculatable; - - // If all inputs are static and the output shape (last operand) is constant, - // then any relationship between the input and the output shape can be - // verified statically. The shape operand is known to be static due to ODS - // constraints. - bool allInputsStatic = llvm::all_of(op->getOperandTypes(), [](Type t) { - return cast(t).hasStaticShape(); - }); - if (allInputsStatic && - matchPattern(op->getOperand(op->getNumOperands() - 1), m_Constant())) - return mlir::Speculation::Speculatable; - return mlir::Speculation::NotSpeculatable; + return getShapedSpeculatability(op, 1); } }; diff --git a/stablehlo/dialect/StablehloOps.cpp b/stablehlo/dialect/StablehloOps.cpp index 7df4612eaa..fa7c31c9f3 100644 --- a/stablehlo/dialect/StablehloOps.cpp +++ b/stablehlo/dialect/StablehloOps.cpp @@ -812,6 +812,21 @@ LogicalResult DynamicGatherOp::inferReturnTypeComponents( adaptor.getDimensionNumbers().getIndexVectorDim(), inferredReturnShapes); } +mlir::Speculation::Speculatability DynamicGatherOp::getSpeculatability() { + // When indices_are_sorted is true, if the start_indices are not sorted, the + // behavior is undefined. + // A possible improvement would be to check if the start_indices are constant + // and if they are sorted, do not return NotSpeculatable. However, such a + // check could be somewhat costly and has unclear ROI. + if (getIndicesAreSorted()) return mlir::Speculation::NotSpeculatable; + bool allOperandsStatic = llvm::all_of( + this->getOperation()->getOperandTypes(), + [](Type t) { return cast(t).hasStaticShape(); }); + return allOperandsStatic && matchPattern(getSliceSizes(), m_Constant()) + ? mlir::Speculation::Speculatable + : mlir::Speculation::NotSpeculatable; +} + //===----------------------------------------------------------------------===// // GetDimensionSizeOp //===----------------------------------------------------------------------===// @@ -1557,6 +1572,10 @@ LogicalResult RealDynamicSliceOp::reifyReturnTypeShapes( return success(); } +mlir::Speculation::Speculatability RealDynamicSliceOp::getSpeculatability() { + return hlo::getShapedSpeculatability(getOperation(), /*count=*/3); +} + //===----------------------------------------------------------------------===// // InfeedOp //===----------------------------------------------------------------------===// @@ -2085,6 +2104,10 @@ LogicalResult DynamicPadOp::reifyReturnTypeShapes( return success(); } +mlir::Speculation::Speculatability DynamicPadOp::getSpeculatability() { + return hlo::getShapedSpeculatability(getOperation(), /*count=*/3); +} + //===----------------------------------------------------------------------===// // ReshapeOp //===----------------------------------------------------------------------===// diff --git a/stablehlo/dialect/StablehloOps.td b/stablehlo/dialect/StablehloOps.td index 9135751135..b8ca5912a5 100644 --- a/stablehlo/dialect/StablehloOps.td +++ b/stablehlo/dialect/StablehloOps.td @@ -1345,7 +1345,7 @@ def StableHLO_AllGatherOp : StableHLO_Op<"all_gather", let results = (outs HLO_Tensor); let hasVerifier = 1; - let extraClassDeclaration = [{ + let extraClassDeclaration = commonClassDeclaration # [{ /// Interface method for ConditionallySpeculatable. mlir::Speculation::Speculatability getSpeculatability(); }]; @@ -3348,7 +3348,8 @@ def StableHLO_ReducePrecisionOp : StableHLO_Op<"reduce_precision", def StableHLO_RealDynamicSliceOp: StableHLO_ShapedInterfaceOp< "real_dynamic_slice", - [Pure, AllElementTypesMatch<["operand", "result"]>, + [ConditionallySpeculatable, NoMemoryEffect, + AllElementTypesMatch<["operand", "result"]>, AllTypesMatch<["start_indices", "limit_indices", "strides"]>]> { let summary = "RealDynamicSlice operation"; let description = [{ @@ -3376,10 +3377,16 @@ def StableHLO_RealDynamicSliceOp: StableHLO_ShapedInterfaceOp< let hasVerifier = 1; let assemblyFormat = "operands attr-dict `:` functional-type(operands, results)"; + + let extraClassDeclaration = commonClassDeclaration # [{ + /// Interface method for ConditionallySpeculatable. + mlir::Speculation::Speculatability getSpeculatability(); + }]; } def StableHLO_DynamicPadOp: StableHLO_ShapedInterfaceOp<"dynamic_pad", - [Pure, AllElementTypesMatch<["operand", "padding_value", "result"]>, + [ConditionallySpeculatable, NoMemoryEffect, + AllElementTypesMatch<["operand", "padding_value", "result"]>, AllTypesMatch<["edge_padding_low", "edge_padding_high", "interior_padding"]>]> { let summary = "DynamicPad operation"; let description = [{ @@ -3413,10 +3420,15 @@ def StableHLO_DynamicPadOp: StableHLO_ShapedInterfaceOp<"dynamic_pad", let hasVerifier = 1; let assemblyFormat = "operands attr-dict `:` functional-type(operands, results)"; + + let extraClassDeclaration = commonClassDeclaration # [{ + /// Interface method for ConditionallySpeculatable. + mlir::Speculation::Speculatability getSpeculatability(); + }]; } def StableHLO_DynamicGatherOp: StableHLO_Op<"dynamic_gather", - [InferTensorTypeWithReify, Pure]> { + [ConditionallySpeculatable, NoMemoryEffect, InferTensorTypeWithReify]> { let summary = "DynamicGather operation"; let description = [{ This operation is a work in progress, so it is not yet included in @@ -3447,9 +3459,15 @@ def StableHLO_DynamicGatherOp: StableHLO_Op<"dynamic_gather", DefaultValuedOptionalAttr:$indices_are_sorted ); let results = (outs HLO_Tensor); + + let extraClassDeclaration = commonClassDeclaration # [{ + /// Interface method for ConditionallySpeculatable. + mlir::Speculation::Speculatability getSpeculatability(); + }]; } -def StableHLO_DynamicConvOp : StableHLO_Op<"dynamic_conv", [Pure]> { +def StableHLO_DynamicConvOp : StableHLO_Op<"dynamic_conv", + [HLO_SpeculatableIfAllInputsStaticAndShapeConstant, NoMemoryEffect]> { let summary = "DynamicConv operation"; let description = [{ This operation is a work in progress, so it is not yet included in diff --git a/stablehlo/tests/ops_speculatability.mlir b/stablehlo/tests/ops_speculatability.mlir index 235b61282b..b386134ece 100644 --- a/stablehlo/tests/ops_speculatability.mlir +++ b/stablehlo/tests/ops_speculatability.mlir @@ -1717,6 +1717,148 @@ func.func @dynamic_broadcast_in_dim( // ----- +// CHECK-LABEL: func @dynamic_conv +// CHECK-NEXT: return +func.func @dynamic_conv( + %static_input: tensor<100x26x26x32xf64>, %static_kernel: tensor<3x3x1x32xf64>, + %dynamic_input: tensor, %dynamic_kernel: tensor, + %unknown_shape: tensor<2x2xi32> +) { + %constant_shape = stablehlo.constant dense<2> : tensor<2x2xi32> + + // Static inputs, constant shape + %0 = "stablehlo.dynamic_conv"(%static_input, %static_kernel, %constant_shape) { + dimension_numbers = #stablehlo.conv<[b, 0, 1, f]x[0, 1, o, i]->[b, 0, 1, f]>, + window_strides = array, lhs_dilation = array, rhs_dilation = array, + feature_group_count = 1 : i64, batch_group_count = 1 : i64 + } : (tensor<100x26x26x32xf64>, tensor<3x3x1x32xf64>, tensor<2x2xi32>) -> tensor<100x28x28x1xf64> + "hlo_test_speculatability.is_speculatable"(%0) : (tensor<100x28x28x1xf64>) -> () + %1 = "stablehlo.dynamic_conv"(%static_input, %static_kernel, %constant_shape) { + dimension_numbers = #stablehlo.conv<[b, 0, 1, f]x[0, 1, o, i]->[b, 0, 1, f]>, + window_strides = array, lhs_dilation = array, rhs_dilation = array, + feature_group_count = 1 : i64, batch_group_count = 1 : i64 + } : (tensor<100x26x26x32xf64>, tensor<3x3x1x32xf64>, tensor<2x2xi32>) -> tensor + "hlo_test_speculatability.is_speculatable"(%1) : (tensor) -> () + + // Dynamic input, static kernel, constant shape + %2 = "stablehlo.dynamic_conv"(%dynamic_input, %static_kernel, %constant_shape) { + dimension_numbers = #stablehlo.conv<[b, 0, 1, f]x[0, 1, o, i]->[b, 0, 1, f]>, + window_strides = array, lhs_dilation = array, rhs_dilation = array, + feature_group_count = 1 : i64, batch_group_count = 1 : i64 + } : (tensor, tensor<3x3x1x32xf64>, tensor<2x2xi32>) -> tensor<100x28x28x1xf64> + "hlo_test_speculatability.is_not_speculatable"(%2) : (tensor<100x28x28x1xf64>) -> () + %3 = "stablehlo.dynamic_conv"(%dynamic_input, %static_kernel, %constant_shape) { + dimension_numbers = #stablehlo.conv<[b, 0, 1, f]x[0, 1, o, i]->[b, 0, 1, f]>, + window_strides = array, lhs_dilation = array, rhs_dilation = array, + feature_group_count = 1 : i64, batch_group_count = 1 : i64 + } : (tensor, tensor<3x3x1x32xf64>, tensor<2x2xi32>) -> tensor + "hlo_test_speculatability.is_speculatable"(%3) : (tensor) -> () + + // Static input, dynamic kernel, constant shape + %4 = "stablehlo.dynamic_conv"(%static_input, %dynamic_kernel, %constant_shape) { + dimension_numbers = #stablehlo.conv<[b, 0, 1, f]x[0, 1, o, i]->[b, 0, 1, f]>, + window_strides = array, lhs_dilation = array, rhs_dilation = array, + feature_group_count = 1 : i64, batch_group_count = 1 : i64 + } : (tensor<100x26x26x32xf64>, tensor, tensor<2x2xi32>) -> tensor<100x28x28x1xf64> + "hlo_test_speculatability.is_not_speculatable"(%4) : (tensor<100x28x28x1xf64>) -> () + %5 = "stablehlo.dynamic_conv"(%static_input, %dynamic_kernel, %constant_shape) { + dimension_numbers = #stablehlo.conv<[b, 0, 1, f]x[0, 1, o, i]->[b, 0, 1, f]>, + window_strides = array, lhs_dilation = array, rhs_dilation = array, + feature_group_count = 1 : i64, batch_group_count = 1 : i64 + } : (tensor<100x26x26x32xf64>, tensor, tensor<2x2xi32>) -> tensor + "hlo_test_speculatability.is_speculatable"(%5) : (tensor) -> () + + // Static input, static kernel, unknown shape + %6 = "stablehlo.dynamic_conv"(%static_input, %static_kernel, %unknown_shape) { + dimension_numbers = #stablehlo.conv<[b, 0, 1, f]x[0, 1, o, i]->[b, 0, 1, f]>, + window_strides = array, lhs_dilation = array, rhs_dilation = array, + feature_group_count = 1 : i64, batch_group_count = 1 : i64 + } : (tensor<100x26x26x32xf64>, tensor<3x3x1x32xf64>, tensor<2x2xi32>) -> tensor<100x28x28x1xf64> + "hlo_test_speculatability.is_not_speculatable"(%6) : (tensor<100x28x28x1xf64>) -> () + %7 = "stablehlo.dynamic_conv"(%static_input, %static_kernel, %unknown_shape) { + dimension_numbers = #stablehlo.conv<[b, 0, 1, f]x[0, 1, o, i]->[b, 0, 1, f]>, + window_strides = array, lhs_dilation = array, rhs_dilation = array, + feature_group_count = 1 : i64, batch_group_count = 1 : i64 + } : (tensor<100x26x26x32xf64>, tensor<3x3x1x32xf64>, tensor<2x2xi32>) -> tensor + "hlo_test_speculatability.is_speculatable"(%7) : (tensor) -> () + + return +} + +// ----- + +// CHECK-LABEL: func @dynamic_gather +// CHECK-NEXT: return +func.func @dynamic_gather( + %static_input: tensor<3x4x2xi32>, %static_indices: tensor<2x3x2xi64>, + %dynamic_input: tensor, %dynamic_indices: tensor, + %unknown_slice_sizes: tensor<3xi32> +) { + %constant_slice_sizes = stablehlo.constant dense<[1, 2, 2]> : tensor<3xi32> + + // Static inputs, constant shape + %0 = "stablehlo.dynamic_gather"(%static_input, %static_indices, %constant_slice_sizes) { + dimension_numbers = #stablehlo.gather< + offset_dims = [2, 3], + collapsed_slice_dims = [0], + start_index_map = [1, 0], + index_vector_dim = 2>, + slice_sizes = array, + indices_are_sorted = false + } : (tensor<3x4x2xi32>, tensor<2x3x2xi64>, tensor<3xi32>) -> tensor + "hlo_test_speculatability.is_speculatable"(%0) : (tensor) -> () + %1 = "stablehlo.dynamic_gather"(%static_input, %static_indices, %constant_slice_sizes) { + dimension_numbers = #stablehlo.gather< + offset_dims = [2, 3], + collapsed_slice_dims = [0], + start_index_map = [1, 0], + index_vector_dim = 2>, + slice_sizes = array, + indices_are_sorted = true + } : (tensor<3x4x2xi32>, tensor<2x3x2xi64>, tensor<3xi32>) -> tensor + "hlo_test_speculatability.is_not_speculatable"(%1) : (tensor) -> () + + // Dynamic input, static start_indices, constant slice_sizes + %2 = "stablehlo.dynamic_gather"(%dynamic_input, %static_indices, %constant_slice_sizes) { + dimension_numbers = #stablehlo.gather< + offset_dims = [2, 3], + collapsed_slice_dims = [0], + start_index_map = [1, 0], + index_vector_dim = 2>, + slice_sizes = array, + indices_are_sorted = false + } : (tensor, tensor<2x3x2xi64>, tensor<3xi32>) -> tensor + "hlo_test_speculatability.is_not_speculatable"(%2) : (tensor) -> () + + // Static input, dynamic start_indices, constant slice_sizes + %3 = "stablehlo.dynamic_gather"(%static_input, %dynamic_indices, %constant_slice_sizes) { + dimension_numbers = #stablehlo.gather< + offset_dims = [2, 3], + collapsed_slice_dims = [0], + start_index_map = [1, 0], + index_vector_dim = 2>, + slice_sizes = array, + indices_are_sorted = false + } : (tensor<3x4x2xi32>, tensor, tensor<3xi32>) -> tensor + "hlo_test_speculatability.is_not_speculatable"(%3) : (tensor) -> () + + // Static input, static start_indices, unknown slice_sizes + %4 = "stablehlo.dynamic_gather"(%static_input, %static_indices, %unknown_slice_sizes) { + dimension_numbers = #stablehlo.gather< + offset_dims = [2, 3], + collapsed_slice_dims = [0], + start_index_map = [1, 0], + index_vector_dim = 2>, + slice_sizes = array, + indices_are_sorted = false + } : (tensor<3x4x2xi32>, tensor<2x3x2xi64>, tensor<3xi32>) -> tensor + "hlo_test_speculatability.is_not_speculatable"(%4) : (tensor) -> () + + return +} + +// ----- + // CHECK-LABEL: func @dynamic_iota // CHECK-NEXT: return func.func @dynamic_iota(%unknown_shape: tensor<2xi32>) { @@ -1742,6 +1884,65 @@ func.func @dynamic_iota(%unknown_shape: tensor<2xi32>) { // ----- +// CHECK-LABEL: func @dynamic_pad +// CHECK-NEXT: return +func.func @dynamic_pad( + %static_arg: tensor<4xf64>, %dynamic_arg: tensor, + %padding_value: tensor, %unknown_padding: tensor<1xi32> +) { + %constant_padding = stablehlo.constant dense<0> : tensor<1xi32> + + // Static input, constant padding + %0 = stablehlo.dynamic_pad %static_arg, %padding_value, + %constant_padding, %constant_padding, %constant_padding + : (tensor<4xf64>, tensor, tensor<1xi32>, tensor<1xi32>, tensor<1xi32>) -> tensor<4xf64> + "hlo_test_speculatability.is_speculatable"(%0) : (tensor<4xf64>) -> () + %1 = stablehlo.dynamic_pad %static_arg, %padding_value, + %constant_padding, %constant_padding, %constant_padding + : (tensor<4xf64>, tensor, tensor<1xi32>, tensor<1xi32>, tensor<1xi32>) -> tensor + "hlo_test_speculatability.is_speculatable"(%1) : (tensor) -> () + + // Dynamic input, constant padding + %2 = stablehlo.dynamic_pad %dynamic_arg, %padding_value, + %unknown_padding, %unknown_padding, %unknown_padding + : (tensor, tensor, tensor<1xi32>, tensor<1xi32>, tensor<1xi32>) -> tensor<4xf64> + "hlo_test_speculatability.is_not_speculatable"(%2) : (tensor<4xf64>) -> () + %3 = stablehlo.dynamic_pad %dynamic_arg, %padding_value, + %unknown_padding, %unknown_padding, %unknown_padding + : (tensor, tensor, tensor<1xi32>, tensor<1xi32>, tensor<1xi32>) -> tensor + "hlo_test_speculatability.is_speculatable"(%3) : (tensor) -> () + + // Static input, unknown paddings + %4 = stablehlo.dynamic_pad %static_arg, %padding_value, + %unknown_padding, %constant_padding, %constant_padding + : (tensor<4xf64>, tensor, tensor<1xi32>, tensor<1xi32>, tensor<1xi32>) -> tensor<4xf64> + "hlo_test_speculatability.is_not_speculatable"(%4) : (tensor<4xf64>) -> () + %5 = stablehlo.dynamic_pad %static_arg, %padding_value, + %unknown_padding, %constant_padding, %constant_padding + : (tensor<4xf64>, tensor, tensor<1xi32>, tensor<1xi32>, tensor<1xi32>) -> tensor + "hlo_test_speculatability.is_speculatable"(%5) : (tensor) -> () + %6 = stablehlo.dynamic_pad %static_arg, %padding_value, + %constant_padding, %unknown_padding, %constant_padding + : (tensor<4xf64>, tensor, tensor<1xi32>, tensor<1xi32>, tensor<1xi32>) -> tensor<4xf64> + "hlo_test_speculatability.is_not_speculatable"(%6) : (tensor<4xf64>) -> () + %7 = stablehlo.dynamic_pad %static_arg, %padding_value, + %constant_padding, %unknown_padding, %constant_padding + : (tensor<4xf64>, tensor, tensor<1xi32>, tensor<1xi32>, tensor<1xi32>) -> tensor + "hlo_test_speculatability.is_speculatable"(%7) : (tensor) -> () + %8 = stablehlo.dynamic_pad %static_arg, %padding_value, + %constant_padding, %constant_padding, %unknown_padding + : (tensor<4xf64>, tensor, tensor<1xi32>, tensor<1xi32>, tensor<1xi32>) -> tensor<4xf64> + "hlo_test_speculatability.is_not_speculatable"(%8) : (tensor<4xf64>) -> () + %9 = stablehlo.dynamic_pad %static_arg, %padding_value, + %constant_padding, %constant_padding, %unknown_padding + : (tensor<4xf64>, tensor, tensor<1xi32>, tensor<1xi32>, tensor<1xi32>) -> tensor + "hlo_test_speculatability.is_speculatable"(%9) : (tensor) -> () + + return +} + +// ----- + // CHECK-LABEL: func @dynamic_reshape // CHECK-NEXT: return func.func @dynamic_reshape( @@ -1773,6 +1974,55 @@ func.func @dynamic_reshape( // ----- +// CHECK-LABEL: func @real_dynamic_slice +// CHECK-NEXT: return +func.func @real_dynamic_slice( + %static_arg: tensor<4xf64>, %dynamic_arg: tensor, + %unknown_value: tensor<1xi32> +) { + %constant_value = stablehlo.constant dense<1> : tensor<1xi32> + + // Static input, constant values + %0 = stablehlo.real_dynamic_slice %static_arg, %constant_value, %constant_value, %constant_value + : (tensor<4xf64>, tensor<1xi32>, tensor<1xi32>, tensor<1xi32>) -> tensor<0xf64> + "hlo_test_speculatability.is_speculatable"(%0) : (tensor<0xf64>) -> () + %1 = stablehlo.real_dynamic_slice %static_arg, %constant_value, %constant_value, %constant_value + : (tensor<4xf64>, tensor<1xi32>, tensor<1xi32>, tensor<1xi32>) -> tensor + "hlo_test_speculatability.is_speculatable"(%1) : (tensor) -> () + + // Dynamic input, constant values + %2 = stablehlo.real_dynamic_slice %dynamic_arg, %constant_value, %constant_value, %constant_value + : (tensor, tensor<1xi32>, tensor<1xi32>, tensor<1xi32>) -> tensor<0xf64> + "hlo_test_speculatability.is_not_speculatable"(%2) : (tensor<0xf64>) -> () + %3 = stablehlo.real_dynamic_slice %dynamic_arg, %constant_value, %constant_value, %constant_value + : (tensor, tensor<1xi32>, tensor<1xi32>, tensor<1xi32>) -> tensor + "hlo_test_speculatability.is_speculatable"(%3) : (tensor) -> () + + // Static input, unknown paddings + %4 = stablehlo.real_dynamic_slice %static_arg, %unknown_value, %constant_value, %constant_value + : (tensor<4xf64>, tensor<1xi32>, tensor<1xi32>, tensor<1xi32>) -> tensor<0xf64> + "hlo_test_speculatability.is_not_speculatable"(%4) : (tensor<0xf64>) -> () + %5 = stablehlo.real_dynamic_slice %static_arg, %unknown_value, %constant_value, %constant_value + : (tensor<4xf64>, tensor<1xi32>, tensor<1xi32>, tensor<1xi32>) -> tensor + "hlo_test_speculatability.is_speculatable"(%5) : (tensor) -> () + %6 = stablehlo.real_dynamic_slice %static_arg, %constant_value, %unknown_value, %constant_value + : (tensor<4xf64>, tensor<1xi32>, tensor<1xi32>, tensor<1xi32>) -> tensor<0xf64> + "hlo_test_speculatability.is_not_speculatable"(%6) : (tensor<0xf64>) -> () + %7 = stablehlo.real_dynamic_slice %static_arg, %constant_value, %unknown_value, %constant_value + : (tensor<4xf64>, tensor<1xi32>, tensor<1xi32>, tensor<1xi32>) -> tensor + "hlo_test_speculatability.is_speculatable"(%7) : (tensor) -> () + %8 = stablehlo.real_dynamic_slice %static_arg, %constant_value, %constant_value, %unknown_value + : (tensor<4xf64>, tensor<1xi32>, tensor<1xi32>, tensor<1xi32>) -> tensor<0xf64> + "hlo_test_speculatability.is_not_speculatable"(%8) : (tensor<0xf64>) -> () + %9 = stablehlo.real_dynamic_slice %static_arg, %constant_value, %constant_value, %unknown_value + : (tensor<4xf64>, tensor<1xi32>, tensor<1xi32>, tensor<1xi32>) -> tensor + "hlo_test_speculatability.is_speculatable"(%9) : (tensor) -> () + + return +} + +// ----- + // Recursively speculatable ops // -----