diff --git a/tensorflow/compiler/mlir/tf2xla/tests/legalize-tf.mlir b/tensorflow/compiler/mlir/tf2xla/tests/legalize-tf.mlir index be0bb7167f2863..a4d4f6a6ced441 100644 --- a/tensorflow/compiler/mlir/tf2xla/tests/legalize-tf.mlir +++ b/tensorflow/compiler/mlir/tf2xla/tests/legalize-tf.mlir @@ -2708,6 +2708,17 @@ func.func @neg_dynamic(%arg0: tensor) -> tensor { // ----- +// CHECK-LABEL: @random_shuffle_float_vector +func.func @random_shuffle_float_vector(%arg0: tensor<2xf32>) -> tensor<2xf32> { + // CHECK: "mhlo.sort" + // CHECK: ^bb0([[KEY_LHS:%.*]]: tensor, [[KEY_RHS:%.*]]: tensor, {{.*}}: tensor, {{.*}}: tensor): + // CHECK: mhlo.compare LT, [[KEY_LHS]], [[KEY_RHS]], TOTALORDER : (tensor, tensor) -> tensor + %0 = "tf.RandomShuffle"(%arg0) {seed = 1 : i64, seed2 = 2 : i64} : (tensor<2xf32>) -> tensor<2xf32> + func.return %0 : tensor<2xf32> +} + +// ----- + // CHECK-LABEL: @sigmoid func.func @sigmoid(%arg0: tensor<2xf32>) -> tensor<2xf32> { // CHECK: mhlo.logistic @@ -2798,4 +2809,3 @@ func.func @func_xla_sharding_consistent(%arg0: tensor<4x8xi32>) -> (tensor<4x8xi %1 = "tf.A"(%0) : (tensor<4x8xi32>) -> (tensor<4x8xi32>) func.return %1 : tensor<4x8xi32> } - diff --git a/tensorflow/python/eager/pywrap_gradient_exclusions.cc b/tensorflow/python/eager/pywrap_gradient_exclusions.cc index 6e39edcad17aa1..3ff8a8ce366a1c 100644 --- a/tensorflow/python/eager/pywrap_gradient_exclusions.cc +++ b/tensorflow/python/eager/pywrap_gradient_exclusions.cc @@ -50,7 +50,7 @@ auto OpGradientInfoInit(const T &a) { absl::optional> OpGradientUnusedInputIndices( const tensorflow::string &op_name) { - static std::array a = {{ + static std::array a = {{ {"Acosh"}, {"AllToAll", 1, {0}}, {"ApproximateEqual"}, @@ -292,7 +292,6 @@ absl::optional> OpGradientUnusedInputIndices( {"SdcaFprint"}, {"SegmentSum", 1, {0}}, {"Select", 1, {2}}, - {"Selu"}, {"SerializeTensor"}, {"SetSize"}, {"Shape"}, @@ -429,7 +428,7 @@ absl::optional> OpGradientUnusedInputIndices( absl::optional> OpGradientUnusedOutputIndices( const tensorflow::string &op_name) { - static std::array a = {{ + static std::array a = {{ {"Abs"}, {"AccumulateNV2"}, {"Acos"}, @@ -760,6 +759,7 @@ absl::optional> OpGradientUnusedOutputIndices( {"SegmentMean"}, {"SegmentSum"}, {"Select"}, + {"Selu"}, {"SeluGrad"}, {"SerializeTensor"}, {"SetSize"}, diff --git a/tensorflow/python/ops/BUILD b/tensorflow/python/ops/BUILD index 16fe6b93b648d0..22b37b8bd77ffd 100644 --- a/tensorflow/python/ops/BUILD +++ b/tensorflow/python/ops/BUILD @@ -2208,6 +2208,7 @@ py_library( ":array_ops_stack", ":math_ops", ":nn_ops_gen", + "//tensorflow/python/framework:constant_op", "//tensorflow/python/framework:dtypes", "//tensorflow/python/framework:ops", ], diff --git a/tensorflow/python/ops/nn_grad.py b/tensorflow/python/ops/nn_grad.py index 1b7dcc0ac9314a..e0c7a1cae72730 100644 --- a/tensorflow/python/ops/nn_grad.py +++ b/tensorflow/python/ops/nn_grad.py @@ -18,6 +18,7 @@ import itertools import operator +from tensorflow.python.framework import constant_op from tensorflow.python.framework import dtypes from tensorflow.python.framework import ops from tensorflow.python.ops import array_ops @@ -466,7 +467,18 @@ def _EluGrad(op: ops.Operation, grad): @ops.RegisterGradient("Selu") def _SeluGrad(op: ops.Operation, grad): - return gen_nn_ops.selu_grad(grad, op.outputs[0]) + x = op.inputs[0] + scale = constant_op.constant(1.0507009873554804934193349852946, dtype=x.dtype) + scale_alpha = constant_op.constant( + 1.7580993408473768599402175208123, dtype=x.dtype + ) + # Reconstructing the negative-branch derivative from the SELU output loses + # precision when the output rounds to -scale_alpha. Compute it from x. + derivative = array_ops.where_v2( + x < 0.0, scale_alpha * math_ops.exp(math_ops.minimum(x, 0.0)), scale + ) + derivative = array_ops.where_v2(math_ops.is_nan(x), x, derivative) + return grad * derivative @ops.RegisterGradient("Softplus") diff --git a/tensorflow/python/ops/nn_grad_test.py b/tensorflow/python/ops/nn_grad_test.py index 1f2b82ecbe9dea..3b84b2cea8345a 100644 --- a/tensorflow/python/ops/nn_grad_test.py +++ b/tensorflow/python/ops/nn_grad_test.py @@ -237,6 +237,57 @@ def testEluGradGradWRTinputs(self): class SeluGradOpTest(test.TestCase): + @test_util.run_in_graph_and_eager_modes + def testSeluGradPreservesSmallNegativeGradients(self): + scale_alpha = 1.7580993408473768599402175208123 + test_cases = ( + (dtypes.float32, -20.0, 1e-6), + (dtypes.float64, -700.0, 1e-14), + ) + for dtype, value, rtol in test_cases: + with self.subTest(dtype=dtype.name): + inputs = constant_op.constant(value, dtype=dtype) + with backprop.GradientTape() as tape: + tape.watch(inputs) + selu = gen_nn_ops.selu(inputs) + + selu_grad = tape.gradient(selu, inputs) + np_dtype = dtype.as_numpy_dtype + expected = np_dtype(scale_alpha) * np.exp(np_dtype(value)) + self.assertAllClose( + expected, self.evaluate(selu_grad), rtol=rtol, atol=0 + ) + + @test_util.run_in_graph_and_eager_modes + def testSeluGradEdgeCases(self): + scale = 1.0507009873554804934193349852946 + scale_alpha = 1.7580993408473768599402175208123 + for dtype in (dtypes.float32, dtypes.float64): + with self.subTest(dtype=dtype.name): + np_dtype = dtype.as_numpy_dtype + values = np.array( + [-np.inf, -2.0, -0.0, 0.0, 2.0, np.inf, np.nan], dtype=np_dtype + ) + inputs = constant_op.constant(values, dtype=dtype) + with backprop.GradientTape() as tape: + tape.watch(inputs) + selu = gen_nn_ops.selu(inputs) + + selu_grad = self.evaluate(tape.gradient(selu, inputs)) + expected = np.array( + [ + 0.0, + np_dtype(scale_alpha) * np.exp(np_dtype(-2.0)), + scale, + scale, + scale, + scale, + np.nan, + ], + dtype=np_dtype, + ) + self.assertAllClose(expected, selu_grad) + @test_util.run_deprecated_v1 def testSeluGradGradWRTgrad_ys(self): inputs = constant_op.constant( diff --git a/tensorflow/python/ops/numpy_ops/np_array_ops.py b/tensorflow/python/ops/numpy_ops/np_array_ops.py index 85d12f133b612e..5b7d87b83138af 100644 --- a/tensorflow/python/ops/numpy_ops/np_array_ops.py +++ b/tensorflow/python/ops/numpy_ops/np_array_ops.py @@ -100,7 +100,7 @@ def ones_like(a, dtype=None): def eye(N, M=None, k=0, dtype=float): # pylint: disable=invalid-name,missing-docstring if dtype: dtype = np_utils.result_type(dtype) - if not M: + if M is None: M = N # Making sure N, M and k are `int` N = int(N) diff --git a/tensorflow/python/ops/numpy_ops/np_array_ops_test.py b/tensorflow/python/ops/numpy_ops/np_array_ops_test.py index 1045b154b4d36e..21b055b2e701e7 100644 --- a/tensorflow/python/ops/numpy_ops/np_array_ops_test.py +++ b/tensorflow/python/ops/numpy_ops/np_array_ops_test.py @@ -230,6 +230,13 @@ def testEye(self): np_array_ops.eye(n, m, k, dtype=dtype), np.eye(n, m, k, dtype=dtype)) + # Test M=0 and N=0 zero-dimension edge cases + for n in (0, 1, 3): + for m in (0, 1, 3): + self.match(np_array_ops.eye(n, m), np.eye(n, m)) + for k in range(-n - 1, m + 2): + self.match(np_array_ops.eye(n, m, k), np.eye(n, m, k)) + def testIdentity(self): n_max = 3 diff --git a/tensorflow/python/training/training_ops_test.py b/tensorflow/python/training/training_ops_test.py index 158ace06a26d54..fe06bd08e8faa2 100644 --- a/tensorflow/python/training/training_ops_test.py +++ b/tensorflow/python/training/training_ops_test.py @@ -603,6 +603,43 @@ def testSparseApplyOpsRejectLowerRankGrad(self): with self.assertRaises(errors.InvalidArgumentError): self.evaluate(apply_op()) + @test_util.run_v2_only + def testResourceSparseApplyAdagradDAInvalidGradRank(self): + # A scalar `grad` with a higher-rank `var` used to hit a fatal CHECK + # instead of raising InvalidArgumentError (see GitHub issue #94130). + # ResourceSparseApplyAdagradDA only has a CPU kernel, so pin the whole + # test to CPU rather than relying on default placement. + with ops.device("/cpu:0"): + var = variables.Variable([[0.0, 0.0]] * 10, dtype=dtypes.float32) + gradient_accumulator = variables.Variable( + [[0.0, 0.0]] * 10, dtype=dtypes.float32 + ) + gradient_squared_accumulator = variables.Variable( + [[0.0, 0.0]] * 10, dtype=dtypes.float32 + ) + self.evaluate(variables.global_variables_initializer()) + + grad = constant_op.constant(0.0, dtype=dtypes.float32) # wrong rank + indices = constant_op.constant([0, 0], dtype=dtypes.int32) + + with self.assertRaisesRegex( + errors.InvalidArgumentError, + "grad must have the same number of dimensions as var", + ): + self.evaluate( + gen_training_ops.resource_sparse_apply_adagrad_da( + var.handle, + gradient_accumulator.handle, + gradient_squared_accumulator.handle, + grad, + indices, + constant_op.constant(0.0, dtype=dtypes.float32), + constant_op.constant(0.0, dtype=dtypes.float32), + constant_op.constant(0.0, dtype=dtypes.float32), + constant_op.constant(1, dtype=dtypes.int64), + ) + ) + if __name__ == '__main__': googletest.main() diff --git a/third_party/xla/third_party/stablehlo/temporary.patch b/third_party/xla/third_party/stablehlo/temporary.patch index ebb9dda1528d7c..b1e2e8b7ff4ea0 100644 --- a/third_party/xla/third_party/stablehlo/temporary.patch +++ b/third_party/xla/third_party/stablehlo/temporary.patch @@ -15,863 +15,4 @@ # # This file is automatically generated by generate_patch tool. # Do not edit directly. -diff --ruN a/stablehlo/docs/spec.md b/stablehlo/docs/spec.md ---- stablehlo/docs/spec.md -+++ stablehlo/docs/spec.md -@@ -2493,7 +2493,9 @@ - * `num_windows = is_empty_window[lhs_dim] ? 0 : floor((padded_input_shape[lhs_dim] - dilated_window_shape[lhs_dim]) / window_strides[spatial_dim]) + 1`. - * (C26) `rank(result) = N`. - * If the operation uses non-quantized tensors: -- * (C27) `element_type(lhs) = element_type(rhs) = element_type(result)`. -+ * (C27) `element_type(lhs) = element_type(rhs) = element_type(result)` or -+ (`is_fp8(element_type(lhs))` and `is_fp8(element_type(rhs))` and -+ `element_type(result) = element_type(lhs)`). - * If the operation uses quantized tensors: - * (C28) `is_quantized(lhs) = is_quantized(result) and is_quantized(rhs)`. - * (C29) If `is_per_axis_quantized(rhs)`, -@@ -3163,7 +3165,9 @@ - * `num_windows = is_empty_window[lhs_dim] ? 0 : floor((padded_input_shape[lhs_dim] - dilated_window_shape[lhs_dim]) / window_strides[spatial_dim]) + 1`. - * (C26) `rank(result) = N`. - * If the operation uses non-quantized tensors: -- * (C27) `element_type(lhs) = element_type(rhs) = element_type(result)`. -+ * (C27) `element_type(lhs) = element_type(rhs) = element_type(result)` or -+ (`is_fp8(element_type(lhs))` and `is_fp8(element_type(rhs))` and -+ `element_type(result) = element_type(lhs)`). - * If the operation uses quantized tensors: - * (C28) `is_quantized(lhs) = is_quantized(result) and is_quantized(rhs)`. - * (C29) If `is_per_axis_quantized(rhs)`, -@@ -7654,6 +7658,12 @@ - * `is_quantized(x: Value | Placeholder | Type) -> Value` is a shortcut for - `is_quantized_tensor_element_type(x)`. - -+* `is_fp8(x: Value | Placeholder | Type) -> Value` returns `true` if `x` is one -+of `Float8E4M3Type`, `Float8E4M3FNType`, `Float8E4M3B11FNUZType`, -+`Float8E4M3FNUZType`, `Float8E5M2Type`, `Float8E5M2FNUZType`, `Float8E3M4Type`, -+or `Float8E8M0FNUType`. If `x` is a value or placeholder, this function is a -+shortcut for `is_fp8(type(x))`. -+ - * `is_type_name(x: Value | Placeholder | Type) -> Value`. Available for all - types. For example, `is_float(x)` returns `true` if `x` is a `FloatType`. - If `x` is a value or placeholder, this function is a shortcut for -diff --ruN a/stablehlo/stablehlo/dialect/Serialization.cpp b/stablehlo/stablehlo/dialect/Serialization.cpp ---- stablehlo/stablehlo/dialect/Serialization.cpp -+++ stablehlo/stablehlo/dialect/Serialization.cpp -@@ -24,6 +24,7 @@ - #include "mlir/IR/Location.h" - #include "mlir/IR/MLIRContext.h" - #include "mlir/IR/OwningOpRef.h" -+#include "mlir/IR/Verifier.h" - #include "mlir/Parser/Parser.h" - #include "mlir/Pass/PassManager.h" - #include "mlir/Support/LLVM.h" -@@ -43,11 +44,18 @@ - bool allowOtherDialects) { - MLIRContext* context = module.getContext(); - -+ // Only enable verifier in debug builds. -+ bool enableVerifier = false; -+#ifndef NDEBUG -+ enableVerifier = true; -+#endif -+ - // Convert StableHLO --> VHLO. - // If allowOtherDialects is true, we will allow other dialects to be present - // in the module, otherwise will fail if there are any other dialects present. - { - PassManager pm(context); -+ pm.enableVerifier(enableVerifier); - StablehloLegalizeToVhloPassOptions options; - options.allowOtherDialects = allowOtherDialects; - pm.addPass(stablehlo::createStablehloLegalizeToVhloPass(options)); -@@ -61,6 +69,7 @@ - // target version failures. - { - PassManager pm(context); -+ pm.enableVerifier(enableVerifier); - pm.addPass(stablehlo::createVhloToVersionPass({targetVersion.str()})); - if (!succeeded(pm.run(module))) { - return failure(); -@@ -81,7 +90,15 @@ - OwningOpRef deserializePortableArtifact(StringRef sourceStr, - MLIRContext* context) { - context->loadDialect(); -- auto module = parseSourceString(sourceStr, context); -+ -+ // Only enable verifier in debug builds. -+ bool enableVerifier = false; -+#ifndef NDEBUG -+ enableVerifier = true; -+#endif -+ -+ ParserConfig config(context, /*verifyAfterParse=*/enableVerifier); -+ auto module = parseSourceString(sourceStr, config); - if (!module) { - emitError(UnknownLoc::get(context)) - << "failed to deserialize portable artifact using StableHLO_v" -@@ -91,8 +108,13 @@ - - // Convert VHLO --> VHLO(current) --> StableHLO - PassManager pm(context); -+ pm.enableVerifier(enableVerifier); - createStablehloDeserializePipeline(pm); - if (!succeeded(pm.run(*module))) { -+ return nullptr; -+ } -+ -+ if (failed(verify(*module))) { - return nullptr; - } - -diff --ruN a/stablehlo/stablehlo/dialect/TypeInference.cpp b/stablehlo/stablehlo/dialect/TypeInference.cpp ---- stablehlo/stablehlo/dialect/TypeInference.cpp -+++ stablehlo/stablehlo/dialect/TypeInference.cpp -@@ -86,6 +86,13 @@ - // Utils for quantization specific verifications - //===----------------------------------------------------------------------===// - -+bool isFp8Type(mlir::Type type) { -+ return llvm::isa(type); -+} -+ - template - bool allQuantized(ArrayRef typeRange) { - return llvm::all_of( -@@ -2221,11 +2228,14 @@ - // convolution_c27 - if (!anyQuantized({rankedLhsType, rankedRhsType}) && - !isCompatibleForHloTypeInference(rankedLhsType.getElementType(), -- rankedRhsType.getElementType())) -+ rankedRhsType.getElementType()) && -+ !(isFp8Type(rankedLhsType.getElementType()) && -+ isFp8Type(rankedRhsType.getElementType()))) { - return emitOptionalError( - location, "expects lhs and rhs to have compatible element type. Got: ", - rankedLhsType.getElementType(), " and ", - rankedRhsType.getElementType()); -+ } - - if (failed(verifyConvolutionAttributes( - location, lhsType, rhsType, inputBatchDimension, -@@ -2548,7 +2558,9 @@ - // dynamic_conv_c27 - if (!anyQuantized({rankedLhsType, rankedRhsType}) && - !isCompatibleForHloTypeInference(rankedLhsType.getElementType(), -- rankedRhsType.getElementType())) -+ rankedRhsType.getElementType()) && -+ !(isFp8Type(rankedLhsType.getElementType()) && -+ isFp8Type(rankedRhsType.getElementType()))) - return emitOptionalError( - location, "expects lhs and rhs to have compatible element type. Got: ", - rankedLhsType.getElementType(), " and ", -diff --ruN a/stablehlo/stablehlo/dialect/Version.h b/stablehlo/stablehlo/dialect/Version.h ---- stablehlo/stablehlo/dialect/Version.h -+++ stablehlo/stablehlo/dialect/Version.h -@@ -38,7 +38,7 @@ - static FailureOr fromString(llvm::StringRef versionRef); - - /// Return a Version representing the current VHLO dialect version. -- static Version getCurrentVersion() { return Version(1, 19, 0); } -+ static Version getCurrentVersion() { return Version(1, 20, 0); } - - /// Return a Version representing the minimum supported VHLO dialect version. - static Version getMinimumVersion() { return Version(0, 9, 0); } -diff --ruN a/stablehlo/stablehlo/dialect/VhloDialect.td b/stablehlo/stablehlo/dialect/VhloDialect.td ---- stablehlo/stablehlo/dialect/VhloDialect.td -+++ stablehlo/stablehlo/dialect/VhloDialect.td -@@ -58,6 +58,7 @@ - 1.17.0: Add `future` type support to `custom_call` op. - 1.18.0: Add `result_tilings` attribute to `custom_call` op. - 1.19.0: Add CollectiveReduceOp. -+ 1.20.0: Allow mixed fp8 operands in `convolution` and `dynamic_conv` ops. - }]; - - let useDefaultAttributePrinterParser = 0; -diff --ruN a/stablehlo/stablehlo/dialect/VhloOps.cpp b/stablehlo/stablehlo/dialect/VhloOps.cpp ---- stablehlo/stablehlo/dialect/VhloOps.cpp -+++ stablehlo/stablehlo/dialect/VhloOps.cpp -@@ -359,11 +359,44 @@ - return success(); - } - -+bool isVhloFp8Type(Type type) { -+ return isa(type); -+} -+ -+LogicalResult verifyConstraint_1_20_0(mlir::Operation* op, -+ Version targetVersion) { -+ if (targetVersion < Version(1, 20, 0)) { -+ if (op->getNumOperands() < 2) { -+ return failure(); -+ } -+ Type lhsElementType = getVhloElementType(op->getOperand(0).getType()); -+ Type rhsElementType = getVhloElementType(op->getOperand(1).getType()); -+ if (lhsElementType != rhsElementType && isVhloFp8Type(lhsElementType) && -+ isVhloFp8Type(rhsElementType)) { -+ return failure(); -+ } -+ } -+ return success(); -+} -+ - } // namespace - - LogicalResult AllReduceOpV1::validateConstraint(mlir::Operation* op, - Version targetVersion) { - return verifyConstraint_0_17_0(op, targetVersion); -+} -+ -+LogicalResult ConvolutionOpV1::validateConstraint(mlir::Operation* op, -+ Version targetVersion) { -+ return verifyConstraint_1_20_0(op, targetVersion); -+} -+ -+LogicalResult DynamicConvOpV2::validateConstraint(mlir::Operation* op, -+ Version targetVersion) { -+ return verifyConstraint_1_20_0(op, targetVersion); - } - - LogicalResult ReduceOpV1::validateConstraint(mlir::Operation* op, -diff --ruN a/stablehlo/stablehlo/dialect/VhloOps.td b/stablehlo/stablehlo/dialect/VhloOps.td ---- stablehlo/stablehlo/dialect/VhloOps.td -+++ stablehlo/stablehlo/dialect/VhloOps.td -@@ -390,7 +390,8 @@ - let results = (outs VHLO_AnyType:$result); - } - --def VHLO_ConvolutionOpV1 : VHLO_Op<"convolution_v1", "0.9.0", "current"> { -+def VHLO_ConvolutionOpV1 : VHLO_Op<"convolution_v1", "0.9.0", "current", -+ [DeclareOpInterfaceMethods]> { - let arguments = (ins - VHLO_AnyType:$lhs, - VHLO_AnyType:$rhs, -@@ -572,7 +573,8 @@ - - // Padding should be specified as an operand only, not an attribute. - // Remove `d_padding` and convert `padding` to an operand. --def VHLO_DynamicConvOpV2 : VHLO_Op<"dynamic_conv_v2", "0.20.0", "current"> { -+def VHLO_DynamicConvOpV2 : VHLO_Op<"dynamic_conv_v2", "0.20.0", "current", -+ [DeclareOpInterfaceMethods]> { - let arguments = (ins - VHLO_AnyType:$lhs, - VHLO_AnyType:$rhs, -diff --ruN a/stablehlo/stablehlo/tests/TestUtils.cpp b/stablehlo/stablehlo/tests/TestUtils.cpp ---- stablehlo/stablehlo/tests/TestUtils.cpp -+++ stablehlo/stablehlo/tests/TestUtils.cpp -@@ -68,6 +68,32 @@ - } - }; - -+struct BroadcastIfNeededPattern : public RewritePattern { -+ explicit BroadcastIfNeededPattern(MLIRContext* context) -+ : RewritePattern("hlo_test_broadcast.broadcast_if_needed", 1, context) {} -+ LogicalResult matchAndRewrite(Operation* op, -+ PatternRewriter& rewriter) const override { -+ if (op->getNumOperands() < 1) return failure(); -+ Value input = op->getOperand(0); -+ -+ SmallVector broadcastDimensions; -+ if (auto bcastDimsAttr = -+ op->getAttrOfType("broadcast_dimensions")) { -+ broadcastDimensions = llvm::to_vector(bcastDimsAttr.asArrayRef()); -+ } -+ -+ auto targetShape = stablehlo::getDimensions(op->getResult(0)); -+ if (failed(targetShape)) return failure(); -+ -+ auto broadcastedVal = stablehlo::broadcastIfNeeded( -+ rewriter, input, *targetShape, broadcastDimensions); -+ if (failed(broadcastedVal)) return failure(); -+ -+ rewriter.replaceOp(op, *broadcastedVal); -+ return success(); -+ } -+}; -+ - struct InferReturnTypesPattern : public RewritePattern { - explicit InferReturnTypesPattern(MLIRContext* context) - : RewritePattern("hlo_test_infer.get_return_types", 1, context) {} -@@ -222,6 +248,7 @@ - LogicalResult initialize(MLIRContext* context) override { - RewritePatternSet patterns(context); - patterns.add(context); -+ patterns.add(context); - patterns_ = std::move(patterns); - return success(); - } -diff --ruN a/stablehlo/stablehlo/tests/chlo/chlo_legalize_to_stablehlo.mlir b/stablehlo/stablehlo/tests/chlo/chlo_legalize_to_stablehlo.mlir ---- stablehlo/stablehlo/tests/chlo/chlo_legalize_to_stablehlo.mlir -+++ stablehlo/stablehlo/tests/chlo/chlo_legalize_to_stablehlo.mlir -@@ -2141,13 +2141,16 @@ - // CHECK: %[[DIVIDE_19:.*]] = stablehlo.divide %[[MULTIPLY_9]], %[[SINE_0]] : tensor - // CHECK: %[[SUBTRACT_10:.*]] = stablehlo.subtract %[[SUBTRACT_9]], %[[DIVIDE_19]] : tensor - // CHECK: %[[SELECT_1:.*]] = stablehlo.select %[[COMPARE_0]], %[[SUBTRACT_10]], %[[SUBTRACT_9]] : tensor, tensor --// CHECK: %[[COMPARE_1:.*]] = stablehlo.compare LE, %[[ARG0]], %[[CONSTANT_2]] : (tensor, tensor) -> tensor -+// CHECK: %[[COMPARE_1:.*]] = stablehlo.compare EQ, %[[ARG0]], %[[CONSTANT_2]] : (tensor, tensor) -> tensor -+// CHECK: %[[COMPARE_2:.*]] = stablehlo.compare LT, %[[ARG0]], %[[CONSTANT_2]] : (tensor, tensor) -> tensor - // CHECK: %[[FLOOR_1:.*]] = stablehlo.floor %[[ARG0]] : tensor --// CHECK: %[[COMPARE_2:.*]] = stablehlo.compare EQ, %[[ARG0]], %[[FLOOR_1]] : (tensor, tensor) -> tensor --// CHECK: %[[AND_0:.*]] = stablehlo.and %[[COMPARE_1]], %[[COMPARE_2]] : tensor -+// CHECK: %[[COMPARE_3:.*]] = stablehlo.compare EQ, %[[ARG0]], %[[FLOOR_1]] : (tensor, tensor) -> tensor -+// CHECK: %[[AND_0:.*]] = stablehlo.and %[[COMPARE_2]], %[[COMPARE_3]] : tensor - // CHECK: %[[CONSTANT_25:.*]] = stablehlo.constant dense<0x7FF8000000000000> : tensor - // CHECK: %[[SELECT_2:.*]] = stablehlo.select %[[AND_0]], %[[CONSTANT_25]], %[[SELECT_1]] : tensor, tensor --// CHECK: return %[[SELECT_2]] : tensor -+// CHECK: %[[CONSTANT_26:.*]] = stablehlo.constant dense<0xFFF0000000000000> : tensor -+// CHECK: %[[SELECT_3:.*]] = stablehlo.select %[[COMPARE_1]], %[[CONSTANT_26]], %[[SELECT_2]] : tensor, tensor -+// CHECK: return %[[SELECT_3]] : tensor - // CHECK: } - func.func @digamma_f64(%arg : tensor) -> tensor { - %1 = chlo.digamma %arg : tensor -> tensor -@@ -2254,13 +2257,16 @@ - // CHECK: %[[DIVIDE_19:.*]] = stablehlo.divide %[[MULTIPLY_9]], %[[SINE_0]] : tensor - // CHECK: %[[SUBTRACT_10:.*]] = stablehlo.subtract %[[SUBTRACT_9]], %[[DIVIDE_19]] : tensor - // CHECK: %[[SELECT_1:.*]] = stablehlo.select %[[COMPARE_0]], %[[SUBTRACT_10]], %[[SUBTRACT_9]] : tensor, tensor --// CHECK: %[[COMPARE_1:.*]] = stablehlo.compare LE, %[[ARG0]], %[[CONSTANT_2]] : (tensor, tensor) -> tensor -+// CHECK: %[[COMPARE_1:.*]] = stablehlo.compare EQ, %[[ARG0]], %[[CONSTANT_2]] : (tensor, tensor) -> tensor -+// CHECK: %[[COMPARE_2:.*]] = stablehlo.compare LT, %[[ARG0]], %[[CONSTANT_2]] : (tensor, tensor) -> tensor - // CHECK: %[[FLOOR_1:.*]] = stablehlo.floor %[[ARG0]] : tensor --// CHECK: %[[COMPARE_2:.*]] = stablehlo.compare EQ, %[[ARG0]], %[[FLOOR_1]] : (tensor, tensor) -> tensor --// CHECK: %[[AND_0:.*]] = stablehlo.and %[[COMPARE_1]], %[[COMPARE_2]] : tensor -+// CHECK: %[[COMPARE_3:.*]] = stablehlo.compare EQ, %[[ARG0]], %[[FLOOR_1]] : (tensor, tensor) -> tensor -+// CHECK: %[[AND_0:.*]] = stablehlo.and %[[COMPARE_2]], %[[COMPARE_3]] : tensor - // CHECK: %[[CONSTANT_25:.*]] = stablehlo.constant dense<0x7FC00000> : tensor - // CHECK: %[[SELECT_2:.*]] = stablehlo.select %[[AND_0]], %[[CONSTANT_25]], %[[SELECT_1]] : tensor, tensor --// CHECK: return %[[SELECT_2]] : tensor -+// CHECK: %[[CONSTANT_26:.*]] = stablehlo.constant dense<0xFF800000> : tensor -+// CHECK: %[[SELECT_3:.*]] = stablehlo.select %[[COMPARE_1]], %[[CONSTANT_26]], %[[SELECT_2]] : tensor, tensor -+// CHECK: return %[[SELECT_3]] : tensor - // CHECK: } - func.func @digamma_f32(%arg : tensor) -> tensor { - %1 = chlo.digamma %arg : tensor -> tensor -@@ -2368,13 +2374,16 @@ - // CHECK: %[[DIVIDE_19:.*]] = stablehlo.divide %[[MULTIPLY_9]], %[[SINE_0]] : tensor - // CHECK: %[[SUBTRACT_10:.*]] = stablehlo.subtract %[[SUBTRACT_9]], %[[DIVIDE_19]] : tensor - // CHECK: %[[SELECT_1:.*]] = stablehlo.select %[[COMPARE_0]], %[[SUBTRACT_10]], %[[SUBTRACT_9]] : tensor, tensor --// CHECK: %[[COMPARE_1:.*]] = stablehlo.compare LE, %[[CONVERT_0]], %[[CONSTANT_2]] : (tensor, tensor) -> tensor -+// CHECK: %[[COMPARE_1:.*]] = stablehlo.compare EQ, %[[CONVERT_0]], %[[CONSTANT_2]] : (tensor, tensor) -> tensor -+// CHECK: %[[COMPARE_2:.*]] = stablehlo.compare LT, %[[CONVERT_0]], %[[CONSTANT_2]] : (tensor, tensor) -> tensor - // CHECK: %[[FLOOR_1:.*]] = stablehlo.floor %[[CONVERT_0]] : tensor --// CHECK: %[[COMPARE_2:.*]] = stablehlo.compare EQ, %[[CONVERT_0]], %[[FLOOR_1]] : (tensor, tensor) -> tensor --// CHECK: %[[AND_0:.*]] = stablehlo.and %[[COMPARE_1]], %[[COMPARE_2]] : tensor -+// CHECK: %[[COMPARE_3:.*]] = stablehlo.compare EQ, %[[CONVERT_0]], %[[FLOOR_1]] : (tensor, tensor) -> tensor -+// CHECK: %[[AND_0:.*]] = stablehlo.and %[[COMPARE_2]], %[[COMPARE_3]] : tensor - // CHECK: %[[CONSTANT_25:.*]] = stablehlo.constant dense<0x7FC00000> : tensor - // CHECK: %[[SELECT_2:.*]] = stablehlo.select %[[AND_0]], %[[CONSTANT_25]], %[[SELECT_1]] : tensor, tensor --// CHECK: %[[CONVERT_1:.*]] = stablehlo.convert %[[SELECT_2]] : (tensor) -> tensor -+// CHECK: %[[CONSTANT_26:.*]] = stablehlo.constant dense<0xFF800000> : tensor -+// CHECK: %[[SELECT_3:.*]] = stablehlo.select %[[COMPARE_1]], %[[CONSTANT_26]], %[[SELECT_2]] : tensor, tensor -+// CHECK: %[[CONVERT_1:.*]] = stablehlo.convert %[[SELECT_3]] : (tensor) -> tensor - // CHECK: return %[[CONVERT_1]] : tensor - // CHECK: } - func.func @digamma_f16(%arg : tensor) -> tensor { -@@ -2942,17 +2951,20 @@ - // CHECK: %[[DIVIDE_32:.*]] = stablehlo.divide %[[MULTIPLY_52]], %[[SINE_1]] : tensor - // CHECK: %[[SUBTRACT_18:.*]] = stablehlo.subtract %[[SUBTRACT_17]], %[[DIVIDE_32]] : tensor - // CHECK: %[[SELECT_12:.*]] = stablehlo.select %[[COMPARE_12]], %[[SUBTRACT_18]], %[[SUBTRACT_17]] : tensor, tensor --// CHECK: %[[COMPARE_13:.*]] = stablehlo.compare LE, %[[ARG1]], %[[CONSTANT_73]] : (tensor, tensor) -> tensor -+// CHECK: %[[COMPARE_13:.*]] = stablehlo.compare EQ, %[[ARG1]], %[[CONSTANT_73]] : (tensor, tensor) -> tensor -+// CHECK: %[[COMPARE_14:.*]] = stablehlo.compare LT, %[[ARG1]], %[[CONSTANT_73]] : (tensor, tensor) -> tensor - // CHECK: %[[FLOOR_5:.*]] = stablehlo.floor %[[ARG1]] : tensor --// CHECK: %[[COMPARE_14:.*]] = stablehlo.compare EQ, %[[ARG1]], %[[FLOOR_5]] : (tensor, tensor) -> tensor --// CHECK: %[[AND_3:.*]] = stablehlo.and %[[COMPARE_13]], %[[COMPARE_14]] : tensor -+// CHECK: %[[COMPARE_15:.*]] = stablehlo.compare EQ, %[[ARG1]], %[[FLOOR_5]] : (tensor, tensor) -> tensor -+// CHECK: %[[AND_3:.*]] = stablehlo.and %[[COMPARE_14]], %[[COMPARE_15]] : tensor - // CHECK: %[[CONSTANT_96:.*]] = stablehlo.constant dense<0x7FC00000> : tensor - // CHECK: %[[SELECT_13:.*]] = stablehlo.select %[[AND_3]], %[[CONSTANT_96]], %[[SELECT_12]] : tensor, tensor --// CHECK: %[[SELECT_14:.*]] = stablehlo.select %[[COMPARE_11]], %[[SELECT_13]], %[[MULTIPLY_42]] : tensor, tensor -+// CHECK: %[[CONSTANT_98:.*]] = stablehlo.constant dense<0xFF800000> : tensor -+// CHECK: %[[SELECT_16:.*]] = stablehlo.select %[[COMPARE_13]], %[[CONSTANT_98]], %[[SELECT_13]] : tensor, tensor -+// CHECK: %[[SELECT_14:.*]] = stablehlo.select %[[COMPARE_11]], %[[SELECT_16]], %[[MULTIPLY_42]] : tensor, tensor - // CHECK: %[[FLOOR_6:.*]] = stablehlo.floor %[[ARG0]] : tensor --// CHECK: %[[COMPARE_15:.*]] = stablehlo.compare NE, %[[ARG0]], %[[FLOOR_6]] : (tensor, tensor) -> tensor --// CHECK: %[[COMPARE_16:.*]] = stablehlo.compare LT, %[[ARG0]], %[[CONSTANT_70]] : (tensor, tensor) -> tensor --// CHECK: %[[OR_0:.*]] = stablehlo.or %[[COMPARE_15]], %[[COMPARE_16]] : tensor -+// CHECK: %[[COMPARE_16:.*]] = stablehlo.compare NE, %[[ARG0]], %[[FLOOR_6]] : (tensor, tensor) -> tensor -+// CHECK: %[[COMPARE_17:.*]] = stablehlo.compare LT, %[[ARG0]], %[[CONSTANT_70]] : (tensor, tensor) -> tensor -+// CHECK: %[[OR_0:.*]] = stablehlo.or %[[COMPARE_16]], %[[COMPARE_17]] : tensor - // CHECK: %[[CONSTANT_97:.*]] = stablehlo.constant dense<0x7FC00000> : tensor - // CHECK: %[[SELECT_15:.*]] = stablehlo.select %[[OR_0]], %[[CONSTANT_97]], %[[SELECT_14]] : tensor, tensor - // CHECK: return %[[SELECT_15]] : tensor -@@ -3332,17 +3344,20 @@ - // CHECK: %[[DIVIDE_32:.*]] = stablehlo.divide %[[MULTIPLY_52]], %[[SINE_1]] : tensor - // CHECK: %[[SUBTRACT_18:.*]] = stablehlo.subtract %[[SUBTRACT_17]], %[[DIVIDE_32]] : tensor - // CHECK: %[[SELECT_12:.*]] = stablehlo.select %[[COMPARE_12]], %[[SUBTRACT_18]], %[[SUBTRACT_17]] : tensor, tensor --// CHECK: %[[COMPARE_13:.*]] = stablehlo.compare LE, %[[ARG1]], %[[CONSTANT_73]] : (tensor, tensor) -> tensor -+// CHECK: %[[COMPARE_13:.*]] = stablehlo.compare EQ, %[[ARG1]], %[[CONSTANT_73]] : (tensor, tensor) -> tensor -+// CHECK: %[[COMPARE_14:.*]] = stablehlo.compare LT, %[[ARG1]], %[[CONSTANT_73]] : (tensor, tensor) -> tensor - // CHECK: %[[FLOOR_5:.*]] = stablehlo.floor %[[ARG1]] : tensor --// CHECK: %[[COMPARE_14:.*]] = stablehlo.compare EQ, %[[ARG1]], %[[FLOOR_5]] : (tensor, tensor) -> tensor --// CHECK: %[[AND_3:.*]] = stablehlo.and %[[COMPARE_13]], %[[COMPARE_14]] : tensor -+// CHECK: %[[COMPARE_15:.*]] = stablehlo.compare EQ, %[[ARG1]], %[[FLOOR_5]] : (tensor, tensor) -> tensor -+// CHECK: %[[AND_3:.*]] = stablehlo.and %[[COMPARE_14]], %[[COMPARE_15]] : tensor - // CHECK: %[[CONSTANT_96:.*]] = stablehlo.constant dense<0x7FF8000000000000> : tensor - // CHECK: %[[SELECT_13:.*]] = stablehlo.select %[[AND_3]], %[[CONSTANT_96]], %[[SELECT_12]] : tensor, tensor --// CHECK: %[[SELECT_14:.*]] = stablehlo.select %[[COMPARE_11]], %[[SELECT_13]], %[[MULTIPLY_42]] : tensor, tensor -+// CHECK: %[[CONSTANT_98:.*]] = stablehlo.constant dense<0xFFF0000000000000> : tensor -+// CHECK: %[[SELECT_16:.*]] = stablehlo.select %[[COMPARE_13]], %[[CONSTANT_98]], %[[SELECT_13]] : tensor, tensor -+// CHECK: %[[SELECT_14:.*]] = stablehlo.select %[[COMPARE_11]], %[[SELECT_16]], %[[MULTIPLY_42]] : tensor, tensor - // CHECK: %[[FLOOR_6:.*]] = stablehlo.floor %[[ARG0]] : tensor --// CHECK: %[[COMPARE_15:.*]] = stablehlo.compare NE, %[[ARG0]], %[[FLOOR_6]] : (tensor, tensor) -> tensor --// CHECK: %[[COMPARE_16:.*]] = stablehlo.compare LT, %[[ARG0]], %[[CONSTANT_70]] : (tensor, tensor) -> tensor --// CHECK: %[[OR_0:.*]] = stablehlo.or %[[COMPARE_15]], %[[COMPARE_16]] : tensor -+// CHECK: %[[COMPARE_16:.*]] = stablehlo.compare NE, %[[ARG0]], %[[FLOOR_6]] : (tensor, tensor) -> tensor -+// CHECK: %[[COMPARE_17:.*]] = stablehlo.compare LT, %[[ARG0]], %[[CONSTANT_70]] : (tensor, tensor) -> tensor -+// CHECK: %[[OR_0:.*]] = stablehlo.or %[[COMPARE_16]], %[[COMPARE_17]] : tensor - // CHECK: %[[CONSTANT_97:.*]] = stablehlo.constant dense<0x7FF8000000000000> : tensor - // CHECK: %[[SELECT_15:.*]] = stablehlo.select %[[OR_0]], %[[CONSTANT_97]], %[[SELECT_14]] : tensor, tensor - // CHECK: return %[[SELECT_15]] : tensor -@@ -3724,17 +3739,20 @@ - // CHECK: %[[DIVIDE_32:.*]] = stablehlo.divide %[[MULTIPLY_52]], %[[SINE_1]] : tensor - // CHECK: %[[SUBTRACT_18:.*]] = stablehlo.subtract %[[SUBTRACT_17]], %[[DIVIDE_32]] : tensor - // CHECK: %[[SELECT_12:.*]] = stablehlo.select %[[COMPARE_12]], %[[SUBTRACT_18]], %[[SUBTRACT_17]] : tensor, tensor --// CHECK: %[[COMPARE_13:.*]] = stablehlo.compare LE, %[[CONVERT_1]], %[[CONSTANT_73]] : (tensor, tensor) -> tensor -+// CHECK: %[[COMPARE_13:.*]] = stablehlo.compare EQ, %[[CONVERT_1]], %[[CONSTANT_73]] : (tensor, tensor) -> tensor -+// CHECK: %[[COMPARE_14:.*]] = stablehlo.compare LT, %[[CONVERT_1]], %[[CONSTANT_73]] : (tensor, tensor) -> tensor - // CHECK: %[[FLOOR_5:.*]] = stablehlo.floor %[[CONVERT_1]] : tensor --// CHECK: %[[COMPARE_14:.*]] = stablehlo.compare EQ, %[[CONVERT_1]], %[[FLOOR_5]] : (tensor, tensor) -> tensor --// CHECK: %[[AND_3:.*]] = stablehlo.and %[[COMPARE_13]], %[[COMPARE_14]] : tensor -+// CHECK: %[[COMPARE_15:.*]] = stablehlo.compare EQ, %[[CONVERT_1]], %[[FLOOR_5]] : (tensor, tensor) -> tensor -+// CHECK: %[[AND_3:.*]] = stablehlo.and %[[COMPARE_14]], %[[COMPARE_15]] : tensor - // CHECK: %[[CONSTANT_96:.*]] = stablehlo.constant dense<0x7FC00000> : tensor - // CHECK: %[[SELECT_13:.*]] = stablehlo.select %[[AND_3]], %[[CONSTANT_96]], %[[SELECT_12]] : tensor, tensor --// CHECK: %[[SELECT_14:.*]] = stablehlo.select %[[COMPARE_11]], %[[SELECT_13]], %[[MULTIPLY_42]] : tensor, tensor -+// CHECK: %[[CONSTANT_98:.*]] = stablehlo.constant dense<0xFF800000> : tensor -+// CHECK: %[[SELECT_16:.*]] = stablehlo.select %[[COMPARE_13]], %[[CONSTANT_98]], %[[SELECT_13]] : tensor, tensor -+// CHECK: %[[SELECT_14:.*]] = stablehlo.select %[[COMPARE_11]], %[[SELECT_16]], %[[MULTIPLY_42]] : tensor, tensor - // CHECK: %[[FLOOR_6:.*]] = stablehlo.floor %[[CONVERT_0]] : tensor --// CHECK: %[[COMPARE_15:.*]] = stablehlo.compare NE, %[[CONVERT_0]], %[[FLOOR_6]] : (tensor, tensor) -> tensor --// CHECK: %[[COMPARE_16:.*]] = stablehlo.compare LT, %[[CONVERT_0]], %[[CONSTANT_70]] : (tensor, tensor) -> tensor --// CHECK: %[[OR_0:.*]] = stablehlo.or %[[COMPARE_15]], %[[COMPARE_16]] : tensor -+// CHECK: %[[COMPARE_16:.*]] = stablehlo.compare NE, %[[CONVERT_0]], %[[FLOOR_6]] : (tensor, tensor) -> tensor -+// CHECK: %[[COMPARE_17:.*]] = stablehlo.compare LT, %[[CONVERT_0]], %[[CONSTANT_70]] : (tensor, tensor) -> tensor -+// CHECK: %[[OR_0:.*]] = stablehlo.or %[[COMPARE_16]], %[[COMPARE_17]] : tensor - // CHECK: %[[CONSTANT_97:.*]] = stablehlo.constant dense<0x7FC00000> : tensor - // CHECK: %[[SELECT_15:.*]] = stablehlo.select %[[OR_0]], %[[CONSTANT_97]], %[[SELECT_14]] : tensor, tensor - // CHECK: %[[CONVERT_2:.*]] = stablehlo.convert %[[SELECT_15]] : (tensor) -> tensor -@@ -5282,6 +5300,32 @@ - - // ----- - -+// CHECK-LABEL: func.func @ragged_dot_mode_2_rank_lhs_lt_rhs( -+// CHECK-SAME: %[[ARG0:.*]]: tensor<2x3xf32>, -+// CHECK-SAME: %[[ARG1:.*]]: tensor<2x3x4xf32>, -+// CHECK-SAME: %[[ARG2:.*]]: tensor<2xi64>) -> tensor<2x3x3x4xf32> { -+// CHECK: %[[DOT_GENERAL_0:.*]] = stablehlo.dot_general %{{.*}}, %[[ARG1]], contracting_dims = [0] x [0], precision = [DEFAULT, DEFAULT] : (tensor<2x3xf32>, tensor<2x3x4xf32>) -> tensor<3x3x4xf32> -+// CHECK: %[[DOT_GENERAL_1:.*]] = stablehlo.dot_general %{{.*}}, %[[ARG1]], contracting_dims = [0] x [0], precision = [DEFAULT, DEFAULT] : (tensor<2x3xf32>, tensor<2x3x4xf32>) -> tensor<3x3x4xf32> -+// CHECK: %[[CONCATENATE_0:.*]] = stablehlo.concatenate{{.*}}dim = 0 : (tensor<1x3x3x4xf32>, tensor<1x3x3x4xf32>) -> tensor<2x3x3x4xf32> -+// CHECK: return %[[CONCATENATE_0]] : tensor<2x3x3x4xf32> -+// CHECK: } -+func.func @ragged_dot_mode_2_rank_lhs_lt_rhs(%lhs : tensor<2x3xf32>, %rhs : tensor<2x3x4xf32>, %group_sizes : tensor<2xi64>) -> tensor<2x3x3x4xf32> { -+ %0 = "chlo.ragged_dot"(%lhs, %rhs, %group_sizes) { -+ ragged_dot_dimension_numbers = #chlo.ragged_dot< -+ lhs_batching_dimensions = [], -+ rhs_batching_dimensions = [], -+ lhs_contracting_dimensions = [0], -+ rhs_contracting_dimensions = [0], -+ lhs_ragged_dimensions = [0], -+ rhs_group_dimensions = [] -+ >, -+ precision_config = [#chlo, #chlo] -+ } : (tensor<2x3xf32>, tensor<2x3x4xf32>, tensor<2xi64>) -> tensor<2x3x3x4xf32> -+ func.return %0 : tensor<2x3x3x4xf32> -+} -+ -+// ----- -+ - // CHECK-LABEL: func.func @ragged_dot_mode_3( - // CHECK-SAME: %[[ARG0:.*]]: tensor<2x3x5xf32>, - // CHECK-SAME: %[[ARG1:.*]]: tensor<2x5x7xf32>, -diff --ruN a/stablehlo/stablehlo/tests/ops_broadcasting.mlir b/stablehlo/stablehlo/tests/ops_broadcasting.mlir ---- stablehlo/stablehlo/tests/ops_broadcasting.mlir -+++ stablehlo/stablehlo/tests/ops_broadcasting.mlir -@@ -320,3 +320,14 @@ - return %0 : !stablehlo.token - } - -+// ----- -+ -+// Non-numpy explicit broadcast_dimensions: [3, 1] -> [3, 4, 5] -+// CHECK-LABEL: func @explicit_broadcast_dims -+func.func @explicit_broadcast_dims(%arg0: tensor<3x1xf64>) -> tensor<3x4x5xf64> { -+ // CHECK: %[[BCAST:.+]] = stablehlo.broadcast_in_dim %arg0, dims = [0, 1] : (tensor<3x1xf64>) -> tensor<3x4x5xf64> -+ // CHECK-NEXT: return %[[BCAST]] : tensor<3x4x5xf64> -+ %0 = "hlo_test_broadcast.broadcast_if_needed"(%arg0) {broadcast_dimensions = array} : (tensor<3x1xf64>) -> tensor<3x4x5xf64> -+ return %0 : tensor<3x4x5xf64> -+} -+ -diff --ruN a/stablehlo/stablehlo/tests/verify_convolution.mlir b/stablehlo/stablehlo/tests/verify_convolution.mlir ---- stablehlo/stablehlo/tests/verify_convolution.mlir -+++ stablehlo/stablehlo/tests/verify_convolution.mlir -@@ -78,6 +78,60 @@ - - // ----- - -+// CHECK-LABEL: func @convolution_mixed_fp8 -+func.func @convolution_mixed_fp8(%arg0 : tensor<100x26x26x32xf8E5M2>, -+ %arg1 : tensor<3x3x1x32xf8E4M3FN>) -> tensor<100x28x28x1xf8E5M2> { -+ %result = "stablehlo.convolution"(%arg0, %arg1) { -+ batch_group_count = 1 : i64, -+ dimension_numbers = #stablehlo.conv, -+ feature_group_count = 1 : i64, -+ lhs_dilation = array, -+ padding = dense<2> : tensor<2x2xi64>, -+ rhs_dilation = array, -+ window_strides = array -+ } : (tensor<100x26x26x32xf8E5M2>, tensor<3x3x1x32xf8E4M3FN>) -> -+ tensor<100x28x28x1xf8E5M2> -+ func.return %result : tensor<100x28x28x1xf8E5M2> -+} -+ -+// ----- -+ -+func.func @convolution_mismatched_element_types(%arg0: tensor<100x26x26x32xf32>, %arg1: tensor<3x3x1x32xf16>) -> tensor<100x28x28x1xf32> { -+ // expected-error@+1{{expects lhs and rhs to have compatible element type. Got: 'f32' and 'f16'}} -+ %result = "stablehlo.convolution"(%arg0, %arg1) { -+ batch_group_count = 1 : i64, -+ dimension_numbers = #stablehlo.conv, -+ feature_group_count = 1 : i64, -+ lhs_dilation = array, -+ padding = dense<2> : tensor<2x2xi64>, -+ rhs_dilation = array, -+ window_strides = array -+ } : (tensor<100x26x26x32xf32>, tensor<3x3x1x32xf16>) -> tensor<100x28x28x1xf32> -+ func.return %result : tensor<100x28x28x1xf32> -+} -+ -+// ----- -+ - func.func @convolution(%arg0: tensor<2x2x3x4xf32>, %arg1: tensor<3x5x5x3xf32>) -> tensor<3x5x5x4xf32> { - // expected-error@+3{{Unexpected keyword stide}} - %0 = stablehlo.convolution(%arg0, %arg1) -diff --ruN a/stablehlo/stablehlo/tests/verify_dynamic_conv.mlir b/stablehlo/stablehlo/tests/verify_dynamic_conv.mlir ---- stablehlo/stablehlo/tests/verify_dynamic_conv.mlir -+++ stablehlo/stablehlo/tests/verify_dynamic_conv.mlir -@@ -48,6 +48,36 @@ - - // ----- - -+// CHECK-LABEL: func @dynamic_conv_mixed_fp8 -+func.func @dynamic_conv_mixed_fp8(%arg0 : tensor<100x26x26x32xf8E5M2>, -+ %arg1 : tensor<3x3x1x32xf8E4M3FN>) -> tensor<100x28x28x1xf8E5M2> { -+ %padding = stablehlo.constant dense<2> : tensor<2x2xi64> -+ %result = "stablehlo.dynamic_conv"(%arg0, %arg1, %padding) { -+ dimension_numbers = #stablehlo.conv<[b, 0, 1, f]x[0, 1, o, i]->[b, 0, 1, f]>, -+ feature_group_count = 1 : i64, -+ batch_group_count = 1 : i64 -+ } : (tensor<100x26x26x32xf8E5M2>, tensor<3x3x1x32xf8E4M3FN>, tensor<2x2xi64>) -> -+ tensor<100x28x28x1xf8E5M2> -+ func.return %result : tensor<100x28x28x1xf8E5M2> -+} -+ -+// ----- -+ -+func.func @dynamic_conv_mismatched_element_types(%arg0: tensor<100x26x26x32xf32>, -+ %arg1: tensor<3x3x1x32xf16>) -> tensor<100x28x28x1xf32> { -+ // expected-error@+2 {{expects lhs and rhs to have compatible element type. Got: 'f32' and 'f16'}} -+ %padding = stablehlo.constant dense<2> : tensor<2x2xi64> -+ %result = "stablehlo.dynamic_conv"(%arg0, %arg1, %padding) { -+ dimension_numbers = #stablehlo.conv<[b, 0, 1, f]x[0, 1, o, i]->[b, 0, 1, f]>, -+ feature_group_count = 1 : i64, -+ batch_group_count = 1 : i64 -+ } : (tensor<100x26x26x32xf32>, tensor<3x3x1x32xf16>, tensor<2x2xi64>) -> -+ tensor<100x28x28x1xf32> -+ func.return %result : tensor<100x28x28x1xf32> -+} -+ -+// ----- -+ - func.func @dynamic_conv_c1(%arg0: tensor<1x8x8x207xf32>, - %arg1: tensor<3x3x207xf32>) -> tensor<1x8x8x16xf32> { - // expected-error@+2 {{expects convolution arguments to have same number of dimensions. Got: 'tensor<1x8x8x207xf32>' and 'tensor<3x3x207xf32>'.}} -diff --ruN a/stablehlo/stablehlo/tests/vhlo/vhlo_to_version_downgrade_invalid.1_19_0.mlir b/stablehlo/stablehlo/tests/vhlo/vhlo_to_version_downgrade_invalid.1_19_0.mlir ---- stablehlo/stablehlo/tests/vhlo/vhlo_to_version_downgrade_invalid.1_19_0.mlir -+++ stablehlo/stablehlo/tests/vhlo/vhlo_to_version_downgrade_invalid.1_19_0.mlir -@@ -0,0 +1,56 @@ -+// RUN: stablehlo-opt --stablehlo-legalize-to-vhlo --vhlo-to-version='target=1.19.0' --verify-diagnostics --split-input-file %s -+ -+// expected-error @+1 {{failed to convert VHLO to v1.19.0}} -+module { -+ func.func @convolution_mixed_fp8(%arg0: tensor<100x26x26x32xf8E5M2>, %arg1: tensor<3x3x1x32xf8E4M3FN>) -> tensor<100x28x28x1xf8E5M2> { -+ // expected-error @+1 {{failed to legalize operation 'vhlo.convolution_v1' that was explicitly marked illegal}} -+ %result = "stablehlo.convolution"(%arg0, %arg1) { -+ batch_group_count = 1 : i64, -+ dimension_numbers = #stablehlo.conv, -+ feature_group_count = 1 : i64, -+ lhs_dilation = array, -+ padding = dense<2> : tensor<2x2xi64>, -+ rhs_dilation = array, -+ window_strides = array -+ } : (tensor<100x26x26x32xf8E5M2>, tensor<3x3x1x32xf8E4M3FN>) -> tensor<100x28x28x1xf8E5M2> -+ func.return %result : tensor<100x28x28x1xf8E5M2> -+ } -+} -+ -+// ----- -+ -+// expected-error @+1 {{failed to convert VHLO to v1.19.0}} -+module { -+ func.func @dynamic_conv_mixed_fp8(%arg0: tensor<100x26x26x32xf8E5M2>, %arg1: tensor<3x3x1x32xf8E4M3FN>, %arg2: tensor<2x2xi64>) -> tensor<100x28x28x1xf8E5M2> { -+ // expected-error @+1 {{failed to legalize operation 'vhlo.dynamic_conv_v2' that was explicitly marked illegal}} -+ %result = "stablehlo.dynamic_conv"(%arg0, %arg1, %arg2) { -+ batch_group_count = 1 : i64, -+ dimension_numbers = #stablehlo.conv, -+ feature_group_count = 1 : i64, -+ lhs_dilation = array, -+ rhs_dilation = array, -+ window_strides = array -+ } : (tensor<100x26x26x32xf8E5M2>, tensor<3x3x1x32xf8E4M3FN>, tensor<2x2xi64>) -> tensor<100x28x28x1xf8E5M2> -+ func.return %result : tensor<100x28x28x1xf8E5M2> -+ } -+} -diff --ruN a/stablehlo/stablehlo/transforms/ChloLegalizeToStablehlo.cpp b/stablehlo/stablehlo/transforms/ChloLegalizeToStablehlo.cpp ---- stablehlo/stablehlo/transforms/ChloLegalizeToStablehlo.cpp -+++ stablehlo/stablehlo/transforms/ChloLegalizeToStablehlo.cpp -@@ -1668,18 +1668,27 @@ - digamma = mlir::stablehlo::SelectOp::create(rewriter, loc, needToReflect, - reflection, digamma); - -- // Digamma has poles at negative integers and zero; return nan for those. -- Value isLeZero = mlir::stablehlo::CompareOp::create( -- rewriter, loc, x, zero, mlir::stablehlo::ComparisonDirection::LE); -+ // Digamma has poles at negative integers and zero; return nan for -+ // negative integers, and -inf for zero. -+ Value isZero = mlir::stablehlo::CompareOp::create( -+ rewriter, loc, x, zero, mlir::stablehlo::ComparisonDirection::EQ); -+ Value isLtZero = mlir::stablehlo::CompareOp::create( -+ rewriter, loc, x, zero, mlir::stablehlo::ComparisonDirection::LT); - Value isInt = mlir::stablehlo::CompareOp::create( - rewriter, loc, x, mlir::stablehlo::FloorOp::create(rewriter, loc, x), - mlir::stablehlo::ComparisonDirection::EQ); -- Value isPole = mlir::stablehlo::AndOp::create(rewriter, loc, isLeZero, isInt); -- return mlir::stablehlo::SelectOp::create( -- rewriter, loc, isPole, -+ Value isNegativeInteger = -+ mlir::stablehlo::AndOp::create(rewriter, loc, isLtZero, isInt); -+ Value resultWithNan = mlir::stablehlo::SelectOp::create( -+ rewriter, loc, isNegativeInteger, - getConstantLike(rewriter, loc, std::numeric_limits::quiet_NaN(), - x), - digamma); -+ return mlir::stablehlo::SelectOp::create( -+ rewriter, loc, isZero, -+ getConstantLike(rewriter, loc, -std::numeric_limits::infinity(), -+ x), -+ resultWithNan); - } - - namespace { -@@ -2752,13 +2761,25 @@ - LogicalResult matchAndRewrite( - mlir::chlo::RaggedDotOp op, OpAdaptor, - ConversionPatternRewriter& rewriter) const override { -- if (op.getLhs().getType().getRank() < op.getRhs().getType().getRank()) { -+ chlo::RaggedDotDimensionNumbersAttr raggedDotDimensionNumbers = -+ op.getRaggedDotDimensionNumbers(); -+ ArrayRef lhsRaggedDimensions = -+ raggedDotDimensionNumbers.getLhsRaggedDimensions(); -+ if (lhsRaggedDimensions.empty()) { -+ return rewriter.notifyMatchFailure( -+ op, "lhs_ragged_dimensions must not be empty"); -+ } -+ const int64_t lhsRaggedDim = lhsRaggedDimensions[0]; -+ if (llvm::is_contained( -+ raggedDotDimensionNumbers.getLhsContractingDimensions(), -+ lhsRaggedDim)) { -+ return handleRaggedDotMode2(op, rewriter); -+ } else if (llvm::is_contained( -+ raggedDotDimensionNumbers.getLhsBatchingDimensions(), -+ lhsRaggedDim)) { -+ return handleRaggedDotMode3(op, rewriter); -+ } else { - return handleRaggedDotMode1(op, rewriter); -- } else if (op.getLhs().getType().getRank() < -- op.getResult().getType().getRank()) { -- return handleRaggedDotMode2(op, rewriter); -- } else { -- return handleRaggedDotMode3(op, rewriter); - } - } - }; -diff --ruN a/stablehlo/stablehlo/transforms/StablehloBroadcastLowering.cpp b/stablehlo/stablehlo/transforms/StablehloBroadcastLowering.cpp ---- stablehlo/stablehlo/transforms/StablehloBroadcastLowering.cpp -+++ stablehlo/stablehlo/transforms/StablehloBroadcastLowering.cpp -@@ -224,23 +224,51 @@ - mlir::RankedTensorType outputType = - getRankedTensorType(shape, inputType.getElementType()); - -- // Short circuit if no broadcasting is needed. -- if (inputType == outputType) return input; -- - int64_t inputRank = inputType.getRank(); - int64_t outputRank = outputType.getRank(); - if (inputRank > outputRank) - return emitError(loc, "input rank must be <= output rank, got ") - << inputRank << " vs " << outputRank; - -- size_t rankDiff = outputRank - inputRank; -+ // Construct broadcast dimensions (right-aligned for NumPy-style -+ // broadcasting). -+ auto broadcastDimensions = -+ llvm::to_vector(llvm::seq(outputRank - inputRank, outputRank)); -+ -+ return broadcastIfNeeded(builder, input, shape, broadcastDimensions); -+} -+ -+FailureOr broadcastIfNeeded(OpBuilder& builder, Value input, -+ const Dimensions& shape, -+ ArrayRef broadcastDimensions) { -+ LLVM_DEBUG(llvm::dbgs() << "[broadcastIfNeeded] Broadcasting input " -+ << input.getType() << " => " << toString(shape) -+ << "\n"); -+ auto loc = input.getLoc(); -+ mlir::RankedTensorType inputType = -+ dyn_cast(input.getType()); -+ if (!inputType) -+ return emitError(loc, "expected ranked tensor type for broadcast inputs"); -+ mlir::RankedTensorType outputType = -+ getRankedTensorType(shape, inputType.getElementType()); -+ -+ // Short circuit if no broadcasting is needed. -+ if (inputType == outputType) return input; -+ -+ int64_t inputRank = inputType.getRank(); -+ int64_t outputRank = outputType.getRank(); -+ if (inputRank > outputRank) -+ return emitError(loc, "input rank must be <= output rank, got ") -+ << inputRank << " vs " << outputRank; -+ -+ if (static_cast(broadcastDimensions.size()) != inputRank) -+ return emitError(loc, "broadcast_dimensions size (") -+ << broadcastDimensions.size() << ") must match input rank (" -+ << inputRank << ")"; -+ - auto inputShapeOrFail = getDimensions(input); - if (failed(inputShapeOrFail)) return failure(); - Dimensions inputShape = std::move(*inputShapeOrFail); -- -- // Construct broadcast dimensions. -- auto broadcastDimensions = -- llvm::to_vector(llvm::seq(outputRank - inputRank, outputRank)); - - // Construct the result type of the broadcast - // - If input is static and target shape is static, use static shape. -@@ -248,14 +276,30 @@ - // - If input is not bounded, but target shape is bounded, broadcast to - // the padded shape then call SetDimensionSize to make dynamic. - auto bcastShape = shape; -+ llvm::SmallVector isMapped(outputRank, false); -+ - for (int64_t i = 0; i < inputRank; ++i) { -+ int64_t resultIdx = broadcastDimensions[i]; -+ if (resultIdx < 0 || resultIdx >= outputRank) -+ return emitError(loc, "broadcast_dimensions index ") -+ << resultIdx << " out of bounds for output rank " << outputRank; -+ -+ isMapped[resultIdx] = true; -+ - int64_t inputDimSize = inputShape[i].size; -- int64_t resultIdx = i + rankDiff; - int64_t resultDimSize = shape[resultIdx].size; - if (inputDimSize != 1 && inputDimSize != resultDimSize) - return emitError(loc, "Cannot broadcast input: ") - << inputType << " to target shape " << toString(shape); - -+ if (inputShape[i].boundOp.has_value() && -+ !shape[resultIdx].boundOp.has_value()) { -+ return emitError( -+ loc, "cannot mix bounded and static dimensions in broadcast: ") -+ << "input dimension " << i << " is bounded, but target dimension " -+ << resultIdx << " is static"; -+ } -+ - if (!inputShape[i].boundOp.has_value() && - shape[resultIdx].boundOp.has_value()) { - // Use padded shape in broadcast. -@@ -263,17 +307,18 @@ - } - } - -- // Broadcast to padded size for remaining dimensions. -- for (size_t i = 0; i < rankDiff; ++i) { -- bcastShape[i] = DimensionInfo{shape[i].size}; -+ // Broadcast to padded size for remaining unmapped dimensions. -+ for (int64_t i = 0; i < outputRank; ++i) { -+ if (!isMapped[i]) { -+ bcastShape[i] = DimensionInfo{shape[i].size}; -+ } - } - - // Insert broadcast ops - mlir::RankedTensorType bcastType = - getRankedTensorType(bcastShape, inputType.getElementType()); -- LLVM_DEBUG( -- llvm::dbgs() << "[numpyBroadcastIfNeeded] Broadcast to padded type " -- << bcastType << "\n"); -+ LLVM_DEBUG(llvm::dbgs() << "[broadcastIfNeeded] Broadcast to padded type " -+ << bcastType << "\n"); - Value bcastOp = stablehlo::BroadcastInDimOp::create( - builder, loc, bcastType, input, broadcastDimensions); - if (bcastOp.getType() == outputType) return bcastOp; -diff --ruN a/stablehlo/stablehlo/transforms/StablehloBroadcastLowering.h b/stablehlo/stablehlo/transforms/StablehloBroadcastLowering.h ---- stablehlo/stablehlo/transforms/StablehloBroadcastLowering.h -+++ stablehlo/stablehlo/transforms/StablehloBroadcastLowering.h -@@ -70,6 +70,12 @@ - FailureOr numpyBroadcastIfNeeded(OpBuilder& builder, Value input, - const Dimensions& shape); - -+// Apply broadcasting to the given operand using the specified -+// broadcast_dimensions, returning an error if the operand is not broadcastable. -+FailureOr broadcastIfNeeded(OpBuilder& builder, Value input, -+ const Dimensions& shape, -+ ArrayRef broadcastDimensions); -+ - } // namespace stablehlo - } // namespace mlir - diff --git a/third_party/xla/third_party/stablehlo/workspace.bzl b/third_party/xla/third_party/stablehlo/workspace.bzl index 5f1d45c57e8e8b..c8508914e189ec 100644 --- a/third_party/xla/third_party/stablehlo/workspace.bzl +++ b/third_party/xla/third_party/stablehlo/workspace.bzl @@ -19,8 +19,8 @@ load("//third_party:repo.bzl", "tf_http_archive", "tf_mirror_urls") def repo(): # LINT.IfChange - STABLEHLO_COMMIT = "7b1b15781ccbd770f50c7eef4b0c3e03834649fd" - STABLEHLO_SHA256 = "d498ed1ba288eaf508369f76912e14d52d7f1f1a31ddaddd8ad4be40c8b7b4dc" + STABLEHLO_COMMIT = "639335932274930617d9724fb25ec9f2f52c5e40" + STABLEHLO_SHA256 = "dd590949798e26838d9ac3d6016def9b508caa60d0d7a69535ed2f5454ad30b8" # LINT.ThenChange(Google-internal path) tf_http_archive( diff --git a/third_party/xla/xla/hlo/ir/hlo_original_value_test.cc b/third_party/xla/xla/hlo/ir/hlo_original_value_test.cc index 3ebc8b62570638..6b8553aa87af37 100644 --- a/third_party/xla/xla/hlo/ir/hlo_original_value_test.cc +++ b/third_party/xla/xla/hlo/ir/hlo_original_value_test.cc @@ -500,5 +500,34 @@ TEST_F(OriginalValueHloTest, CopyOriginalValueWithMap) { Optional(Eq(OriginalArray{"instA", {0}}))); } +TEST_F(OriginalValueHloTest, CopyOriginalValueWithVector) { + const char* hlo_string = R"( +HloModule test + +ENTRY main { + p0 = f32[] parameter(0) + p1 = f32[] parameter(1) + p2 = f32[] parameter(2) + ROOT tuple = (f32[], f32[], f32[]) tuple(p0, p1, p2), origin={({"p0"}, {"p1"}, {"p2"})} +} +)"; + ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); + HloInstruction* root = module->entry_computation()->root_instruction(); + HloInstruction* dest = + module->entry_computation()->AddInstruction(HloInstruction::CreateTuple( + {root->mutable_operand(2), root->mutable_operand(0)})); + + // Index 0 -> 1, Index 1 is pruned (-1), Index 2 -> 0. + std::vector old_to_new_tuple_idx = {1, -1, 0}; + + CopyOriginalValue(root, dest, old_to_new_tuple_idx); + + ASSERT_NE(dest->original_value(), nullptr); + EXPECT_THAT(dest->original_value()->original_array({0}), + Optional(Eq(OriginalArray{"p2"}))); + EXPECT_THAT(dest->original_value()->original_array({1}), + Optional(Eq(OriginalArray{"p0"}))); +} + } // namespace } // namespace xla diff --git a/third_party/xla/xla/hlo/ir/hlo_original_value_util.h b/third_party/xla/xla/hlo/ir/hlo_original_value_util.h index c3e4740e538176..671127ce011072 100644 --- a/third_party/xla/xla/hlo/ir/hlo_original_value_util.h +++ b/third_party/xla/xla/hlo/ir/hlo_original_value_util.h @@ -16,9 +16,11 @@ limitations under the License. #ifndef XLA_HLO_IR_HLO_ORIGINAL_VALUE_UTIL_H_ #define XLA_HLO_IR_HLO_ORIGINAL_VALUE_UTIL_H_ +#include #include #include #include +#include #include "absl/container/flat_hash_map.h" #include "xla/hlo/ir/hlo_instruction.h" @@ -28,8 +30,12 @@ limitations under the License. namespace xla { // Checks if the type of the map is a matching integer map. +template +struct is_matching_integer_map : std::false_type {}; + template -struct is_matching_integer_map { +struct is_matching_integer_map< + T, std::void_t> { static constexpr bool value = std::is_integral::value && std::is_same::value; @@ -81,6 +87,26 @@ CopyOriginalValue(const HloInstruction* src_instruction, dest_instruction->set_original_value(new_original_value); } +// Copies the original value of the source to the destination instruction. +// Original arrays in the source original value are rearranged in the new +// original value according to the given vector of old to new tuple indices. +// Elements with negative values in the vector are treated as pruned/unused. +template +std::enable_if_t> CopyOriginalValue( + const HloInstruction* src_instruction, HloInstruction* dest_instruction, + const std::vector& old_to_new_tuple_idx) { + absl::flat_hash_map mapping; + for (size_t old_idx = 0; old_idx < old_to_new_tuple_idx.size(); ++old_idx) { + if constexpr (std::is_signed_v) { + if (old_to_new_tuple_idx[old_idx] < 0) { + continue; + } + } + mapping[static_cast(old_idx)] = old_to_new_tuple_idx[old_idx]; + } + CopyOriginalValue(src_instruction, dest_instruction, mapping); +} + // Copies the original value of the source to the destination instruction if the // shapes of the source and destination are compatible. This performs a deep // copy if clone is set to true. Otherwise, it performs a shallow copy. Print a diff --git a/third_party/xla/xla/hlo/transforms/simplifiers/algebraic_simplifier.cc b/third_party/xla/xla/hlo/transforms/simplifiers/algebraic_simplifier.cc index a72a150261dba4..2a3ef2c1cb5aef 100644 --- a/third_party/xla/xla/hlo/transforms/simplifiers/algebraic_simplifier.cc +++ b/third_party/xla/xla/hlo/transforms/simplifiers/algebraic_simplifier.cc @@ -54,6 +54,7 @@ limitations under the License. #include "xla/hlo/ir/hlo_instructions.h" #include "xla/hlo/ir/hlo_opcode.h" #include "xla/hlo/ir/hlo_original_value.h" +#include "xla/hlo/ir/hlo_original_value_util.h" #include "xla/hlo/ir/hlo_sharding.h" #include "xla/hlo/transforms/simplifiers/conv_operand_swapper.h" #include "xla/hlo/utils/hlo_sharding_util.h" @@ -5430,6 +5431,8 @@ absl::Status AlgebraicSimplifierVisitor::HandleOptimizationBarrier( operand->AddInstruction(HloInstruction::CreateTuple(operands)); ABSL_RETURN_IF_ERROR(barrier->ReplaceOperandWithDifferentShape(0, new_operand)); *barrier->mutable_shape() = new_operand->shape(); + CopyOriginalValue(barrier, barrier, index_map); + CopyOriginalValue(operand, new_operand, index_map); for (auto use : barrier->users()) { CHECK_EQ(use->opcode(), HloOpcode::kGetTupleElement); use->set_tuple_index(index_map[use->tuple_index()]); diff --git a/third_party/xla/xla/hlo/transforms/simplifiers/algebraic_simplifier_test.cc b/third_party/xla/xla/hlo/transforms/simplifiers/algebraic_simplifier_test.cc index 5c827aec5f682e..3d80a30755bf1a 100644 --- a/third_party/xla/xla/hlo/transforms/simplifiers/algebraic_simplifier_test.cc +++ b/third_party/xla/xla/hlo/transforms/simplifiers/algebraic_simplifier_test.cc @@ -12669,6 +12669,36 @@ TEST_F(AlgebraicSimplifierTest, SimplifyOptimizationBarrier) { 2); } +TEST_F(AlgebraicSimplifierTest, SimplifyOptimizationBarrierWithOriginalValue) { + constexpr absl::string_view kModuleStr = R"( + HloModule m + + ENTRY entry { + param.0 = f32[] parameter(0) + param.1 = f32[] parameter(1) + add.0 = f32[] add(param.0, param.1) + sub.0 = f32[] subtract(param.0, param.1) + mul.0 = f32[] multiply(param.0, param.1) + tuple.0 = (f32[], f32[], f32[]) tuple(mul.0, sub.0, add.0), origin={({"mul.0"}, {"sub.0"}, {"add.0"})} + b = (f32[], f32[], f32[]) opt-barrier(tuple.0), origin={({"b" {0}}, {"b" {1}}, {"b" {2}})} + gte.0 = f32[] get-tuple-element(b), index=1 + ROOT t = (f32[]) tuple(gte.0) + } + )"; + ASSERT_OK_AND_ASSIGN(auto m, ParseAndReturnVerifiedModule(kModuleStr)); + ASSERT_OK_AND_ASSIGN(bool changed, + AlgebraicSimplifier(default_options_).Run(m.get())); + EXPECT_TRUE(changed); + ASSERT_THAT(verifier().Run(m.get()), absl_testing::IsOk()); + const HloInstruction* b = FindInstruction(m.get(), "b"); + ASSERT_NE(b, nullptr); + EXPECT_EQ(b->shape().tuple_shapes().size(), 1); + ASSERT_NE(b->original_value(), nullptr); + EXPECT_EQ(b->original_value()->original_array({0})->instruction_name, "b"); + EXPECT_EQ(b->original_value()->original_array({0})->shape_index, + ShapeIndex({1})); +} + TEST_F(AlgebraicSimplifierTest, DoNotSimplifyOptimizationBarrierSideEffects) { constexpr absl::string_view kModuleStr = R"( HloModule m diff --git a/third_party/xla/xla/hlo/translate/mhlo_to_hlo/mlir_hlo_to_hlo.cc b/third_party/xla/xla/hlo/translate/mhlo_to_hlo/mlir_hlo_to_hlo.cc index 85852dba54c504..1db8afcf410d30 100644 --- a/third_party/xla/xla/hlo/translate/mhlo_to_hlo/mlir_hlo_to_hlo.cc +++ b/third_party/xla/xla/hlo/translate/mhlo_to_hlo/mlir_hlo_to_hlo.cc @@ -2481,8 +2481,8 @@ LogicalResult ExportXlaOp(BitcastConvertOp op, OpLoweringContext ctx) { LogicalResult ExportXlaOp(CollectiveBroadcastOp op, OpLoweringContext ctx) { auto& value_map = *ctx.values; - xla::XlaOp operand; - if (failed(GetXlaOp(op.getOperand(), value_map, &operand, op))) { + SmallVector operands; + if (failed(GetTuple(op.getOperation(), op.getOperands(), ctx, operands))) { return failure(); } auto replica_groups = Convert_replica_groups(op.getReplicaGroups(), op); @@ -2490,8 +2490,20 @@ LogicalResult ExportXlaOp(CollectiveBroadcastOp op, OpLoweringContext ctx) { return op.emitOpError(replica_groups.status().ToString()); } auto result = xla::CollectiveBroadcastWithDeviceList( - operand, **replica_groups, Convert_channel_handle(op.getChannelHandle())); - value_map[op->getResult(0)] = result; + operands, **replica_groups, Convert_channel_handle(op.getChannelHandle()), + op.getHasDynamicRoot()); + + // A collective_broadcast with more than one data operand produces a tuple. + mlir::FailureOr shape_or = + xla::ExtractXlaShape(op.getOperation()); + if (failed(shape_or)) { + return failure(); + } + if (shape_or->IsTuple()) { + BuildGetTupleElementsForTupleResults(op, result, ctx); + } else { + value_map[op->getResult(0)] = result; + } return success(); } @@ -3616,7 +3628,7 @@ LogicalResult ExportXlaOp(AsyncStartOp op, OpLoweringContext ctx) { xla::Shape input_shape = xla::ShapeUtil::MakeTupleShape( {xla::TypeToShape(op.getOperand(0).getType())}); - xla::Shape output_shape = xla::TypeToShape(collective_broadcast.getType()); + xla::Shape output_shape = xla::TypeToShape(collective_broadcast.getType(0)); xla::Shape start_shape = xla::ShapeUtil::MakeTupleShape({input_shape, output_shape}); (*ctx.values)[op.getResult()] = @@ -3698,7 +3710,7 @@ LogicalResult ExportXlaOp(AsyncDoneOp op, OpLoweringContext ctx) { (*ctx.values)[op.getResult()] = xla::internal::XlaBuilderFriend::BuildAsyncDone( ctx.builder, operand, - xla::TypeToShape(collective_broadcast.getType())); + xla::TypeToShape(collective_broadcast.getType(0))); return success(); } diff --git a/third_party/xla/xla/hlo/utils/BUILD b/third_party/xla/xla/hlo/utils/BUILD index c5aa2126276acc..c4b0ff8a532a9c 100644 --- a/third_party/xla/xla/hlo/utils/BUILD +++ b/third_party/xla/xla/hlo/utils/BUILD @@ -78,8 +78,10 @@ cc_library( ":hlo_stack_trace", "//xla:shape_util", "//xla/hlo/analysis:hlo_alias_analysis", + "//xla/hlo/analysis:hlo_dataflow_analysis", "//xla/hlo/ir:hlo", "//xla/hlo/ir:hlo_instruction_utils", + "//xla/service:buffer_value", "//xla/service:hlo_buffer", "//xla/service:hlo_value", "@com_google_absl//absl/algorithm:container", @@ -110,6 +112,7 @@ xla_cc_test( "//xla/hlo/ir:hlo", "//xla/hlo/parser:hlo_parser", "//xla/hlo/testlib:hlo_hardware_independent_test_base", + "//xla/service:buffer_value", "//xla/service:hlo_value", "//xla/tsl/lib/core:status_test_util", "//xla/tsl/platform:statusor", diff --git a/third_party/xla/xla/hlo/utils/hlo_live_range.cc b/third_party/xla/xla/hlo/utils/hlo_live_range.cc index 6288d09d8723c5..1db60423a53b3e 100644 --- a/third_party/xla/xla/hlo/utils/hlo_live_range.cc +++ b/third_party/xla/xla/hlo/utils/hlo_live_range.cc @@ -37,12 +37,14 @@ limitations under the License. #include "absl/strings/string_view.h" #include "absl/types/span.h" #include "xla/hlo/analysis/hlo_alias_analysis.h" +#include "xla/hlo/analysis/hlo_dataflow_analysis.h" #include "xla/hlo/ir/dfs_hlo_visitor.h" #include "xla/hlo/ir/hlo_computation.h" #include "xla/hlo/ir/hlo_instruction_utils.h" #include "xla/hlo/ir/hlo_opcode.h" #include "xla/hlo/ir/hlo_schedule.h" #include "xla/hlo/utils/hlo_stack_trace.h" +#include "xla/service/buffer_value.h" #include "xla/service/hlo_buffer.h" #include "xla/service/hlo_value.h" #include "xla/shape.h" @@ -63,6 +65,119 @@ absl::StatusOr> HloLiveRange::Run( return hlo_live_range; } +/*static*/ +std::vector HloLiveRange::GetValuesDefined( + const HloInstruction* instruction, const HloDataflowAnalysis& dataflow) { + std::vector values; + const auto& value_set_tree = dataflow.GetInstructionValueSet(instruction); + for (const auto& entry : value_set_tree) { + if (dataflow.ValueIsDefinedAt(instruction, entry.first)) { + values.push_back(&dataflow.GetValueDefinedAt(instruction, entry.first)); + } + } + return values; +} + +/*static*/ +std::vector HloLiveRange::GetBuffersDefined( + const HloInstruction* instruction, const HloAliasAnalysis& alias_analysis) { + std::vector buffers; + absl::flat_hash_set seen; + for (const HloValue* value : + GetValuesDefined(instruction, alias_analysis.dataflow_analysis())) { + const HloBuffer* buffer = &alias_analysis.GetBufferContainingValue(*value); + if (seen.insert(buffer).second) { + buffers.push_back(buffer); + } + } + return buffers; +} + +/*static*/ +int64_t HloLiveRange::GetBytesDefined( + const HloInstruction* instruction, const HloAliasAnalysis& alias_analysis, + const BufferValue::SizeFunction& size_fn) { + if (instruction->opcode() == HloOpcode::kParameter) { + return 0; + } + int64_t bytes = 0; + for (const HloBuffer* buffer : + GetBuffersDefined(instruction, alias_analysis)) { + if (buffer->IsHeapPressureImpacting()) { + bytes += buffer->ComputeSize(size_fn); + } + } + return bytes; +} + +/*static*/ +std::vector HloLiveRange::GetBuffersUsed( + const HloInstruction* instruction, const HloAliasAnalysis& alias_analysis) { + std::vector buffers; + absl::flat_hash_set seen; + const auto& dataflow = alias_analysis.dataflow_analysis(); + for (const HloInstruction* operand : instruction->operands()) { + HloValueSet value_set = dataflow.GetFlattenedValueSet(operand); + for (const HloValue* value : value_set.values()) { + const HloBuffer* buffer = + &alias_analysis.GetBufferContainingValue(*value); + if (seen.insert(buffer).second) { + buffers.push_back(buffer); + } + } + } + return buffers; +} + +/*static*/ +int32_t HloLiveRange::GetTotalUsers(const HloBuffer& buffer, + const HloComputation* computation) { + int32_t total = 0; + for (const HloValue* value : buffer.values()) { + for (const HloUse& use : value->GetUses()) { + if (computation == nullptr || use.instruction->parent() == computation) { + ++total; + } + } + } + return total; +} + +/*static*/ +int64_t HloLiveRange::GetParameterBytesAtStart( + const HloComputation& computation, const HloAliasAnalysis& alias_analysis, + const BufferValue::SizeFunction& size_fn) { + int64_t bytes = 0; + absl::flat_hash_set seen; + for (const HloInstruction* param : computation.parameter_instructions()) { + for (const HloBuffer* buffer : GetBuffersDefined(param, alias_analysis)) { + if (seen.insert(buffer).second && buffer->IsHeapPressureImpacting()) { + bytes += buffer->ComputeSize(size_fn); + } + } + } + return bytes; +} + +/*static*/ +bool HloLiveRange::BufferLivesOut(const HloBuffer& buffer, + const HloAliasAnalysis& alias_analysis, + const HloComputation* computation) { + if (alias_analysis.BufferLivesOut(buffer)) { + return true; + } + if (computation != nullptr) { + for (const HloValue* value : buffer.values()) { + for (const HloUse& use : value->GetUses()) { + if (use.instruction->parent() != computation) { + return true; + } + } + } + } + return false; +} + void HloLiveRange::NormalizeAliasedBuffers() { absl::flat_hash_map>> @@ -288,36 +403,26 @@ void HloLiveRange::CalculateBufferStartEndMap() { << definition_end_time; } - const InstructionValueSet& value_set_tree = - alias_analysis_.dataflow_analysis().GetInstructionValueSet( - &instruction); - - for (const auto& entry : value_set_tree) { - for (const HloValue* value : entry.second.values()) { - // The start time is only correct for the defining instruction. - if (value->defining_instruction() != &instruction) { - continue; - } - - auto [end_time, end_position] = - ComputeValueLiveRangeEnd(*value, definition_end_time); - LiveRangeBounds live_range{start_time, end_time, end_position}; - - // Readonly entry parameters (parameters that don't alias) live across - // whole computation. - const HloModule& module = *computation->parent(); - if (instruction.opcode() == HloOpcode::kParameter && - computation == module.entry_computation() && - !module.input_output_alias_config().ParameterHasAlias( - instruction.parameter_number(), value->index())) { - live_range.end = schedule_end_time(); - } else { - live_range.end = std::max(live_range.end, GetLastUsageTime(*value)); - } - - CHECK_LE(live_range.start, live_range.end) << instruction.ToString(); - CHECK(buffer_live_ranges_.insert({value, live_range}).second); + for (const HloValue* value : + GetValuesDefined(&instruction, alias_analysis_.dataflow_analysis())) { + auto [end_time, end_position] = + ComputeValueLiveRangeEnd(*value, definition_end_time); + LiveRangeBounds live_range{start_time, end_time, end_position}; + + // Readonly entry parameters (parameters that don't alias) live across + // whole computation. + const HloModule& module = *computation->parent(); + if (instruction.opcode() == HloOpcode::kParameter && + computation == module.entry_computation() && + !module.input_output_alias_config().ParameterHasAlias( + instruction.parameter_number(), value->index())) { + live_range.end = schedule_end_time(); + } else { + live_range.end = std::max(live_range.end, GetLastUsageTime(*value)); } + + CHECK_LE(live_range.start, live_range.end) << instruction.ToString(); + CHECK(buffer_live_ranges_.insert({value, live_range}).second); } } } diff --git a/third_party/xla/xla/hlo/utils/hlo_live_range.h b/third_party/xla/xla/hlo/utils/hlo_live_range.h index cc407ced4e42bd..92758777416b47 100644 --- a/third_party/xla/xla/hlo/utils/hlo_live_range.h +++ b/third_party/xla/xla/hlo/utils/hlo_live_range.h @@ -19,6 +19,7 @@ the License. #include #include #include +#include #include "absl/container/flat_hash_map.h" #include "absl/container/flat_hash_set.h" @@ -26,9 +27,12 @@ the License. #include "absl/status/statusor.h" #include "absl/strings/string_view.h" #include "xla/hlo/analysis/hlo_alias_analysis.h" +#include "xla/hlo/analysis/hlo_dataflow_analysis.h" #include "xla/hlo/ir/dfs_hlo_visitor.h" #include "xla/hlo/ir/hlo_instruction.h" #include "xla/hlo/ir/hlo_schedule.h" +#include "xla/service/buffer_value.h" +#include "xla/service/hlo_buffer.h" #include "xla/service/hlo_value.h" namespace xla { @@ -46,6 +50,46 @@ class HloLiveRange { const HloComputation* computation, bool module_scoped_analysis = true, absl::flat_hash_set execution_threads = {}); + // Returns all HloValues defined by this instruction. + static std::vector GetValuesDefined( + const HloInstruction* instruction, const HloDataflowAnalysis& dataflow); + + // Returns the distinct physical HloBuffers newly allocated/defined by this + // instruction (excluding buffers forwarded or aliased from operands). + static std::vector GetBuffersDefined( + const HloInstruction* instruction, + const HloAliasAnalysis& alias_analysis); + + // Returns the total bytes defined by this instruction according to size_fn. + // Returns 0 for parameter instructions (parameters are attributed to + // computation start). + static int64_t GetBytesDefined(const HloInstruction* instruction, + const HloAliasAnalysis& alias_analysis, + const BufferValue::SizeFunction& size_fn); + + // Returns the distinct physical HloBuffers read by the operands of this + // instruction. + static std::vector GetBuffersUsed( + const HloInstruction* instruction, + const HloAliasAnalysis& alias_analysis); + + // Returns the total number of instruction reads across the computation for + // all values contained in this buffer. If computation is null, counts across + // all computations. + static int32_t GetTotalUsers(const HloBuffer& buffer, + const HloComputation* computation = nullptr); + + // Returns the total parameter bytes allocated at the start of the computation + // (matching HloLiveRange parameter attribution). + static int64_t GetParameterBytesAtStart( + const HloComputation& computation, const HloAliasAnalysis& alias_analysis, + const BufferValue::SizeFunction& size_fn); + + // Returns true if any value in this buffer lives out of the computation. + static bool BufferLivesOut(const HloBuffer& buffer, + const HloAliasAnalysis& alias_analysis, + const HloComputation* computation = nullptr); + // LogicalTime represents the time in a virtual clock. Each instruction has // one monotonically increasing logical time assigned according to the // schedule. diff --git a/third_party/xla/xla/hlo/utils/hlo_live_range_test.cc b/third_party/xla/xla/hlo/utils/hlo_live_range_test.cc index cc2a4a9963364e..e3bb28c339bff7 100644 --- a/third_party/xla/xla/hlo/utils/hlo_live_range_test.cc +++ b/third_party/xla/xla/hlo/utils/hlo_live_range_test.cc @@ -37,6 +37,7 @@ limitations under the License. #include "xla/hlo/parser/hlo_parser.h" #include "xla/hlo/testlib/hlo_hardware_independent_test_base.h" #include "xla/literal_util.h" +#include "xla/service/buffer_value.h" #include "xla/service/hlo_value.h" #include "xla/shape.h" #include "xla/shape_util.h" @@ -1140,5 +1141,66 @@ ENTRY %entry { EXPECT_TRUE(hlo_live_range_->instruction_schedule().contains(neg0)); } +TEST_F(HloLiveRangeTest, HelpersTest) { + const char* hlo_string = R"( +HloModule module + +ENTRY entry { + param0 = f32[100] parameter(0) + param1 = f32[200] parameter(1) + const = f32[10] constant({0, 1, 2, 3, 4, 5, 6, 7, 8, 9}) + add0 = f32[100] add(param0, param0) + add1 = f32[200] add(param1, param1) + ROOT root = (f32[100], f32[200]) tuple(add0, add1) +} +)"; + + ASSERT_OK_AND_ASSIGN(module_, ParseAndReturnVerifiedModule(hlo_string)); + ASSERT_OK_AND_ASSIGN(std::unique_ptr aa, + HloAliasAnalysis::Run(module_.get(), &alias_info_)); + + HloComputation* entry = module_->entry_computation(); + const HloInstruction* param0 = entry->parameter_instruction(0); + const HloInstruction* param1 = entry->parameter_instruction(1); + const HloInstruction* const_inst = entry->GetInstructionWithName("const"); + const HloInstruction* add0 = entry->GetInstructionWithName("add0"); + const HloInstruction* add1 = entry->GetInstructionWithName("add1"); + const HloInstruction* root = entry->root_instruction(); + + auto size_fn = [](const BufferValue& buffer) { + return ShapeUtil::ByteSizeOf(buffer.shape(), 8); + }; + + // 1. GetBuffersDefined + auto param0_buffers = HloLiveRange::GetBuffersDefined(param0, *aa); + EXPECT_EQ(param0_buffers.size(), 1); + auto param1_buffers = HloLiveRange::GetBuffersDefined(param1, *aa); + EXPECT_EQ(param1_buffers.size(), 1); + auto add0_buffers = HloLiveRange::GetBuffersDefined(add0, *aa); + EXPECT_EQ(add0_buffers.size(), 1); + + // 2. GetBytesDefined + EXPECT_EQ(HloLiveRange::GetBytesDefined(param0, *aa, size_fn), 0); + EXPECT_EQ(HloLiveRange::GetBytesDefined(add0, *aa, size_fn), 400); + EXPECT_EQ(HloLiveRange::GetBytesDefined(add1, *aa, size_fn), 800); + EXPECT_EQ(HloLiveRange::GetBytesDefined(const_inst, *aa, size_fn), 0); + + // 3. GetBuffersUsed + auto add0_used = HloLiveRange::GetBuffersUsed(add0, *aa); + EXPECT_EQ(add0_used.size(), 1); + EXPECT_EQ(add0_used[0], param0_buffers[0]); + + // 4. GetParameterBytesAtStart + EXPECT_EQ(HloLiveRange::GetParameterBytesAtStart(*entry, *aa, size_fn), 1200); + + // 5. GetTotalUsers + EXPECT_EQ(HloLiveRange::GetTotalUsers(*param0_buffers[0], entry), 2); + + // 6. BufferLivesOut + EXPECT_FALSE(HloLiveRange::BufferLivesOut(*param0_buffers[0], *aa, entry)); + auto root_buffers = HloLiveRange::GetBuffersDefined(root, *aa); + EXPECT_FALSE(root_buffers.empty()); + EXPECT_TRUE(HloLiveRange::BufferLivesOut(*root_buffers[0], *aa, entry)); +} } // namespace } // namespace xla diff --git a/third_party/xla/xla/mlir_hlo/mhlo/IR/hlo_ops.cc b/third_party/xla/xla/mlir_hlo/mhlo/IR/hlo_ops.cc index 3fee8629835b1f..26ff56452437ea 100644 --- a/third_party/xla/xla/mlir_hlo/mhlo/IR/hlo_ops.cc +++ b/third_party/xla/xla/mlir_hlo/mhlo/IR/hlo_ops.cc @@ -1949,7 +1949,9 @@ void CollectiveBroadcastOp::build(OpBuilder& odsBuilder, } LogicalResult CollectiveBroadcastOp::verify() { - return hlo::verifyCollectiveBroadcastOp(getLoc(), getReplicaGroups()); + return hlo::verifyCollectiveBroadcastOp(getLoc(), (*this)->getOperands(), + getReplicaGroups(), + /*hasDynamicRoot=*/false); } //===----------------------------------------------------------------------===// diff --git a/third_party/xla/xla/mosaic/dialect/tpu/tpu_dialect.h b/third_party/xla/xla/mosaic/dialect/tpu/tpu_dialect.h index a87e9cd7426b5a..ea72eb35a3be98 100644 --- a/third_party/xla/xla/mosaic/dialect/tpu/tpu_dialect.h +++ b/third_party/xla/xla/mosaic/dialect/tpu/tpu_dialect.h @@ -53,6 +53,11 @@ class TPUDialect; namespace mlir { namespace tpu { +inline constexpr StringRef kTpuVectorSubcoreGlobalBarrierName = + "__tpu_vector_subcore_global_barrier"; +inline constexpr StringRef kTpuScalarSubcoreGlobalBarrierName = + "__tpu_scalar_subcore_global_barrier"; + DEFINE_ABSL_STRINGIFY_FOR_ENUMS(); struct TpuTilingFlags { diff --git a/third_party/xla/xla/pjrt/c_api_client/pjrt_c_api_client.cc b/third_party/xla/xla/pjrt/c_api_client/pjrt_c_api_client.cc index 9e08a8071358bb..7ea1ed00f02ba1 100644 --- a/third_party/xla/xla/pjrt/c_api_client/pjrt_c_api_client.cc +++ b/third_party/xla/xla/pjrt/c_api_client/pjrt_c_api_client.cc @@ -747,6 +747,15 @@ InitializeArgsAndCompileAot(const PJRT_Api* c_api, PjRtClient* client, } // namespace +absl::StatusOr> PjRtCApiClient::Compile( + const XlaComputation& computation, CompileOptions options) { + tsl::profiler::TraceMe traceme("PjRtCApiClient::Compile(XlaComputation)"); + ABSL_ASSIGN_OR_RETURN(const PjRtTopologyDescription* const topology, + GetTopologyDescription()); + return InitializeArgsAndCompileAot(c_api_, this, &computation, options, + *topology); +} + absl::StatusOr> PjRtCApiClient::Compile( MaybeOwningMlirModule module, CompileOptions options) { tsl::profiler::TraceMe traceme([&module]() { diff --git a/third_party/xla/xla/pjrt/c_api_client/pjrt_c_api_client.h b/third_party/xla/xla/pjrt/c_api_client/pjrt_c_api_client.h index e9fdfd162fc88f..4b4da58f2ffdcc 100644 --- a/third_party/xla/xla/pjrt/c_api_client/pjrt_c_api_client.h +++ b/third_party/xla/xla/pjrt/c_api_client/pjrt_c_api_client.h @@ -422,6 +422,9 @@ class PjRtCApiClient : public PjRtClient { absl::StatusOr GetDefaultLayout( PrimitiveType element_type, absl::Span dims) override; + absl::StatusOr> Compile( + const XlaComputation& computation, CompileOptions options) override; + absl::StatusOr> CompileAndLoad( const XlaComputation& computation, CompileOptions options) override; diff --git a/third_party/xla/xla/pjrt/c_api_client/pjrt_c_api_client_test.cc b/third_party/xla/xla/pjrt/c_api_client/pjrt_c_api_client_test.cc index 3c47f0ea8e8a07..b7efe0487da9d0 100644 --- a/third_party/xla/xla/pjrt/c_api_client/pjrt_c_api_client_test.cc +++ b/third_party/xla/xla/pjrt/c_api_client/pjrt_c_api_client_test.cc @@ -573,6 +573,23 @@ TEST(PjRtClientTest, CompileMlirModule) { EXPECT_NE(executable.get(), nullptr); } +TEST(PjRtClientTest, CompileXlaComputation) { + SetUpCpuPjRtApi(); + ASSERT_OK_AND_ASSIGN(std::unique_ptr client, + GetCApiClient("cpu")); + Shape shape = ShapeUtil::MakeShape(S32, {4}); + XlaBuilder builder("add_one"); + auto input = Parameter(&builder, 0, shape, "input"); + auto one = ConstantR0(&builder, 1); + auto add = Add(input, one); + ASSERT_OK_AND_ASSIGN(XlaComputation computation, builder.Build(add)); + + CompileOptions options; + ASSERT_OK_AND_ASSIGN(std::unique_ptr executable, + client->Compile(computation, options)); + EXPECT_NE(executable.get(), nullptr); +} + TEST(PjRtCApiClientTest, LoadExecutable) { SetUpCpuPjRtApi(); ASSERT_OK_AND_ASSIGN(std::unique_ptr client, diff --git a/third_party/xla/xla/python/ifrt/BUILD b/third_party/xla/xla/python/ifrt/BUILD index 7a5c8592aa348b..23bc8a98a5d733 100644 --- a/third_party/xla/xla/python/ifrt/BUILD +++ b/third_party/xla/xla/python/ifrt/BUILD @@ -126,6 +126,7 @@ cc_library( "//xla/tsl/platform:logging", "//xla/tsl/platform:statusor", "@com_google_absl//absl/algorithm:container", + "@com_google_absl//absl/base", "@com_google_absl//absl/base:core_headers", "@com_google_absl//absl/base:no_destructor", "@com_google_absl//absl/base:nullability", diff --git a/third_party/xla/xla/python/ifrt/sharding.cc b/third_party/xla/xla/python/ifrt/sharding.cc index dd16a79114abae..a641d51f29e315 100644 --- a/third_party/xla/xla/python/ifrt/sharding.cc +++ b/third_party/xla/xla/python/ifrt/sharding.cc @@ -31,6 +31,7 @@ limitations under the License. #include "absl/log/check.h" #include "absl/status/status.h" #include "absl/status/status_macros.h" +#include "absl/status/statusor.h" #include "absl/strings/str_cat.h" #include "absl/strings/str_format.h" #include "absl/strings/str_join.h" @@ -75,6 +76,16 @@ bool Sharding::operator==(const Sharding& other) const { *devices() == *other.devices(); } +absl::StatusOr> +Sharding::UniqueIndexDomains(const Shape& shape) const { + return sharding_spec()->UniqueIndexDomains(shape); +} + +absl::StatusOr> Sharding::ShardToUniqueIndexDomainIndex() + const { + return sharding_spec()->ShardToUniqueIndexDomainIndex(); +} + absl::StatusOr Sharding::FromProto( Client* client, const ShardingProto& sharding_proto) { return Deserialize( diff --git a/third_party/xla/xla/python/ifrt/sharding.h b/third_party/xla/xla/python/ifrt/sharding.h index 0b658604161553..7490a43a8d4656 100644 --- a/third_party/xla/xla/python/ifrt/sharding.h +++ b/third_party/xla/xla/python/ifrt/sharding.h @@ -21,13 +21,17 @@ limitations under the License. #include #include #include -#include #include +#include "absl/base/attributes.h" #include "absl/base/nullability.h" +#include "absl/container/inlined_vector.h" #include "absl/hash/hash.h" #include "absl/log/check.h" +#include "absl/status/status.h" #include "absl/status/status_macros.h" +#include "absl/status/statusor.h" +#include "absl/types/span.h" #include "xla/python/ifrt/device.h" #include "xla/python/ifrt/device_list.h" #include "xla/python/ifrt/index_domain.h" @@ -149,6 +153,25 @@ class Sharding : public RTTIExtends { const Shape& shape, SingleDeviceShardSemantics single_device_shard_semantics) const = 0; + using IndexDomainAndShardIndices = ShardingSpec::IndexDomainAndShardIndices; + + // Breaks a shape up into unique `IndexDomain`s and the shard indices mapped + // to it. The result is calculated for all shards. + // + // The result is valid for the lifetime of this `Sharding`. + absl::StatusOr> + UniqueIndexDomains(const Shape& shape) const ABSL_ATTRIBUTE_LIFETIME_BOUND; + + // Inverse of `UniqueIndexDomains()` for `shard_indices`. Does not take + // `shape` because the result is independent of `shape`. + // + // Suppose `j` be `unique_index_domain_indices[shard_i]`. Then, + // `unique_index_domains[j].shard_indices` contains `shard_i`. + // + // The result is valid for the lifetime of this `Sharding`. + absl::StatusOr> ShardToUniqueIndexDomainIndex() const + ABSL_ATTRIBUTE_LIFETIME_BOUND; + template friend H AbslHashValue(H h, const Sharding& value) { value.Hash(absl::HashState::Create(&h)); diff --git a/third_party/xla/xla/python/ifrt/sharding_spec.cc b/third_party/xla/xla/python/ifrt/sharding_spec.cc index 72f3746521e526..a0d3b02d82f454 100644 --- a/third_party/xla/xla/python/ifrt/sharding_spec.cc +++ b/third_party/xla/xla/python/ifrt/sharding_spec.cc @@ -15,6 +15,7 @@ limitations under the License. #include "xla/python/ifrt/sharding_spec.h" +#include #include #include #include @@ -26,6 +27,8 @@ limitations under the License. #include #include "absl/algorithm/container.h" +#include "absl/base/call_once.h" +#include "absl/container/flat_hash_map.h" #include "absl/container/inlined_vector.h" #include "absl/hash/hash.h" #include "absl/log/check.h" @@ -244,6 +247,23 @@ absl::StatusOr> SingleDeviceShardingSpec::IndexDomains( return std::vector{IndexDomain(shape)}; } +absl::StatusOr> +SingleDeviceShardingSpec::UniqueIndexDomains(const Shape& shape) const { + static constexpr std::array kShardIndices({0}); + return absl::InlinedVector{ + IndexDomainAndShardIndices{ + /*index_domain=*/IndexDomain(shape), + /*shard_indices=*/absl::MakeConstSpan(kShardIndices), + }, + }; +} + +absl::StatusOr> +SingleDeviceShardingSpec::ShardToUniqueIndexDomainIndex() const { + static constexpr std::array kShardToUniqueIndexDomainIndex({0}); + return absl::MakeConstSpan(kShardToUniqueIndexDomainIndex); +} + std::string SingleDeviceShardingSpec::DebugString() const { return "SingleDeviceShardingSpec()"; } @@ -311,6 +331,18 @@ absl::StatusOr> OpaqueShardingSpec::IndexDomains( "OpaqueShardingSpec does not have index domain information"); } +absl::StatusOr> +OpaqueShardingSpec::UniqueIndexDomains(const Shape& shape) const { + return absl::InvalidArgumentError( + "OpaqueShardingSpec does not support UniqueIndexDomains"); +} + +absl::StatusOr> +OpaqueShardingSpec::ShardToUniqueIndexDomainIndex() const { + return absl::InvalidArgumentError( + "OpaqueShardingSpec does not support ShardToUniqueIndexDomainIndex"); +} + std::string OpaqueShardingSpec::DebugString() const { return absl::StrFormat("OpaqueShardingSpec(num_shards: %d)", num_shards_); } @@ -370,6 +402,13 @@ ConcreteShardingSpec::ConcreteShardingSpec( shape_(std::move(dynamic_shape)), shard_shapes_(std::move(shard_dynamic_shapes)) {} +ConcreteShardingSpec::ConcreteShardingSpec(const ConcreteShardingSpec& other) + : RTTIExtends(other), + shape_(other.shape_), + shard_shapes_(other.shard_shapes_), + shard_shape_(other.shard_shape_), + index_domains_(other.index_domains_) {} + absl::StatusOr ConcreteShardingSpec::ToSharding( DeviceListRef devices, MemoryKind memory_kind) const { if (devices->size() != num_shards()) { @@ -474,6 +513,76 @@ absl::StatusOr> ConcreteShardingSpec::IndexDomains( return *index_domains_; } +absl::StatusOr> +ConcreteShardingSpec::UniqueIndexDomains(const Shape& shape) const { + if (!index_domains_.has_value()) { + return absl::InvalidArgumentError( + "ConcreteShardingSpec does not have index domain information"); + } + if (has_static_shape() && this->shape() != shape) { + return absl::InvalidArgumentError(absl::StrFormat( + "ConcreteShardingSpec has index domains for shape %v, but was asked " + "to get unique index domains for shape %v", + this->shape(), shape)); + } + absl::call_once(unique_shard_indices_once_, [this] { + absl::flat_hash_map index_domain_to_unique_idx; + std::vector> shard_indices; + for (int i = 0; i < index_domains_->size(); ++i) { + const IndexDomain& domain = (*index_domains_)[i]; + auto [it, inserted] = + index_domain_to_unique_idx.try_emplace(domain, shard_indices.size()); + if (inserted) { + shard_indices.emplace_back(); + } + shard_indices[it->second].push_back(i); + } + cached_shard_indices_.reserve(index_domains_->size()); + cached_shard_indices_offsets_.reserve(shard_indices.size() + 1); + for (const auto& indices : shard_indices) { + cached_shard_indices_offsets_.push_back(cached_shard_indices_.size()); + cached_shard_indices_.insert(cached_shard_indices_.end(), indices.begin(), + indices.end()); + } + cached_shard_indices_offsets_.push_back(cached_shard_indices_.size()); + }); + + const int num_unique = cached_shard_indices_offsets_.size() - 1; + absl::InlinedVector unique_domains; + unique_domains.reserve(num_unique); + for (int i = 0; i < num_unique; ++i) { + const int offset = cached_shard_indices_offsets_[i]; + const int count = cached_shard_indices_offsets_[i + 1] - offset; + const int first_shard = cached_shard_indices_[offset]; + unique_domains.push_back(IndexDomainAndShardIndices{ + /*index_domain=*/(*index_domains_)[first_shard], + /*shard_indices=*/ + absl::MakeConstSpan(cached_shard_indices_).subspan(offset, count), + }); + } + return unique_domains; +} + +absl::StatusOr> +ConcreteShardingSpec::ShardToUniqueIndexDomainIndex() const { + if (!index_domains_.has_value()) { + return absl::InvalidArgumentError( + "ConcreteShardingSpec does not have index domain information"); + } + absl::call_once(shard_to_unique_index_domain_index_once_, [this] { + absl::flat_hash_map domain_to_unique_idx; + cached_shard_to_unique_index_domain_index_.reserve(index_domains_->size()); + for (int i = 0; i < index_domains_->size(); ++i) { + const IndexDomain& domain = (*index_domains_)[i]; + auto it = + domain_to_unique_idx.try_emplace(domain, domain_to_unique_idx.size()) + .first; + cached_shard_to_unique_index_domain_index_.push_back(it->second); + } + }); + return absl::MakeConstSpan(cached_shard_to_unique_index_domain_index_); +} + std::string ConcreteShardingSpec::DebugString() const { return std::visit( [this](const auto& shape, const auto& shard_shapes) { @@ -508,6 +617,12 @@ ConcreteEvenShardingSpec::ConcreteEvenShardingSpec(int num_shards, Shape shape, shape_(std::move(shape)), shard_shape_(std::move(shard_shape)) {} +ConcreteEvenShardingSpec::ConcreteEvenShardingSpec( + const ConcreteEvenShardingSpec& other) + : RTTIExtends(other), + shape_(other.shape_), + shard_shape_(other.shard_shape_) {} + absl::StatusOr ConcreteEvenShardingSpec::ToSharding( DeviceListRef devices, MemoryKind memory_kind) const { if (devices->size() != num_shards()) { @@ -584,6 +699,39 @@ absl::StatusOr> ConcreteEvenShardingSpec::IndexDomains( "ConcreteEvenShardingSpec does not have index domain information"); } +absl::StatusOr> +ConcreteEvenShardingSpec::UniqueIndexDomains(const Shape& shape) const { + if (!IsFullyReplicated() || this->shape() != shard_shape() || + this->shape() != shape) { + return absl::InvalidArgumentError( + "ConcreteEvenShardingSpec does not have index domain information"); + } + absl::call_once(unique_shard_indices_once_, [this] { + cached_shard_indices_.reserve(num_shards_); + for (int i = 0; i < num_shards_; ++i) { + cached_shard_indices_.push_back(i); + } + }); + return absl::InlinedVector{ + IndexDomainAndShardIndices{ + /*index_domain=*/IndexDomain(shape), + /*shard_indices=*/absl::MakeConstSpan(cached_shard_indices_), + }, + }; +} + +absl::StatusOr> +ConcreteEvenShardingSpec::ShardToUniqueIndexDomainIndex() const { + if (!IsFullyReplicated() || this->shape() != shard_shape()) { + return absl::InvalidArgumentError( + "ConcreteEvenShardingSpec does not have index domain information"); + } + absl::call_once(shard_to_unique_index_domain_index_once_, [this] { + cached_shard_to_unique_index_domain_index_.assign(num_shards_, 0); + }); + return absl::MakeConstSpan(cached_shard_to_unique_index_domain_index_); +} + std::string ConcreteEvenShardingSpec::DebugString() const { return absl::StrFormat( "ConcreteEvenShardingSpec(num_shards: %d, shape: %v, " @@ -610,6 +758,11 @@ ShardingParamShardingSpec::ShardingParamShardingSpec( num_shards, ComputeIsFullyReplicated(sharding_param)), sharding_param_(std::move(sharding_param)) {} +ShardingParamShardingSpec::ShardingParamShardingSpec( + const ShardingParamShardingSpec& other) + : RTTIExtends(other), + sharding_param_(other.sharding_param_) {} + absl::StatusOr ShardingParamShardingSpec::ToSharding( DeviceListRef devices, MemoryKind memory_kind) const { if (devices->size() != num_shards()) { @@ -722,6 +875,90 @@ ShardingParamShardingSpec::IndexDomains(const Shape& shape) const { return result; } +absl::StatusOr> +ShardingParamShardingSpec::UniqueIndexDomains(const Shape& shape) const { + ABSL_ASSIGN_OR_RETURN(Shape local_shape, GetShardShape(shape)); + + absl::call_once(unique_shard_indices_once_, [this] { + absl::InlinedVector device_list; + sharding_param_.minor_to_major().ToDeviceList(device_list); + if (device_list.size() != num_shards_) { + cached_shard_indices_ = absl::InvalidArgumentError(absl::StrFormat( + "ShardingParamShardingSpec has %d shards, but sharding param has %d " + "shards", + num_shards_, device_list.size())); + return; + } + cached_shard_indices_ = + std::vector(device_list.begin(), device_list.end()); + }); + ABSL_RETURN_IF_ERROR(cached_shard_indices_.status()); + + std::vector tile_indices = + GetTileIndices(sharding_param_.dim_shards()); + const int num_unique_tiles = tile_indices.size(); + if (num_shards_ % num_unique_tiles != 0) { + return absl::InvalidArgumentError(absl::StrFormat( + "ShardingParamShardingSpec has %d shards, but sharding param has %d " + "unique tiles, which is not a divisor of the number of shards", + num_shards_, num_unique_tiles)); + } + const int replication = num_shards_ / num_unique_tiles; + + absl::InlinedVector unique_domains; + unique_domains.reserve(num_unique_tiles); + for (int tile_idx = 0; tile_idx < num_unique_tiles; ++tile_idx) { + const Index& tile_index = tile_indices[tile_idx]; + unique_domains.push_back(IndexDomainAndShardIndices{ + /*index_domain=*/IndexDomain(tile_index * local_shape.dims(), + local_shape), + /*shard_indices=*/ + absl::MakeConstSpan(*cached_shard_indices_) + .subspan(tile_idx * replication, replication), + }); + } + + return unique_domains; +} + +absl::StatusOr> +ShardingParamShardingSpec::ShardToUniqueIndexDomainIndex() const { + absl::call_once(shard_to_unique_index_domain_index_once_, [this] { + std::vector tile_indices = + GetTileIndices(sharding_param_.dim_shards()); + const int num_unique_tiles = tile_indices.size(); + absl::InlinedVector device_list; + sharding_param_.minor_to_major().ToDeviceList(device_list); + if (device_list.size() != num_shards_) { + cached_shard_to_unique_index_domain_index_ = absl::InvalidArgumentError( + absl::StrFormat("ShardingParamShardingSpec has %d shards, but " + "sharding param has %d shards", + num_shards_, device_list.size())); + return; + } + if (device_list.size() % num_unique_tiles != 0) { + cached_shard_to_unique_index_domain_index_ = + absl::InvalidArgumentError(absl::StrFormat( + "ShardingParamShardingSpec has %d shards, but sharding param has " + "%d unique tiles, which is not a divisor of the number of shards", + num_shards_, num_unique_tiles)); + return; + } + const int replication = device_list.size() / num_unique_tiles; + + std::vector shard_to_unique_index_domain_index(num_shards_); + for (int i = 0; i < device_list.size(); ++i) { + const int device_idx = device_list[i]; + const int tile_idx = i / replication; + shard_to_unique_index_domain_index[device_idx] = tile_idx; + } + cached_shard_to_unique_index_domain_index_ = + std::move(shard_to_unique_index_domain_index); + }); + ABSL_RETURN_IF_ERROR(cached_shard_to_unique_index_domain_index_.status()); + return absl::MakeConstSpan(*cached_shard_to_unique_index_domain_index_); +} + std::string ShardingParamShardingSpec::DebugString() const { return absl::StrFormat("ShardingParamShardingSpec(num_shards: %d, %s)", num_shards_, sharding_param_.DebugString()); diff --git a/third_party/xla/xla/python/ifrt/sharding_spec.h b/third_party/xla/xla/python/ifrt/sharding_spec.h index bafdff441ba469..38b6941546ae24 100644 --- a/third_party/xla/xla/python/ifrt/sharding_spec.h +++ b/third_party/xla/xla/python/ifrt/sharding_spec.h @@ -24,11 +24,15 @@ limitations under the License. #include #include +#include "absl/base/attributes.h" +#include "absl/base/call_once.h" #include "absl/base/nullability.h" +#include "absl/container/inlined_vector.h" #include "absl/hash/hash.h" #include "absl/log/check.h" #include "absl/status/status_macros.h" #include "absl/status/statusor.h" +#include "absl/types/span.h" #include "xla/python/ifrt/device_list.h" #include "xla/python/ifrt/index_domain.h" #include "xla/python/ifrt/ir/sharding_param.h" @@ -116,6 +120,40 @@ class ShardingSpec : public RTTIExtends, virtual absl::StatusOr> IndexDomains( const Shape& shape) const = 0; + struct IndexDomainAndShardIndices { + // The index domain mapped from shards. + IndexDomain index_domain; + + // All shard indices that map to this index domain. + absl::Span shard_indices; + + bool operator==(const IndexDomainAndShardIndices& other) const { + return index_domain == other.index_domain && + shard_indices == other.shard_indices; + } + bool operator!=(const IndexDomainAndShardIndices& other) const { + return !(*this == other); + } + }; + + // Breaks a shape up into unique `IndexDomain`s and the shard indices mapped + // to it. + // + // The result is valid for the lifetime of this `ShardingSpec`. + virtual absl::StatusOr> + UniqueIndexDomains(const Shape& shape) const + ABSL_ATTRIBUTE_LIFETIME_BOUND = 0; + + // Inverse of `UniqueIndexDomains()` for `shard_indices`. Does not take + // `shape` because the result is independent of `shape`. + // + // Suppose `j` be `unique_index_domain_indices[shard_i]`. Then, + // `unique_index_domains[j].shard_indices` contains `shard_i`. + // + // The result is valid for the lifetime of this `ShardingSpec`. + virtual absl::StatusOr> ShardToUniqueIndexDomainIndex() + const ABSL_ATTRIBUTE_LIFETIME_BOUND = 0; + template friend H AbslHashValue(H h, const ShardingSpec& value) { value.Hash(absl::HashState::Create(&h)); @@ -212,6 +250,12 @@ class SingleDeviceShardingSpec final absl::StatusOr> IndexDomains( const Shape& shape) const override; + absl::StatusOr> + UniqueIndexDomains(const Shape& shape) const override; + + absl::StatusOr> ShardToUniqueIndexDomainIndex() + const override; + static char ID; // NOLINT private: @@ -250,6 +294,12 @@ class OpaqueShardingSpec absl::StatusOr> IndexDomains( const Shape& shape) const override; + absl::StatusOr> + UniqueIndexDomains(const Shape& shape) const override; + + absl::StatusOr> ShardToUniqueIndexDomainIndex() + const override; + static char ID; // NOLINT private: @@ -340,8 +390,16 @@ class ConcreteShardingSpec absl::StatusOr> IndexDomains( const Shape& shape) const override; + absl::StatusOr> + UniqueIndexDomains(const Shape& shape) const override; + + absl::StatusOr> ShardToUniqueIndexDomainIndex() + const override; + static char ID; // NOLINT + ConcreteShardingSpec(const ConcreteShardingSpec& other); + private: ConcreteShardingSpec( int num_shards, Shape shape, std::vector shard_shapes, @@ -358,6 +416,12 @@ class ConcreteShardingSpec std::variant, std::vector> shard_shapes_; std::optional shard_shape_; std::optional> index_domains_; + + mutable absl::once_flag unique_shard_indices_once_; + mutable std::vector cached_shard_indices_; + mutable std::vector cached_shard_indices_offsets_; + mutable absl::once_flag shard_to_unique_index_domain_index_once_; + mutable std::vector cached_shard_to_unique_index_domain_index_; }; // Opaque sharding spec that does not define a fixed semantics for conversion @@ -402,8 +466,16 @@ class ConcreteEvenShardingSpec absl::StatusOr> IndexDomains( const Shape& shape) const override; + absl::StatusOr> + UniqueIndexDomains(const Shape& shape) const override; + + absl::StatusOr> ShardToUniqueIndexDomainIndex() + const override; + static char ID; // NOLINT + ConcreteEvenShardingSpec(const ConcreteEvenShardingSpec& other); + private: ConcreteEvenShardingSpec(int num_shards, Shape shape, Shape shard_shape, bool is_fully_replicated); @@ -414,6 +486,11 @@ class ConcreteEvenShardingSpec Shape shape_; Shape shard_shape_; + + mutable absl::once_flag unique_shard_indices_once_; + mutable std::vector cached_shard_indices_; + mutable absl::once_flag shard_to_unique_index_domain_index_once_; + mutable std::vector cached_shard_to_unique_index_domain_index_; }; // Sharding spec derived from an IR ShardingParam. @@ -441,8 +518,16 @@ class ShardingParamShardingSpec absl::StatusOr> IndexDomains( const Shape& shape) const override; + absl::StatusOr> + UniqueIndexDomains(const Shape& shape) const override; + + absl::StatusOr> ShardToUniqueIndexDomainIndex() + const override; + static char ID; // NOLINT + ShardingParamShardingSpec(const ShardingParamShardingSpec& other); + private: ShardingParamShardingSpec(int num_shards, ShardingParam sharding_param); @@ -451,6 +536,12 @@ class ShardingParamShardingSpec void Hash(absl::HashState state) const override; ShardingParam sharding_param_; + + mutable absl::once_flag unique_shard_indices_once_; + mutable absl::StatusOr> cached_shard_indices_; + mutable absl::once_flag shard_to_unique_index_domain_index_once_; + mutable absl::StatusOr> + cached_shard_to_unique_index_domain_index_; }; } // namespace ifrt diff --git a/third_party/xla/xla/python/ifrt/sharding_spec_test.cc b/third_party/xla/xla/python/ifrt/sharding_spec_test.cc index a1c7f2969b9338..86d45252180f5b 100644 --- a/third_party/xla/xla/python/ifrt/sharding_spec_test.cc +++ b/third_party/xla/xla/python/ifrt/sharding_spec_test.cc @@ -43,6 +43,7 @@ namespace { using ::absl_testing::StatusIs; using ::testing::ElementsAre; using ::testing::ElementsAreArray; +using ::testing::FieldsAre; using ::testing::HasSubstr; using ::testing::SizeIs; @@ -101,6 +102,17 @@ TEST_P(SingleDeviceShardingSpecTest, IndexDomains) { EXPECT_THAT(index_domains, ElementsAre(IndexDomain(shape))); } +TEST_P(SingleDeviceShardingSpecTest, UniqueIndexDomains) { + ShardingSpecRef sharding = SingleDeviceShardingSpec::Create(); + + Shape shape({10, 20}); + EXPECT_THAT(sharding->UniqueIndexDomains(shape), + absl_testing::IsOkAndHolds( + ElementsAre(FieldsAre(IndexDomain(shape), ElementsAre(0))))); + EXPECT_THAT(sharding->ShardToUniqueIndexDomainIndex(), + absl_testing::IsOkAndHolds(ElementsAre(0))); +} + TEST_P(SingleDeviceShardingSpecTest, Disassemble) { ShardingSpecRef sharding = SingleDeviceShardingSpec::Create(); @@ -183,6 +195,21 @@ TEST_P(OpaqueShardingSpecTest, IndexDomainsFails) { "OpaqueShardingSpec does not have index domain information"))); } +TEST_P(OpaqueShardingSpecTest, UniqueIndexDomains) { + ShardingSpecRef sharding = OpaqueShardingSpec::Create(num_shards()); + + EXPECT_THAT( + sharding->UniqueIndexDomains(Shape({30})), + absl_testing::StatusIs( + tsl::error::INVALID_ARGUMENT, + HasSubstr("OpaqueShardingSpec does not support UniqueIndexDomains"))); + EXPECT_THAT( + sharding->ShardToUniqueIndexDomainIndex(), + absl_testing::StatusIs(tsl::error::INVALID_ARGUMENT, + HasSubstr("OpaqueShardingSpec does not support " + "ShardToUniqueIndexDomainIndex"))); +} + TEST_P(OpaqueShardingSpecTest, Hash) { EXPECT_TRUE(absl::VerifyTypeImplementsAbslHashCorrectly({ *OpaqueShardingSpec::Create(2), @@ -371,6 +398,44 @@ TEST_P(ConcreteShardingSpecTest, IndexDomainsMissing) { "ConcreteShardingSpec does not have index domain information"))); } +TEST_P(ConcreteShardingSpecTest, UniqueIndexDomains) { + std::vector shard_shapes{Shape({10}), Shape({10}), Shape({20})}; + std::vector index_domains{ + IndexDomain(Index({0}), Shape({10})), + IndexDomain(Index({0}), Shape({10})), + IndexDomain(Index({10}), Shape({20})), + }; + ShardingSpecRef sharding = + ConcreteShardingSpec::Create(Shape({30}), shard_shapes, index_domains); + + EXPECT_THAT( + sharding->UniqueIndexDomains(Shape({30})), + absl_testing::IsOkAndHolds(ElementsAre( + FieldsAre(IndexDomain(Index({0}), Shape({10})), ElementsAre(0, 1)), + FieldsAre(IndexDomain(Index({10}), Shape({20})), ElementsAre(2))))); + EXPECT_THAT(sharding->ShardToUniqueIndexDomainIndex(), + absl_testing::IsOkAndHolds(ElementsAre(0, 0, 1))); +} + +TEST_P(ConcreteShardingSpecTest, UniqueIndexDomainsMissing) { + std::vector shard_shapes{Shape({10}), Shape({20})}; + ShardingSpecRef sharding = + ConcreteShardingSpec::Create(Shape({30}), shard_shapes); + + EXPECT_THAT( + sharding->UniqueIndexDomains(Shape({30})), + absl_testing::StatusIs( + tsl::error::INVALID_ARGUMENT, + HasSubstr( + "ConcreteShardingSpec does not have index domain information"))); + EXPECT_THAT( + sharding->ShardToUniqueIndexDomainIndex(), + absl_testing::StatusIs( + tsl::error::INVALID_ARGUMENT, + HasSubstr( + "ConcreteShardingSpec does not have index domain information"))); +} + TEST_P(ConcreteShardingSpecTest, Hash) { ASSERT_OK_AND_ASSIGN( auto dynamic_shape, @@ -505,6 +570,38 @@ TEST_P(ConcreteEvenShardingSpecTest, IndexDomainsFails) { "index domain information"))); } +TEST_P(ConcreteEvenShardingSpecTest, UniqueIndexDomains) { + { + // Fully replicated. + ShardingSpecRef sharding = ConcreteEvenShardingSpec::Create( + /*num_shards=*/2, Shape({30}), Shape({30}), + /*is_fully_replicated=*/true); + + EXPECT_THAT(sharding->UniqueIndexDomains(Shape({30})), + absl_testing::IsOkAndHolds(ElementsAre( + FieldsAre(IndexDomain(Shape({30})), ElementsAre(0, 1))))); + EXPECT_THAT(sharding->ShardToUniqueIndexDomainIndex(), + absl_testing::IsOkAndHolds(ElementsAre(0, 0))); + } + { + // Not fully replicated. + ShardingSpecRef sharding = ConcreteEvenShardingSpec::Create( + /*num_shards=*/2, Shape({30}), Shape({15}), + /*is_fully_replicated=*/false); + + EXPECT_THAT(sharding->UniqueIndexDomains(Shape({30})), + absl_testing::StatusIs( + tsl::error::INVALID_ARGUMENT, + HasSubstr("ConcreteEvenShardingSpec does not have index " + "domain information"))); + EXPECT_THAT(sharding->ShardToUniqueIndexDomainIndex(), + absl_testing::StatusIs( + tsl::error::INVALID_ARGUMENT, + HasSubstr("ConcreteEvenShardingSpec does not have index " + "domain information"))); + } +} + TEST_P(ConcreteEvenShardingSpecTest, Hash) { EXPECT_TRUE(absl::VerifyTypeImplementsAbslHashCorrectly({ *ConcreteEvenShardingSpec::Create( @@ -679,6 +776,49 @@ TEST_P(ShardingParamShardingSpecTest, IndexDomainWithReplication) { } } +TEST_P(ShardingParamShardingSpecTest, UniqueIndexDomains) { + { + // 2x3 tiled with replication 1. + ShardingParam param{/*dim_shards=*/{2, 3}, + {/*permutation=*/{0, 1}, /*axis_sizes=*/{2, 3}}}; + ShardingSpecRef sharding = ShardingParamShardingSpec::Create(param); + + Shape shape({4, 9}); + EXPECT_THAT(sharding->UniqueIndexDomains(shape), + absl_testing::IsOkAndHolds(ElementsAre( + FieldsAre(IndexDomain(Index({0, 0}), Shape({2, 3})), + ElementsAre(0)), + FieldsAre(IndexDomain(Index({0, 3}), Shape({2, 3})), + ElementsAre(1)), + FieldsAre(IndexDomain(Index({0, 6}), Shape({2, 3})), + ElementsAre(2)), + FieldsAre(IndexDomain(Index({2, 0}), Shape({2, 3})), + ElementsAre(3)), + FieldsAre(IndexDomain(Index({2, 3}), Shape({2, 3})), + ElementsAre(4)), + FieldsAre(IndexDomain(Index({2, 6}), Shape({2, 3})), + ElementsAre(5))))); + EXPECT_THAT(sharding->ShardToUniqueIndexDomainIndex(), + absl_testing::IsOkAndHolds(ElementsAre(0, 1, 2, 3, 4, 5))); + } + { + // 2x1 tiled with replication 3. + ShardingParam param{/*dim_shards=*/{2, 1}, + {/*permutation=*/{0, 1}, /*axis_sizes=*/{2, 3}}}; + ShardingSpecRef sharding = ShardingParamShardingSpec::Create(param); + + Shape shape({4, 9}); + EXPECT_THAT(sharding->UniqueIndexDomains(shape), + absl_testing::IsOkAndHolds(ElementsAre( + FieldsAre(IndexDomain(Index({0, 0}), Shape({2, 9})), + ElementsAre(0, 1, 2)), + FieldsAre(IndexDomain(Index({2, 0}), Shape({2, 9})), + ElementsAre(3, 4, 5))))); + EXPECT_THAT(sharding->ShardToUniqueIndexDomainIndex(), + absl_testing::IsOkAndHolds(ElementsAre(0, 0, 0, 1, 1, 1))); + } +} + TEST_P(ShardingParamShardingSpecTest, IndexDomainZeroRank) { ShardingParam param{/*dim_shards=*/{}, {/*permutation=*/{0}, /*axis_sizes=*/{6}}}; diff --git a/third_party/xla/xla/python/ifrt/sharding_test.cc b/third_party/xla/xla/python/ifrt/sharding_test.cc index 9ddcdaffdc1b06..0c403771a16579 100644 --- a/third_party/xla/xla/python/ifrt/sharding_test.cc +++ b/third_party/xla/xla/python/ifrt/sharding_test.cc @@ -44,6 +44,7 @@ using ::absl_testing::IsOkAndHolds; using ::absl_testing::StatusIs; using ::testing::ElementsAre; using ::testing::ElementsAreArray; +using ::testing::FieldsAre; using ::testing::HasSubstr; using ::testing::SizeIs; @@ -142,6 +143,19 @@ TEST_P(SingleDeviceShardingTest, IndexDomains) { } } +TEST_P(SingleDeviceShardingTest, UniqueIndexDomains) { + DeviceListRef device_list = GetDevices({0}); + ShardingRef sharding = SingleDeviceSharding::Create( + device_list->devices().front(), MemoryKind()); + + Shape shape({10, 20}); + EXPECT_THAT( + sharding->UniqueIndexDomains(shape), + IsOkAndHolds(ElementsAre(FieldsAre(IndexDomain(shape), ElementsAre(0))))); + EXPECT_THAT(sharding->ShardToUniqueIndexDomainIndex(), + IsOkAndHolds(ElementsAre(0))); +} + TEST_P(SingleDeviceShardingTest, Disassemble) { auto device_list = GetDevices({0}); ShardingRef sharding = SingleDeviceSharding::Create( @@ -298,6 +312,19 @@ TEST_P(OpaqueShardingTest, IndexDomainsFails) { HasSubstr("OpaqueSharding does not have index domain information"))); } +TEST_P(OpaqueShardingTest, UniqueIndexDomainsFails) { + DeviceListRef device_list = GetDevices({0, 1}); + ShardingRef sharding = OpaqueSharding::Create(device_list, MemoryKind()); + + EXPECT_THAT(sharding->UniqueIndexDomains(Shape({30})), + StatusIs(tsl::error::INVALID_ARGUMENT, + HasSubstr("does not support UniqueIndexDomains"))); + EXPECT_THAT( + sharding->ShardToUniqueIndexDomainIndex(), + StatusIs(tsl::error::INVALID_ARGUMENT, + HasSubstr("does not support ShardToUniqueIndexDomainIndex"))); +} + TEST_P(OpaqueShardingTest, Hash) { EXPECT_TRUE(absl::VerifyTypeImplementsAbslHashCorrectly({ *OpaqueSharding::Create(GetDevices({0, 1}), MemoryKind()), @@ -668,6 +695,35 @@ TEST_P(ConcreteShardingTest, IndexDomainsFails) { "of index domains and addressable devices"))); } +TEST_P(ConcreteShardingTest, UniqueIndexDomains) { + DeviceListRef device_list = GetDevices({0, 1, 2, 3, 4, 5}); + // devices 0..3 are addressable, 4..5 are non-addressable. + std::vector shard_shapes = { + Shape({10}), Shape({10}), Shape({10}), + Shape({10}), Shape({10}), Shape({10}), + }; + std::vector index_domains{ + IndexDomain(Index({0}), Shape({10})), + IndexDomain(Index({0}), Shape({10})), + IndexDomain(Index({10}), Shape({10})), + IndexDomain(Index({10}), Shape({10})), + IndexDomain(Index({20}), Shape({10})), + IndexDomain(Index({20}), Shape({10})), + }; + ShardingRef sharding = ConcreteSharding::Create( + device_list, MemoryKind(), Shape({30}), shard_shapes, index_domains); + + EXPECT_THAT( + sharding->UniqueIndexDomains(Shape({30})), + IsOkAndHolds(ElementsAre( + FieldsAre(IndexDomain(Index({0}), Shape({10})), ElementsAre(0, 1)), + FieldsAre(IndexDomain(Index({10}), Shape({10})), ElementsAre(2, 3)), + FieldsAre(IndexDomain(Index({20}), Shape({10})), + ElementsAre(4, 5))))); + EXPECT_THAT(sharding->ShardToUniqueIndexDomainIndex(), + IsOkAndHolds(ElementsAre(0, 0, 1, 1, 2, 2))); +} + TEST_P(ConcreteShardingTest, Hash) { ASSERT_OK_AND_ASSIGN( auto dynamic_shape, @@ -894,6 +950,38 @@ TEST_P(ConcreteEvenShardingTest, IndexDomainsFailsForNonFullyReplicated) { "ConcreteEvenSharding does not have index domain information"))); } +TEST_P(ConcreteEvenShardingTest, UniqueIndexDomains) { + Shape shape({10, 20}); + + auto device_list = GetDevices({0, 4}); + ASSERT_TRUE(device_list->devices()[0]->IsAddressable()); + ASSERT_FALSE(device_list->devices()[1]->IsAddressable()); + + ShardingRef sharding = ConcreteEvenSharding::Create( + device_list, MemoryKind(), /*shape=*/shape, /*shard_shape=*/shape, + /*is_fully_replicated=*/true); + + EXPECT_THAT(sharding->UniqueIndexDomains(shape), + IsOkAndHolds(ElementsAre( + FieldsAre(IndexDomain(shape), ElementsAre(0, 1))))); + EXPECT_THAT(sharding->ShardToUniqueIndexDomainIndex(), + IsOkAndHolds(ElementsAre(0, 0))); +} + +TEST_P(ConcreteEvenShardingTest, UniqueIndexDomainsFailsForNonFullyReplicated) { + auto device_list = GetDevices({0, 1}); + ShardingRef sharding = + ConcreteEvenSharding::Create(device_list, MemoryKind(), Shape({30}), + Shape({5}), /*is_fully_replicated=*/false); + + EXPECT_THAT(sharding->UniqueIndexDomains(Shape({30})), + StatusIs(tsl::error::INVALID_ARGUMENT, + HasSubstr("does not have index domain information"))); + EXPECT_THAT(sharding->ShardToUniqueIndexDomainIndex(), + StatusIs(tsl::error::INVALID_ARGUMENT, + HasSubstr("does not have index domain information"))); +} + TEST_P(ConcreteEvenShardingTest, Hash) { EXPECT_TRUE(absl::VerifyTypeImplementsAbslHashCorrectly({ *ConcreteEvenSharding::Create(GetDevices({0, 1}), MemoryKind(), @@ -1237,6 +1325,27 @@ TEST_P(ShardingParamShardingTest, IndexDomainWithReplication) { } } +TEST_P(ShardingParamShardingTest, UniqueIndexDomains) { + auto device_list = GetDevices({0, 1, 2, 3, 4, 5}); + // devices 0..3 are addressable, 4..5 are non-addressable. + // 2x1 tiled with replication 3. + ShardingParam param{/*dim_shards=*/{2, 1}, + {/*permutation=*/{0, 1}, /*axis_sizes=*/{2, 3}}}; + ASSERT_OK_AND_ASSIGN( + ShardingRef sharding, + ShardingParamSharding::Create(param, device_list, MemoryKind())); + + Shape shape({6, 6}); + EXPECT_THAT(sharding->UniqueIndexDomains(shape), + IsOkAndHolds(ElementsAre( + FieldsAre(IndexDomain(Index({0, 0}), Shape({3, 6})), + ElementsAre(0, 1, 2)), + FieldsAre(IndexDomain(Index({3, 0}), Shape({3, 6})), + ElementsAre(3, 4, 5))))); + EXPECT_THAT(sharding->ShardToUniqueIndexDomainIndex(), + IsOkAndHolds(ElementsAre(0, 0, 0, 1, 1, 1))); +} + TEST_P(ShardingParamShardingTest, IndexDomainZeroRank) { auto device_list = GetDevices({0, 1, 2, 3, 4, 5}); ShardingParam param{/*dim_shards=*/{}, diff --git a/third_party/xla/xla/python/pjrt_ifrt/BUILD b/third_party/xla/xla/python/pjrt_ifrt/BUILD index abc818cfac4811..13bf8eb35a761c 100644 --- a/third_party/xla/xla/python/pjrt_ifrt/BUILD +++ b/third_party/xla/xla/python/pjrt_ifrt/BUILD @@ -47,7 +47,9 @@ cc_library( "//xla/python/ifrt:serdes", "//xla/python/ifrt:serdes_version", "//xla/service:device_assignment", + "@com_google_absl//absl/base", "@com_google_absl//absl/base:core_headers", + "@com_google_absl//absl/container:inlined_vector", "@com_google_absl//absl/hash", "@com_google_absl//absl/log", "@com_google_absl//absl/log:check", diff --git a/third_party/xla/xla/python/pjrt_ifrt/xla_sharding.cc b/third_party/xla/xla/python/pjrt_ifrt/xla_sharding.cc index 29091deb0a6a74..9bc5bc00053bd1 100644 --- a/third_party/xla/xla/python/pjrt_ifrt/xla_sharding.cc +++ b/third_party/xla/xla/python/pjrt_ifrt/xla_sharding.cc @@ -26,6 +26,7 @@ limitations under the License. #include "absl/log/check.h" #include "absl/status/status.h" #include "absl/status/status_macros.h" +#include "absl/status/statusor.h" #include "absl/strings/str_format.h" #include "absl/types/span.h" #include "xla/hlo/ir/hlo_sharding.h" diff --git a/third_party/xla/xla/python/pjrt_ifrt/xla_sharding.h b/third_party/xla/xla/python/pjrt_ifrt/xla_sharding.h index 38140cfa0ea491..13d26a44a3752b 100644 --- a/third_party/xla/xla/python/pjrt_ifrt/xla_sharding.h +++ b/third_party/xla/xla/python/pjrt_ifrt/xla_sharding.h @@ -16,8 +16,6 @@ limitations under the License. #ifndef XLA_PYTHON_PJRT_IFRT_XLA_SHARDING_H_ #define XLA_PYTHON_PJRT_IFRT_XLA_SHARDING_H_ -#include -#include #include #include #include @@ -26,7 +24,6 @@ limitations under the License. #include "absl/hash/hash.h" #include "absl/status/statusor.h" -#include "absl/types/span.h" #include "xla/hlo/ir/hlo_sharding.h" #include "xla/python/ifrt/device_list.h" #include "xla/python/ifrt/index_domain.h" diff --git a/third_party/xla/xla/python/pjrt_ifrt/xla_sharding_spec.cc b/third_party/xla/xla/python/pjrt_ifrt/xla_sharding_spec.cc index cc3a02453c152d..d92d8a643426c6 100644 --- a/third_party/xla/xla/python/pjrt_ifrt/xla_sharding_spec.cc +++ b/third_party/xla/xla/python/pjrt_ifrt/xla_sharding_spec.cc @@ -23,12 +23,15 @@ limitations under the License. #include #include +#include "absl/base/call_once.h" #include "absl/base/optimization.h" +#include "absl/container/inlined_vector.h" #include "absl/hash/hash.h" #include "absl/log/check.h" #include "absl/log/log.h" #include "absl/status/status.h" #include "absl/status/status_macros.h" +#include "absl/status/statusor.h" #include "absl/strings/str_format.h" #include "absl/types/span.h" #include "xla/hlo/ir/hlo_sharding.h" @@ -40,6 +43,7 @@ limitations under the License. #include "xla/python/ifrt/shape.h" #include "xla/python/ifrt/sharding_spec.h" #include "xla/python/pjrt_ifrt/xla_sharding.h" +#include "xla/shape.h" #include "xla/shape_util.h" #include "xla/util.h" #include "xla/xla_data.pb.h" @@ -53,7 +57,6 @@ char HloShardingSpec::ID = 0; // NOLINT namespace { // Generates IndexDomains for an HloShardingSpec, using XLA HloSharding APIs. -// Note that this is O(N^2) where N is the number of devices (shards). std::vector IndexDomainsSlowPath( const xla::HloSharding& hlo_sharding, int num_shards, const Shape& shape) { // Only shape dimensions are used. @@ -108,6 +111,11 @@ HloShardingSpec::HloShardingSpec(int num_shards, num_shards_ == 1); } +HloShardingSpec::HloShardingSpec(const HloShardingSpec& other) + : RTTIExtends(other), + xla_hlo_sharding_(other.xla_hlo_sharding_), + hash_(other.hash_.load(std::memory_order_relaxed)) {} + absl::StatusOr HloShardingSpec::ToSharding( DeviceListRef devices, MemoryKind memory_kind) const { if (devices->size() != num_shards()) { @@ -281,6 +289,139 @@ absl::StatusOr> HloShardingSpec::IndexDomains( return result; } +absl::StatusOr> +HloShardingSpec::UniqueIndexDomains(const Shape& shape) const { + if (xla_hlo_sharding_.IsManual()) { + return absl::InvalidArgumentError( + "Manual sharding does not support UniqueIndexDomains"); + } + if (xla_hlo_sharding_.IsUnreduced()) { + return absl::InvalidArgumentError( + "Unreduced sharding does not support UniqueIndexDomains"); + } + if (xla_hlo_sharding_.HasNonReplicatedSubgroup()) { + return absl::InvalidArgumentError( + "Non-replicated subgroup (e.g., manual or unreduced subgroup) sharding " + "does not support UniqueIndexDomains"); + } + if (xla_hlo_sharding_.IsReplicatedOrSingleDevice()) { + absl::call_once(unique_shard_indices_once_, [this] { + cached_shard_indices_.reserve(num_shards_); + for (int i = 0; i < num_shards_; ++i) { + cached_shard_indices_.push_back(i); + } + }); + return absl::InlinedVector{ + IndexDomainAndShardIndices{ + /*index_domain=*/IndexDomain(shape), + /*shard_indices=*/absl::MakeConstSpan(cached_shard_indices_), + }, + }; + } + + const int64_t tiled_data_rank = xla_hlo_sharding_.TiledDataRank(); + if (shape.dims().size() != tiled_data_rank) { + return absl::InvalidArgumentError( + absl::StrFormat("shape must have %d dimensions, but has %d dimensions: " + "shape=%v, sharding=%s", + tiled_data_rank, shape.dims().size(), shape, + xla_hlo_sharding_.ToString())); + } + + absl::call_once(unique_shard_indices_once_, [this] { + const int64_t* flat_tile_assignment = + xla_hlo_sharding_.tile_assignment().array().data(); + cached_shard_indices_.reserve(num_shards_); + for (int64_t i = 0; i < num_shards_; ++i) { + cached_shard_indices_.push_back( + static_cast(flat_tile_assignment[i])); + } + }); + + const int64_t num_unique_tiles = xla_hlo_sharding_.NumTiles(); + if (num_shards_ % num_unique_tiles != 0) { + return absl::InvalidArgumentError(absl::StrFormat( + "HloShardingSpec has %d shards, but HloSharding has %d unique tiles, " + "which is not a divisor of the number of shards", + num_shards_, num_unique_tiles)); + } + const int64_t num_replicas = num_shards_ / num_unique_tiles; + + xla::Shape xla_shape = xla::ShapeUtil::MakeShapeWithDescendingLayout( + xla::PrimitiveType::S32, shape.dims()); + absl::InlinedVector unique_domains; + unique_domains.reserve(num_unique_tiles); + for (int64_t tile_idx = 0; tile_idx < num_unique_tiles; ++tile_idx) { + const int first_shard = cached_shard_indices_[tile_idx * num_replicas]; + std::vector tile_offset = + xla_hlo_sharding_.TileOffsetForDevice(xla_shape, first_shard); + std::vector tile_limit = + xla_hlo_sharding_.TileLimitForDevice(xla_shape, first_shard); + Index::Elements origin(shape.dims().size()); + Shape::Dimensions shard_shape(shape.dims().size()); + for (int i = 0; i < shape.dims().size(); ++i) { + origin[i] = tile_offset[i]; + shard_shape[i] = tile_limit[i] - tile_offset[i]; + } + unique_domains.push_back(IndexDomainAndShardIndices{ + /*index_domain=*/ + IndexDomain(Index(std::move(origin)), Shape(std::move(shard_shape))), + /*shard_indices=*/ + absl::MakeConstSpan(cached_shard_indices_) + .subspan(tile_idx * num_replicas, num_replicas), + }); + } + + return unique_domains; +} + +absl::StatusOr> +HloShardingSpec::ShardToUniqueIndexDomainIndex() const { + if (xla_hlo_sharding_.IsManual()) { + return absl::InvalidArgumentError( + "Manual sharding does not support ShardToUniqueIndexDomainIndex"); + } + if (xla_hlo_sharding_.IsUnreduced()) { + return absl::InvalidArgumentError( + "Unreduced sharding does not support ShardToUniqueIndexDomainIndex"); + } + if (xla_hlo_sharding_.HasNonReplicatedSubgroup()) { + return absl::InvalidArgumentError( + "Non-replicated subgroup (e.g., manual or unreduced subgroup) sharding " + "does not support ShardToUniqueIndexDomainIndex"); + } + if (xla_hlo_sharding_.IsReplicatedOrSingleDevice()) { + absl::call_once(shard_to_unique_index_domain_index_once_, [this] { + cached_shard_to_unique_index_domain_index_.assign(num_shards_, 0); + }); + return absl::MakeConstSpan(cached_shard_to_unique_index_domain_index_); + } + + const int64_t num_unique_tiles = xla_hlo_sharding_.NumTiles(); + if (num_shards_ % num_unique_tiles != 0) { + return absl::InvalidArgumentError(absl::StrFormat( + "HloShardingSpec has %d shards, but HloSharding has %d unique tiles, " + "which is not a divisor of the number of shards", + num_shards_, num_unique_tiles)); + } + const int64_t num_replicas = num_shards_ / num_unique_tiles; + + absl::call_once(shard_to_unique_index_domain_index_once_, [&, this] { + cached_shard_to_unique_index_domain_index_.resize(num_shards_); + const int64_t* flat_tile_assignment = + xla_hlo_sharding_.tile_assignment().array().data(); + for (int64_t tile_idx = 0; tile_idx < num_unique_tiles; ++tile_idx) { + const int64_t offset = tile_idx * num_replicas; + for (int64_t i = 0; i < num_replicas; ++i) { + const int device_idx = + static_cast(flat_tile_assignment[offset + i]); + cached_shard_to_unique_index_domain_index_[device_idx] = tile_idx; + } + } + }); + return absl::MakeConstSpan(cached_shard_to_unique_index_domain_index_); +} + std::string HloShardingSpec::DebugString() const { return absl::StrFormat("HloShardingSpec(num_shards: %d, hlo_sharding: %s)", num_shards_, xla_hlo_sharding_.ToString()); diff --git a/third_party/xla/xla/python/pjrt_ifrt/xla_sharding_spec.h b/third_party/xla/xla/python/pjrt_ifrt/xla_sharding_spec.h index b2685077274f3a..4d5fc96ad59d53 100644 --- a/third_party/xla/xla/python/pjrt_ifrt/xla_sharding_spec.h +++ b/third_party/xla/xla/python/pjrt_ifrt/xla_sharding_spec.h @@ -23,6 +23,7 @@ limitations under the License. #include #include +#include "absl/base/call_once.h" #include "absl/hash/hash.h" #include "absl/status/statusor.h" #include "xla/hlo/ir/hlo_sharding.h" @@ -79,8 +80,16 @@ class HloShardingSpec final absl::StatusOr> IndexDomains( const Shape& shape) const override; + absl::StatusOr> + UniqueIndexDomains(const Shape& shape) const override; + + absl::StatusOr> ShardToUniqueIndexDomainIndex() + const override; + static char ID; // NOLINT + HloShardingSpec(const HloShardingSpec& other); + private: HloShardingSpec(int num_shards, xla::HloSharding xla_hlo_sharding); @@ -94,6 +103,11 @@ class HloShardingSpec final // May be written multiple times with the same non-zero value. static constexpr uint64_t kUnsetHash = 0; mutable std::atomic hash_ = kUnsetHash; + + mutable absl::once_flag unique_shard_indices_once_; + mutable std::vector cached_shard_indices_; + mutable absl::once_flag shard_to_unique_index_domain_index_once_; + mutable std::vector cached_shard_to_unique_index_domain_index_; }; // Test only: returns `HloShardingSpec::IndexDomains()`, using diff --git a/third_party/xla/xla/python/pjrt_ifrt/xla_sharding_spec_test.cc b/third_party/xla/xla/python/pjrt_ifrt/xla_sharding_spec_test.cc index 99cc2e0411c5c2..214bff8445f9e3 100644 --- a/third_party/xla/xla/python/pjrt_ifrt/xla_sharding_spec_test.cc +++ b/third_party/xla/xla/python/pjrt_ifrt/xla_sharding_spec_test.cc @@ -42,6 +42,7 @@ namespace { using ::testing::ElementsAre; using ::testing::ElementsAreArray; +using ::testing::FieldsAre; using ::testing::HasSubstr; using ::testing::SizeIs; @@ -206,6 +207,21 @@ TEST_F(HloShardingSpecTest, IndexDomainsWithReplication) { ElementsAreArray(TEST_HloShardingSpecIndexDomainsSlowPath(*spec, shape))); } +TEST_F(HloShardingSpecTest, UniqueIndexDomainsWithReplication) { + int num_shards = 6; + // Fully replicated. + auto xla_hlo_sharding = xla::HloSharding::Replicate(); + std::shared_ptr spec = + HloShardingSpec::Create(num_shards, xla_hlo_sharding); + + Shape shape({10, 20}); + EXPECT_THAT(spec->UniqueIndexDomains(shape), + absl_testing::IsOkAndHolds(ElementsAre(FieldsAre( + IndexDomain(shape), ElementsAre(0, 1, 2, 3, 4, 5))))); + EXPECT_THAT(spec->ShardToUniqueIndexDomainIndex(), + absl_testing::IsOkAndHolds(ElementsAre(0, 0, 0, 0, 0, 0))); +} + TEST_F(HloShardingSpecTest, DisassembleWithReplication) { int num_shards = 6; // Fully replicated. @@ -244,6 +260,28 @@ TEST_F(HloShardingSpecTest, IndexDomainsWithTile) { ElementsAreArray(TEST_HloShardingSpecIndexDomainsSlowPath(*spec, shape))); } +TEST_F(HloShardingSpecTest, UniqueIndexDomainsWithTile) { + int num_shards = 6; + // 6-way sharded along axis 0, 1-way sharded along axis 1. + auto xla_hlo_sharding = xla::HloSharding::Tile(xla::TileAssignment({6, 1})); + std::shared_ptr spec = + HloShardingSpec::Create(num_shards, xla_hlo_sharding); + + Shape shape({12, 20}); + EXPECT_THAT( + spec->UniqueIndexDomains(shape), + absl_testing::IsOkAndHolds(ElementsAre( + FieldsAre(IndexDomain(Index({0, 0}), Shape({2, 20})), ElementsAre(0)), + FieldsAre(IndexDomain(Index({2, 0}), Shape({2, 20})), ElementsAre(1)), + FieldsAre(IndexDomain(Index({4, 0}), Shape({2, 20})), ElementsAre(2)), + FieldsAre(IndexDomain(Index({6, 0}), Shape({2, 20})), ElementsAre(3)), + FieldsAre(IndexDomain(Index({8, 0}), Shape({2, 20})), ElementsAre(4)), + FieldsAre(IndexDomain(Index({10, 0}), Shape({2, 20})), + ElementsAre(5))))); + EXPECT_THAT(spec->ShardToUniqueIndexDomainIndex(), + absl_testing::IsOkAndHolds(ElementsAre(0, 1, 2, 3, 4, 5))); +} + TEST_F(HloShardingSpecTest, DisassembleWithTile) { int num_shards = 6; // 6-way sharded along axis 0, 1-way sharded along axis 1. @@ -282,6 +320,28 @@ TEST_F(HloShardingSpecTest, IndexDomainsWithUnevenTile) { ElementsAreArray(TEST_HloShardingSpecIndexDomainsSlowPath(*spec, shape))); } +TEST_F(HloShardingSpecTest, UniqueIndexDomainsWithUnevenTile) { + int num_shards = 6; + // 6-way sharded along axis 0, 1-way sharded along axis 1. + auto xla_hlo_sharding = xla::HloSharding::Tile(xla::TileAssignment({6, 1})); + std::shared_ptr spec = + HloShardingSpec::Create(num_shards, xla_hlo_sharding); + + Shape shape({11, 20}); + EXPECT_THAT( + spec->UniqueIndexDomains(shape), + absl_testing::IsOkAndHolds(ElementsAre( + FieldsAre(IndexDomain(Index({0, 0}), Shape({2, 20})), ElementsAre(0)), + FieldsAre(IndexDomain(Index({2, 0}), Shape({2, 20})), ElementsAre(1)), + FieldsAre(IndexDomain(Index({4, 0}), Shape({2, 20})), ElementsAre(2)), + FieldsAre(IndexDomain(Index({6, 0}), Shape({2, 20})), ElementsAre(3)), + FieldsAre(IndexDomain(Index({8, 0}), Shape({2, 20})), ElementsAre(4)), + FieldsAre(IndexDomain(Index({10, 0}), Shape({1, 20})), + ElementsAre(5))))); + EXPECT_THAT(spec->ShardToUniqueIndexDomainIndex(), + absl_testing::IsOkAndHolds(ElementsAre(0, 1, 2, 3, 4, 5))); +} + TEST_F(HloShardingSpecTest, DisassembleWithUnevenTile) { int num_shards = 6; // 6-way sharded along axis 0, 1-way sharded along axis 1. @@ -326,6 +386,45 @@ TEST_F(HloShardingSpecTest, IndexDomainsWithPartialTile) { ElementsAreArray(TEST_HloShardingSpecIndexDomainsSlowPath(*spec, shape))); } +TEST_F(HloShardingSpecTest, UniqueIndexDomainsWithPartialTile) { + int num_shards = 6; + // 2-way sharded along axis 0, 1-way sharded along axis 1, each shard + // replicated by 3 times. + auto xla_hlo_sharding = + xla::HloSharding::PartialTile(xla::TileAssignment({2, 1, 3})); + std::shared_ptr spec = + HloShardingSpec::Create(num_shards, xla_hlo_sharding); + + Shape shape({10, 20}); + EXPECT_THAT(spec->UniqueIndexDomains(shape), + absl_testing::IsOkAndHolds(ElementsAre( + FieldsAre(IndexDomain(Index({0, 0}), Shape({5, 20})), + ElementsAre(0, 1, 2)), + FieldsAre(IndexDomain(Index({5, 0}), Shape({5, 20})), + ElementsAre(3, 4, 5))))); + EXPECT_THAT(spec->ShardToUniqueIndexDomainIndex(), + absl_testing::IsOkAndHolds(ElementsAre(0, 0, 0, 1, 1, 1))); +} + +TEST_F(HloShardingSpecTest, UniqueIndexDomainsManualFails) { + int num_shards = 6; + auto xla_hlo_sharding = xla::HloSharding::Manual(); + std::shared_ptr spec = + HloShardingSpec::Create(num_shards, xla_hlo_sharding); + + Shape shape({10, 20}); + EXPECT_THAT( + spec->UniqueIndexDomains(shape), + absl_testing::StatusIs( + tsl::error::INVALID_ARGUMENT, + HasSubstr("Manual sharding does not support UniqueIndexDomains"))); + EXPECT_THAT( + spec->ShardToUniqueIndexDomainIndex(), + absl_testing::StatusIs(tsl::error::INVALID_ARGUMENT, + HasSubstr("Manual sharding does not support " + "ShardToUniqueIndexDomainIndex"))); +} + TEST_F(HloShardingSpecTest, DisassembleWithPartialTile) { int num_shards = 6; // 2-way sharded along axis 0, 1-way sharded along axis 1, each shard diff --git a/third_party/xla/xla/python/pjrt_ifrt/xla_sharding_test.cc b/third_party/xla/xla/python/pjrt_ifrt/xla_sharding_test.cc index ca24cd1d6aa397..d22c8afd621538 100644 --- a/third_party/xla/xla/python/pjrt_ifrt/xla_sharding_test.cc +++ b/third_party/xla/xla/python/pjrt_ifrt/xla_sharding_test.cc @@ -46,6 +46,7 @@ namespace { using ::testing::ElementsAre; using ::testing::ElementsAreArray; +using ::testing::FieldsAre; using ::testing::HasSubstr; using ::testing::SizeIs; @@ -277,6 +278,21 @@ TEST_P(HloShardingTest, IndexDomainsWithReplication) { } } +TEST_P(HloShardingTest, UniqueIndexDomainsWithReplication) { + auto device_list = GetDevices({0, 1, 2, 3, 4, 5}); + // Fully replicated. + auto xla_hlo_sharding = xla::HloSharding::Replicate(); + std::shared_ptr sharding = + HloSharding::Create(device_list, MemoryKind(), xla_hlo_sharding); + + Shape shape({10, 20}); + EXPECT_THAT(sharding->UniqueIndexDomains(shape), + absl_testing::IsOkAndHolds(ElementsAre(FieldsAre( + IndexDomain(shape), ElementsAre(0, 1, 2, 3, 4, 5))))); + EXPECT_THAT(sharding->ShardToUniqueIndexDomainIndex(), + absl_testing::IsOkAndHolds(ElementsAre(0, 0, 0, 0, 0, 0))); +} + TEST_P(HloShardingTest, DisassembleWithReplication) { auto device_list = GetDevices({0, 1, 2, 3, 4, 5}); // Fully replicated. @@ -354,6 +370,28 @@ TEST_P(HloShardingTest, IndexDomainsWithTile) { } } +TEST_P(HloShardingTest, UniqueIndexDomainsWithTile) { + auto device_list = GetDevices({0, 1, 2, 3, 4, 5}); + // 6-way sharded along axis 0, 1-way sharded along axis 1. + auto xla_hlo_sharding = xla::HloSharding::Tile(xla::TileAssignment({6, 1})); + std::shared_ptr sharding = + HloSharding::Create(device_list, MemoryKind(), xla_hlo_sharding); + + Shape shape({12, 20}); + EXPECT_THAT( + sharding->UniqueIndexDomains(shape), + absl_testing::IsOkAndHolds(ElementsAre( + FieldsAre(IndexDomain(Index({0, 0}), Shape({2, 20})), ElementsAre(0)), + FieldsAre(IndexDomain(Index({2, 0}), Shape({2, 20})), ElementsAre(1)), + FieldsAre(IndexDomain(Index({4, 0}), Shape({2, 20})), ElementsAre(2)), + FieldsAre(IndexDomain(Index({6, 0}), Shape({2, 20})), ElementsAre(3)), + FieldsAre(IndexDomain(Index({8, 0}), Shape({2, 20})), ElementsAre(4)), + FieldsAre(IndexDomain(Index({10, 0}), Shape({2, 20})), + ElementsAre(5))))); + EXPECT_THAT(sharding->ShardToUniqueIndexDomainIndex(), + absl_testing::IsOkAndHolds(ElementsAre(0, 1, 2, 3, 4, 5))); +} + TEST_P(HloShardingTest, DisassembleWithTile) { auto device_list = GetDevices({0, 1, 2, 3, 4, 5}); // 6-way sharded along axis 0, 1-way sharded along axis 1. @@ -474,7 +512,7 @@ TEST_P(HloShardingTest, DisassembleWithUnevenTile) { TEST_P(HloShardingTest, IndexDomainsWithPartialTile) { auto device_list = GetDevices({0, 1, 2, 3, 4, 5}); // 2-way sharded along axis 0, 1-way sharded along axis 1, each shard - // replicated by 3 times. + // replicated 3 times. auto xla_hlo_sharding = xla::HloSharding::PartialTile(xla::TileAssignment({2, 1, 3})); std::shared_ptr sharding = @@ -514,10 +552,30 @@ TEST_P(HloShardingTest, IndexDomainsWithPartialTile) { } } +TEST_P(HloShardingTest, UniqueIndexDomainsWithPartialTile) { + auto device_list = GetDevices({0, 1, 2, 3, 4, 5}); + // 2-way sharded along axis 0, 1-way sharded along axis 1, each shard + // replicated 3 times. + auto xla_hlo_sharding = + xla::HloSharding::PartialTile(xla::TileAssignment({2, 1, 3})); + std::shared_ptr sharding = + HloSharding::Create(device_list, MemoryKind(), xla_hlo_sharding); + + Shape shape({10, 20}); + EXPECT_THAT(sharding->UniqueIndexDomains(shape), + absl_testing::IsOkAndHolds(ElementsAre( + FieldsAre(IndexDomain(Index({0, 0}), Shape({5, 20})), + ElementsAre(0, 1, 2)), + FieldsAre(IndexDomain(Index({5, 0}), Shape({5, 20})), + ElementsAre(3, 4, 5))))); + EXPECT_THAT(sharding->ShardToUniqueIndexDomainIndex(), + absl_testing::IsOkAndHolds(ElementsAre(0, 0, 0, 1, 1, 1))); +} + TEST_P(HloShardingTest, DisassembleWithPartialTile) { auto device_list = GetDevices({0, 1, 2, 3, 4, 5}); // 2-way sharded along axis 0, 1-way sharded along axis 1, each shard - // replicated by 3 times. + // replicated 3 times. auto xla_hlo_sharding = xla::HloSharding::PartialTile(xla::TileAssignment({2, 1, 3})); std::shared_ptr sharding = @@ -555,7 +613,7 @@ TEST_P(HloShardingTest, DisassembleWithPartialTile) { TEST_P(HloShardingTest, IndexDomainsWithSubgroupReplicated) { auto device_list = GetDevices({0, 1, 2, 3, 4, 5}); // 2-way sharded along axis 0, 1-way sharded along axis 1, each shard - // replicated by 3 times. + // replicated 3 times. auto xla_hlo_sharding = xla::HloSharding::Subgroup( xla::TileAssignment({2, 1, 3}), {xla::OpSharding::REPLICATED}); std::shared_ptr sharding = @@ -598,7 +656,7 @@ TEST_P(HloShardingTest, IndexDomainsWithSubgroupReplicated) { TEST_P(HloShardingTest, DisassembleWithSubgroupReplicated) { auto device_list = GetDevices({0, 1, 2, 3, 4, 5}); // 2-way sharded along axis 0, 1-way sharded along axis 1, each shard - // replicated by 3 times. + // replicated 3 times. auto xla_hlo_sharding = xla::HloSharding::Subgroup( xla::TileAssignment({2, 1, 3}), {xla::OpSharding::REPLICATED}); std::shared_ptr sharding = @@ -636,7 +694,7 @@ TEST_P(HloShardingTest, DisassembleWithSubgroupReplicated) { TEST_P(HloShardingTest, IndexDomainsWithSubgroupMaximalSlowPath) { auto device_list = GetDevices({0, 1, 2, 3, 4, 5}); // 2-way sharded along axis 0, 1-way sharded along axis 1, each shard - // maximal-replicated by 3 times, device#0 in each replication is maximal. + // maximal-replicated 3 times, device#0 in each replication is maximal. auto xla_hlo_sharding = xla::HloSharding::Subgroup( xla::TileAssignment({2, 1, 3}), {xla::OpSharding::MAXIMAL}); std::shared_ptr sharding = @@ -679,7 +737,7 @@ TEST_P(HloShardingTest, IndexDomainsWithSubgroupMaximalSlowPath) { TEST_P(HloShardingTest, DisassembleWithSubgroupMaximalSlowPath) { auto device_list = GetDevices({0, 1, 2, 3, 4, 5}); // 2-way sharded along axis 0, 1-way sharded along axis 1, each shard - // maximal-replicated by 3 times, device#0 in each replication is maximal. + // maximal-replicated 3 times, device#0 in each replication is maximal. auto xla_hlo_sharding = xla::HloSharding::Subgroup( xla::TileAssignment({2, 1, 3}), {xla::OpSharding::MAXIMAL}); std::shared_ptr sharding = diff --git a/third_party/xla/xla/python/version.h b/third_party/xla/xla/python/version.h index a2ae7e4250a98f..5f71f9ef3ecc56 100644 --- a/third_party/xla/xla/python/version.h +++ b/third_party/xla/xla/python/version.h @@ -18,6 +18,7 @@ limitations under the License. // An increasing version number to protect jax code against breaking changes. // In JAX, reference this via jax._src.lib.ifrt_version. -#define JAX_IFRT_VERSION_NUMBER 67 // Scan dimension in TPU MLIR dialect. +#define JAX_IFRT_VERSION_NUMBER \ + 68 // StableHlo CollectiveBroadcastOp Compatibility #endif // XLA_PYTHON_VERSION_H_ diff --git a/third_party/xla/xla/service/BUILD b/third_party/xla/xla/service/BUILD index f1a7f7134f2de0..9e402b8148fed0 100644 --- a/third_party/xla/xla/service/BUILD +++ b/third_party/xla/xla/service/BUILD @@ -4059,7 +4059,9 @@ cc_library( ":buffer_value", ":hlo_value", "//xla:xla_data_proto_cc", + "//xla/hlo/ir:hlo", "@com_google_absl//absl/algorithm:container", + "@com_google_absl//absl/container:flat_hash_set", "@com_google_absl//absl/log:check", "@com_google_absl//absl/status", "@com_google_absl//absl/status:statusor", diff --git a/third_party/xla/xla/service/buffer_assignment.h b/third_party/xla/xla/service/buffer_assignment.h index bb898e2f1fb823..6d69deac524334 100644 --- a/third_party/xla/xla/service/buffer_assignment.h +++ b/third_party/xla/xla/service/buffer_assignment.h @@ -701,11 +701,7 @@ class BufferAssignment { int64_t HloBufferSize(const HloBuffer& buffer) { auto [it, inserted] = cached_buffer_sizes_.try_emplace(buffer.id()); if (inserted) { - int64_t result = 0; - for (const HloValue* value : buffer.values()) { - result = std::max(result, buffer_size_(*value)); - } - it->second = result; + it->second = buffer.ComputeSize(buffer_size_); } return it->second; } diff --git a/third_party/xla/xla/service/heap_simulator/heap_simulator.cc b/third_party/xla/xla/service/heap_simulator/heap_simulator.cc index 999ffacb383873..847bcd49250dc4 100644 --- a/third_party/xla/xla/service/heap_simulator/heap_simulator.cc +++ b/third_party/xla/xla/service/heap_simulator/heap_simulator.cc @@ -347,7 +347,7 @@ absl::Status HeapSimulator::RunComputation( if (!buffer_live_ranges.contains(value)) { continue; } - if (IgnoreBuffer(value)) { + if (!IsHeapPressureImpacting(value)) { continue; } @@ -380,10 +380,7 @@ absl::Status HeapSimulator::RunComputation( // Populate buffer sizes with the maximum size of the constituent HloValues. for (const HloBuffer& buffer : alias_analysis.buffers()) { - int64_t size = 0; - for (const HloValue* value : buffer.values()) { - size = std::max(size, (*size_fn_)(*value)); - } + int64_t size = buffer.ComputeSize(*size_fn_); const HloValue* first_value = nullptr; for (const HloValue* value : buffer.values()) { buffer_groups_.emplace(value, size); @@ -443,7 +440,7 @@ absl::Status HeapSimulator::RunComputation( continue; } - if (IgnoreBuffer(operand_value)) { + if (!IsHeapPressureImpacting(operand_value)) { continue; } @@ -519,17 +516,9 @@ HeapSimulator::HeapSimulator( HeapSimulator::~HeapSimulator() {} -bool HeapSimulator::IgnoreBuffer(const HloValue* buffer) const { - // Buffers for constants are ignored unless the alloc_constants option is - // set. Also ignore buffers that we're not meant to assign. - // - // TODO(b/32248867): For consistency, constants should get allocations. - if (!options_.alloc_constants && - buffer->instruction()->opcode() == HloOpcode::kConstant) { - return true; - } - return options_.buffers_to_assign != nullptr && - !options_.buffers_to_assign->contains(buffer); +bool HeapSimulator::IsHeapPressureImpacting(const HloValue* buffer) const { + return HloBuffer::IsHeapPressureImpacting(*buffer, options_.alloc_constants, + options_.buffers_to_assign); } // Alloc always calls the underlying heap algorithm. diff --git a/third_party/xla/xla/service/heap_simulator/heap_simulator.h b/third_party/xla/xla/service/heap_simulator/heap_simulator.h index beb04f68c64c2f..187e06dd731fb5 100644 --- a/third_party/xla/xla/service/heap_simulator/heap_simulator.h +++ b/third_party/xla/xla/service/heap_simulator/heap_simulator.h @@ -214,7 +214,10 @@ class HeapSimulator { const HloAliasAnalysis& alias_analysis, const AliasInfo* alias_info, HloLiveRange* live_range); - bool IgnoreBuffer(const HloValue* buffer) const; + // Returns whether the buffer should be allocated space in the heap simulation + // (excludes constants unless alloc_constants is set, and respects the + // buffers_to_assign filter). + bool IsHeapPressureImpacting(const HloValue* buffer) const; void Alloc(const HloValue* buffer, const HloInstruction* instruction); void Free(const HloValue* buffer, const HloInstruction* instruction); // ShareBuffer indicates that a new buffer is defined and it has to be the diff --git a/third_party/xla/xla/service/hlo_buffer.cc b/third_party/xla/xla/service/hlo_buffer.cc index 7e2ba451f6a87f..7fe0f5041828d2 100644 --- a/third_party/xla/xla/service/hlo_buffer.cc +++ b/third_party/xla/xla/service/hlo_buffer.cc @@ -16,14 +16,19 @@ limitations under the License. #include "xla/service/hlo_buffer.h" #include +#include #include #include #include #include "absl/algorithm/container.h" +#include "absl/container/flat_hash_set.h" #include "absl/log/check.h" #include "absl/strings/str_cat.h" #include "absl/strings/str_join.h" +#include "xla/hlo/ir/hlo_instruction.h" +#include "xla/hlo/ir/hlo_opcode.h" +#include "xla/service/buffer_value.h" #include "xla/service/hlo_value.h" namespace xla { @@ -73,4 +78,33 @@ std::ostream& operator<<(std::ostream& out, const HloBuffer& buffer) { return out; } +int64_t HloBuffer::ComputeSize(const BufferValue::SizeFunction& size_fn) const { + int64_t max_size = 0; + for (const HloValue* value : values_) { + max_size = std::max(max_size, size_fn(*value)); + } + return max_size; +} + +bool HloBuffer::IsHeapPressureImpacting( + const HloValue& value, bool alloc_constants, + const absl::flat_hash_set* buffers_to_assign) { + if (!alloc_constants && + value.instruction()->opcode() == HloOpcode::kConstant) { + return false; + } + if (buffers_to_assign != nullptr && !buffers_to_assign->contains(&value)) { + return false; + } + return true; +} + +bool HloBuffer::IsHeapPressureImpacting( + bool alloc_constants, + const absl::flat_hash_set* buffers_to_assign) const { + return absl::c_any_of(values_, [&](const HloValue* value) { + return IsHeapPressureImpacting(*value, alloc_constants, buffers_to_assign); + }); +} + } // namespace xla diff --git a/third_party/xla/xla/service/hlo_buffer.h b/third_party/xla/xla/service/hlo_buffer.h index 2ee0887ee76a08..4be60c9097d251 100644 --- a/third_party/xla/xla/service/hlo_buffer.h +++ b/third_party/xla/xla/service/hlo_buffer.h @@ -22,6 +22,7 @@ limitations under the License. #include #include +#include "absl/container/flat_hash_set.h" #include "absl/log/check.h" #include "absl/status/status.h" #include "absl/status/statusor.h" @@ -97,6 +98,25 @@ class HloBuffer { // Return all values contained in this buffer. const std::vector& values() const { return values_; } + // Computes the physical size of the buffer as the maximum size of its + // constituent HloValues according to the given size function. + int64_t ComputeSize(const BufferValue::SizeFunction& size_fn) const; + + // Returns whether this value impacts dynamic heap allocation pressure (e.g. + // not an embedded constant or a buffer excluded by allocation filter). + static bool IsHeapPressureImpacting( + const HloValue& value, bool alloc_constants = false, + const absl::flat_hash_set* buffers_to_assign = nullptr); + + // Returns whether this buffer impacts dynamic heap allocation pressure. + // Matches HeapSimulator's allocation logic: a buffer impacts heap pressure if + // any of its constituent values impacts heap pressure (i.e. is not ignored + // by HeapSimulator). A buffer is only exempt if all of its constituent values + // are ignored (e.g. purely constants or outside buffers_to_assign). + bool IsHeapPressureImpacting(bool alloc_constants = false, + const absl::flat_hash_set* + buffers_to_assign = nullptr) const; + // Memory space color. Used to indicate the memory space that the hlo buffer // needs to live in. absl::StatusOr color() const { diff --git a/third_party/xla/xla/service/spmd/shardy/stablehlo_round_trip/shard_map_export.cc b/third_party/xla/xla/service/spmd/shardy/stablehlo_round_trip/shard_map_export.cc index 6a072cab73678b..71b6c94c5a3e40 100644 --- a/third_party/xla/xla/service/spmd/shardy/stablehlo_round_trip/shard_map_export.cc +++ b/third_party/xla/xla/service/spmd/shardy/stablehlo_round_trip/shard_map_export.cc @@ -184,16 +184,24 @@ void setOpManualAxes(Operation* op, ManualAxesAttr manualAxes, op->setAttr(kManualAxes, manualAxes); } -void setFuncManualAxesRecursively(FuncOp funcOp, ManualAxesAttr manualAxes, - Attribute meshOrRef, - const mlir::SymbolTable& symbolTable); - -mlir::WalkResult setManualAxes(Operation* op, ManualAxesAttr manualAxes, - Attribute meshOrRef, - const mlir::SymbolTable& symbolTable) { - if (mlir::isa(op)) { - // Skip `ManualComputationOp`s and their nested operations, they will - // be handled separately. +void setManualAxesForOpsInBody( + ManualComputationOp op, const mlir::SymbolTable& symbolTable, + ManualComputationToParentManualAxes& parentManualCompAxes); + +void setFuncManualAxesRecursively( + FuncOp funcOp, ManualAxesAttr manualAxes, Attribute meshOrRef, + const mlir::SymbolTable& symbolTable, + ManualComputationToParentManualAxes& parentManualCompAxes); + +mlir::WalkResult setManualAxes( + Operation* op, ManualAxesAttr manualAxes, Attribute meshOrRef, + const mlir::SymbolTable& symbolTable, + ManualComputationToParentManualAxes& parentManualCompAxes) { + if (auto manualCompOp = mlir::dyn_cast(op)) { + // Record parent manual axes for this manualCompOp and process its body. + parentManualCompAxes[manualCompOp].assign(manualAxes.getValue().begin(), + manualAxes.getValue().end()); + setManualAxesForOpsInBody(manualCompOp, symbolTable, parentManualCompAxes); return mlir::WalkResult::skip(); } if (!mlir::isa(op)) { @@ -202,14 +210,16 @@ mlir::WalkResult setManualAxes(Operation* op, ManualAxesAttr manualAxes, if (CallOp callOp = mlir::dyn_cast(op)) { FuncOp funcOp = symbolTable.lookup(callOp.getCallee()); CHECK(funcOp) << "Failed to lookup function: " << callOp.getCallee().str(); - setFuncManualAxesRecursively(funcOp, manualAxes, meshOrRef, symbolTable); + setFuncManualAxesRecursively(funcOp, manualAxes, meshOrRef, symbolTable, + parentManualCompAxes); } return mlir::WalkResult::advance(); } -void setFuncManualAxesRecursively(FuncOp funcOp, ManualAxesAttr manualAxes, - Attribute meshOrRef, - const mlir::SymbolTable& symbolTable) { +void setFuncManualAxesRecursively( + FuncOp funcOp, ManualAxesAttr manualAxes, Attribute meshOrRef, + const mlir::SymbolTable& symbolTable, + ManualComputationToParentManualAxes& parentManualCompAxes) { llvm::SmallVector funcArgAttrs; funcArgAttrs.reserve(funcOp.getNumArguments()); for (int argNum = 0; argNum < funcOp.getNumArguments(); argNum++) { @@ -250,15 +260,15 @@ void setFuncManualAxesRecursively(FuncOp funcOp, ManualAxesAttr manualAxes, // Walk in preorder of blocks in order to stop walks on manual computations. funcOp->walk([&](Operation* op) { - return setManualAxes(op, manualAxes, meshOrRef, symbolTable); + return setManualAxes(op, manualAxes, meshOrRef, symbolTable, + parentManualCompAxes); }); } // Sets the manual axes of all operations in `op`'s body. void setManualAxesForOpsInBody( - ManualComputationOp op, - const ManualComputationToParentManualAxes& parentManualCompAxes, - const mlir::SymbolTable& symbolTable) { + ManualComputationOp op, const mlir::SymbolTable& symbolTable, + ManualComputationToParentManualAxes& parentManualCompAxes) { TensorShardingAttr sharding = getFirstSharding(op); if (!sharding) { // If there are no in/out shardings, op.getManualAxes() must be empty. We do @@ -280,7 +290,8 @@ void setManualAxesForOpsInBody( // Set the manual axes of all operations in the body. op.getBody().front().walk( [&](Operation* opInBody) { - return setManualAxes(opInBody, manualAxesAttr, meshOrRef, symbolTable); + return setManualAxes(opInBody, manualAxesAttr, meshOrRef, symbolTable, + parentManualCompAxes); }); } @@ -483,12 +494,16 @@ class ShardMapExportPass // walk. module->walk([&](ManualComputationOp op) { if (auto parentOp = op->getParentOfType()) { - SmallVector& parentAxes = parentManualCompAxes[op]; - parentAxes = parentManualCompAxes[parentOp]; - parentAxes.insert(parentAxes.end(), parentOp.getManualAxes().begin(), + SmallVector parentAxes; + if (auto it = parentManualCompAxes.find(parentOp); + it != parentManualCompAxes.end()) { + parentAxes = it->second; + } + parentAxes.append(parentOp.getManualAxes().begin(), parentOp.getManualAxes().end()); + parentManualCompAxes[op] = std::move(parentAxes); } - setManualAxesForOpsInBody(op, parentManualCompAxes, symbolTable); + setManualAxesForOpsInBody(op, symbolTable, parentManualCompAxes); }); // Need to do a separate post order walk to inline the diff --git a/third_party/xla/xla/service/spmd/shardy/test/stablehlo_round_trip_shard_map_export.mlir b/third_party/xla/xla/service/spmd/shardy/test/stablehlo_round_trip_shard_map_export.mlir index 8a7d5aa3cbd6b4..5ca34525922de5 100644 --- a/third_party/xla/xla/service/spmd/shardy/test/stablehlo_round_trip_shard_map_export.mlir +++ b/third_party/xla/xla/service/spmd/shardy/test/stablehlo_round_trip_shard_map_export.mlir @@ -575,3 +575,43 @@ func.func private @bar(%arg0: tensor<8xi32>) -> tensor<8xi32> { // CHECK-NEXT: %1 = call @bar(%0) {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{}]>]>, xla.sdy.manual_axes = #sdy} : (tensor<8xi32>) -> tensor<8xi32> // CHECK-NEXT: return %arg0 : tensor<4xi32> // CHECK-NEXT: } + +// ----- +sdy.mesh @mesh = <["a"=2, "b"=4]> + +// Tests that when a function containing a `ManualComputationOp` (callee) is +// called from within another `ManualComputationOp` (caller), the callee's +// `ManualComputationOp` correctly inherits parent manual axes from the caller. +// +// 1. Caller passes parent manual axes {"a"} to callee's inputs/outputs: +// CHECK-LABEL: func private @called_func_with_manual_comp( +// CHECK-SAME: %arg0: tensor<8x8xf32> {sdy.sharding = #sdy.sharding<@mesh, [{}, {}]>, xla.sdy.manual_axes = #sdy}) +// CHECK-SAME: -> (tensor<8x8xf32> {sdy.sharding = #sdy.sharding<@mesh, [{}, {}]>, xla.sdy.manual_axes = #sdy}) { +// CHECK-NEXT: %[[COPY:.*]] = mhlo.copy %arg0 {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{}, {"b"}]>]>, xla.sdy.manual_axes = #sdy} : tensor<8x8xf32> +// +// 2. Full-to-shard inside callee inherits caller axis "a" + callee axis "b" -> manual axes {"a", "b"}: +// CHECK-NEXT: %[[FULL_TO_SHARD:.*]] = stablehlo.custom_call @SPMDFullToShardShape(%[[COPY]]) {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{}, {}]>]>, xla.sdy.manual_axes = #sdy} : (tensor<8x8xf32>) -> tensor<8x2xf32> +// +// 3. Inner body call runs with both axes manual: +// CHECK-NEXT: %[[CALL:.*]] = call @xla.sdy.inlinable_manual_computation_body(%[[FULL_TO_SHARD]]) {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{}, {}]>]>, xla.sdy.manual_axes = #sdy} : (tensor<8x2xf32>) -> tensor<8x2xf32> +// CHECK-NEXT: %[[COPY_1:.*]] = mhlo.copy %[[CALL]] {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{}, {}]>]>, xla.sdy.manual_axes = #sdy} : tensor<8x2xf32> +// +// 4. Shard-to-full returns to caller scope, with manual axis {"a"} preserved: +// CHECK-NEXT: %[[SHARD_TO_FULL:.*]] = stablehlo.custom_call @SPMDShardToFullShape(%[[COPY_1]]) {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{}, {"b"}]>]>, xla.sdy.manual_axes = #sdy} : (tensor<8x2xf32>) -> tensor<8x8xf32> +// CHECK-NEXT: return %[[SHARD_TO_FULL]] : tensor<8x8xf32> +func.func private @called_func_with_manual_comp(%arg0: tensor<8x8xf32>) -> tensor<8x8xf32> { + %0 = sdy.manual_computation(%arg0) in_shardings=[<@mesh, [{}, {"b"}]>] out_shardings=[<@mesh, [{}, {"b"}]>] manual_axes={"b"} (%arg1: tensor<8x2xf32>) { + %1 = stablehlo.add %arg1, %arg1 : tensor<8x2xf32> + sdy.return %1 : tensor<8x2xf32> + } : (tensor<8x8xf32>) -> tensor<8x8xf32> + return %0 : tensor<8x8xf32> +} + +// CHECK-LABEL: func @manual_comp_calls_func_with_manual_comp( +func.func @manual_comp_calls_func_with_manual_comp(%arg0: tensor<16x8xf32>) -> tensor<16x8xf32> { + %0 = sdy.manual_computation(%arg0) in_shardings=[<@mesh, [{"a"}, {}]>] out_shardings=[<@mesh, [{"a"}, {}]>] manual_axes={"a"} (%arg1: tensor<8x8xf32>) { + %1 = func.call @called_func_with_manual_comp(%arg1) : (tensor<8x8xf32>) -> tensor<8x8xf32> + sdy.return %1 : tensor<8x8xf32> + } : (tensor<16x8xf32>) -> tensor<16x8xf32> + return %0 : tensor<16x8xf32> +}