diff --git a/tensorflow/compiler/mlir/tensorflow/transforms/BUILD b/tensorflow/compiler/mlir/tensorflow/transforms/BUILD index bc47f98a1ba4a3..d19b0e3a8c45a5 100644 --- a/tensorflow/compiler/mlir/tensorflow/transforms/BUILD +++ b/tensorflow/compiler/mlir/tensorflow/transforms/BUILD @@ -717,6 +717,7 @@ cc_library( "//tensorflow/core/protobuf:for_core_protos_cc", "@com_google_absl//absl/base:core_headers", "@com_google_absl//absl/log", + "@com_google_absl//absl/log:vlog_is_on", "@com_google_absl//absl/status", "@llvm-project//llvm:Support", "@llvm-project//mlir:FuncDialect", diff --git a/tensorflow/compiler/mlir/tensorflow/transforms/bridge.cc b/tensorflow/compiler/mlir/tensorflow/transforms/bridge.cc index d5284e1ae5736b..1fdd6d46ed3296 100644 --- a/tensorflow/compiler/mlir/tensorflow/transforms/bridge.cc +++ b/tensorflow/compiler/mlir/tensorflow/transforms/bridge.cc @@ -16,6 +16,7 @@ limitations under the License. #include "tensorflow/compiler/mlir/tensorflow/transforms/bridge.h" #include "absl/log/log.h" +#include "absl/log/vlog_is_on.h" #include "llvm/ADT/StringRef.h" #include "mlir/IR/BuiltinOps.h" // from @llvm-project #include "mlir/Pass/PassManager.h" // from @llvm-project diff --git a/tensorflow/compiler/mlir/tfrt/transforms/ifrt/BUILD b/tensorflow/compiler/mlir/tfrt/transforms/ifrt/BUILD index 7ae680f60668c2..3fc654d38b506d 100644 --- a/tensorflow/compiler/mlir/tfrt/transforms/ifrt/BUILD +++ b/tensorflow/compiler/mlir/tfrt/transforms/ifrt/BUILD @@ -241,10 +241,9 @@ tf_cc_test( "@xla//xla/pjrt:pjrt_compiler", "@xla//xla/pjrt/plugin/xla_cpu:cpu_topology_description", "@xla//xla/python/ifrt", - "@xla//xla/python/ifrt:mock", "@xla//xla/python/ifrt:test_util", "@xla//xla/python/pjrt_ifrt", - "@xla//xla/python/pjrt_ifrt:tfrt_cpu_client_test_lib", + "@xla//xla/python/pjrt_ifrt:pjrt_cpu_client_test_lib", "@xla//xla/service:device_assignment", "@xla//xla/tsl/platform:statusor", ], diff --git a/tensorflow/compiler/mlir/tfrt/transforms/ifrt/tf2hlo_test.cc b/tensorflow/compiler/mlir/tfrt/transforms/ifrt/tf2hlo_test.cc index d5ac5184661058..780e361f06f2dd 100644 --- a/tensorflow/compiler/mlir/tfrt/transforms/ifrt/tf2hlo_test.cc +++ b/tensorflow/compiler/mlir/tfrt/transforms/ifrt/tf2hlo_test.cc @@ -43,7 +43,6 @@ limitations under the License. #include "xla/pjrt/pjrt_compiler.h" #include "xla/pjrt/plugin/xla_cpu/cpu_topology_description.h" #include "xla/python/ifrt/client.h" -#include "xla/python/ifrt/mock.h" #include "xla/python/ifrt/test_util.h" #include "xla/python/pjrt_ifrt/pjrt_topology.h" #include "xla/service/device_assignment.h" @@ -514,18 +513,15 @@ TEST_F(Tf2HloTest, GpuCompile) { ASSERT_TRUE(mlir_module); ASSERT_TRUE(mlir_module.get() != nullptr); - xla::ifrt::MockClient mock_client; - ON_CALL(mock_client, GetDefaultDeviceAssignment) - .WillByDefault([]() -> absl::StatusOr { - return xla::DeviceAssignment(1, 1); - }); + TF_ASSERT_OK_AND_ASSIGN(std::shared_ptr client, + xla::ifrt::test_util::GetClient()); std::vector dtype_and_shapes; dtype_and_shapes.push_back(DtypeAndShape{DT_FLOAT, {}}); TF_ASSERT_OK_AND_ASSIGN( tensorflow::tpu::TPUCompileMetadataProto compile_metadata, - GetCompileMetadata(mlir_module.get(), mock_client)); + GetCompileMetadata(mlir_module.get(), *client)); TF_ASSERT_OK(UpdateCompileMetadata(compile_metadata, dtype_and_shapes)); std::vector variable_arg_indices; diff --git a/tensorflow/core/kernels/rnn/lstm_ops.cc b/tensorflow/core/kernels/rnn/lstm_ops.cc index 6307a62c91d3d0..fae4e36733e5fc 100644 --- a/tensorflow/core/kernels/rnn/lstm_ops.cc +++ b/tensorflow/core/kernels/rnn/lstm_ops.cc @@ -593,6 +593,35 @@ class LSTMBlockCellGradOp : public OpKernel { const Tensor* h_grad_tensor = nullptr; OP_REQUIRES_OK(ctx, ctx->input("h_grad", &h_grad_tensor)); + // The dimensions read below, and the sanity checks that follow them, index + // dimension 1 of these tensors. That is a fatal check rather than an error + // when a tensor has a lower rank, so validate the ranks first, the way + // LSTMBlockCellOp does before using the same inputs. + const auto check_rank = [](const Tensor* t, const char* name, + int expected) -> absl::Status { + if (t->dims() != expected) { + return absl::InvalidArgumentError(absl::StrCat( + name, " must be rank ", expected, " but is rank ", t->dims())); + } + return absl::OkStatus(); + }; + OP_REQUIRES_OK(ctx, check_rank(x_tensor, "x", 2)); + OP_REQUIRES_OK(ctx, check_rank(cs_prev_tensor, "cs_prev", 2)); + OP_REQUIRES_OK(ctx, check_rank(h_prev_tensor, "h_prev", 2)); + OP_REQUIRES_OK(ctx, check_rank(w_tensor, "w", 2)); + OP_REQUIRES_OK(ctx, check_rank(wci_tensor, "wci", 1)); + OP_REQUIRES_OK(ctx, check_rank(wcf_tensor, "wcf", 1)); + OP_REQUIRES_OK(ctx, check_rank(wco_tensor, "wco", 1)); + OP_REQUIRES_OK(ctx, check_rank(b_tensor, "b", 1)); + OP_REQUIRES_OK(ctx, check_rank(i_tensor, "i", 2)); + OP_REQUIRES_OK(ctx, check_rank(cs_tensor, "cs", 2)); + OP_REQUIRES_OK(ctx, check_rank(f_tensor, "f", 2)); + OP_REQUIRES_OK(ctx, check_rank(o_tensor, "o", 2)); + OP_REQUIRES_OK(ctx, check_rank(ci_tensor, "ci", 2)); + OP_REQUIRES_OK(ctx, check_rank(co_tensor, "co", 2)); + OP_REQUIRES_OK(ctx, check_rank(cs_grad_tensor, "cs_grad", 2)); + OP_REQUIRES_OK(ctx, check_rank(h_grad_tensor, "h_grad", 2)); + const int64_t batch_size = x_tensor->dim_size(0); const int64_t input_size = x_tensor->dim_size(1); const int64_t cell_size = cs_prev_tensor->dim_size(1); @@ -1292,6 +1321,50 @@ class BlockLSTMGradOp : public OpKernel { const Tensor* h_grad = nullptr; OP_REQUIRES_OK(ctx, ctx->input("h_grad", &h_grad)); + // `timelen`, `batch_size` and `cell_size` are taken from x and w, and every + // other input is sliced with them below. A tensor whose dimensions + // disagree is read out of bounds rather than failing the op: a zero batch + // on one side and a non-zero batch on the other crashes, and a shorter + // time dimension trips a fatal check while slicing. LSTMBlockCellGradOp + // already validates its own inputs this way. The rank is checked here as + // well so that the dimension lookups below are valid for any caller. + const auto check_dims = [&](const Tensor* t, const char* name, + bool has_time) -> absl::Status { + const int expected_dims = has_time ? 3 : 2; + if (t->dims() != expected_dims) { + return absl::InvalidArgumentError(absl::StrCat( + name, " must be rank ", expected_dims, " but is rank ", t->dims())); + } + const int batch_dim = has_time ? 1 : 0; + if (has_time && t->dim_size(0) != timelen) { + return absl::InvalidArgumentError( + absl::StrCat(name, ".dim_size(0) != timelen: ", t->dim_size(0), + " vs. ", timelen)); + } + if (t->dim_size(batch_dim) != batch_size) { + return absl::InvalidArgumentError(absl::StrCat( + name, ".dim_size(", batch_dim, + ") != batch_size: ", t->dim_size(batch_dim), " vs. ", batch_size)); + } + if (t->dim_size(batch_dim + 1) != cell_size) { + return absl::InvalidArgumentError( + absl::StrCat(name, ".dim_size(", batch_dim + 1, ") != cell_size: ", + t->dim_size(batch_dim + 1), " vs. ", cell_size)); + } + return absl::OkStatus(); + }; + OP_REQUIRES_OK(ctx, check_dims(cs_prev_tensor, "cs_prev", false)); + OP_REQUIRES_OK(ctx, check_dims(h_prev_tensor, "h_prev", false)); + OP_REQUIRES_OK(ctx, check_dims(i_out, "i", true)); + OP_REQUIRES_OK(ctx, check_dims(cs_out, "cs", true)); + OP_REQUIRES_OK(ctx, check_dims(f_out, "f", true)); + OP_REQUIRES_OK(ctx, check_dims(o_out, "o", true)); + OP_REQUIRES_OK(ctx, check_dims(ci_out, "ci", true)); + OP_REQUIRES_OK(ctx, check_dims(co_out, "co", true)); + OP_REQUIRES_OK(ctx, check_dims(h_out, "h", true)); + OP_REQUIRES_OK(ctx, check_dims(cs_grad, "cs_grad", true)); + OP_REQUIRES_OK(ctx, check_dims(h_grad, "h_grad", true)); + TensorShape batch_input_shape({timelen, batch_size, input_size}); Tensor* x_grad; OP_REQUIRES_OK(ctx, diff --git a/tensorflow/python/kernel_tests/nn_ops/rnn_cell_test.py b/tensorflow/python/kernel_tests/nn_ops/rnn_cell_test.py index 9309f178ee9122..34da26ab30d746 100644 --- a/tensorflow/python/kernel_tests/nn_ops/rnn_cell_test.py +++ b/tensorflow/python/kernel_tests/nn_ops/rnn_cell_test.py @@ -1468,6 +1468,102 @@ def testLSTMBlockInvalidArgument(self): ) ) + def _blockLSTMGradInputs( + self, timelen, batch_size, input_size, cell_size, x_batch_size=None + ): + """Returns a consistent set of BlockLSTMGrad inputs, x batch overridable.""" + if x_batch_size is None: + x_batch_size = batch_size + + def rand(*shape): + return random_ops.random_uniform(list(shape), dtype=dtypes.float32) + + def per_step(): + return rand(timelen, batch_size, cell_size) + + return dict( + seq_len_max=constant_op.constant(timelen, dtype=dtypes.int64), + x=rand(timelen, x_batch_size, input_size), + cs_prev=rand(batch_size, cell_size), + h_prev=rand(batch_size, cell_size), + w=rand(input_size + cell_size, 4 * cell_size), + wci=rand(cell_size), + wcf=rand(cell_size), + wco=rand(cell_size), + b=rand(4 * cell_size), + i=per_step(), + cs=per_step(), + f=per_step(), + o=per_step(), + ci=per_step(), + co=per_step(), + h=per_step(), + cs_grad=per_step(), + h_grad=per_step(), + use_peephole=True, + ) + + def testBlockLSTMGradBatchSizeMismatch(self): + # Test case for GitHub issue 112746. x carried a different batch size from + # the other inputs, which were then sliced with the batch size taken from + # x and read out of bounds, crashing the process. + kwargs = self._blockLSTMGradInputs(3, 2, 4, 5, x_batch_size=0) + with self.assertRaisesRegex(errors_impl.InvalidArgumentError, "batch_size"): + self.evaluate(gen_rnn_ops.BlockLSTMGrad(**kwargs)) + + def testBlockLSTMGradTimeLenMismatch(self): + # A shorter time dimension on one of the per step inputs was sliced past + # its end. + kwargs = self._blockLSTMGradInputs(3, 2, 4, 5) + kwargs["i"] = random_ops.random_uniform([1, 2, 5], dtype=dtypes.float32) + with self.assertRaisesRegex(errors_impl.InvalidArgumentError, "timelen"): + self.evaluate(gen_rnn_ops.BlockLSTMGrad(**kwargs)) + + def testBlockLSTMGradConsistentShapesStillWork(self): + # Consistent shapes must keep working, including a zero batch size. + for batch_size in (0, 2): + kwargs = self._blockLSTMGradInputs(3, batch_size, 4, 5) + self.evaluate(gen_rnn_ops.BlockLSTMGrad(**kwargs)) + + def testLSTMBlockCellGradInvalidRank(self): + # Test case for GitHub issue 112746. A rank 1 x was indexed at dimension 1 + # before any shape check ran, which aborted the process. Eager execution + # reaches the kernel check ("x must be rank 2"), while graph construction + # is stopped earlier by the op's shape function ("Shape must be rank 2"), + # which surfaces as a ValueError, so both are accepted. + batch_size, input_size, cell_size = 2, 4, 5 + + def rand(*shape): + return random_ops.random_uniform(list(shape), dtype=dtypes.float32) + + def per_step(): + return rand(batch_size, cell_size) + + with self.assertRaisesRegex( + (ValueError, errors_impl.InvalidArgumentError), "must be rank 2" + ): + self.evaluate( + gen_rnn_ops.LSTMBlockCellGrad( + x=constant_op.constant([], dtype=dtypes.float32), + cs_prev=per_step(), + h_prev=per_step(), + w=rand(input_size + cell_size, 4 * cell_size), + wci=rand(cell_size), + wcf=rand(cell_size), + wco=rand(cell_size), + b=rand(4 * cell_size), + i=per_step(), + cs=per_step(), + f=per_step(), + o=per_step(), + ci=per_step(), + co=per_step(), + cs_grad=per_step(), + h_grad=per_step(), + use_peephole=True, + ) + ) + @test_util.run_in_graph_and_eager_modes def testBlockLSTMSeqLenMaxBounds(self): timelen, batch_size, input_size, cell_size = 4, 2, 3, 5 diff --git a/third_party/xla/build_tools/lint/tags.py b/third_party/xla/build_tools/lint/tags.py index a909af93153f8e..6da264a0827519 100644 --- a/third_party/xla/build_tools/lint/tags.py +++ b/third_party/xla/build_tools/lint/tags.py @@ -104,6 +104,7 @@ "xla_interpreter": "Uses interpreter backend.", "xla_cpu": "Uses CPU backend.", "xla_amdgpu_any": "Uses ROCm backend.", + "xla_mi350": "Runs on MI350 backend.", "xla_nvgpu_any": "Uses NVIDIA GPU backend.", "xla_intelgpu_any": "Uses Intel GPU backend.", # Below tags are emitted alongside `requires-gpu-x` tags, which is what the diff --git a/third_party/xla/third_party/cuda_tile/cuda_tile.BUILD b/third_party/xla/third_party/cuda_tile/cuda_tile.BUILD index 72a4d67e2d0297..a2765e58eeb66a 100644 --- a/third_party/xla/third_party/cuda_tile/cuda_tile.BUILD +++ b/third_party/xla/third_party/cuda_tile/cuda_tile.BUILD @@ -156,6 +156,35 @@ gentbl_cc_library( deps = [":CudaTileTdFiles"], ) +gentbl_cc_library( + name = "CudaTileAttrCodeGenIncGen", + tbl_outs = [ + ( + ["-gen-enum-attr-list"], + "include/cuda_tile/Dialect/CudaTile/IR/CudaTileEnumAttrs.inc", + ), + ( + ["-gen-gpu-arch-def"], + "include/cuda_tile/Dialect/CudaTile/IR/GpuArchitectures.inc", + ), + ( + ["-gen-gpu-arch-constants"], + "include/cuda_tile/Dialect/CudaTile/IR/GpuArchConstants.inc", + ), + ( + ["-gen-hint-key-accessors"], + "include/cuda_tile/Dialect/CudaTile/IR/HintKeyAccessors.inc", + ), + ( + ["-gen-hint-key-impl"], + "include/cuda_tile/Dialect/CudaTile/IR/HintKeyImpl.inc", + ), + ], + tblgen = ":cuda-tile-tblgen", + td_file = "include/cuda_tile/Dialect/CudaTile/IR/AttrDefs.td", + deps = [":CudaTileTdFiles"], +) + gentbl_cc_library( name = "CudaTileRemarksIncGen", tbl_outs = [ @@ -219,12 +248,15 @@ cc_library( ], visibility = ["//visibility:public"], deps = [ + ":CudaTileAttrCodeGenIncGen", ":CudaTileAttrDefsIncGen", ":CudaTileDialectIncGen", ":CudaTileInterfacesIncGen", ":CudaTileOpsCanonicalizationIncGen", ":CudaTileOpsIncGen", + ":CudaTileRemarksIncGen", ":CudaTileTypesIncGen", + "@llvm-project//llvm:Core", "@llvm-project//llvm:Support", "@llvm-project//mlir:ArithDialect", "@llvm-project//mlir:BytecodeOpInterface", @@ -276,6 +308,7 @@ cc_library( "@llvm-project//llvm:Support", "@llvm-project//mlir:IR", "@llvm-project//mlir:Pass", + "@llvm-project//mlir:SideEffectInterfaces", "@llvm-project//mlir:TransformUtils", ], ) @@ -363,7 +396,9 @@ cc_library( ":CudaTileBytecodeTypeIncGen", ":CudaTileDialect", "@llvm-project//llvm:Support", + "@llvm-project//mlir:FunctionInterfaces", "@llvm-project//mlir:IR", + "@llvm-project//mlir:Support", "@llvm-project//mlir:TranslateLib", ], ) diff --git a/third_party/xla/third_party/cuda_tile/patches/constructor.patch b/third_party/xla/third_party/cuda_tile/patches/constructor.patch index 98b074b4be7358..06374ab761ad3d 100644 --- a/third_party/xla/third_party/cuda_tile/patches/constructor.patch +++ b/third_party/xla/third_party/cuda_tile/patches/constructor.patch @@ -3,7 +3,7 @@ Remove redundant defaulted copy/move constructors to preserve aggregate status. diff --git a/tools/cuda-tile-tblgen/CudaTileOp.h b/tools/cuda-tile-tblgen/CudaTileOp.h --- a/tools/cuda-tile-tblgen/CudaTileOp.h +++ b/tools/cuda-tile-tblgen/CudaTileOp.h -@@ -55,11 +55,6 @@ +@@ -56,11 +56,6 @@ std::optional> selectedVariants; ParameterType type; diff --git a/third_party/xla/third_party/cuda_tile/patches/pow_op_alias.patch b/third_party/xla/third_party/cuda_tile/patches/pow_op_alias.patch new file mode 100644 index 00000000000000..14bb426cb6d73b --- /dev/null +++ b/third_party/xla/third_party/cuda_tile/patches/pow_op_alias.patch @@ -0,0 +1,13 @@ +diff --git a/include/cuda_tile/Dialect/CudaTile/IR/Ops.h b/include/cuda_tile/Dialect/CudaTile/IR/Ops.h +--- a/include/cuda_tile/Dialect/CudaTile/IR/Ops.h ++++ b/include/cuda_tile/Dialect/CudaTile/IR/Ops.h +@@ -59,6 +59,9 @@ + + namespace mlir::cuda_tile { + ++/// Backwards compatibility alias for PowOp which was renamed to FPowFOp in v13.4.0. ++using PowOp = FPowFOp; ++ + // Helper function to extract cuda_tile::ModuleOp + cuda_tile::ModuleOp extractCudaTileModuleOp(Operation *op); + diff --git a/third_party/xla/third_party/cuda_tile/patches/series b/third_party/xla/third_party/cuda_tile/patches/series index f26d3afb3502f8..9134fcfde4c7bd 100644 --- a/third_party/xla/third_party/cuda_tile/patches/series +++ b/third_party/xla/third_party/cuda_tile/patches/series @@ -1,2 +1,3 @@ constructor.patch -symbol_op_interface.patch +unused_variable.patch +pow_op_alias.patch diff --git a/third_party/xla/third_party/cuda_tile/patches/symbol_op_interface.patch b/third_party/xla/third_party/cuda_tile/patches/symbol_op_interface.patch deleted file mode 100644 index d48d6ab5b4408a..00000000000000 --- a/third_party/xla/third_party/cuda_tile/patches/symbol_op_interface.patch +++ /dev/null @@ -1,163 +0,0 @@ -Fix SymbolOpInterface usages for LLVM r632326. - -diff --git a/include/cuda_tile/Dialect/CudaTile/IR/Ops.td b/include/cuda_tile/Dialect/CudaTile/IR/Ops.td ---- a/include/cuda_tile/Dialect/CudaTile/IR/Ops.td -+++ b/include/cuda_tile/Dialect/CudaTile/IR/Ops.td -@@ -2178,7 +2178,7 @@ def CudaTile_FToIOp : CudaTileConversion - //===----------------------------------------------------------------------===// - - def CudaTile_EntryOp : CudaTileCoreOpDef<"entry", "13.1", [ -- FunctionOpInterface, IsolatedFromAbove, OpAsmOpInterface, SingleBlock, -+ SymbolName, FunctionOpInterface, IsolatedFromAbove, OpAsmOpInterface, SingleBlock, - SingleBlockImplicitTerminator<"ReturnOp"> - ]> { - let summary = "Define a tile kernel"; -@@ -2217,6 +2217,12 @@ def CudaTile_EntryOp : CudaTileCoreOpDef - /// Returns the result types of this function. - ArrayRef getResultTypes() { return getFunctionType().getResults(); } - -+ // SymbolOpInterface Methods -+ mlir::SymbolTable::Visibility getVisibility() { -+ return mlir::SymbolTable::Visibility::Public; -+ } -+ void setVisibility(mlir::SymbolTable::Visibility vis) {} -+ - static void build(::mlir::OpBuilder &odsBuilder, - ::mlir::OperationState &odsState, - ::mlir::StringAttr sym_name, -@@ -2350,7 +2356,7 @@ def CudaTile_GetGlobalOp : CudaTileCore - // GlobalOp - //===----------------------------------------------------------------------===// - --def CudaTile_GlobalOp : CudaTileCoreOpDef<"global", "13.1", [Symbol]> { -+def CudaTile_GlobalOp : CudaTileCoreOpDef<"global", "13.1", [Symbol, SymbolName]> { - let summary = "Allocate static global memory"; - - let description = [{ -@@ -2404,6 +2410,23 @@ def CudaTile_GlobalOp : CudaTileCoreOpDe - - let assemblyFormat = "custom($symbol_visibility) (`constant` $constant^)? $sym_name (`alignment` `=` $alignment^)? attr-dict custom($value)"; - let hasVerifier = 1; -+ -+ let extraClassDeclaration = [{ -+ // SymbolOpInterface Methods -+ mlir::SymbolTable::Visibility getVisibility() { -+ if (getSymbolVisibility() == SymbolVisibility::Public) { -+ return mlir::SymbolTable::Visibility::Public; -+ } -+ return mlir::SymbolTable::Visibility::Private; -+ } -+ void setVisibility(mlir::SymbolTable::Visibility vis) { -+ if (vis == mlir::SymbolTable::Visibility::Public) { -+ setSymbolVisibility(SymbolVisibility::Public); -+ } else { -+ setSymbolVisibility(SymbolVisibility::Private); -+ } -+ } -+ }]; - } - - //===----------------------------------------------------------------------===// -diff --git a/include/cuda_tile/Dialect/CudaTile/IR/SharedFuncParserAndPrinter.h b/include/cuda_tile/Dialect/CudaTile/IR/SharedFuncParserAndPrinter.h ---- a/include/cuda_tile/Dialect/CudaTile/IR/SharedFuncParserAndPrinter.h -+++ b/include/cuda_tile/Dialect/CudaTile/IR/SharedFuncParserAndPrinter.h -@@ -24,9 +24,9 @@ template - ParseResult parseFuncOp(OpAsmParser &parser, OperationState &result) { - // Parse the name as a symbol. - StringAttr nameAttr; -- if (parser.parseSymbolName(nameAttr, SymbolTable::getSymbolAttrName(), -- result.attributes)) -+ if (parser.parseSymbolName(nameAttr, "sym_name", result.attributes)) { - return failure(); -+ } - - // Parse the function signature using custom parsing that supports both - // short form (tile>) and long form (!cuda_tile.tile>) types -@@ -80,10 +80,9 @@ ParseResult parseFuncOp(OpAsmParser &par - template - void printFuncOp(OpTy op, OpAsmPrinter &printer) { - // Print the operation and the function name. -- auto funcName = -- op.getOperation() -- ->template getAttrOfType(SymbolTable::getSymbolAttrName()) -- .getValue(); -+ auto funcName = op.getOperation() -+ ->template getAttrOfType("sym_name") -+ .getValue(); - printer << ' '; - printer.printSymbolName(funcName); - auto fnType = op.getFunctionType(); -@@ -95,8 +94,8 @@ void printFuncOp(OpTy op, OpAsmPrinter & - /*printEmptyBlock=*/false); - printer.printOptionalAttrDict( - op.getOperation()->getAttrs(), -- {op.getArgAttrsAttrName(), op.getFunctionTypeAttrName(), -- SymbolTable::getSymbolAttrName(), op.getResAttrsAttrName()}); -+ {op.getArgAttrsAttrName(), op.getFunctionTypeAttrName(), "sym_name", -+ op.getResAttrsAttrName()}); - } - - } // end namespace mlir::cuda_tile. -diff --git a/include/cuda_tile/Dialect/CudaTile/IR/TestingOps.td b/include/cuda_tile/Dialect/CudaTile/IR/TestingOps.td ---- a/include/cuda_tile/Dialect/CudaTile/IR/TestingOps.td -+++ b/include/cuda_tile/Dialect/CudaTile/IR/TestingOps.td -@@ -31,7 +31,7 @@ include "cuda_tile/Dialect/CudaTile/IR/T - //===----------------------------------------------------------------------===// - - def CudaTile_Test_FuncOp : CudaTileTestingOpDef<"func", "250.0", [ -- IsolatedFromAbove, FunctionOpInterface, SingleBlock, OpAsmOpInterface, -+ IsolatedFromAbove, SymbolName, FunctionOpInterface, SingleBlock, OpAsmOpInterface, - SingleBlockImplicitTerminator<"ReturnOp"> - ]> { - -@@ -56,6 +56,12 @@ def CudaTile_Test_FuncOp : CudaTileTesti - - /// Returns the result types of this function. - ArrayRef getResultTypes() { return getFunctionType().getResults(); } -+ -+ // SymbolOpInterface Methods -+ mlir::SymbolTable::Visibility getVisibility() { -+ return mlir::SymbolTable::Visibility::Public; -+ } -+ void setVisibility(mlir::SymbolTable::Visibility vis) {} - }]; - } - -diff --git a/lib/Dialect/CudaTile/IR/CudaTile.cpp b/lib/Dialect/CudaTile/IR/CudaTile.cpp ---- a/lib/Dialect/CudaTile/IR/CudaTile.cpp -+++ b/lib/Dialect/CudaTile/IR/CudaTile.cpp -@@ -2669,9 +2669,9 @@ constexpr char kOptimizationHintsAttr[] - ParseResult EntryOp::parse(OpAsmParser &parser, OperationState &result) { - // Parse the name as a symbol. - StringAttr nameAttr; -- if (parser.parseSymbolName(nameAttr, SymbolTable::getSymbolAttrName(), -- result.attributes)) -+ if (parser.parseSymbolName(nameAttr, "sym_name", result.attributes)) { - return failure(); -+ } - - // Parse the function signature using custom parsing that supports both - // short form (tile>) and long form (!cuda_tile.tile>) types -@@ -2736,9 +2736,7 @@ ParseResult EntryOp::parse(OpAsmParser & - void EntryOp::print(OpAsmPrinter &printer) { - // Print the operation and the function name. - auto funcName = -- getOperation() -- ->getAttrOfType(SymbolTable::getSymbolAttrName()) -- .getValue(); -+ getOperation()->getAttrOfType("sym_name").getValue(); - printer << ' '; - printer.printSymbolName(funcName); - auto fnType = getFunctionType(); -@@ -2754,9 +2752,8 @@ void EntryOp::print(OpAsmPrinter &printe - /*printEmptyBlock=*/false); - printer.printOptionalAttrDict( - getOperation()->getAttrs(), -- {getArgAttrsAttrName(), getFunctionTypeAttrName(), -- SymbolTable::getSymbolAttrName(), getResAttrsAttrName(), -- kOptimizationHintsAttr}); -+ {getArgAttrsAttrName(), getFunctionTypeAttrName(), "sym_name", -+ getResAttrsAttrName(), kOptimizationHintsAttr}); - } - - LogicalResult EntryOp::verify() { diff --git a/third_party/xla/third_party/cuda_tile/patches/unused_variable.patch b/third_party/xla/third_party/cuda_tile/patches/unused_variable.patch new file mode 100644 index 00000000000000..3258c90744e083 --- /dev/null +++ b/third_party/xla/third_party/cuda_tile/patches/unused_variable.patch @@ -0,0 +1,58 @@ +diff --git a/tools/cuda-tile-tblgen/BytecodeTypeAnalysis.cpp b/tools/cuda-tile-tblgen/BytecodeTypeAnalysis.cpp +--- a/tools/cuda-tile-tblgen/BytecodeTypeAnalysis.cpp ++++ b/tools/cuda-tile-tblgen/BytecodeTypeAnalysis.cpp +@@ -68,7 +68,7 @@ + return Kind::OptionalEnum; + } + +- if (const auto *defInit = dyn_cast_if_present(param.getDef())) { ++ if (const auto* defInit = dyn_cast_if_present(param.getDef())) { + if (defInit->getDef()->isSubClassOf("CudaTileAttrInterfaceParam")) { + return Kind::PolymorphicNestedAttr; + } +@@ -110,7 +110,6 @@ + cppType(param.getCppType().str()), + cppStorageType(param.getCppStorageType().str()), + isOptional(param.isOptional()), kind(classifyParameter(param)) { +- + // Extract version from wrapped parameters (CudaTileTypeParam adds + // sinceVersion). Use getValue() to check if field exists before accessing - + // raw parameters like ArrayRefParameter don't have this field in their class +@@ -153,8 +152,6 @@ + "::" + typeDef.getCppClassName().str()), + typeTagValue(tagValue), sinceVersion(version.str()), + needsReverseOrder(typeDef.getCppClassName() == "TileType") { +- const Record *typeRecord = typeDef.getDef(); +- + // Analyze and validate all parameters. + for (const auto &attrParam : typeDef.getParameters()) { + BytecodeTypeParameter param(attrParam); +@@ -239,7 +236,7 @@ + // through it. Both qualified forms ("ns::Foo" and "::ns::Foo") are indexed + // so lookups don't depend on how OptionalParameter spells the name. + llvm::StringSet<> enumAttrCppTypes; +- for (const Record *r : records.getAllDerivedDefinitions("CudaTileEnumAttr")) { ++ for (const Record* r : records.getAllDerivedDefinitions("CudaTileEnumAttr")) { + AttrOrTypeDef def(r); + std::string qualified = + (def.getDialect().getCppNamespace() + "::" + def.getCppClassName()) +@@ -250,7 +247,7 @@ + + // Build map of CudaTileTypeDef for matching. + StringMap cudaTileTypeDefRecords; +- for (const Record *typeRecord : ++ for (const Record* typeRecord : + records.getAllDerivedDefinitions("CudaTileTypeDef")) { + cudaTileTypeDefRecords[AttrOrTypeDef(typeRecord).getCppClassName()] = + typeRecord; +@@ -327,8 +324,8 @@ + + // Validate every OptionalEnum-classified parameter is actually a + // CudaTileEnumAttr-derived enum-attribute wrapper. +- for (const CudaTileType &type : structure.cudaTileTypes) { +- for (const BytecodeTypeParameter &p : type.parameters) { ++ for (const CudaTileType& type : structure.cudaTileTypes) { ++ for (const BytecodeTypeParameter& p : type.parameters) { + if (p.kind != BytecodeTypeParameter::Kind::OptionalEnum) { + continue; + } diff --git a/third_party/xla/third_party/cuda_tile/workspace.bzl b/third_party/xla/third_party/cuda_tile/workspace.bzl index 72f49935a0ca06..c04963da4a4130 100644 --- a/third_party/xla/third_party/cuda_tile/workspace.bzl +++ b/third_party/xla/third_party/cuda_tile/workspace.bzl @@ -5,8 +5,8 @@ load("//third_party:repo.bzl", "tf_mirror_urls") def repo(): """Imports CUDA Tile IR.""" - CUDA_TILE_COMMIT = "af2417041cc939b87ef56d92cfdcf61737c5457e" - CUDA_TILE_SHA256 = "81597e49469171bf8fa7319fbd44ebe133001521f484589e3dd3fb3fad282dc0" + CUDA_TILE_COMMIT = "7e8e2e68fa219716103824c01f7303367cf7df8d" + CUDA_TILE_SHA256 = "2398f88dd0dbcda3732934c8527485c3941d3a6453c98b6c249461e6309ea44c" quilt_http_archive( name = "cuda_tile", diff --git a/third_party/xla/third_party/gpus/crosstool/hipcc_cc_toolchain_config.bzl.tpl b/third_party/xla/third_party/gpus/crosstool/hipcc_cc_toolchain_config.bzl.tpl index e6c3f6e17bd28d..5451024a08a09d 100644 --- a/third_party/xla/third_party/gpus/crosstool/hipcc_cc_toolchain_config.bzl.tpl +++ b/third_party/xla/third_party/gpus/crosstool/hipcc_cc_toolchain_config.bzl.tpl @@ -1166,12 +1166,6 @@ def _impl(ctx): flag_sets = [ flag_set( actions = all_compile_actions + all_link_actions, - flag_groups = [ - flag_group( - flags = ["--rocm-path=" + _rocm_path] + - ["--offload-arch=" + arch for arch in _hipcc_config.gpu_architectures], - ), - ], ), ], ) diff --git a/third_party/xla/xla/backends/gpu/libraries/native_custom_call_thunks/native_custom_call_handler_utils.h b/third_party/xla/xla/backends/gpu/libraries/native_custom_call_thunks/native_custom_call_handler_utils.h index ea80082b1c0447..37bf1330934fee 100644 --- a/third_party/xla/xla/backends/gpu/libraries/native_custom_call_thunks/native_custom_call_handler_utils.h +++ b/third_party/xla/xla/backends/gpu/libraries/native_custom_call_thunks/native_custom_call_handler_utils.h @@ -34,8 +34,7 @@ limitations under the License. #include "xla/stream_executor/launch_dim.h" // Helpers for writing custom-call thunk-folding handlers. See -// `native_custom_call_handler_registry.h` for the mechanism itself and -// `g3doc/native_custom_call_handlers.md` for a tutorial. +// `native_custom_call_handler_registry.h` for the mechanism itself. namespace xla::gpu { diff --git a/third_party/xla/xla/backends/gpu/runtime/BUILD b/third_party/xla/xla/backends/gpu/runtime/BUILD index 97f21128b1829e..bfe3b67572bec5 100644 --- a/third_party/xla/xla/backends/gpu/runtime/BUILD +++ b/third_party/xla/xla/backends/gpu/runtime/BUILD @@ -4779,6 +4779,7 @@ cc_library( srcs = ["command_buffer_conversion_pass.cc"], hdrs = ["command_buffer_conversion_pass.h"], deps = [ + ":async_execution", ":async_thunk", ":collective_group_thunk", ":command_buffer_cmd_emitter", @@ -4788,6 +4789,7 @@ cc_library( ":custom_call_thunk", ":device_to_device_copy_thunk", ":dynamic_slice_fusion_v2_thunk", + ":execution_stream_id", ":ragged_all_to_all_thunk", ":sequential_thunk", ":thunk", @@ -4802,8 +4804,10 @@ cc_library( "//xla/stream_executor:semantic_version", "@com_google_absl//absl/algorithm:container", "@com_google_absl//absl/base:nullability", + "@com_google_absl//absl/container:flat_hash_map", "@com_google_absl//absl/container:flat_hash_set", "@com_google_absl//absl/log", + "@com_google_absl//absl/log:check", "@com_google_absl//absl/log:vlog_is_on", "@com_google_absl//absl/status", "@com_google_absl//absl/status:status_macros", @@ -4817,6 +4821,31 @@ cc_library( ], ) +xla_cc_test( + name = "command_buffer_async_conversion_test", + srcs = ["command_buffer_async_conversion_test.cc"], + deps = [ + ":async_thunk", + ":command_buffer_conversion_pass", + ":command_buffer_thunk", + ":execution_stream_id", + ":memset_thunk", + ":replica_id_thunk", + ":thunk", + ":thunk_executor", + ":thunk_id", + ":thunk_pass_pipeline", + ":while_thunk", + "//xla:xla_proto_cc", + "//xla/service:buffer_assignment", + "//xla/service/gpu:gpu_device_info_for_tests", + "//xla/stream_executor:device_description", + "@com_google_absl//absl/status", + "@com_google_absl//absl/status:statusor", + "@com_google_googletest//:gtest_main", + ], +) + xla_test( name = "command_buffer_conversion_pass_test", srcs = ["command_buffer_conversion_pass_test.cc"], diff --git a/third_party/xla/xla/backends/gpu/runtime/async_execution.cc b/third_party/xla/xla/backends/gpu/runtime/async_execution.cc index a5b45b78de9a32..26823b190a8800 100644 --- a/third_party/xla/xla/backends/gpu/runtime/async_execution.cc +++ b/third_party/xla/xla/backends/gpu/runtime/async_execution.cc @@ -138,9 +138,11 @@ absl::StatusOr AsyncExecution::Start( // Wait for all prior operations on `stream` before launching operations on // `async_stream`. We use a stream-level wait (not the shared event) so that // the event remains exclusively used for the async→main completion signal. - // This is critical for pipelined send/recv where multiple Start() calls can - // happen before Done() (the event is safely overwritten on the async stream - // because the stream is ordered). + // This matters for pipelined send/recv, where several start thunks share one + // execution and Start()/Done() alternate across loop iterations: each Start() + // re-records the shared event on the ordered async stream after the previous + // Done() has waited on it. A second Start() before the matching Done() is + // rejected by the counter check above. ABSL_RETURN_IF_ERROR(async_stream->WaitFor(stream)); return ExecutionGuard(event, async_stream); diff --git a/third_party/xla/xla/backends/gpu/runtime/command_buffer_async_conversion_test.cc b/third_party/xla/xla/backends/gpu/runtime/command_buffer_async_conversion_test.cc new file mode 100644 index 00000000000000..333dcaacb39701 --- /dev/null +++ b/third_party/xla/xla/backends/gpu/runtime/command_buffer_async_conversion_test.cc @@ -0,0 +1,434 @@ +/* Copyright 2026 The OpenXLA Authors. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#include +#include +#include +#include + +#include +#include +#include "absl/status/status.h" +#include "absl/status/statusor.h" +#include "xla/backends/gpu/runtime/async_thunk.h" +#include "xla/backends/gpu/runtime/command_buffer_conversion_pass.h" +#include "xla/backends/gpu/runtime/command_buffer_thunk.h" +#include "xla/backends/gpu/runtime/execution_stream_id.h" +#include "xla/backends/gpu/runtime/memset_thunk.h" +#include "xla/backends/gpu/runtime/replica_id_thunk.h" +#include "xla/backends/gpu/runtime/thunk.h" +#include "xla/backends/gpu/runtime/thunk_executor.h" +#include "xla/backends/gpu/runtime/thunk_id.h" +#include "xla/backends/gpu/runtime/thunk_pass_pipeline.h" +#include "xla/backends/gpu/runtime/while_thunk.h" +#include "xla/service/buffer_assignment.h" +#include "xla/service/gpu/gpu_device_info_for_tests.h" +#include "xla/stream_executor/device_description.h" +#include "xla/xla.pb.h" + +namespace xla::gpu { +namespace { + +using ::testing::ElementsAre; +using ::testing::Pointee; + +MATCHER_P(ThunkKindIs, kind, "") { return arg.kind() == kind; } + +template +auto ThunkKindsAre(Kinds... kinds) { + return ElementsAre(Pointee(ThunkKindIs(kinds))...); +} + +class RejectingAllocator : public ThunkPassBufferAllocator { + public: + absl::StatusOr NewEmptyAllocation(int64_t size) override { + return absl::InternalError("Unexpected allocation during conversion"); + } +}; + +class CommandBufferAsyncConversionTest : public testing::Test { + protected: + Thunk::ThunkInfo Info() { + Thunk::ThunkInfo info; + info.thunk_id = ids_.GetNextThunkId(); + return info; + } + + void AddCommand(ThunkSequence& thunks) { + thunks.Emplace(Info(), slice_); + } + + AsyncStartThunk* Start(ThunkSequence& thunks, bool convertible = true, + ExecutionStreamId stream = ComputationStreamId(0)) { + ThunkSequence body; + if (convertible) { + AddCommand(body); + } else { + // Memset thunks are not eligible for conversion in this pass. + body.Emplace(Info(), 0, slice_); + } + auto start = + std::make_unique(Info(), stream, std::move(body)); + auto* result = start.get(); + thunks.push_back(std::move(start)); + return result; + } + + void Done(ThunkSequence& thunks, AsyncStartThunk* start) { + thunks.Emplace(Info(), start->async_execution()); + } + + WhileThunk* Loop(ThunkSequence& thunks, ThunkSequence body) { + auto loop = std::make_unique( + Info(), BufferAllocation::Slice(&allocation_, 0, 1), ThunkSequence{}, + std::move(body), /*trip_count=*/1); + auto* result = loop.get(); + thunks.push_back(std::move(loop)); + return result; + } + + const ThunkSequence& CommandBufferThunks(const Thunk& thunk) { + return static_cast(thunk).thunks()->thunks(); + } + + absl::StatusOr Convert(ThunkSequence& thunks, + bool enable_while = false) { + DebugOptions options; + options.add_xla_gpu_enable_command_buffer(DebugOptions::FUSION); + if (enable_while) { + options.add_xla_gpu_enable_command_buffer(DebugOptions::WHILE); + } + options.set_xla_gpu_graph_min_graph_size(1); + options.set_xla_gpu_command_buffer_scheduling_mode(DebugOptions::LHS); + se::DeviceDescription device = TestGpuDeviceInfo::RTXA6000DeviceInfo(); + RejectingAllocator allocator; + return CommandBufferConversionPass("test").Run( + &thunks, options, /*hlo_module=*/nullptr, device, allocator); + } + + BufferAllocation allocation_{0, sizeof(int32_t), 0}; + BufferAllocation::Slice slice_{&allocation_, 0, sizeof(int32_t)}; + ThunkIdGenerator ids_; +}; + +TEST_F(CommandBufferAsyncConversionTest, KeepsUnsupportedRegionsIntact) { + for (bool crossed : {false, true}) { + SCOPED_TRACE(crossed); + ThunkSequence thunks; + AddCommand(thunks); + auto* outer = Start(thunks, /*convertible=*/false); + auto* inner = Start(thunks); + Done(thunks, crossed ? outer : inner); + Done(thunks, crossed ? inner : outer); + AddCommand(thunks); + + ASSERT_OK_AND_ASSIGN(bool changed, Convert(thunks)); + EXPECT_TRUE(changed); + // Capturing the inner start/done alone would drop its stream ordering + // against the unsupported outer operation. Capture resumes after both + // joins, and ordinary commands before the region can still be captured. + EXPECT_THAT(thunks, + ThunkKindsAre(Thunk::kCommandBuffer, Thunk::kAsyncStart, + Thunk::kAsyncStart, Thunk::kAsyncDone, + Thunk::kAsyncDone, Thunk::kCommandBuffer)); + } +} + +TEST_F(CommandBufferAsyncConversionTest, ConvertsPlainCommandsInOpenRegion) { + ThunkSequence thunks; + auto* start = Start(thunks, /*convertible=*/false); + AddCommand(thunks); + AddCommand(thunks); + Done(thunks, start); + AddCommand(thunks); + + ASSERT_OK_AND_ASSIGN(bool changed, Convert(thunks)); + EXPECT_TRUE(changed); + // Commands between the unsupported start and its done execute on the main + // stream, exactly as their thunks would, so they can still be captured. + EXPECT_THAT(thunks, ThunkKindsAre(Thunk::kAsyncStart, Thunk::kCommandBuffer, + Thunk::kAsyncDone, Thunk::kCommandBuffer)); + EXPECT_THAT(CommandBufferThunks(*thunks[1]), + ThunkKindsAre(Thunk::kReplicaId, Thunk::kReplicaId)); +} + +TEST_F(CommandBufferAsyncConversionTest, KeepsNestedRegionsInOpenRegion) { + ThunkSequence thunks; + auto* outer = Start(thunks, /*convertible=*/false); + auto* inner = Start(thunks); + AddCommand(thunks); + Done(thunks, inner); + AddCommand(thunks); + Done(thunks, outer); + + ASSERT_OK_AND_ASSIGN(bool changed, Convert(thunks)); + EXPECT_TRUE(changed); + // Capturing the inner pair would run its body in a graph on the main stream + // and lose stream order with the outstanding outer operation. Its start and + // done stay thunks while the plain commands around them are captured. + EXPECT_THAT(thunks, ThunkKindsAre(Thunk::kAsyncStart, Thunk::kAsyncStart, + Thunk::kCommandBuffer, Thunk::kAsyncDone, + Thunk::kCommandBuffer, Thunk::kAsyncDone)); +} + +TEST_F(CommandBufferAsyncConversionTest, ConvertsLoopBodiesInOpenRegion) { + for (bool body_has_async : {false, true}) { + SCOPED_TRACE(body_has_async); + ThunkSequence thunks; + auto* start = Start(thunks); + ThunkSequence body; + if (body_has_async) { + auto* inner = Start(body); + Done(body, inner); + } + AddCommand(body); + // WHILE conversion is disabled, so the loop keeps the region open and its + // body is converted recursively under the open-region rules. + WhileThunk* loop = Loop(thunks, std::move(body)); + Done(thunks, start); + AddCommand(thunks); + + ASSERT_OK_AND_ASSIGN(bool changed, Convert(thunks)); + EXPECT_TRUE(changed); + EXPECT_THAT(thunks, + ThunkKindsAre(Thunk::kAsyncStart, Thunk::kWhile, + Thunk::kAsyncDone, Thunk::kCommandBuffer)); + if (body_has_async) { + EXPECT_THAT(loop->body_executor().thunks(), + ThunkKindsAre(Thunk::kAsyncStart, Thunk::kAsyncDone, + Thunk::kCommandBuffer)); + } else { + EXPECT_THAT(loop->body_executor().thunks(), + ThunkKindsAre(Thunk::kCommandBuffer)); + } + } +} + +TEST_F(CommandBufferAsyncConversionTest, + DoesNotCaptureLoopWithAsyncInOpenRegion) { + for (bool open_region : {false, true}) { + SCOPED_TRACE(open_region); + ThunkSequence thunks; + AsyncStartThunk* start = nullptr; + if (open_region) { + start = Start(thunks, /*convertible=*/false); + } + ThunkSequence body; + auto* inner = Start(body); + Done(body, inner); + WhileThunk* loop = Loop(thunks, std::move(body)); + if (open_region) { + Done(thunks, start); + } + + ASSERT_OK_AND_ASSIGN(bool changed, Convert(thunks, /*enable_while=*/true)); + if (open_region) { + // A convertible loop whose body starts async work on a stream with an + // outstanding operation must not become a graph on the main stream. + EXPECT_FALSE(changed); + EXPECT_THAT(thunks, ThunkKindsAre(Thunk::kAsyncStart, Thunk::kWhile, + Thunk::kAsyncDone)); + EXPECT_THAT(loop->body_executor().thunks(), + ThunkKindsAre(Thunk::kAsyncStart, Thunk::kAsyncDone)); + } else { + EXPECT_TRUE(changed); + EXPECT_THAT(thunks, ThunkKindsAre(Thunk::kCommandBuffer)); + } + } +} + +TEST_F(CommandBufferAsyncConversionTest, ConvertsTailAfterUnmatchedStart) { + ThunkSequence thunks; + AddCommand(thunks); + Start(thunks); // Its done can belong to a different pipelined computation. + AddCommand(thunks); + auto* inner = Start(thunks); + Done(thunks, inner); + AddCommand(thunks); + + ASSERT_OK_AND_ASSIGN(bool changed, Convert(thunks)); + EXPECT_TRUE(changed); + // The stream stays open until the end of the sequence: later pairs on the + // same stream remain thunks while plain commands are still captured. + EXPECT_THAT(thunks, ThunkKindsAre(Thunk::kCommandBuffer, Thunk::kAsyncStart, + Thunk::kCommandBuffer, Thunk::kAsyncStart, + Thunk::kAsyncDone, Thunk::kCommandBuffer)); +} + +TEST_F(CommandBufferAsyncConversionTest, + ConvertsTailAfterUnmatchedStartOnOtherStream) { + ThunkSequence thunks; + AddCommand(thunks); + Start(thunks, /*convertible=*/true, ComputationStreamId(0)); + AddCommand(thunks); + auto* inner = Start(thunks, /*convertible=*/true, ComputationStreamId(1)); + Done(thunks, inner); + AddCommand(thunks); + + ASSERT_OK_AND_ASSIGN(bool changed, Convert(thunks)); + EXPECT_TRUE(changed); + // The unmatched start only keeps its own stream open. A pair on another + // stream never observed that stream in thunk mode either, so capturing it + // loses no ordering. + EXPECT_THAT(thunks, ThunkKindsAre(Thunk::kCommandBuffer, Thunk::kAsyncStart, + Thunk::kCommandBuffer)); + EXPECT_THAT(CommandBufferThunks(*thunks[2]), + ThunkKindsAre(Thunk::kReplicaId, Thunk::kAsyncStart, + Thunk::kAsyncDone, Thunk::kReplicaId)); +} + +TEST_F(CommandBufferAsyncConversionTest, ConvertsNestedRegionOnOtherStream) { + ThunkSequence thunks; + auto* outer = Start(thunks, /*convertible=*/false, ComputationStreamId(0)); + auto* inner = Start(thunks, /*convertible=*/true, ComputationStreamId(1)); + AddCommand(thunks); + Done(thunks, inner); + AddCommand(thunks); + Done(thunks, outer); + + ASSERT_OK_AND_ASSIGN(bool changed, Convert(thunks)); + EXPECT_TRUE(changed); + EXPECT_THAT(thunks, ThunkKindsAre(Thunk::kAsyncStart, Thunk::kCommandBuffer, + Thunk::kAsyncDone)); + EXPECT_THAT(CommandBufferThunks(*thunks[1]), + ThunkKindsAre(Thunk::kAsyncStart, Thunk::kReplicaId, + Thunk::kAsyncDone, Thunk::kReplicaId)); +} + +TEST_F(CommandBufferAsyncConversionTest, ClosesStreamAtMatchingDone) { + ThunkSequence thunks; + auto* a = Start(thunks, /*convertible=*/false, ComputationStreamId(0)); + auto* b = Start(thunks, /*convertible=*/false, ComputationStreamId(1)); + Done(thunks, a); + auto* c = Start(thunks, /*convertible=*/true, ComputationStreamId(0)); + Done(thunks, c); + Done(thunks, b); + AddCommand(thunks); + + ASSERT_OK_AND_ASSIGN(bool changed, Convert(thunks)); + EXPECT_TRUE(changed); + // Stream 0 is joined at `a`'s done, so `c` can be captured even though + // stream 1 is still open. + EXPECT_THAT(thunks, ThunkKindsAre(Thunk::kAsyncStart, Thunk::kAsyncStart, + Thunk::kAsyncDone, Thunk::kCommandBuffer, + Thunk::kAsyncDone, Thunk::kCommandBuffer)); + EXPECT_THAT(CommandBufferThunks(*thunks[3]), + ThunkKindsAre(Thunk::kAsyncStart, Thunk::kAsyncDone)); +} + +TEST_F(CommandBufferAsyncConversionTest, KeepsRegionWithStartOnOpenStream) { + ThunkSequence thunks; + auto* a = Start(thunks, /*convertible=*/false, ComputationStreamId(0)); + auto* b = Start(thunks, /*convertible=*/true, ComputationStreamId(1)); + auto* c = Start(thunks, /*convertible=*/true, ComputationStreamId(0)); + Done(thunks, c); + Done(thunks, b); + Done(thunks, a); + + ASSERT_OK_AND_ASSIGN(bool changed, Convert(thunks)); + // `b`'s region is closed and convertible, but it contains `c`, which starts + // work on the stream still occupied by `a`. Capturing the region would run + // `c`'s body in a graph on the main stream. + EXPECT_FALSE(changed); + EXPECT_THAT(thunks, ThunkKindsAre(Thunk::kAsyncStart, Thunk::kAsyncStart, + Thunk::kAsyncStart, Thunk::kAsyncDone, + Thunk::kAsyncDone, Thunk::kAsyncDone)); +} + +TEST_F(CommandBufferAsyncConversionTest, ConvertsLoopBodyRegionOnOtherStream) { + ThunkSequence thunks; + auto* start = Start(thunks, /*convertible=*/false, ComputationStreamId(0)); + ThunkSequence body; + auto* inner = Start(body, /*convertible=*/true, ComputationStreamId(1)); + Done(body, inner); + AddCommand(body); + WhileThunk* loop = Loop(thunks, std::move(body)); + Done(thunks, start); + + ASSERT_OK_AND_ASSIGN(bool changed, Convert(thunks)); + EXPECT_TRUE(changed); + EXPECT_THAT(thunks, ThunkKindsAre(Thunk::kAsyncStart, Thunk::kWhile, + Thunk::kAsyncDone)); + // Only stream 0 is open inside the loop body; the pair on stream 1 is + // captured together with the plain command. + EXPECT_THAT(loop->body_executor().thunks(), + ThunkKindsAre(Thunk::kCommandBuffer)); + EXPECT_THAT( + CommandBufferThunks(*loop->body_executor().thunks()[0]), + ThunkKindsAre(Thunk::kAsyncStart, Thunk::kAsyncDone, Thunk::kReplicaId)); +} + +TEST_F(CommandBufferAsyncConversionTest, DoesNotCaptureLoopWithLoneDone) { + ThunkSequence thunks; + auto* start = Start(thunks, /*convertible=*/false); + ThunkSequence body; + Done(body, start); + AddCommand(body); + WhileThunk* loop = Loop(thunks, std::move(body)); + + ASSERT_OK_AND_ASSIGN(bool changed, Convert(thunks, /*enable_while=*/true)); + EXPECT_TRUE(changed); + // The loop body joins an operation started outside the loop. Capturing the + // loop whole would drop that join, so only the plain command in the body is + // captured. + EXPECT_THAT(thunks, ThunkKindsAre(Thunk::kAsyncStart, Thunk::kWhile)); + EXPECT_THAT(loop->body_executor().thunks(), + ThunkKindsAre(Thunk::kAsyncDone, Thunk::kCommandBuffer)); +} + +TEST_F(CommandBufferAsyncConversionTest, MatchesCanonicalAsyncExecution) { + ThunkSequence thunks; + auto* canonical = Start(thunks); + Done(thunks, canonical); + ThunkSequence body; + AddCommand(body); + thunks.push_back(std::make_unique( + Info(), ComputationStreamId(0), std::move(body), + canonical->async_execution())); + Done(thunks, canonical); + + ASSERT_OK_AND_ASSIGN(bool changed, Convert(thunks)); + EXPECT_TRUE(changed); + EXPECT_THAT(thunks, ThunkKindsAre(Thunk::kCommandBuffer)); +} + +TEST_F(CommandBufferAsyncConversionTest, DuplicateOutstandingStart) { + ThunkSequence thunks; + auto* canonical = Start(thunks); + ThunkSequence body; + AddCommand(body); + thunks.push_back(std::make_unique( + Info(), ComputationStreamId(0), std::move(body), + canonical->async_execution())); + Done(thunks, canonical); + Done(thunks, canonical); + AddCommand(thunks); + + // The runtime permits only one outstanding start per AsyncExecution, so this + // sequence can never execute. Debug builds catch it in the pass; optimized + // builds leave the region as thunks and still capture the trailing command. +#ifndef NDEBUG + EXPECT_DEATH(Convert(thunks).IgnoreError(), "started twice"); +#else + ASSERT_OK_AND_ASSIGN(bool changed, Convert(thunks)); + EXPECT_TRUE(changed); + EXPECT_THAT(thunks, ThunkKindsAre(Thunk::kAsyncStart, Thunk::kAsyncStart, + Thunk::kAsyncDone, Thunk::kAsyncDone, + Thunk::kCommandBuffer)); +#endif +} + +} // namespace +} // namespace xla::gpu diff --git a/third_party/xla/xla/backends/gpu/runtime/command_buffer_conversion_pass.cc b/third_party/xla/xla/backends/gpu/runtime/command_buffer_conversion_pass.cc index 3730af6a1cd1db..295772501d1968 100644 --- a/third_party/xla/xla/backends/gpu/runtime/command_buffer_conversion_pass.cc +++ b/third_party/xla/xla/backends/gpu/runtime/command_buffer_conversion_pass.cc @@ -29,7 +29,9 @@ limitations under the License. #include "absl/algorithm/container.h" #include "absl/base/nullability.h" +#include "absl/container/flat_hash_map.h" #include "absl/container/flat_hash_set.h" +#include "absl/log/check.h" #include "absl/log/log.h" #include "absl/log/vlog_is_on.h" #include "absl/status/status.h" @@ -38,6 +40,7 @@ limitations under the License. #include "absl/strings/str_cat.h" #include "absl/strings/str_join.h" #include "absl/types/span.h" +#include "xla/backends/gpu/runtime/async_execution.h" #include "xla/backends/gpu/runtime/async_thunk.h" #include "xla/backends/gpu/runtime/collective_group_thunk.h" #include "xla/backends/gpu/runtime/command_buffer_cmd_emitter.h" @@ -47,6 +50,7 @@ limitations under the License. #include "xla/backends/gpu/runtime/custom_call_thunk.h" #include "xla/backends/gpu/runtime/device_to_device_copy_thunk.h" #include "xla/backends/gpu/runtime/dynamic_slice_fusion_v2_thunk.h" +#include "xla/backends/gpu/runtime/execution_stream_id.h" #include "xla/backends/gpu/runtime/ragged_all_to_all_thunk.h" #include "xla/backends/gpu/runtime/sequential_thunk.h" #include "xla/backends/gpu/runtime/thunk.h" @@ -459,19 +463,23 @@ bool ThunkSequenceIsConvertible(const ThunkSequence& thunks, return false; } i += region_size - 1; + } else if (thunk->kind() == Thunk::kAsyncDone) { + // Every done that belongs to a start in this sequence was consumed above + // as part of its region. This done joins an operation started outside + // the sequence; capturing it would drop the join (the top-level loop in + // `RunImpl` keeps such thunks in place for the same reason). + return false; } } return true; } // Collects and returns the size of the shortest non-empty sequence of thunks -// that form a valid async region. -// The sequence considered as a valid async region if each start thunk has a -// corresponding done thunk and vice versa, and all thunks in between are -// convertible. If there is another start thunk between the original start and -// done, we may potentially extend the sequence to include its corresponding -// done thunk. For example, if we call this function on async-start_a in the -// following sequence: +// that form a closed async region: each start thunk has a corresponding done +// thunk and vice versa. If there is another start thunk between the original +// start and done, we may potentially extend the sequence to include its +// corresponding done thunk. For example, if we call this function on +// async-start_a in the following sequence: // // async_start_a // async_start_b @@ -480,48 +488,113 @@ bool ThunkSequenceIsConvertible(const ThunkSequence& thunks, // // The returned sequence will contain async_done_b. So that all async pairs // are captured by the same command buffer. -size_t CheckAsyncRegion(absl::Span> thunks, - const CommandBufferConfig& config) { - absl::flat_hash_set unpaired_ids; +// Find the boundary independently of command-buffer eligibility. If any thunk +// is unsupported, the whole region must remain outside the command buffer: +// capturing an inner region could lose ordering with an outstanding operation +// on the same async stream. +// +// Returns 0 when no closed region starts here: either a start in the scanned +// range has no done in this sequence, or a done in the range joins a start +// from outside it. The caller handles both the same way, by leaving the start +// as a thunk and keeping its stream open until it is joined. +size_t AsyncRegionSize(absl::Span> thunks) { + absl::flat_hash_set unpaired_executions; for (size_t i = 0; i < thunks.size(); ++i) { auto& thunk = thunks[i]; - // Check if thunk is convertible - if (!IsConvertible(*thunk, config)) { - return 0; // All thunks in the region must be convertible. - } - - // Track AsyncStartThunk/AsyncDoneThunk pairs via AsyncExecutionId. if (thunk->kind() == Thunk::kAsyncStart) { - unpaired_ids.insert(static_cast(*thunk) - .async_execution_id() - .value()); + // Pipelined starts can share the canonical start's execution state, but + // AsyncExecution::Start rejects a second start before the matching done, + // so a valid sequence never inserts the same execution twice. Optimized + // builds fall back to leaving such a region uncaptured. + bool inserted = unpaired_executions + .insert(static_cast(*thunk) + .async_execution() + .get()) + .second; + DCHECK(inserted) << "Async execution started twice before its done: " + << thunk->profile_annotation(); + if (!inserted) { + return 0; + } } if (thunk->kind() == Thunk::kAsyncDone) { - auto id = static_cast(*thunk) - .async_execution_id() - .value(); - auto it = unpaired_ids.find(id); - if (it == unpaired_ids.end()) { + auto* execution = + static_cast(*thunk).async_execution().get(); + auto it = unpaired_executions.find(execution); + if (it == unpaired_executions.end()) { return 0; // Done without matching start in the region. } - unpaired_ids.erase(it); + unpaired_executions.erase(it); } - if (unpaired_ids.empty()) { + if (unpaired_executions.empty()) { return i + 1; // All start/done pairs are matched. } } - return 0; // error didn't find an end for some start + return 0; // At least one start has no matching done in this sequence. } -// Returns the shortest non-empty sequence of thunks that form a valid async -// region as a span. If no such region is found, an empty span is returned. -absl::Span> CollectAndCheckAsyncRegion( - absl::Span> thunks, - const CommandBufferConfig& config) { - return thunks.subspan(0, CheckAsyncRegion(thunks, config)); +// Returns the size of a closed region only if every thunk can be converted. +size_t CheckAsyncRegion(absl::Span> thunks, + const CommandBufferConfig& config) { + size_t size = AsyncRegionSize(thunks); + for (const std::unique_ptr& thunk : thunks.first(size)) { + if (!IsConvertible(*thunk, config)) { + return 0; + } + } + return size; +} + +// Returns true if `thunk` or any thunk nested in it starts an async region on +// one of `streams`. DynamicSliceFusionV2Thunk hides its embedded thunks from +// Thunk::Walk, so they are inspected explicitly. +bool ContainsAsyncStartOnStreams( + const Thunk& thunk, const absl::flat_hash_set& streams) { + if (streams.empty()) { + return false; + } + // A non-OK status stops the walk at the first match. + absl::Status walk = thunk.Walk([&](const Thunk* nested) -> absl::Status { + if (nested->kind() == Thunk::kAsyncStart) { + if (streams.contains(static_cast(*nested) + .execution_stream_id())) { + return absl::CancelledError(); + } + } else if (nested->kind() == Thunk::kDynamicSliceFusion) { + const auto& fusion = + static_cast(*nested); + for (const std::unique_ptr& embedded : fusion.thunks()) { + if (ContainsAsyncStartOnStreams(*embedded, streams)) { + return absl::CancelledError(); + } + } + } + return absl::OkStatus(); + }); + return !walk.ok(); +} + +// Returns the index one past the done that joins the start at +// `thunks[start_index]`, or `thunks.size()` if no thunk in this sequence joins +// it. Until that index, work started on the start's stream is outstanding. +size_t AsyncJoinEnd(absl::Span> thunks, + size_t start_index) { + const AsyncExecution* execution = + static_cast(*thunks[start_index]) + .async_execution() + .get(); + for (size_t i = start_index + 1; i < thunks.size(); ++i) { + if (thunks[i]->kind() == Thunk::kAsyncDone && + static_cast(*thunks[i]) + .async_execution() + .get() == execution) { + return i + 1; + } + } + return thunks.size(); } absl::StatusOr GetSynchronizationMode( @@ -671,6 +744,16 @@ absl::StatusOr CommandBufferConversionPass::Run( const HloModule* absl_nullable hlo_module, const se::DeviceDescription& device_info, ThunkPassBufferAllocator& allocator) { + return RunImpl(thunk_sequence, debug_options, hlo_module, device_info, + allocator, /*open_async_streams=*/{}); +} + +absl::StatusOr CommandBufferConversionPass::RunImpl( + ThunkSequence* thunk_sequence, const DebugOptions& debug_options, + const HloModule* absl_nullable hlo_module, + const se::DeviceDescription& device_info, + ThunkPassBufferAllocator& allocator, + const absl::flat_hash_set& open_async_streams) { tsl::profiler::TraceMe traceme("CommandBufferConversionPass"); CommandBufferConfig config = @@ -694,54 +777,103 @@ absl::StatusOr CommandBufferConversionPass::Run( auto& original_thunks = *thunk_sequence; + // An async region is "open" when it cannot be captured as a whole: some + // thunk in it is not convertible, or its start has no matching done in this + // sequence. Its start stays a thunk and enqueues work on its async stream + // that a command buffer launched on the main stream cannot observe. That + // stream is "open" until the done that joins it, so while it is open: + // * plain commands still execute on the main stream and can be captured, + // exactly as their thunks never observed the async stream either; + // * async regions on other streams can still be captured: their bodies + // only ever waited on the main stream, so a graph loses no ordering; + // * async starts on the open stream stay in place, because capturing them + // would move their bodies into a graph on the main stream and lose + // stream order with the outstanding operation; + // * control flow containing such starts is not captured whole; its bodies + // are converted recursively with the same streams open throughout. + // `open_stream_ends` maps each open stream to the index one past the thunk + // that joins it, or the end of the sequence for an unmatched start. + absl::flat_hash_map open_stream_ends; + for (ExecutionStreamId stream : open_async_streams) { + open_stream_ends[stream] = original_thunks.size(); + } + auto open_streams_at = [&](size_t index) { + absl::flat_hash_set streams; + for (const auto& [stream, end] : open_stream_ends) { + if (index < end) { + streams.insert(stream); + } + } + return streams; + }; + for (size_t i = 0; i < original_thunks.size(); ++i) { auto& thunk = original_thunks[i]; + const absl::flat_hash_set open_streams = + open_streams_at(i); - // We always have to capture both corresponding start and done events in the - // same command buffer. if (thunk->kind() == Thunk::kAsyncStart) { - // Collect and check async region - absl::Span> region = CollectAndCheckAsyncRegion( - absl::MakeSpan(original_thunks).subspan(i), config); - - if (!region.empty()) { - // If a valid region is found, add the whole region to the current - // sequence and continue processing. - i += region.size() - 1; - absl::c_move(region, std::back_inserter(current_command_buffer_thunks)); - continue; + const auto& start = static_cast(*thunk); + // We always have to capture both corresponding start and done events in + // the same command buffer. + if (!open_streams.contains(start.execution_stream_id())) { + absl::Span> tail = + absl::MakeSpan(original_thunks).subspan(i); + absl::Span> region = + tail.first(CheckAsyncRegion(tail, config)); + if (!region.empty() && + absl::c_none_of(region, [&](const std::unique_ptr& nested) { + return ContainsAsyncStartOnStreams(*nested, open_streams); + })) { + // If a valid region is found, add the whole region to the current + // sequence and continue processing. + i += region.size() - 1; + absl::c_move(region, + std::back_inserter(current_command_buffer_thunks)); + continue; + } } - } else if (IsConvertible(*thunk.get(), config) && - thunk->kind() != Thunk::kAsyncDone) { - // Check if thunk is convertible and not an async done: async done thunks - // can be only added to the current_command_buffer_thunks as part of a - // valid async regions. + // This start stays a thunk so that its body keeps executing on its async + // stream; the stream is open until the done that joins it. + size_t& end = open_stream_ends[start.execution_stream_id()]; + end = std::max(end, AsyncJoinEnd(original_thunks, i)); + } + + // Async start and done thunks are only captured as part of a valid async + // region above; on their own they stay in place. + const bool is_async_boundary = thunk->kind() == Thunk::kAsyncStart || + thunk->kind() == Thunk::kAsyncDone; + if (!is_async_boundary && IsConvertible(*thunk, config) && + !ContainsAsyncStartOnStreams(*thunk, open_streams)) { current_command_buffer_thunks.push_back(std::move(thunk)); continue; } + if (thunk->kind() == Thunk::kWhile) { - // If a `WhileThunk` itself is not eligible for conversion into a - // command buffer, we attempt to convert thunks within its body + // If a `WhileThunk` itself is not captured into a command buffer, we + // attempt to convert thunks within its body. auto while_thunk = static_cast(thunk.get()); - ABSL_ASSIGN_OR_RETURN(bool changed_in_body, - Run(&while_thunk->body_executor().thunks(), - debug_options, hlo_module, device_info, allocator)); + ABSL_ASSIGN_OR_RETURN( + bool changed_in_body, + RunImpl(&while_thunk->body_executor().thunks(), debug_options, + hlo_module, device_info, allocator, open_streams)); changed |= changed_in_body; } else if (thunk->kind() == Thunk::kConditional) { - // If a `ConditionalThunk` itself is not eligible for conversion into a - // command buffer, we attempt to convert thunks within its branches. + // If a `ConditionalThunk` itself is not captured into a command buffer, + // we attempt to convert thunks within its branches. auto conditional_thunk = static_cast(thunk.get()); for (auto& branch_executor : conditional_thunk->branch_executors()) { - ABSL_ASSIGN_OR_RETURN(bool changed_in_branch, - Run(&branch_executor.thunks(), debug_options, - hlo_module, device_info, allocator)); + ABSL_ASSIGN_OR_RETURN( + bool changed_in_branch, + RunImpl(&branch_executor.thunks(), debug_options, hlo_module, + device_info, allocator, open_streams)); changed |= changed_in_branch; } } - // If the current thunk is not convertible, flush collected eligible thunk + // If the current thunk is not captured, flush collected eligible thunks // to a command buffer thunk and add it to the processed sequence. Then add - // non-convertible thunk to the sequence. + // the thunk itself to the sequence. ABSL_RETURN_IF_ERROR(flush_command_buffer()); new_thunks.push_back(std::move(thunk)); } diff --git a/third_party/xla/xla/backends/gpu/runtime/command_buffer_conversion_pass.h b/third_party/xla/xla/backends/gpu/runtime/command_buffer_conversion_pass.h index c7cff299b8cb16..30e74f21fd2885 100644 --- a/third_party/xla/xla/backends/gpu/runtime/command_buffer_conversion_pass.h +++ b/third_party/xla/xla/backends/gpu/runtime/command_buffer_conversion_pass.h @@ -24,6 +24,7 @@ limitations under the License. #include "absl/container/flat_hash_set.h" #include "absl/status/statusor.h" #include "absl/strings/string_view.h" +#include "xla/backends/gpu/runtime/execution_stream_id.h" #include "xla/backends/gpu/runtime/thunk.h" #include "xla/backends/gpu/runtime/thunk_pass_pipeline.h" #include "xla/hlo/ir/hlo_module.h" @@ -66,6 +67,17 @@ class CommandBufferConversionPass : public ThunkPassInterface { }; private: + // Implementation of `Run`. `open_async_streams` are the async streams with + // an outstanding operation for the whole of `thunk_sequence`: it is the body + // of a loop or branch nested in an async region that is not captured as a + // whole; see the comment in `RunImpl`. + absl::StatusOr RunImpl( + ThunkSequence* thunk_sequence, const DebugOptions& debug_options, + const HloModule* absl_nullable hlo_module, + const se::DeviceDescription& device_info, + ThunkPassBufferAllocator& allocator, + const absl::flat_hash_set& open_async_streams); + std::string module_name_; }; diff --git a/third_party/xla/xla/backends/gpu/target_config/specs/mi300.txtpb b/third_party/xla/xla/backends/gpu/target_config/specs/mi300.txtpb new file mode 100644 index 00000000000000..869d97f8f69d01 --- /dev/null +++ b/third_party/xla/xla/backends/gpu/target_config/specs/mi300.txtpb @@ -0,0 +1,40 @@ +# Copyright 2026 The OpenXLA Authors. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +gpu_device_info { + threads_per_block_limit: 1024 + threads_per_warp: 64 + shared_memory_per_block: 65536 + shared_memory_per_block_optin: 65536 + shared_memory_per_core: 65536 + threads_per_core_limit: 2048 + core_count: 304 + fpus_per_core: 128 + block_dim_limit_x: 2147483647 + block_dim_limit_y: 65536 + block_dim_limit_z: 65536 + memory_bandwidth: 5300000000000 + l2_cache_size: 4194304 + clock_rate_ghz: 2.1 + device_memory_size: 206158430208 + registers_per_core_limit: 131072 + registers_per_block_limit: 131072 + rocm_compute_capability { + gcn_arch_name: "gfx942:sramecc+:xnack-" + } +} +platform_name: "ROCM" +dnn_version_info { +} +device_description_str: "MI300X" diff --git a/third_party/xla/xla/backends/gpu/target_config/specs/gfx1250.txtpb b/third_party/xla/xla/backends/gpu/target_config/specs/mi450.txtpb similarity index 100% rename from third_party/xla/xla/backends/gpu/target_config/specs/gfx1250.txtpb rename to third_party/xla/xla/backends/gpu/target_config/specs/mi450.txtpb diff --git a/third_party/xla/xla/backends/gpu/target_config/target_config_map.bzl b/third_party/xla/xla/backends/gpu/target_config/target_config_map.bzl index d658aa86dcef6e..fac887ef139acf 100644 --- a/third_party/xla/xla/backends/gpu/target_config/target_config_map.bzl +++ b/third_party/xla/xla/backends/gpu/target_config/target_config_map.bzl @@ -38,7 +38,7 @@ target_config_map = { "bmg_g21": "//xla/backends/gpu/target_config:specs/bmg_g21.txtpb", "gb200": "//xla/backends/gpu/target_config:specs/gb200.txtpb", "gb300": "//xla/backends/gpu/target_config:specs/gb300.txtpb", - "gfx1250": "//xla/backends/gpu/target_config:specs/gfx1250.txtpb", + "mi450": "//xla/backends/gpu/target_config:specs/mi450.txtpb", "h100_pcie": "//xla/backends/gpu/target_config:specs/h100_pcie.txtpb", "h100_sxm": "//xla/backends/gpu/target_config:specs/h100_sxm.txtpb", "h100": "//xla/backends/gpu/target_config:specs/h100_sxm.txtpb", diff --git a/third_party/xla/xla/backends/gpu/tests/BUILD b/third_party/xla/xla/backends/gpu/tests/BUILD index 680c99c6556244..7bc17e886098fe 100644 --- a/third_party/xla/xla/backends/gpu/tests/BUILD +++ b/third_party/xla/xla/backends/gpu/tests/BUILD @@ -769,14 +769,21 @@ lit_test_suite_for_gpus( "vectorization.hlo", "pdl.hlo", ], + "mi300": [ + "memory_bandwidth_fusion.hlo", + "dot_bf16.hlo", + "single_instruction.hlo", + "vectorization.hlo", + "pdl.hlo", + ], "mi350": [ "dot_bf16.hlo", "single_instruction.hlo", "vectorization.hlo", "pdl.hlo", ], - "gfx1250": [ - # TODO: re-enable once gfx1250 memory_bandwidth is finalized. + "mi450": [ + # TODO: re-enable once mi450 memory_bandwidth is finalized. "memory_bandwidth_fusion.hlo", "single_instruction.hlo", "vectorization.hlo", @@ -787,10 +794,11 @@ lit_test_suite_for_gpus( "a100_pcie_80", "a6000", "b200", - "gfx1250", "h100_sxm", "mi200", + "mi300", "mi350", + "mi450", ], tags = ["no-oneapi"], tools = [ diff --git a/third_party/xla/xla/backends/gpu/tests/triton_tdm_elemwise.hlo b/third_party/xla/xla/backends/gpu/tests/triton_tdm_elemwise.hlo index 3e21447232700d..22c085a0638500 100644 --- a/third_party/xla/xla/backends/gpu/tests/triton_tdm_elemwise.hlo +++ b/third_party/xla/xla/backends/gpu/tests/triton_tdm_elemwise.hlo @@ -1,18 +1,22 @@ // RUN: hlo-opt %s --platform=gpu --stage=llvm-before-optimizations --xla_gpu_target_config_filename=%S/../target_config/specs/%{GPU}.txtpb | FileCheck %s --check-prefixes=CHECK-%{GPU},CHECK-%{PTX} --allow-unused-prefixes -// Verify that element-wise Triton fusions targeting gfx1250 use TDM +// Verify that element-wise Triton fusions targeting mi450 family use TDM // (Tensor Descriptor Memory) intrinsics instead of pointer arithmetic. // On non-TDM ROCm targets (e.g. mi200), buffer load/store ops are used. // On CUDA targets, the AMD-specific TDM intrinsics must not appear. -// CHECK-gfx1250: @llvm.amdgcn.tensor.load.to.lds -// CHECK-gfx1250: @llvm.amdgcn.s.wait.tensorcnt -// CHECK-gfx1250: @llvm.amdgcn.tensor.store.from.lds +// CHECK-mi450: @llvm.amdgcn.tensor.load.to.lds +// CHECK-mi450: @llvm.amdgcn.s.wait.tensorcnt +// CHECK-mi450: @llvm.amdgcn.tensor.store.from.lds // CHECK-mi200-NOT: @llvm.amdgcn.{{(tensor\.|s\.wait\.tensorcnt)}} // CHECK-mi200: @llvm.amdgcn.raw.ptr.buffer.load // CHECK-mi200-NOT: @llvm.amdgcn.{{(tensor\.|s\.wait\.tensorcnt)}} +// CHECK-mi300-NOT: @llvm.amdgcn.{{(tensor\.|s\.wait\.tensorcnt)}} +// CHECK-mi300: @llvm.amdgcn.raw.ptr.buffer.load +// CHECK-mi300-NOT: @llvm.amdgcn.{{(tensor\.|s\.wait\.tensorcnt)}} + // CHECK-mi350-NOT: @llvm.amdgcn.{{(tensor\.|s\.wait\.tensorcnt)}} // CHECK-mi350: @llvm.amdgcn.raw.ptr.buffer.load // CHECK-mi350-NOT: @llvm.amdgcn.{{(tensor\.|s\.wait\.tensorcnt)}} diff --git a/third_party/xla/xla/comparison_util.cc b/third_party/xla/xla/comparison_util.cc index 9d7ffcd7012ad7..aed2072aa2d5d5 100644 --- a/third_party/xla/xla/comparison_util.cc +++ b/third_party/xla/xla/comparison_util.cc @@ -223,42 +223,19 @@ Comparison::Order Comparison::DefaultOrdering(PrimitiveType type) { LOG(FATAL) << "Unsupported type: " << PrimitiveType_Name(type); } -namespace { -Comparison::Type ComparisonTypeFromPrimitiveTypeAndOrder( - PrimitiveType type, Comparison::Order order) { - if (primitive_util::IsFloatingPointType(type) || - primitive_util::IsComplexType(type)) { - return order == Comparison::Order::kTotal - ? Comparison::Type::kFloatTotalOrder - : Comparison::Type::kFloat; - } - if (primitive_util::IsSignedIntegralType(type)) { - return Comparison::Type::kSigned; - } - if (primitive_util::IsUnsignedIntegralType(type) || type == PRED) { - return Comparison::Type::kUnsigned; - } - LOG(FATAL) << "Unexpected: " << PrimitiveType_Name(type); -} -} // namespace - Comparison::Comparison(Direction dir, PrimitiveType type, Order order) - : dir_(dir), - primitive_type_(type), - order_(order), - type_(ComparisonTypeFromPrimitiveTypeAndOrder(type, order)) {} + : dir_(dir), primitive_type_(type), order_(order) { + CHECK(primitive_util::IsArrayType(type)) + << "Unsupported type: " << PrimitiveType_Name(type); +} Comparison::Comparison(Direction dir, PrimitiveType type) - : dir_(dir), - primitive_type_(type), - order_(DefaultOrdering(type)), - type_(DefaultComparisonType(type)) {} + : dir_(dir), primitive_type_(type), order_(DefaultOrdering(type)) {} Comparison::Comparison(Direction dir, Type type) : dir_(dir), primitive_type_(DefaultPrimitiveType(type)), - order_(DefaultOrdering(type)), - type_(type) {} + order_(DefaultOrdering(type)) {} Comparison Comparison::Converse() const { return Comparison(xla::Converse(dir_), primitive_type_, order_); diff --git a/third_party/xla/xla/comparison_util.h b/third_party/xla/xla/comparison_util.h index 1b852d21c98040..a626a8d772dcab 100644 --- a/third_party/xla/xla/comparison_util.h +++ b/third_party/xla/xla/comparison_util.h @@ -115,11 +115,6 @@ class Comparison { inline PrimitiveType GetPrimitiveType() const { return primitive_type_; } inline Order GetOrder() const { return order_; } - [[deprecated("Use GetPrimitiveType() and GetOrder()")]] inline Type GetType() - const { - return type_; - } - inline bool IsEq() const { return dir_ == Direction::kEq; } inline bool IsNe() const { return dir_ == Direction::kNe; } inline bool IsGe() const { return dir_ == Direction::kGe; } @@ -233,9 +228,6 @@ class Comparison { const PrimitiveType primitive_type_; // The ordering of the Comparison, e.g., kPartial. const Order order_; - // The Type of the Comparison. This tries to mesh together the ordering and - // the numerical data classification. - [[deprecated]] const Type type_; }; using ComparisonDirection = Comparison::Direction; diff --git a/third_party/xla/xla/ffi/api/collectives_api.h b/third_party/xla/xla/ffi/api/collectives_api.h index 4d832aa9616781..c76f4686f34863 100644 --- a/third_party/xla/xla/ffi/api/collectives_api.h +++ b/third_party/xla/xla/ffi/api/collectives_api.h @@ -16,6 +16,7 @@ limitations under the License. #ifndef XLA_FFI_API_COLLECTIVES_API_H_ #define XLA_FFI_API_COLLECTIVES_API_H_ +#include #include #include @@ -33,6 +34,17 @@ enum class GroupMode { kFlattenedId = XLA_FFI_GROUP_FLATTENED_ID, }; +struct CollectiveMemoryRegion { + const void* buffer; + size_t byte_size; + uint64_t flags = 0; +}; + +struct WindowLookup { + XLA_FFI_Window* window; + size_t offset; +}; + namespace internal { // C++ wrapper for the XLA FFI Collectives extension API. @@ -87,6 +99,62 @@ class CommunicatorContextBase { return args.communicator; } + //===--------------------------------------------------------------------===// + // Collective memory window + //===--------------------------------------------------------------------===// + // + // Request window registration for a batch of already-allocated buffers in + // Prepare; look up an opaque backend-defined window handle per buffer in + // Init/Execute. The handler reinterprets the window and calls the backend's + // collective device APIs directly to obtain local, peer, and multicast + // pointers. Allocation stays on the JAX-side. + + Status RequestWindow(GroupMode group_mode, + const std::vector>& groups, + int64_t communication_id, + const std::vector& regions) { + std::vector raw_groups = ToRawGroups(groups); + std::vector raw_regions; + raw_regions.reserve(regions.size()); + for (const CollectiveMemoryRegion& r : regions) { + raw_regions.push_back( + XLA_FFI_CollectiveMemoryRegion{r.buffer, r.byte_size, r.flags}); + } + XLA_FFI_Window_Request_Args args; + args.struct_size = XLA_FFI_Window_Request_Args_STRUCT_SIZE; + args.extension_start = nullptr; + args.group_mode = static_cast(group_mode); + args.groups = raw_groups.data(); + args.num_groups = raw_groups.size(); + args.communication_id = communication_id; + args.regions = raw_regions.data(); + args.num_regions = raw_regions.size(); + if (XLA_FFI_Error* err = ext_->request_window(ext_, &args)) { + return ErrorPolicy::TakeError(api_, err); + } + return ErrorPolicy::Ok(); + } + + StatusOr GetWindow( + GroupMode group_mode, const std::vector>& groups, + int64_t communication_id, const void* buffer) { + std::vector raw_groups = ToRawGroups(groups); + XLA_FFI_Window_Get_Args args; + args.struct_size = XLA_FFI_Window_Get_Args_STRUCT_SIZE; + args.extension_start = nullptr; + args.group_mode = static_cast(group_mode); + args.groups = raw_groups.data(); + args.num_groups = raw_groups.size(); + args.communication_id = communication_id; + args.buffer = buffer; + args.window = nullptr; + args.window_offset = 0; + if (XLA_FFI_Error* err = ext_->get_window(ext_, &args)) { + return StatusOr(ErrorPolicy::TakeError(api_, err)); + } + return WindowLookup{args.window, args.window_offset}; + } + private: // Converts a vector of replica groups to a vector of `XLA_FFI_ReplicaGroup`. // The results reference the id storage in `groups`, which must outlive them. diff --git a/third_party/xla/xla/ffi/api/collectives_c_api.h b/third_party/xla/xla/ffi/api/collectives_c_api.h index b635a92df6c44d..9a2e4719169d77 100644 --- a/third_party/xla/xla/ffi/api/collectives_c_api.h +++ b/third_party/xla/xla/ffi/api/collectives_c_api.h @@ -29,18 +29,22 @@ extern "C" { // XLA FFI Collectives API //===----------------------------------------------------------------------===// -// Exposes the XLA-owned host collective communicator to FFI handlers. +// Exposes the XLA-owned host collective communicator and collective memory +// window to FFI handlers. // `request_communicator` requests a clique in the Prepare stage; -// `get_communicator` returns the handle once cliques are acquired. The -// communicator is backend-defined and passed as an opaque `void*` (a handler on -// XLA:GPU reinterprets it as `ncclComm_t`, XLA:CPU as its own comm type, etc.). +// `get_communicator` returns the handle once cliques are acquired. +// `request_window` requests window registration for a batch of already- +// allocated buffers with the clique in Prepare; `get_window` returns the +// corresponding window handle once collective memory is acquired. +// The communicator and window are backend-defined and passed as opaque +// pointers (a handler on XLA:GPU reinterprets `XLA_FFI_Communicator*` as +// `ncclComm_t` and `XLA_FFI_Window*` as `ncclWindow_t`). // This API is backend agnostic; backend wiring lives in the runtime that // attaches the extension. -// Version macros for the collectives extension. #define XLA_FFI_Extension_Collectives 129 #define XLA_FFI_Extension_Collectives_MajorVersion 0 -#define XLA_FFI_Extension_Collectives_MinorVersion 1 +#define XLA_FFI_Extension_Collectives_MinorVersion 2 // Mirrors `xla::CollectiveOpGroupMode`. typedef enum XLA_FFI_CollectiveGroupMode { @@ -65,6 +69,10 @@ typedef struct XLA_FFI_CollectivesState XLA_FFI_CollectivesState; // type; callers reinterpret it (e.g. as `ncclComm_t`). typedef struct XLA_FFI_Communicator XLA_FFI_Communicator; +//===----------------------------------------------------------------------===// +// Communicator: request in Prepare, get in Initialize/Execute +//===----------------------------------------------------------------------===// + typedef struct XLA_FFI_Communicator_Request_Args { size_t struct_size; XLA_FFI_InternalExtension* extension_start; @@ -97,13 +105,81 @@ typedef struct XLA_FFI_Communicator_Get_Args { XLA_FFI_DEFINE_STRUCT_TRAITS(XLA_FFI_Communicator_Get_Args, communicator); -// Returns the non-owning communicator handle for the clique. Valid once cliques -// are acquired (Initialize/Execute stages). +// Returns the non-owning communicator handle for the clique. Valid once +// cliques are acquired (Initialize/Execute stages). typedef XLA_FFI_Error* XLA_FFI_Communicator_Get( const XLA_FFI_Collectives_Extension* self, XLA_FFI_Communicator_Get_Args* args); -// Main extension struct for the collectives API. +//===----------------------------------------------------------------------===// +// Collective memory window +//===----------------------------------------------------------------------===// +// +// Handlers request window registration for a batch of already-allocated +// buffers with the clique via `request_window` in Prepare, then look up an +// opaque `XLA_FFI_Window` per buffer via `get_window` in Init/Execute. +// Handlers reinterpret the window and call the backend's collective device APIs +// directly to obtain local, peer, and multicast pointers. + +// Opaque, non-owning collective memory window handle. The backend defines the +// concrete type; callers reinterpret it (e.g. as `ncclWindow_t`). +typedef struct XLA_FFI_Window XLA_FFI_Window; + +typedef struct XLA_FFI_CollectiveMemoryRegion { + const void* buffer; + size_t byte_size; + uint64_t flags; +} XLA_FFI_CollectiveMemoryRegion; + +typedef struct XLA_FFI_Window_Request_Args { + size_t struct_size; + XLA_FFI_InternalExtension* extension_start; + + XLA_FFI_CollectiveGroupMode group_mode; + const XLA_FFI_ReplicaGroup* groups; + size_t num_groups; + int64_t communication_id; + + const XLA_FFI_CollectiveMemoryRegion* regions; + size_t num_regions; +} XLA_FFI_Window_Request_Args; + +XLA_FFI_DEFINE_STRUCT_TRAITS(XLA_FFI_Window_Request_Args, num_regions); + +// Requests window registration for a batch of already-allocated buffers with +// the clique. Actual registration happens later once the clique is acquired; +// look up the resulting window handle with `get_window`. Prepare stage only. +typedef XLA_FFI_Error* XLA_FFI_Window_Request( + const XLA_FFI_Collectives_Extension* self, + XLA_FFI_Window_Request_Args* args); + +typedef struct XLA_FFI_Window_Get_Args { + size_t struct_size; + XLA_FFI_InternalExtension* extension_start; + + XLA_FFI_CollectiveGroupMode group_mode; + const XLA_FFI_ReplicaGroup* groups; + size_t num_groups; + int64_t communication_id; + + const void* buffer; // registered address (in) + XLA_FFI_Window* window; // out + size_t window_offset; // byte offset of `buffer` within `window` (out) +} XLA_FFI_Window_Get_Args; + +XLA_FFI_DEFINE_STRUCT_TRAITS(XLA_FFI_Window_Get_Args, window_offset); + +// Returns the non-owning collective memory window handle for a buffer whose +// registration was previously requested via `request_window`, along with the +// byte offset of `buffer` within that window. Valid once collective memory is +// acquired (Initialize/Execute stages). +typedef XLA_FFI_Error* XLA_FFI_Window_Get( + const XLA_FFI_Collectives_Extension* self, XLA_FFI_Window_Get_Args* args); + +//===----------------------------------------------------------------------===// +// Extension struct +//===----------------------------------------------------------------------===// + struct XLA_FFI_Collectives_Extension { XLA_FFI_Extension extension_base; @@ -111,9 +187,12 @@ struct XLA_FFI_Collectives_Extension { XLA_FFI_Communicator_Request* request_communicator; XLA_FFI_Communicator_Get* get_communicator; + + XLA_FFI_Window_Request* request_window; + XLA_FFI_Window_Get* get_window; }; -XLA_FFI_DEFINE_STRUCT_TRAITS(XLA_FFI_Collectives_Extension, get_communicator); +XLA_FFI_DEFINE_STRUCT_TRAITS(XLA_FFI_Collectives_Extension, get_window); #ifdef __cplusplus } // extern "C" diff --git a/third_party/xla/xla/ffi/api/collectives_ffi_test.cc b/third_party/xla/xla/ffi/api/collectives_ffi_test.cc index 397f24323a8372..95c487c54fd3a6 100644 --- a/third_party/xla/xla/ffi/api/collectives_ffi_test.cc +++ b/third_party/xla/xla/ffi/api/collectives_ffi_test.cc @@ -37,6 +37,9 @@ static const XLA_FFI_Api* Api() { return GetXlaFfiApi(); } static XLA_FFI_Communicator* const kFakeCommunicator = reinterpret_cast(0xC0FFEE); +static XLA_FFI_Window* const kFakeWindow = + reinterpret_cast(0xDEC0DE); +static constexpr size_t kFakeWindowOffset = 42; // Fake backend: implements the extension callbacks and records what the handler // requested so the test can check the C++ wrapper builds the C args correctly. @@ -47,6 +50,15 @@ struct FakeBackend { size_t requested_num_groups = 0; int64_t requested_communication_id = -1; std::vector requested_ids; + + bool request_window_called = false; + XLA_FFI_CollectiveGroupMode requested_window_group_mode = + XLA_FFI_GROUP_CROSS_REPLICA; + size_t requested_window_num_groups = 0; + int64_t requested_window_communication_id = -1; + std::vector requested_window_ids; + std::vector requested_regions; + const void* get_window_buffer = nullptr; }; static XLA_FFI_Error* FakeRequestCommunicator( @@ -73,18 +85,52 @@ static XLA_FFI_Error* FakeGetCommunicator( return nullptr; } +static XLA_FFI_Error* FakeRequestWindow( + const XLA_FFI_Collectives_Extension* self, + XLA_FFI_Window_Request_Args* args) { + auto* backend = reinterpret_cast(self->state); + backend->request_window_called = true; + backend->requested_window_group_mode = args->group_mode; + backend->requested_window_num_groups = args->num_groups; + backend->requested_window_communication_id = args->communication_id; + backend->requested_window_ids.clear(); + for (size_t i = 0; i < args->num_groups; ++i) { + for (size_t j = 0; j < args->groups[i].size; ++j) { + backend->requested_window_ids.push_back(args->groups[i].ids[j]); + } + } + backend->requested_regions.clear(); + for (size_t i = 0; i < args->num_regions; ++i) { + backend->requested_regions.push_back(args->regions[i]); + } + return nullptr; +} + +static XLA_FFI_Error* FakeGetWindow(const XLA_FFI_Collectives_Extension* self, + XLA_FFI_Window_Get_Args* args) { + auto* backend = reinterpret_cast(self->state); + backend->get_window_buffer = args->buffer; + args->window = kFakeWindow; + args->window_offset = kFakeWindowOffset; + return nullptr; +} + // Builds a collectives extension for the fake backend. Mirrors the small // builder each backend uses to publish the extension (see the GPU backend). static XLA_FFI_Collectives_Extension MakeFakeCollectivesExtension( XLA_FFI_CollectivesState* state, XLA_FFI_Communicator_Request* request_communicator, - XLA_FFI_Communicator_Get* get_communicator) { - XLA_FFI_Collectives_Extension ext; + XLA_FFI_Communicator_Get* get_communicator, + XLA_FFI_Window_Request* request_window = nullptr, + XLA_FFI_Window_Get* get_window = nullptr) { + XLA_FFI_Collectives_Extension ext = {}; ext.extension_base = MakeExtensionHeader>(); ext.state = state; ext.request_communicator = request_communicator; ext.get_communicator = get_communicator; + ext.request_window = request_window; + ext.get_window = get_window; return ext; } @@ -133,5 +179,66 @@ TEST(CollectivesFfiTest, RequestAndGetCommunicator) { EXPECT_THAT(backend.requested_ids, ElementsAre(0, 1)); } +TEST(CollectivesFfiTest, RequestAndGetWindow) { + static uint32_t fake_buffer0 = 0; + static uint32_t fake_buffer1 = 0; + const std::vector regions = { + {&fake_buffer0, sizeof(fake_buffer0)}, + {&fake_buffer1, sizeof(fake_buffer1)}, + }; + + bool called = false; + WindowLookup got_lookup = {nullptr, 0}; + + auto handler = Ffi::Bind().Ctx>().To( + [&](Communicator comm) -> Error { + called = true; + if (Error status = comm.RequestWindow(GroupMode::kFlattenedId, {{0, 1}}, + /*communication_id=*/9, regions); + status.failure()) { + return status; + } + ErrorOr lookup_or = + comm.GetWindow(GroupMode::kFlattenedId, {{0, 1}}, + /*communication_id=*/9, &fake_buffer0); + if (lookup_or.has_error()) { + return lookup_or.error(); + } + got_lookup = lookup_or.value(); + return Error::Success(); + }); + + CallFrameBuilder builder(/*num_args=*/0, /*num_rets=*/0); + auto call_frame = builder.Build(); + + FakeBackend backend; + XLA_FFI_Collectives_Extension ext = MakeFakeCollectivesExtension( + reinterpret_cast(&backend), + FakeRequestCommunicator, FakeGetCommunicator, FakeRequestWindow, + FakeGetWindow); + + InvokeContext context; + context.extension_start = &ext.extension_base; + + auto status = + Invoke(Api(), *handler, call_frame, context, ExecutionStage::kExecute); + + ASSERT_OK(status); + EXPECT_TRUE(called); + EXPECT_TRUE(backend.request_window_called); + EXPECT_EQ(backend.requested_window_group_mode, XLA_FFI_GROUP_FLATTENED_ID); + EXPECT_EQ(backend.requested_window_num_groups, 1u); + EXPECT_EQ(backend.requested_window_communication_id, 9); + EXPECT_THAT(backend.requested_window_ids, ElementsAre(0, 1)); + ASSERT_EQ(backend.requested_regions.size(), 2u); + EXPECT_EQ(backend.requested_regions[0].buffer, &fake_buffer0); + EXPECT_EQ(backend.requested_regions[0].byte_size, sizeof(fake_buffer0)); + EXPECT_EQ(backend.requested_regions[1].buffer, &fake_buffer1); + EXPECT_EQ(backend.requested_regions[1].byte_size, sizeof(fake_buffer1)); + EXPECT_EQ(backend.get_window_buffer, &fake_buffer0); + EXPECT_EQ(got_lookup.window, kFakeWindow); + EXPECT_EQ(got_lookup.offset, kFakeWindowOffset); +} + } // namespace } // namespace xla::ffi diff --git a/third_party/xla/xla/ffi/collectives_ffi_test.cc b/third_party/xla/xla/ffi/collectives_ffi_test.cc index 389b957cf9dbef..af29ebc81ffa23 100644 --- a/third_party/xla/xla/ffi/collectives_ffi_test.cc +++ b/third_party/xla/xla/ffi/collectives_ffi_test.cc @@ -39,6 +39,9 @@ static const XLA_FFI_Api* Api() { return GetXlaFfiApi(); } static XLA_FFI_Communicator* const kFakeCommunicator = reinterpret_cast(0xC0FFEE); +static XLA_FFI_Window* const kFakeWindow = + reinterpret_cast(0xDEC0DE); +static constexpr size_t kFakeWindowOffset = 42; // Fake backend: implements the extension callbacks and records what the handler // requested so the test can check the C++ wrapper builds the C args correctly. @@ -49,6 +52,15 @@ struct FakeBackend { size_t requested_num_groups = 0; int64_t requested_communication_id = -1; std::vector requested_ids; + + bool request_window_called = false; + XLA_FFI_CollectiveGroupMode requested_window_group_mode = + XLA_FFI_GROUP_CROSS_REPLICA; + size_t requested_window_num_groups = 0; + int64_t requested_window_communication_id = -1; + std::vector requested_window_ids; + std::vector requested_regions; + const void* get_window_buffer = nullptr; }; static XLA_FFI_Error* FakeRequestCommunicator( @@ -75,18 +87,52 @@ static XLA_FFI_Error* FakeGetCommunicator( return nullptr; } +static XLA_FFI_Error* FakeRequestWindow( + const XLA_FFI_Collectives_Extension* self, + XLA_FFI_Window_Request_Args* args) { + auto* backend = reinterpret_cast(self->state); + backend->request_window_called = true; + backend->requested_window_group_mode = args->group_mode; + backend->requested_window_num_groups = args->num_groups; + backend->requested_window_communication_id = args->communication_id; + backend->requested_window_ids.clear(); + for (size_t i = 0; i < args->num_groups; ++i) { + for (size_t j = 0; j < args->groups[i].size; ++j) { + backend->requested_window_ids.push_back(args->groups[i].ids[j]); + } + } + backend->requested_regions.clear(); + for (size_t i = 0; i < args->num_regions; ++i) { + backend->requested_regions.push_back(args->regions[i]); + } + return nullptr; +} + +static XLA_FFI_Error* FakeGetWindow(const XLA_FFI_Collectives_Extension* self, + XLA_FFI_Window_Get_Args* args) { + auto* backend = reinterpret_cast(self->state); + backend->get_window_buffer = args->buffer; + args->window = kFakeWindow; + args->window_offset = kFakeWindowOffset; + return nullptr; +} + // Builds a collectives extension for the fake backend. Mirrors the small // builder each backend uses to publish the extension (see the GPU backend). static XLA_FFI_Collectives_Extension MakeFakeCollectivesExtension( XLA_FFI_CollectivesState* state, XLA_FFI_Communicator_Request* request_communicator, - XLA_FFI_Communicator_Get* get_communicator) { - XLA_FFI_Collectives_Extension ext; + XLA_FFI_Communicator_Get* get_communicator, + XLA_FFI_Window_Request* request_window = nullptr, + XLA_FFI_Window_Get* get_window = nullptr) { + XLA_FFI_Collectives_Extension ext = {}; ext.extension_base = MakeExtensionHeader>(); ext.state = state; ext.request_communicator = request_communicator; ext.get_communicator = get_communicator; + ext.request_window = request_window; + ext.get_window = get_window; return ext; } @@ -135,5 +181,67 @@ TEST(CollectivesFfiTest, RequestAndGetCommunicator) { EXPECT_THAT(backend.requested_ids, ElementsAre(0, 1)); } +TEST(CollectivesFfiTest, RequestAndGetWindow) { + static uint32_t fake_buffer0 = 0; + static uint32_t fake_buffer1 = 0; + const std::vector regions = { + {&fake_buffer0, sizeof(fake_buffer0)}, + {&fake_buffer1, sizeof(fake_buffer1)}, + }; + + bool called = false; + WindowLookup got_lookup = {nullptr, 0}; + + auto handler = Ffi::Bind().Ctx>().To( + [&](Communicator comm) -> absl::Status { + called = true; + if (absl::Status status = + comm.RequestWindow(GroupMode::kFlattenedId, {{0, 1}}, + /*communication_id=*/9, regions); + !status.ok()) { + return status; + } + absl::StatusOr lookup_or = + comm.GetWindow(GroupMode::kFlattenedId, {{0, 1}}, + /*communication_id=*/9, &fake_buffer0); + if (!lookup_or.ok()) { + return lookup_or.status(); + } + got_lookup = *lookup_or; + return absl::OkStatus(); + }); + + CallFrameBuilder builder(/*num_args=*/0, /*num_rets=*/0); + auto call_frame = builder.Build(); + + FakeBackend backend; + XLA_FFI_Collectives_Extension ext = MakeFakeCollectivesExtension( + reinterpret_cast(&backend), + FakeRequestCommunicator, FakeGetCommunicator, FakeRequestWindow, + FakeGetWindow); + + InvokeContext context; + context.extension_start = &ext.extension_base; + + auto status = + Invoke(Api(), *handler, call_frame, context, ExecutionStage::kExecute); + + ASSERT_OK(status); + EXPECT_TRUE(called); + EXPECT_TRUE(backend.request_window_called); + EXPECT_EQ(backend.requested_window_group_mode, XLA_FFI_GROUP_FLATTENED_ID); + EXPECT_EQ(backend.requested_window_num_groups, 1u); + EXPECT_EQ(backend.requested_window_communication_id, 9); + EXPECT_THAT(backend.requested_window_ids, ElementsAre(0, 1)); + ASSERT_EQ(backend.requested_regions.size(), 2u); + EXPECT_EQ(backend.requested_regions[0].buffer, &fake_buffer0); + EXPECT_EQ(backend.requested_regions[0].byte_size, sizeof(fake_buffer0)); + EXPECT_EQ(backend.requested_regions[1].buffer, &fake_buffer1); + EXPECT_EQ(backend.requested_regions[1].byte_size, sizeof(fake_buffer1)); + EXPECT_EQ(backend.get_window_buffer, &fake_buffer0); + EXPECT_EQ(got_lookup.window, kFakeWindow); + EXPECT_EQ(got_lookup.offset, kFakeWindowOffset); +} + } // namespace } // namespace xla::ffi diff --git a/third_party/xla/xla/hlo/builder/xla_builder.cc b/third_party/xla/xla/hlo/builder/xla_builder.cc index b96c716823cd80..a672fa261726b7 100644 --- a/third_party/xla/xla/hlo/builder/xla_builder.cc +++ b/third_party/xla/xla/hlo/builder/xla_builder.cc @@ -1400,7 +1400,7 @@ absl::StatusOr BroadcastToOutputShapeWithUnbounded( XlaOp XlaBuilder::BinaryOp(HloOpcode binop, XlaOp lhs, XlaOp rhs, absl::Span broadcast_dimensions, std::optional direction, - std::optional type) { + std::optional order) { return ReportErrorOrReturn([&]() -> absl::StatusOr { ABSL_ASSIGN_OR_RETURN(const Shape* lhs_shape, GetShapePtr(lhs)); ABSL_ASSIGN_OR_RETURN(const Shape* rhs_shape, GetShapePtr(rhs)); @@ -1463,10 +1463,10 @@ XlaOp XlaBuilder::BinaryOp(HloOpcode binop, XlaOp lhs, XlaOp rhs, return InvalidArgument( "kCompare expects a ComparisonDirection, but none provided."); } - if (type == std::nullopt) { + if (order == std::nullopt) { return Compare(shape, updated_lhs, updated_rhs, *direction); } - return Compare(shape, updated_lhs, updated_rhs, *direction, *type); + return Compare(shape, updated_lhs, updated_rhs, *direction, *order); } if (direction.has_value()) { @@ -1491,18 +1491,17 @@ absl::StatusOr XlaBuilder::Compare(const Shape& shape, XlaOp lhs, XlaOp rhs, ComparisonDirection direction) { ABSL_ASSIGN_OR_RETURN(auto operand_shape, GetShape(lhs)); - return Compare( - shape, lhs, rhs, direction, - Comparison::DefaultComparisonType(operand_shape.element_type())); + return Compare(shape, lhs, rhs, direction, + Comparison::DefaultOrdering(operand_shape.element_type())); } absl::StatusOr XlaBuilder::Compare(const Shape& shape, XlaOp lhs, XlaOp rhs, ComparisonDirection direction, - Comparison::Type type) { + ComparisonOrder order) { HloInstructionProto instr; instr.set_comparison_direction(ComparisonDirectionToString(direction)); - instr.set_comparison_type(ComparisonTypeToString(type)); + instr.set_comparison_order(ComparisonOrderToShortString(order)); *instr.mutable_shape() = shape.ToProto(); return AddInstruction(std::move(instr), HloOpcode::kCompare, {lhs, rhs}); } @@ -5696,12 +5695,12 @@ static XlaOp CompareTotalOrder(const XlaOp lhs, const XlaOp rhs, return b->ReportErrorOrReturn([&]() -> absl::StatusOr { ABSL_ASSIGN_OR_RETURN(auto operand_shape, b->GetShape(lhs)); auto operand_element_type = operand_shape.element_type(); - auto compare_type = + auto compare_order = primitive_util::IsFloatingPointType(operand_element_type) - ? Comparison::Type::kFloatTotalOrder - : Comparison::DefaultComparisonType(operand_element_type); + ? ComparisonOrder::kTotal + : Comparison::DefaultOrdering(operand_element_type); return Compare(lhs, rhs, broadcast_dimensions, comparison_direction, - compare_type); + compare_order); }); } @@ -5775,9 +5774,16 @@ XlaOp Compare(const XlaOp lhs, const XlaOp rhs, XlaOp Compare(const XlaOp lhs, const XlaOp rhs, absl::Span broadcast_dimensions, - ComparisonDirection direction, Comparison::Type compare_type) { + ComparisonDirection direction, ComparisonOrder order) { return lhs.builder()->BinaryOp(HloOpcode::kCompare, lhs, rhs, - broadcast_dimensions, direction, compare_type); + broadcast_dimensions, direction, order); +} + +XlaOp Compare(const XlaOp lhs, const XlaOp rhs, + absl::Span broadcast_dimensions, + ComparisonDirection direction, Comparison::Type compare_type) { + return Compare(lhs, rhs, broadcast_dimensions, direction, + Comparison::DefaultOrdering(compare_type)); } XlaOp Compare(const XlaOp lhs, const XlaOp rhs, ComparisonDirection direction) { diff --git a/third_party/xla/xla/hlo/builder/xla_builder.h b/third_party/xla/xla/hlo/builder/xla_builder.h index 065be7eb27763b..3c0e937a8c1088 100644 --- a/third_party/xla/xla/hlo/builder/xla_builder.h +++ b/third_party/xla/xla/hlo/builder/xla_builder.h @@ -1254,7 +1254,7 @@ class XlaBuilder { XlaOp BinaryOp(HloOpcode binop, XlaOp lhs, XlaOp rhs, absl::Span broadcast_dimensions, std::optional direction = std::nullopt, - std::optional type = std::nullopt); + std::optional order = std::nullopt); absl::StatusOr Compare(const Shape& shape, XlaOp lhs, XlaOp rhs, ComparisonDirection direction); @@ -1263,7 +1263,7 @@ class XlaBuilder { virtual absl::StatusOr Compare(const Shape& shape, XlaOp lhs, XlaOp rhs, ComparisonDirection direction, - Comparison::Type type); + ComparisonOrder order); // Internal helper method that does the building for an arbitrary binary op // with same ranked operands that doesn't broadcast. @@ -1519,6 +1519,9 @@ class XlaBuilder { friend XlaOp Compare(XlaOp lhs, XlaOp rhs, absl::Span broadcast_dimensions, ComparisonDirection direction); + friend XlaOp Compare(XlaOp lhs, XlaOp rhs, + absl::Span broadcast_dimensions, + ComparisonDirection direction, ComparisonOrder order); friend XlaOp Compare(XlaOp lhs, XlaOp rhs, absl::Span broadcast_dimensions, ComparisonDirection direction, @@ -2574,6 +2577,9 @@ XlaOp LeTotalOrder(XlaOp lhs, XlaOp rhs, // Enqueues a comparison instruction onto the computation (optionally without // broadcast_dimensions for consistency with others). +XlaOp Compare(XlaOp lhs, XlaOp rhs, + absl::Span broadcast_dimensions, + ComparisonDirection direction, ComparisonOrder order); XlaOp Compare(XlaOp lhs, XlaOp rhs, absl::Span broadcast_dimensions, ComparisonDirection direction, Comparison::Type compare_type); diff --git a/third_party/xla/xla/hlo/evaluator/BUILD b/third_party/xla/xla/hlo/evaluator/BUILD index bcb0565fd8f01b..2ae54c6ac85335 100644 --- a/third_party/xla/xla/hlo/evaluator/BUILD +++ b/third_party/xla/xla/hlo/evaluator/BUILD @@ -120,6 +120,7 @@ cc_library( "@com_google_absl//absl/strings", "@com_google_absl//absl/strings:cord", "@com_google_absl//absl/strings:str_format", + "@com_google_absl//absl/synchronization", "@com_google_absl//absl/types:span", "@eigen_archive//:eigen3", "@tsl//tsl/platform:ml_dtypes", diff --git a/third_party/xla/xla/hlo/evaluator/hlo_evaluator.cc b/third_party/xla/xla/hlo/evaluator/hlo_evaluator.cc index 7200b9f496ed82..69b3cc8cff771c 100644 --- a/third_party/xla/xla/hlo/evaluator/hlo_evaluator.cc +++ b/third_party/xla/xla/hlo/evaluator/hlo_evaluator.cc @@ -3672,10 +3672,10 @@ absl::Status HloEvaluator::HandleCall(const HloInstruction* call) { TF_RET_CHECK(specialization_cache_ != nullptr); - const Literal* cached_result = - specialization_cache_->Find(computation, arg_literals); - if (cached_result != nullptr) { - SetEvaluatedLiteralFor(call, cached_result->Clone()); + if (std::optional cached_result = + specialization_cache_->Find(computation, arg_literals); + cached_result.has_value()) { + SetEvaluatedLiteralFor(call, std::move(*cached_result)); return absl::OkStatus(); } @@ -3686,9 +3686,7 @@ absl::Status HloEvaluator::HandleCall(const HloInstruction* call) { ABSL_ASSIGN_OR_RETURN(Literal result, embedded_evaluator->Evaluate(*computation, arg_literals)); - if (specialization_cache_->Find(computation, arg_literals) == nullptr) { - specialization_cache_->Insert(computation, arg_literals, result.Clone()); - } + specialization_cache_->Insert(computation, arg_literals, result.Clone()); SetEvaluatedLiteralFor(call, std::move(result)); return absl::OkStatus(); diff --git a/third_party/xla/xla/hlo/evaluator/hlo_evaluator.h b/third_party/xla/xla/hlo/evaluator/hlo_evaluator.h index 1453825d8a3579..acc6b37b43dfeb 100644 --- a/third_party/xla/xla/hlo/evaluator/hlo_evaluator.h +++ b/third_party/xla/xla/hlo/evaluator/hlo_evaluator.h @@ -29,6 +29,7 @@ limitations under the License. #include #include +#include "absl/base/thread_annotations.h" #include "absl/container/flat_hash_map.h" #include "absl/container/inlined_vector.h" #include "absl/container/node_hash_map.h" @@ -38,6 +39,7 @@ limitations under the License. #include "absl/status/status_macros.h" #include "absl/status/statusor.h" #include "absl/strings/string_view.h" +#include "absl/synchronization/mutex.h" #include "absl/types/span.h" #include "Eigen/Core" #include "xla/array2d.h" @@ -369,13 +371,13 @@ class HloEvaluator : public ConstDfsHloVisitorWithDefault, // See the description on SpecializationKey for more details. struct SpecializationCache { using EntryMap = absl::flat_hash_map; - using iterator = EntryMap::iterator; - using const_iterator = EntryMap::const_iterator; + + mutable absl::Mutex mu; // A deque never relocates existing elements. It means all existing pointers // and references to previously stored Literals remain valid for the // lifetime of the SpecializationCache. - std::deque arg_storage; - EntryMap entries; + std::deque arg_storage ABSL_GUARDED_BY(mu); + EntryMap entries ABSL_GUARDED_BY(mu); SpecializationCache() = default; SpecializationCache(const SpecializationCache&) = delete; @@ -384,31 +386,39 @@ class HloEvaluator : public ConstDfsHloVisitorWithDefault, SpecializationCache& operator=(SpecializationCache&&) = delete; void Clear() { + absl::MutexLock lock(&mu); entries.clear(); arg_storage.clear(); } - size_t size() const { return entries.size(); } - bool empty() const { return entries.empty(); } - iterator begin() { return entries.begin(); } - iterator end() { return entries.end(); } - const_iterator begin() const { return entries.begin(); } - const_iterator end() const { return entries.end(); } - - const Literal* Find(const SpecializationKey& key) const { + size_t size() const { + absl::MutexLock lock(&mu); + return entries.size(); + } + bool empty() const { + absl::MutexLock lock(&mu); + return entries.empty(); + } + + std::optional Find(const SpecializationKey& key) const { + absl::ReaderMutexLock lock(&mu); auto it = entries.find(key); if (it != entries.end()) { - return &it->second; + return it->second.Clone(); } - return nullptr; + return std::nullopt; } - const Literal* Find(const HloComputation* computation, - absl::Span args) const { + std::optional Find(const HloComputation* computation, + absl::Span args) const { return Find(SpecializationKey(computation, args)); } void Insert(const HloComputation* computation, absl::Span args, Literal result) { + absl::MutexLock lock(&mu); + if (entries.contains(SpecializationKey(computation, args))) { + return; + } std::vector slices; slices.reserve(args.size()); for (const Literal* arg : args) { diff --git a/third_party/xla/xla/hlo/evaluator/hlo_evaluator_test.cc b/third_party/xla/xla/hlo/evaluator/hlo_evaluator_test.cc index 7b813bacaeeac8..e8f47905b40997 100644 --- a/third_party/xla/xla/hlo/evaluator/hlo_evaluator_test.cc +++ b/third_party/xla/xla/hlo/evaluator/hlo_evaluator_test.cc @@ -8724,15 +8724,12 @@ TEST_F(HloEvaluatorTest, NestedCallSharedCache) { const HloComputation* comp_b = m_->GetComputationWithName("comp_b"); ASSERT_NE(comp_b, nullptr); - int comp_b_entries = 0; - for (const auto& [key, val] : *evaluator.specialization_cache()) { - if (key.computation == comp_b) { - ++comp_b_entries; - EXPECT_TRUE( - LiteralTestUtil::Equal(LiteralUtil::CreateR0(15.0f), val)); - } - } - EXPECT_EQ(comp_b_entries, 1); + Literal arg_5 = LiteralUtil::CreateR0(5.0f); + std::optional cached_b = + evaluator.specialization_cache()->Find(comp_b, {&arg_5}); + ASSERT_TRUE(cached_b.has_value()); + EXPECT_TRUE( + LiteralTestUtil::Equal(LiteralUtil::CreateR0(15.0f), *cached_b)); } TEST_F(HloEvaluatorTest, ClearSpecializationCache) { diff --git a/third_party/xla/xla/hlo/ir/hlo_instructions.cc b/third_party/xla/xla/hlo/ir/hlo_instructions.cc index 9f5c773561d98f..ce9d59cefa5a12 100644 --- a/third_party/xla/xla/hlo/ir/hlo_instructions.cc +++ b/third_party/xla/xla/hlo/ir/hlo_instructions.cc @@ -961,7 +961,6 @@ void HloCompareInstruction::ToProto(HloInstructionProto* proto) const { ComparisonDirectionToString(compare_.GetDirection())); proto->set_comparison_order( ComparisonOrderToShortString(compare_.GetOrder())); - proto->set_comparison_type(ComparisonTypeToString(compare_.GetType())); } void HloCompareInstruction::PrintExtraAttributesImpl( diff --git a/third_party/xla/xla/hlo/ir/hlo_instructions.h b/third_party/xla/xla/hlo/ir/hlo_instructions.h index 10ae4cba26b71d..054e036d666b7d 100644 --- a/third_party/xla/xla/hlo/ir/hlo_instructions.h +++ b/third_party/xla/xla/hlo/ir/hlo_instructions.h @@ -454,9 +454,6 @@ class HloCompareInstruction : public HloInstruction { std::optional order = std::nullopt); ComparisonDirection direction() const { return compare_.GetDirection(); } ComparisonOrder order() const { return compare_.GetOrder(); } - [[deprecated("Use order()")]] Comparison::Type type() const { - return compare_.GetType(); - } const Comparison& comparison() const { return compare_; } void ToProto(HloInstructionProto* proto) const override; diff --git a/third_party/xla/xla/hlo/transforms/collectives/BUILD b/third_party/xla/xla/hlo/transforms/collectives/BUILD index d80191f24c2867..51a74a2a0043dd 100644 --- a/third_party/xla/xla/hlo/transforms/collectives/BUILD +++ b/third_party/xla/xla/hlo/transforms/collectives/BUILD @@ -102,6 +102,7 @@ cc_library( "@com_google_absl//absl/log", "@com_google_absl//absl/log:check", "@com_google_absl//absl/log:vlog_is_on", + "@com_google_absl//absl/status:status_macros", "@com_google_absl//absl/status:statusor", "@com_google_absl//absl/strings:string_view", ], @@ -785,6 +786,8 @@ cc_library( deps = [ ":collective_permute_cycle", "//xla:shape_util", + "//xla:xla_data_proto_cc", + "//xla:xla_proto_cc", "//xla/hlo/ir:hlo", "//xla/hlo/pass:hlo_pass", "//xla/service:call_graph", @@ -808,6 +811,8 @@ xla_cc_test( srcs = ["collective_permute_decomposer_test.cc"], deps = [ ":collective_permute_decomposer", + "//xla:xla_data_proto_cc", + "//xla:xla_proto_cc", "//xla/hlo/ir:hlo", "//xla/hlo/testlib:hlo_hardware_independent_test_base", "//xla/service:collective_ops_utils", @@ -865,6 +870,7 @@ xla_cc_test( srcs = ["reduce_scatter_combiner_test.cc"], deps = [ ":reduce_scatter_combiner", + "//xla:xla_data_proto_cc", "//xla/hlo/ir:hlo", "//xla/hlo/testlib:hlo_hardware_independent_test_base", "//xla/hlo/utils:hlo_matchers", @@ -961,6 +967,7 @@ cc_library( hdrs = ["reduce_scatter_decomposer.h"], deps = [ "//xla:shape_util", + "//xla:xla_data_proto_cc", "//xla/hlo/ir:collective_op_group_mode", "//xla/hlo/ir:hlo", "//xla/hlo/pass:hlo_pass", @@ -982,6 +989,7 @@ xla_cc_test( "//xla:literal", "//xla:literal_util", "//xla:shape_util", + "//xla:xla_data_proto_cc", "//xla/hlo/ir:hlo", "//xla/hlo/testlib:hlo_hardware_independent_test_base", "//xla/hlo/utils:hlo_matchers", diff --git a/third_party/xla/xla/hlo/transforms/collectives/all_gather_decomposer.h b/third_party/xla/xla/hlo/transforms/collectives/all_gather_decomposer.h index 84ef994189c11e..55ac3c62864f1a 100644 --- a/third_party/xla/xla/hlo/transforms/collectives/all_gather_decomposer.h +++ b/third_party/xla/xla/hlo/transforms/collectives/all_gather_decomposer.h @@ -30,6 +30,7 @@ limitations under the License. #include "xla/hlo/ir/hlo_module.h" #include "xla/hlo/pass/hlo_pass_interface.h" #include "xla/shape.h" +#include "xla/xla_data.pb.h" namespace xla { diff --git a/third_party/xla/xla/hlo/transforms/collectives/async_collective_replacer.cc b/third_party/xla/xla/hlo/transforms/collectives/async_collective_replacer.cc index 8433b2ee83a818..019c5a8cceca1a 100644 --- a/third_party/xla/xla/hlo/transforms/collectives/async_collective_replacer.cc +++ b/third_party/xla/xla/hlo/transforms/collectives/async_collective_replacer.cc @@ -23,6 +23,7 @@ limitations under the License. #include "absl/log/check.h" #include "absl/log/log.h" #include "absl/log/vlog_is_on.h" +#include "absl/status/status_macros.h" #include "absl/strings/string_view.h" #include "xla/hlo/ir/hlo_computation.h" #include "xla/hlo/ir/hlo_instruction.h" diff --git a/third_party/xla/xla/hlo/transforms/collectives/collective_permute_decomposer.cc b/third_party/xla/xla/hlo/transforms/collectives/collective_permute_decomposer.cc index 780927fbb9e9c9..83991e68a13546 100644 --- a/third_party/xla/xla/hlo/transforms/collectives/collective_permute_decomposer.cc +++ b/third_party/xla/xla/hlo/transforms/collectives/collective_permute_decomposer.cc @@ -45,6 +45,8 @@ limitations under the License. #include "xla/service/source_target_pairs.h" #include "xla/shape.h" #include "xla/shape_util.h" +#include "xla/xla.pb.h" +#include "xla/xla_data.pb.h" namespace xla { diff --git a/third_party/xla/xla/hlo/transforms/collectives/collective_permute_decomposer.h b/third_party/xla/xla/hlo/transforms/collectives/collective_permute_decomposer.h index 1e963f2b451133..ce05ce5a5a936b 100644 --- a/third_party/xla/xla/hlo/transforms/collectives/collective_permute_decomposer.h +++ b/third_party/xla/xla/hlo/transforms/collectives/collective_permute_decomposer.h @@ -23,6 +23,7 @@ limitations under the License. #include "absl/strings/string_view.h" #include "xla/hlo/ir/hlo_module.h" #include "xla/hlo/pass/hlo_pass_interface.h" +#include "xla/xla.pb.h" namespace xla { diff --git a/third_party/xla/xla/hlo/transforms/collectives/collective_permute_decomposer_test.cc b/third_party/xla/xla/hlo/transforms/collectives/collective_permute_decomposer_test.cc index c924facb525c83..40d713710554b2 100644 --- a/third_party/xla/xla/hlo/transforms/collectives/collective_permute_decomposer_test.cc +++ b/third_party/xla/xla/hlo/transforms/collectives/collective_permute_decomposer_test.cc @@ -32,6 +32,8 @@ limitations under the License. #include "xla/hlo/ir/hlo_opcode.h" #include "xla/hlo/testlib/hlo_hardware_independent_test_base.h" #include "xla/service/collective_ops_utils.h" +#include "xla/xla.pb.h" +#include "xla/xla_data.pb.h" namespace xla { namespace { diff --git a/third_party/xla/xla/hlo/transforms/collectives/reduce_scatter_combiner_test.cc b/third_party/xla/xla/hlo/transforms/collectives/reduce_scatter_combiner_test.cc index 20f5a515662506..1f5c7ef3666628 100644 --- a/third_party/xla/xla/hlo/transforms/collectives/reduce_scatter_combiner_test.cc +++ b/third_party/xla/xla/hlo/transforms/collectives/reduce_scatter_combiner_test.cc @@ -32,6 +32,7 @@ limitations under the License. #include "xla/hlo/ir/hlo_opcode.h" #include "xla/hlo/testlib/hlo_hardware_independent_test_base.h" #include "xla/hlo/utils/hlo_matchers.h" +#include "xla/xla_data.pb.h" namespace xla { namespace { diff --git a/third_party/xla/xla/hlo/transforms/collectives/reduce_scatter_decomposer.cc b/third_party/xla/xla/hlo/transforms/collectives/reduce_scatter_decomposer.cc index 09cdb0924c9e2e..57c6c2f20a3539 100644 --- a/third_party/xla/xla/hlo/transforms/collectives/reduce_scatter_decomposer.cc +++ b/third_party/xla/xla/hlo/transforms/collectives/reduce_scatter_decomposer.cc @@ -34,6 +34,7 @@ limitations under the License. #include "xla/hlo/ir/hlo_module.h" #include "xla/hlo/utils/hlo_query.h" #include "xla/service/collective_decomposer_utils.h" +#include "xla/xla_data.pb.h" namespace xla { diff --git a/third_party/xla/xla/hlo/transforms/collectives/reduce_scatter_decomposer_test.cc b/third_party/xla/xla/hlo/transforms/collectives/reduce_scatter_decomposer_test.cc index 41e6bdbf77dab4..0e6dc41f4a2deb 100644 --- a/third_party/xla/xla/hlo/transforms/collectives/reduce_scatter_decomposer_test.cc +++ b/third_party/xla/xla/hlo/transforms/collectives/reduce_scatter_decomposer_test.cc @@ -32,6 +32,7 @@ limitations under the License. #include "xla/literal.h" #include "xla/literal_util.h" #include "xla/shape.h" +#include "xla/xla_data.pb.h" namespace xla { namespace { 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 39f6046fe0549a..97edf369ef7ba6 100644 --- a/third_party/xla/xla/hlo/transforms/simplifiers/algebraic_simplifier.cc +++ b/third_party/xla/xla/hlo/transforms/simplifiers/algebraic_simplifier.cc @@ -8070,19 +8070,26 @@ absl::Status AlgebraicSimplifierVisitor::HandleDynamicSlice( } } - // ds(ds(x,id),inner_id) -> ds(x, id + inner_id) + // ds(ds(x, inner_id), id) -> + // ds(x, clamp(0, id, inner_size - outer_size) + + // clamp(0, inner_id, operand_size - inner_size)) if (operand->opcode() == HloOpcode::kDynamicSlice) { ABSL_RETURN_IF_ERROR(dynamic_slice->ReplaceOperandWithDifferentShape( 0, operand->mutable_operand(0))); for (int64_t i = 1; i < dynamic_slice->operand_count(); ++i) { HloInstruction* index = dynamic_slice->mutable_operand(i); + index = index->AddInstruction(HloInstruction::CreateTernary( + index->shape(), HloOpcode::kClamp, MakeScalarLike(index, 0), index, + MakeScalarLike(index, + operand->dynamic_slice_sizes()[i - 1] - + dynamic_slice->dynamic_slice_sizes()[i - 1]))); HloInstruction* inner_index = operand->mutable_operand(i); inner_index = inner_index->AddInstruction(HloInstruction::CreateTernary( inner_index->shape(), HloOpcode::kClamp, MakeScalarLike(inner_index, 0), inner_index, MakeScalarLike(inner_index, operand->operand(0)->shape().dimensions(i - 1) - - dynamic_slice->dynamic_slice_sizes()[i - 1]))); + operand->dynamic_slice_sizes()[i - 1]))); if (inner_index->shape().element_type() != index->shape().element_type()) { inner_index = inner_index->AddInstruction( 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 757244860a7629..b439bf6ed0d202 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 @@ -15063,6 +15063,30 @@ TEST_F(AlgebraicSimplifierTest, DoNotFoldTransposeIntoScatterWhenDisabled) { absl_testing::IsOkAndHolds(false)); } +TEST_F(AlgebraicSimplifierTest, DynamicSliceOfDynamicSlice) { + constexpr absl::string_view hlo_string = R"( + HloModule module + + ENTRY test { + operand = f32[10] parameter(0) + i = s32[] parameter(1) + j = s32[] parameter(2) + inner_ds = f32[4] dynamic-slice(operand, i), dynamic_slice_sizes={4} + ROOT outer_ds = f32[1] dynamic-slice(inner_ds, j), dynamic_slice_sizes={1} + } + )"; + ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); + AlgebraicSimplifier simplifier(default_options_); + ASSERT_THAT(simplifier.Run(module.get()), absl_testing::IsOkAndHolds(true)); + EXPECT_THAT(module->entry_computation()->root_instruction(), + GmockMatch(m::DynamicSlice( + m::Parameter(0), + m::Add(m::Clamp(m::ConstantScalar(0), m::Parameter(2), + m::ConstantScalar(3)), + m::Clamp(m::ConstantScalar(0), m::Parameter(1), + m::ConstantScalar(6)))))); +} + } // namespace } // namespace xla // Trivial comment to force rebuild diff --git a/third_party/xla/xla/hlo/transforms/simplifiers/hlo_constant_folding.cc b/third_party/xla/xla/hlo/transforms/simplifiers/hlo_constant_folding.cc index 003ef2e539c8a3..4f9f4e286bf1af 100644 --- a/third_party/xla/xla/hlo/transforms/simplifiers/hlo_constant_folding.cc +++ b/third_party/xla/xla/hlo/transforms/simplifiers/hlo_constant_folding.cc @@ -610,7 +610,8 @@ absl::StatusOr HloConstantFolding::RunImpl( // default case. This retains the behavior from before while loop support in // HloEvaluator and may be revised. auto evaluator = std::make_unique( - /*max_loop_iterations=*/options_.level == Level::kAggressive ? -1 : 0); + /*max_loop_iterations=*/options_.level == Level::kAggressive ? -1 : 0, + /*cache_call_computation_evals=*/true); // fast-path lets us e.g. use Eigen for matmuls. evaluator->set_use_fast_path(true); diff --git a/third_party/xla/xla/hlo/transforms/simplifiers/hlo_dce.cc b/third_party/xla/xla/hlo/transforms/simplifiers/hlo_dce.cc index c68f5cdf75aff8..5732a50d4a5456 100644 --- a/third_party/xla/xla/hlo/transforms/simplifiers/hlo_dce.cc +++ b/third_party/xla/xla/hlo/transforms/simplifiers/hlo_dce.cc @@ -228,7 +228,8 @@ bool CanRemoveInstruction( bool remove_cross_partition_collective_ops, const std::function(const HloComputation*)>& computation_callers) { - if (!instruction->IsDead() || HasDisableWhileLoopDceAttr(instruction)) { + if (instruction->parent() == nullptr || !instruction->IsDead() || + HasDisableWhileLoopDceAttr(instruction)) { return false; } @@ -250,18 +251,42 @@ bool CanRemoveInstruction( instruction->operand(0)->user_count() != 1)) { return false; } + auto has_dce_side_effect_attr = [](const HloInstruction* inst, + absl::string_view value) { + if (inst == nullptr) { + return false; + } + auto it = + inst->frontend_attributes().map().find(kDceSideEffectFrontendAttribute); + return it != inst->frontend_attributes().map().end() && it->second == value; + }; + // Reverse postorder visits async-done before async-start. If "false" is + // set on async-start, async-done must inspect async_chain_start(). + // otherwise only async-done will be removed, leaving a dangling async-start + if (has_dce_side_effect_attr(instruction, "false") || + has_dce_side_effect_attr(instruction->async_chain_start(), "false")) { + return false; + } if (instruction->HasSideEffect()) { - auto maybe_collective_op = DynCast(instruction); + auto maybe_collective_op = DynCast( + instruction->async_wrapped_instruction() + ? instruction->async_wrapped_instruction() + : instruction); bool allow_collective = remove_cross_partition_collective_ops && maybe_collective_op && !maybe_collective_op->constrain_layout(); bool allow_while = IsRemovableWhile(instruction, remove_cross_partition_collective_ops); - bool allow_custom_call = instruction->IsCustomCall("tpu_custom_call") && - instruction->frontend_attributes().map().contains( - kDceSideEffectFrontendAttribute) && - instruction->frontend_attributes().map().at( - kDceSideEffectFrontendAttribute) == "true"; + // Reverse postorder visits async-done before async-start. If "true" is + // set on async-start, async-done must inspect + // async_chain_start(). otherwise async-done is kept + // alive, which keeps async-start used (user_count > 0) and prevents + // either from being removed. + bool allow_custom_call = + (instruction->IsCustomCall("tpu_custom_call") || + instruction->IsAsynchronous()) && + (has_dce_side_effect_attr(instruction, "true") || + has_dce_side_effect_attr(instruction->async_chain_start(), "true")); if (!allow_collective && !allow_while && !allow_custom_call) { return false; } @@ -274,27 +299,25 @@ absl::StatusOr RemoveDeadRoots( const std::function(const HloComputation*)>& computation_callers) { bool changed = false; - std::vector dead_roots; - for (auto* instruction : computation->instructions()) { + auto post_order = computation->MakeInstructionPostOrder(); + for (auto it = post_order.rbegin(); it != post_order.rend(); ++it) { + HloInstruction* instruction = *it; if (!CanRemoveInstruction(instruction, remove_cross_partition_collective_ops, computation_callers)) { continue; } - dead_roots.push_back(instruction); - } - - for (HloInstruction* dead_root : dead_roots) { - VLOG(1) << "Removing dead root " << dead_root->ToString() + VLOG(1) << "Removing dead root " << instruction->ToString() << " and its unused operands"; ABSL_RETURN_IF_ERROR(computation->RemoveInstructionAndUnusedOperands( - dead_root, /*cleanup=*/std::nullopt, + instruction, /*cleanup=*/std::nullopt, /*ignore_control_dependencies=*/false, /*computation_callers=*/computation_callers)); changed = true; } return changed; } + absl::Status RemoveDeadParametersFromEntryComputationLayout( HloModule* module, std::vector& dead_parameter_indexes) { if (dead_parameter_indexes.empty()) { diff --git a/third_party/xla/xla/hlo/transforms/simplifiers/hlo_dce_test.cc b/third_party/xla/xla/hlo/transforms/simplifiers/hlo_dce_test.cc index 3646f62f1864dd..fe8911e33aed6c 100644 --- a/third_party/xla/xla/hlo/transforms/simplifiers/hlo_dce_test.cc +++ b/third_party/xla/xla/hlo/transforms/simplifiers/hlo_dce_test.cc @@ -1227,5 +1227,216 @@ ENTRY main { EXPECT_FALSE(changed); } +TEST_F(HloDceTest, DceAlreadyRemovedInstructionInReversePostOrder) { + constexpr absl::string_view kHloString = R"hlo( +HloModule test_already_removed_guard + +ENTRY main { + p0 = f32[8] parameter(0) + dead_pure_producer = f32[8] negate(p0) + dead_consumer = f32[8] custom-call(dead_pure_producer), custom_call_target="tpu_custom_call", custom_call_has_side_effect=true, frontend_attributes={xla_allow_dce_side_effecting_op="true"} + ROOT root = f32[8] copy(p0) +} +)hlo"; + ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(kHloString)); + HloDCE dce; + // Removing `dead_consumer` cascades removal to its pure operand + // `dead_pure_producer` (detaching it from parent). Subsequent reverse + // postorder iteration must skip `dead_pure_producer` rather than crashing. + ASSERT_OK_AND_ASSIGN(bool changed, dce.Run(module.get())); + EXPECT_TRUE(changed); + EXPECT_EQ(FindInstruction(module.get(), "dead_consumer"), nullptr); + EXPECT_EQ(FindInstruction(module.get(), "dead_pure_producer"), nullptr); +} + +TEST_F(HloDceTest, PreserveInstructionWithFrontendAttributeFalse) { + constexpr absl::string_view kHloString = R"hlo( +HloModule test_preserve_self_attr_false + +ENTRY main { + p0 = f32[8] parameter(0) + live_ag = f32[16] all-gather(p0), dimensions={0}, replica_groups={{0,1}}, channel_id=1, frontend_attributes={xla_allow_dce_side_effecting_op="false"} + ROOT root = f32[8] copy(p0) +} +)hlo"; + ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(kHloString)); + HloDCE dce(/*remove_cross_partition_collective_ops=*/true); + ASSERT_OK_AND_ASSIGN(bool changed, dce.Run(module.get())); + EXPECT_FALSE(changed); + EXPECT_NE(FindInstruction(module.get(), "live_ag"), nullptr); +} + +TEST_F(HloDceTest, PreserveAsyncPairWhenStartHasFrontendAttributeFalse) { + constexpr absl::string_view kHloString = R"hlo( +HloModule test_preserve_async_pair_start_attr_false + +async_comp { + p0 = f32[8] parameter(0) + ROOT ag = f32[16] all-gather(p0), dimensions={0}, replica_groups={{0,1}}, channel_id=1 +} + +ENTRY main { + p0 = f32[8] parameter(0) + live_start = ((f32[8]), f32[16]) async-start(p0), calls=async_comp, frontend_attributes={xla_allow_dce_side_effecting_op="false"} + live_done = f32[16] async-done(live_start), calls=async_comp + ROOT root = f32[8] copy(p0) +} +)hlo"; + ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(kHloString)); + // `live_done` must check its `async-start` operand's attribute so `live_done` + // is not removed and `live_start` is not orphaned. + HloDCE dce(/*remove_cross_partition_collective_ops=*/true); + ASSERT_OK_AND_ASSIGN(bool changed, dce.Run(module.get())); + EXPECT_FALSE(changed); + EXPECT_NE(FindInstruction(module.get(), "live_done"), nullptr); + EXPECT_NE(FindInstruction(module.get(), "live_start"), nullptr); +} + +TEST_F(HloDceTest, DceAsyncWrappedCollective) { + constexpr absl::string_view kUnconstrainedHlo = R"hlo( +HloModule test_async_wrapped_collective_unconstrained + +async_comp { + p0 = f32[8] parameter(0) + ROOT ag = f32[16] all-gather(p0), dimensions={0}, replica_groups={{0,1}}, channel_id=1 +} + +ENTRY main { + p0 = f32[8] parameter(0) + dead_start = ((f32[8]), f32[16]) async-start(p0), calls=async_comp + dead_done = f32[16] async-done(dead_start), calls=async_comp + ROOT root = f32[8] copy(p0) +} +)hlo"; + // Preserved when remove_cross_partition_collective_ops=false. + { + ASSERT_OK_AND_ASSIGN(auto module, + ParseAndReturnVerifiedModule(kUnconstrainedHlo)); + HloDCE dce(/*remove_cross_partition_collective_ops=*/false); + ASSERT_OK_AND_ASSIGN(bool changed, dce.Run(module.get())); + EXPECT_FALSE(changed); + EXPECT_NE(FindInstruction(module.get(), "dead_done"), nullptr); + EXPECT_NE(FindInstruction(module.get(), "dead_start"), nullptr); + } + + // Preserved when constrain_layout=true on the wrapped collective. + constexpr absl::string_view kConstrainedHlo = R"hlo( +HloModule test_async_wrapped_collective_constrained + +async_comp { + p0 = f32[8] parameter(0) + ROOT ag = f32[16] all-gather(p0), dimensions={0}, replica_groups={{0,1}}, channel_id=1, constrain_layout=true +} + +ENTRY main { + p0 = f32[8] parameter(0) + dead_start = ((f32[8]), f32[16]) async-start(p0), calls=async_comp + dead_done = f32[16] async-done(dead_start), calls=async_comp + ROOT root = f32[8] copy(p0) +} +)hlo"; + { + ASSERT_OK_AND_ASSIGN(auto module, + ParseAndReturnVerifiedModule(kConstrainedHlo)); + HloDCE dce(/*remove_cross_partition_collective_ops=*/true); + ASSERT_OK_AND_ASSIGN(bool changed, dce.Run(module.get())); + EXPECT_FALSE(changed); + EXPECT_NE(FindInstruction(module.get(), "dead_done"), nullptr); + EXPECT_NE(FindInstruction(module.get(), "dead_start"), nullptr); + } + + // Eliminated when unconstrained and + // remove_cross_partition_collective_ops=true. + { + ASSERT_OK_AND_ASSIGN(auto module, + ParseAndReturnVerifiedModule(kUnconstrainedHlo)); + HloDCE dce(/*remove_cross_partition_collective_ops=*/true); + ASSERT_OK_AND_ASSIGN(bool changed, dce.Run(module.get())); + EXPECT_TRUE(changed); + EXPECT_EQ(FindInstruction(module.get(), "dead_done"), nullptr); + EXPECT_EQ(FindInstruction(module.get(), "dead_start"), nullptr); + } +} + +TEST_F(HloDceTest, DceSelfAttributeSideEffectingCustomCallAndAsync) { + constexpr absl::string_view kHloString = R"hlo( +HloModule test_self_attr_true_custom_call_and_async + +async_comp { + p0 = f32[8] parameter(0) + ROOT inner_cc = f32[8] custom-call(p0), custom_call_target="tpu_custom_call", custom_call_has_side_effect=true +} + +ENTRY main { + p0 = f32[8] parameter(0) + dead_cc = f32[8] custom-call(p0), custom_call_target="tpu_custom_call", custom_call_has_side_effect=true, frontend_attributes={xla_allow_dce_side_effecting_op="true"} + dead_start = ((f32[8]), f32[8]) async-start(p0), calls=async_comp, frontend_attributes={xla_allow_dce_side_effecting_op="true"} + dead_done = f32[8] async-done(dead_start), calls=async_comp, frontend_attributes={xla_allow_dce_side_effecting_op="true"} + ROOT root = f32[8] copy(p0) +} +)hlo"; + ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(kHloString)); + HloDCE dce; + ASSERT_OK_AND_ASSIGN(bool changed, dce.Run(module.get())); + EXPECT_TRUE(changed); + EXPECT_EQ(FindInstruction(module.get(), "dead_cc"), nullptr); + EXPECT_EQ(FindInstruction(module.get(), "dead_done"), nullptr); + EXPECT_EQ(FindInstruction(module.get(), "dead_start"), nullptr); +} + +TEST_F(HloDceTest, DcePairedAsyncDoneInheritingStartAttributeTrue) { + constexpr absl::string_view kHloString = R"hlo( +HloModule test_paired_async_done_inheriting_start_attr_true + +async_comp { + p0 = f32[8] parameter(0) + ROOT inner_cc = f32[8] custom-call(p0), custom_call_target="tpu_custom_call", custom_call_has_side_effect=true +} + +ENTRY main { + p0 = f32[8] parameter(0) + dead_start = ((f32[8]), f32[8]) async-start(p0), calls=async_comp, frontend_attributes={xla_allow_dce_side_effecting_op="true"} + dead_done = f32[8] async-done(dead_start), calls=async_comp + ROOT root = f32[8] copy(p0) +} +)hlo"; + ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(kHloString)); + HloDCE dce; + ASSERT_OK_AND_ASSIGN(bool changed, dce.Run(module.get())); + EXPECT_TRUE(changed); + EXPECT_EQ(FindInstruction(module.get(), "dead_done"), nullptr); + EXPECT_EQ(FindInstruction(module.get(), "dead_start"), nullptr); +} + +TEST_F(HloDceTest, DceSinglePassReversePostOrderSideEffectingChain) { + constexpr absl::string_view kHloString = R"hlo( +HloModule test_single_pass_reverse_postorder_chain + +async_comp { + p0 = f32[8] parameter(0) + ROOT inner_cc = f32[8] custom-call(p0), custom_call_target="tpu_custom_call", custom_call_has_side_effect=true +} + +ENTRY main { + p0 = f32[8] parameter(0) + producer_cc = f32[8] custom-call(p0), custom_call_target="tpu_custom_call", custom_call_has_side_effect=true, frontend_attributes={xla_allow_dce_side_effecting_op="true"} + dead_start = ((f32[8]), f32[8]) async-start(producer_cc), calls=async_comp, frontend_attributes={xla_allow_dce_side_effecting_op="true"} + dead_done = f32[8] async-done(dead_start), calls=async_comp, frontend_attributes={xla_allow_dce_side_effecting_op="true"} + ROOT root = f32[8] copy(p0) +} +)hlo"; + ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(kHloString)); + HloDCE dce; + // Reverse postorder visits consumers before operands (`dead_done` -> + // `dead_start` -> `producer_cc`), allowing a single pass to eliminate the + // entire chain even though `RemoveInstructionAndUnusedOperands` does not + // recurse into side-effecting operands. + ASSERT_OK_AND_ASSIGN(bool changed, dce.Run(module.get())); + EXPECT_TRUE(changed); + EXPECT_EQ(FindInstruction(module.get(), "dead_done"), nullptr); + EXPECT_EQ(FindInstruction(module.get(), "dead_start"), nullptr); + EXPECT_EQ(FindInstruction(module.get(), "producer_cc"), nullptr); +} + } // namespace } // namespace xla diff --git a/third_party/xla/xla/hlo/translate/hlo_to_mhlo/hlo_function_importer.cc b/third_party/xla/xla/hlo/translate/hlo_to_mhlo/hlo_function_importer.cc index 0f11cbe3d2f722..b180a3292b3748 100644 --- a/third_party/xla/xla/hlo/translate/hlo_to_mhlo/hlo_function_importer.cc +++ b/third_party/xla/xla/hlo/translate/hlo_to_mhlo/hlo_function_importer.cc @@ -1294,10 +1294,10 @@ absl::StatusOr HloFunctionImporter::ImportInstructionImpl( case HloOpcode::kCompare: { auto compare = Cast(instruction); attributes.push_back(ConvertComparisonDirection(compare->direction())); - auto default_type = Comparison::DefaultComparisonType( + auto default_order = Comparison::DefaultOrdering( compare->operand(0)->shape().element_type()); - if (compare->type() != default_type) { - attributes.push_back(ConvertComparisonType(compare->type())); + if (compare->order() != default_order) { + attributes.push_back(ConvertComparisonOrder(compare->order())); } return mlir::stablehlo::CompareOp::create(*func_builder, loc, result_type, operands, attributes) @@ -2592,14 +2592,20 @@ mlir::NamedAttribute HloFunctionImporter::ConvertComparisonDirection( .value())); } -mlir::NamedAttribute HloFunctionImporter::ConvertComparisonType( - Comparison::Type type) { +mlir::NamedAttribute HloFunctionImporter::ConvertComparisonOrder( + ComparisonOrder order) { + mlir::stablehlo::ComparisonType type; + switch (order) { + case ComparisonOrder::kPartial: + type = mlir::stablehlo::ComparisonType::FLOAT; + break; + case ComparisonOrder::kTotal: + type = mlir::stablehlo::ComparisonType::TOTALORDER; + break; + } return builder_->getNamedAttr( "compare_type", - mlir::stablehlo::ComparisonTypeAttr::get( - builder_->getContext(), - mlir::stablehlo::symbolizeComparisonType(ComparisonTypeToString(type)) - .value())); + mlir::stablehlo::ComparisonTypeAttr::get(builder_->getContext(), type)); } mlir::DenseIntElementsAttr HloFunctionImporter::ConvertDimensions( diff --git a/third_party/xla/xla/hlo/translate/hlo_to_mhlo/hlo_function_importer.h b/third_party/xla/xla/hlo/translate/hlo_to_mhlo/hlo_function_importer.h index 7b34a32e720f59..1116cff866dbb6 100644 --- a/third_party/xla/xla/hlo/translate/hlo_to_mhlo/hlo_function_importer.h +++ b/third_party/xla/xla/hlo/translate/hlo_to_mhlo/hlo_function_importer.h @@ -201,8 +201,8 @@ class HloFunctionImporter { mlir::NamedAttribute ConvertComparisonDirection( ComparisonDirection direction); - // Converts an XLA Comparison::Type to the corresponding MLIR attribute. - mlir::NamedAttribute ConvertComparisonType(Comparison::Type type); + // Converts an XLA ComparisonOrder to the corresponding MLIR attribute. + mlir::NamedAttribute ConvertComparisonOrder(ComparisonOrder order); // Converts an XLA CustomCallSchedule to the corresponding MLIR attribute. mlir::NamedAttribute ConvertCustomCallSchedule(CustomCallSchedule schedule); diff --git a/third_party/xla/xla/hlo/translate/mhlo_to_hlo/BUILD b/third_party/xla/xla/hlo/translate/mhlo_to_hlo/BUILD index 627898aa1fc3d3..9d83f381c4bab8 100644 --- a/third_party/xla/xla/hlo/translate/mhlo_to_hlo/BUILD +++ b/third_party/xla/xla/hlo/translate/mhlo_to_hlo/BUILD @@ -261,6 +261,7 @@ cc_library( "//xla/tsl/platform:statusor", "//xla/tsl/platform:types", "@com_google_absl//absl/container:flat_hash_map", + "@com_google_absl//absl/log", "@com_google_absl//absl/log:check", "@com_google_absl//absl/status", "@com_google_absl//absl/status:status_macros", 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 1db8afcf410d30..8d225e18d5711f 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 @@ -28,6 +28,7 @@ limitations under the License. #include "mhlo/transforms/passes.h" #include "absl/container/flat_hash_map.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" @@ -2566,10 +2567,21 @@ mlir::LogicalResult ExportXlaOp(mlir::stablehlo::CompareOp op, xla::XlaOp xla_result; if (type_attr && type_attr.getValue() != mlir::stablehlo::ComparisonType::NOTYPE) { - auto type = xla::StringToComparisonType( - stringifyComparisonType(type_attr.getValue()).str()) - .value(); - xla_result = xla::Compare(lhs, rhs, /*broadcast_dimensions=*/{}, dir, type); + xla::ComparisonOrder order; + switch (type_attr.getValue()) { + case mlir::stablehlo::ComparisonType::FLOAT: + order = xla::ComparisonOrder::kPartial; + break; + case mlir::stablehlo::ComparisonType::TOTALORDER: + case mlir::stablehlo::ComparisonType::SIGNED: + case mlir::stablehlo::ComparisonType::UNSIGNED: + order = xla::ComparisonOrder::kTotal; + break; + case mlir::stablehlo::ComparisonType::NOTYPE: + LOG(FATAL) << "Unreachable"; + } + xla_result = + xla::Compare(lhs, rhs, /*broadcast_dimensions=*/{}, dir, order); } else { xla_result = xla::Compare(lhs, rhs, dir); } 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 f3027bb4cff609..b2dfdaca4e8661 100644 --- a/third_party/xla/xla/hlo/utils/hlo_live_range.cc +++ b/third_party/xla/xla/hlo/utils/hlo_live_range.cc @@ -368,11 +368,29 @@ void HloLiveRange::CalculateBufferStartEndMap() { instruction.IsRoot() ? computation_span_times_[computation].end : entry.second; - // If the instruction is in an asynchronous context, extend the live range - // until the end of the async-done instruction. + // If the instruction is in an asynchronous context, adjust its live range + // to cover the async window precisely: + // - end_time is extended to the async-done instruction so the buffer + // remains live for the full async duration. + // - start_time is tightened to the first-fully-bound instruction + // (async-start for standard chains, async-update for late-binding + // chains), because FlattenSchedule inlines the async computation + // immediately before that instruction and the inner buffer is not live + // before it fires. auto async_context_it = computations_in_async_context_.find(computation); if (async_context_it != computations_in_async_context_.end()) { const HloComputation* async_context = async_context_it->second; + // async_context can have multiple callers when the same async-wrapped + // computation is shared by several async-start/update instructions + // (e.g. one outside and one inside a while loop). FlattenSchedule only + // ever inlines the computation once, at whichever caller the schedule + // walk reaches first, so we take the minimum schedule time across all + // first-fully-bound callers below. This is deterministic (unlike + // picking an arbitrary caller from caller_instructions(), which is + // returned in no particular order) and never overshoots the schedule + // position the instructions were actually flattened at, which would + // otherwise push start_time past definition_end_time. + std::optional tightened_start_time; for (const HloInstruction* caller : async_context->caller_instructions()) { if (caller->IsAsynchronous()) { @@ -397,10 +415,39 @@ void HloLiveRange::CalculateBufferStartEndMap() { definition_end_time = std::max( definition_end_time, computation_span_times_[computation].end); } + // Track the earliest first-fully-bound caller's schedule time. + // FlattenSchedule inlines the async computation's instructions + // immediately before the first-fully-bound instruction (async-start + // for standard chains, async-update for late-binding chains) that + // actually triggers the flattening. The inner buffer is not live + // before that point. + absl::StatusOr is_first_fully_bound = + hlo_instruction_utils::async::IsFirstFullyBound(caller); + if (is_first_fully_bound.ok() && *is_first_fully_bound) { + auto first_bound_it = instruction_schedule_.find(caller); + if (first_bound_it != instruction_schedule_.end()) { + tightened_start_time = + tightened_start_time.has_value() + ? std::min(*tightened_start_time, first_bound_it->second) + : first_bound_it->second; + } + } } } + if (tightened_start_time.has_value()) { + // Cap by definition_end_time (computed above): for computations + // reached transitively through a kCall inside the async-wrapped + // computation, FlattenSchedule inlines their bodies even earlier + // than the async-wrapped computation's own instructions, so a + // caller's schedule time can land after such an instruction's own + // (already-correct) end. Capping avoids pushing start past end + // while still tightening whenever it's safe to do so. + start_time = std::min(*tightened_start_time, definition_end_time); + } VLOG(2) << "Setting the definition end time for op in async context: " << definition_end_time; + VLOG(2) << "Setting the definition start time for op in async context: " + << start_time; } for (const HloValue* value : 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 e3bb28c339bff7..c77a682e24a9d7 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 @@ -1202,5 +1202,217 @@ ENTRY entry { EXPECT_FALSE(root_buffers.empty()); EXPECT_TRUE(HloLiveRange::BufferLivesOut(*root_buffers[0], *aa, entry)); } + +TEST_F(HloLiveRangeTest, AsyncInnerBufferStartExtendedToOuterStart) { + const std::string hlo_string = R"hlo( + HloModule AsyncInnerBufferStart, is_scheduled=true + + %async_wrapped (p: f32[4]) -> f32[4] { + %p = f32[4] parameter(0) + ROOT %inner_op = f32[4] negate(%p) + } + + ENTRY %main (a: f32[4]) -> f32[4] { + %a = f32[4] parameter(0) + %outer_op = f32[4] negate(%a) + %async-start = ((f32[4]), f32[4], u32[]) async-start(%outer_op), + calls=%async_wrapped + %async-done = f32[4] async-done(%async-start) + ROOT %result = f32[4] add(%outer_op, %async-done) + } + )hlo"; + + ASSERT_OK_AND_ASSIGN(module_, ParseAndReturnVerifiedModule(hlo_string)); + ASSERT_OK_AND_ASSIGN(alias_analysis_, + HloAliasAnalysis::Run(module_.get(), &alias_info_)); + ASSERT_OK_AND_ASSIGN(hlo_live_range_, + HloLiveRange::Run(module_->schedule(), *alias_analysis_, + module_->entry_computation())); + CheckSchedule(); + + HloComputation* async_wrapped = + module_->GetComputationWithName("async_wrapped"); + ASSERT_NE(async_wrapped, nullptr); + const HloInstruction* inner_op = + async_wrapped->GetInstructionWithName("inner_op"); + ASSERT_NE(inner_op, nullptr); + + const HloInstruction* async_done = + module_->entry_computation()->GetInstructionWithName("async-done"); + ASSERT_NE(async_done, nullptr); + const HloInstruction* result = + module_->entry_computation()->GetInstructionWithName("result"); + ASSERT_NE(result, nullptr); + + const HloInstruction* async_start = + module_->entry_computation()->GetInstructionWithName("async-start"); + ASSERT_NE(async_start, nullptr); + + auto inner_range = LiveRangeAt(inner_op); + auto async_start_time = + hlo_live_range_->instruction_schedule().at(async_start); + auto async_done_time = hlo_live_range_->instruction_schedule().at(async_done); + auto result_time = hlo_live_range_->instruction_schedule().at(result); + + EXPECT_EQ(inner_range.start, async_start_time); + EXPECT_GT(result_time, async_done_time); + EXPECT_EQ(inner_range.end, result_time); +} + +TEST_F(HloLiveRangeTest, AsyncInnerParameterStartExtendedToOuterStart) { + const std::string hlo_string = R"hlo( + HloModule AsyncInnerParamStart, is_scheduled=true + + %async_wrapped (p: f32[4]) -> f32[4] { + %p = f32[4] parameter(0) + ROOT %inner_op = f32[4] negate(%p) + } + + ENTRY %main (a: f32[4]) -> f32[4] { + %a = f32[4] parameter(0) + %async-start = ((f32[4]), f32[4], u32[]) async-start(%a), + calls=%async_wrapped + %async-done = f32[4] async-done(%async-start) + ROOT %result = f32[4] negate(%async-done) + } + )hlo"; + + ASSERT_OK_AND_ASSIGN(module_, ParseAndReturnVerifiedModule(hlo_string)); + ASSERT_OK_AND_ASSIGN(alias_analysis_, + HloAliasAnalysis::Run(module_.get(), &alias_info_)); + ASSERT_OK_AND_ASSIGN(hlo_live_range_, + HloLiveRange::Run(module_->schedule(), *alias_analysis_, + module_->entry_computation())); + CheckSchedule(); + + HloComputation* async_wrapped = + module_->GetComputationWithName("async_wrapped"); + ASSERT_NE(async_wrapped, nullptr); + + const HloInstruction* async_start = + module_->entry_computation()->GetInstructionWithName("async-start"); + ASSERT_NE(async_start, nullptr); + const HloInstruction* inner_param = + async_wrapped->GetInstructionWithName("p"); + const HloInstruction* inner_op = + async_wrapped->GetInstructionWithName("inner_op"); + ASSERT_NE(inner_param, nullptr); + ASSERT_NE(inner_op, nullptr); + + auto async_start_time = + hlo_live_range_->instruction_schedule().at(async_start); + + EXPECT_EQ(LiveRangeAt(inner_param).start, 0); + EXPECT_EQ(LiveRangeAt(inner_op).start, async_start_time); +} + +TEST_F(HloLiveRangeTest, AsyncInnerBufferStartIsAsyncStartTime) { + const std::string hlo_string = R"hlo( + HloModule AsyncInnerOverlap, is_scheduled=true + + %async_wrapped (p: f32[4]) -> f32[4] { + %p = f32[4] parameter(0) + ROOT %inner_op = f32[4] negate(%p) + } + + ENTRY %main (a: f32[4]) -> f32[4] { + %a = f32[4] parameter(0) + %temp = f32[4] negate(%a) + %outer = f32[4] negate(%temp) + %async-start = ((f32[4]), f32[4], u32[]) async-start(%outer), + calls=%async_wrapped + %async-done = f32[4] async-done(%async-start) + ROOT %result = f32[4] negate(%async-done) + } + )hlo"; + + ASSERT_OK_AND_ASSIGN(module_, ParseAndReturnVerifiedModule(hlo_string)); + ASSERT_OK_AND_ASSIGN(alias_analysis_, + HloAliasAnalysis::Run(module_.get(), &alias_info_)); + ASSERT_OK_AND_ASSIGN(hlo_live_range_, + HloLiveRange::Run(module_->schedule(), *alias_analysis_, + module_->entry_computation())); + CheckSchedule(); + + const HloInstruction* async_start = + module_->entry_computation()->GetInstructionWithName("async-start"); + const HloInstruction* temp = + module_->entry_computation()->GetInstructionWithName("temp"); + HloComputation* async_wrapped = + module_->GetComputationWithName("async_wrapped"); + ASSERT_NE(async_start, nullptr); + ASSERT_NE(temp, nullptr); + ASSERT_NE(async_wrapped, nullptr); + + const HloInstruction* inner_op = + async_wrapped->GetInstructionWithName("inner_op"); + ASSERT_NE(inner_op, nullptr); + + auto async_start_time = + hlo_live_range_->instruction_schedule().at(async_start); + auto temp_range = LiveRangeAt(temp); + auto inner_range = LiveRangeAt(inner_op); + + EXPECT_EQ(inner_range.start, async_start_time); + EXPECT_GT(inner_range.start, temp_range.end); +} + +TEST_F(HloLiveRangeTest, AsyncComputationSharedByMultipleCallersUsesMinStart) { + const std::string hlo_string = R"hlo( + HloModule m, is_scheduled=true + + %nested (q: f32[4]) -> f32[4] { + %q = f32[4] parameter(0) + %t = f32[4] negate(%q) + ROOT %n = f32[4] negate(%t) + } + + %async_computation (p: f32[4]) -> f32[4] { + %p = f32[4] parameter(0) + ROOT %c = f32[4] call(%p), to_apply=%nested + } + + %while_cond (wp: ((f32[4]), f32[4], u32[])) -> pred[] { + %wp = ((f32[4]), f32[4], u32[]) parameter(0) + ROOT %cond = pred[] constant(true) + } + + %while_body (bp: ((f32[4]), f32[4], u32[])) -> ((f32[4]), f32[4], u32[]) { + %bp = ((f32[4]), f32[4], u32[]) parameter(0) + %g = f32[4] get-tuple-element(%bp), index=1 + ROOT %s = ((f32[4]), f32[4], u32[]) async-start(%g), calls=%async_computation + } + + ENTRY %main (a: f32[4]) -> ((f32[4]), f32[4], u32[]) { + %a = f32[4] parameter(0) + %init = ((f32[4]), f32[4], u32[]) async-start(%a), calls=%async_computation + ROOT %loop = ((f32[4]), f32[4], u32[]) while(%init), condition=%while_cond, body=%while_body + } + )hlo"; + + ASSERT_OK_AND_ASSIGN(module_, ParseAndReturnVerifiedModule(hlo_string)); + ASSERT_OK_AND_ASSIGN(alias_analysis_, + HloAliasAnalysis::Run(module_.get(), &alias_info_)); + ASSERT_OK_AND_ASSIGN(hlo_live_range_, + HloLiveRange::Run(module_->schedule(), *alias_analysis_, + module_->entry_computation())); + CheckSchedule(); + + HloComputation* nested = module_->GetComputationWithName("nested"); + ASSERT_NE(nested, nullptr); + const HloInstruction* t = nested->GetInstructionWithName("t"); + ASSERT_NE(t, nullptr); + + const HloInstruction* init = + module_->entry_computation()->GetInstructionWithName("init"); + ASSERT_NE(init, nullptr); + + auto init_time = hlo_live_range_->instruction_schedule().at(init); + auto t_range = LiveRangeAt(t); + + EXPECT_LE(t_range.start, t_range.end); + EXPECT_LT(t_range.start, init_time); +} + } // namespace } // namespace xla diff --git a/third_party/xla/xla/lit.bzl b/third_party/xla/xla/lit.bzl index 415c8525921845..d9fa7784bb15c7 100644 --- a/third_party/xla/xla/lit.bzl +++ b/third_party/xla/xla/lit.bzl @@ -223,7 +223,7 @@ def lit_test_suite_for_gpus( # If there are kwargs that need to be passed to only some of the generated # rules, they should be extracted into separate named arguments. - rocm_gpus = ["mi200", "mi350", "gfx1250"] + rocm_gpus = ["mi200", "mi300", "mi350", "mi450"] for gpu in gpus: is_rocm = gpu in rocm_gpus diff --git a/third_party/xla/xla/mosaic/dialect/tpu/tpu_ops.cc b/third_party/xla/xla/mosaic/dialect/tpu/tpu_ops.cc index 70fbf7739effb6..89a0171789b02c 100644 --- a/third_party/xla/xla/mosaic/dialect/tpu/tpu_ops.cc +++ b/third_party/xla/xla/mosaic/dialect/tpu/tpu_ops.cc @@ -1214,6 +1214,31 @@ LogicalResult verifyLoadOp(Op op) { return success(); } +LogicalResult ExpandLoadVregOp::verify() { + MemRefType ref_ty = getBase().getType(); + if (ref_ty.getMemorySpace() && !HasMemorySpace(ref_ty, MemorySpace::kVmem)) { + return emitOpError("Expected base memref to be in VMEM."); + } + VectorType result_ty = getResult().getType(); + if (result_ty.getElementType() != ref_ty.getElementType()) { + return emitOpError("Expected base and result element type to match"); + } + if (llvm::size(getIndices()) != ref_ty.getRank()) { + return emitOpError("Expected ") << ref_ty.getRank() << " indices."; + } + // Note: We deliberately do not go through verifyLoadOp, which rejects masked + // loads of non-32-bit element types. The mask here is mandatory and narrow + // element types are allowed. They are rejected later, in the vector layout + // passes. + VectorType mask_ty = getMask().getType(); + if (result_ty.getShape()[0] != mask_ty.getShape()[0]) { + return emitOpError( + "Expected result dimension 0 to match mask dimension 0: ") + << result_ty.getShape()[0] << " vs " << mask_ty.getShape()[0] << "."; + } + return success(); +} + LogicalResult VectorLoadOp::verify() { const MemRefType ref_ty = getBase().getType(); if (llvm::size(getIndices()) != ref_ty.getRank()) { diff --git a/third_party/xla/xla/mosaic/dialect/tpu/tpu_ops.td b/third_party/xla/xla/mosaic/dialect/tpu/tpu_ops.td index 3d050637762cd5..084deae550db80 100644 --- a/third_party/xla/xla/mosaic/dialect/tpu/tpu_ops.td +++ b/third_party/xla/xla/mosaic/dialect/tpu/tpu_ops.td @@ -377,6 +377,42 @@ def TPU_VectorCompressStoreOp : TPU_Op<"vector_compress_store", [DefaultMemWrite let hasVerifier = 1; } +def TPU_ExpandLoadVregOp : TPU_Op<"expand_load_vreg", [DefaultMemRead]> { + let summary = "Performs a masked, expanded load of a single vreg."; + let description = [{ + `tpu.expand_load_vreg` performs a masked, expanded load of a single vreg + from memory starting at the address given by `base[indices]`. It is the + physical (vreg-level) counterpart of `tpu.vector_expand_load`. + + Masked-out lanes do not consume an address: the enabled lanes are filled + from contiguous, sequential locations in memory starting at + `base[indices]`. The n-th enabled lane reads `base[indices] + n`. + + For example, loading with mask [true, false, true, false] yields + [`base[indices]`, , `base[indices] + 1`, ]. + + Elements in masked-out lanes are undefined. + + base : Memref to load from. + indices: Scalar indices into `base`, identifying the position of the first + loaded element. The length must be the same as the rank of `base`. + mask : 1-D lane mask. + result : The loaded vreg. For packed (sub-32-bit) types, the vreg is 2-D + with shape `(lane_count, packing)`, and a lane loads all of its + `packing` subelements together. + }]; + let arguments = (ins + AnyMemRef:$base, + Variadic:$indices, + VectorOfRankAndType<[1], [I1]>:$mask + ); + let results = (outs TPU_Vreg:$result); + let assemblyFormat = [{ + $base `[` $indices `]` `,` $mask attr-dict `:` type($base) `,` type($result) `,` type($mask) + }]; + let hasVerifier = 1; +} + // tpu.vector_load loads a vector from memory into a register. // // base : Memref to load from. diff --git a/third_party/xla/xla/mosaic/dialect/tpu/tpu_ops_verification_test.cc b/third_party/xla/xla/mosaic/dialect/tpu/tpu_ops_verification_test.cc index c6b31587e3bd02..2fa8f8df4ecae9 100644 --- a/third_party/xla/xla/mosaic/dialect/tpu/tpu_ops_verification_test.cc +++ b/third_party/xla/xla/mosaic/dialect/tpu/tpu_ops_verification_test.cc @@ -377,7 +377,7 @@ TEST_F(TpuOpsVerificationTest, CompressStoreVregValid) { TEST_F(TpuOpsVerificationTest, CompressStoreVregValid2D) { auto c0 = Create(0); Value memref = AllocaI32({8, 2}, MemorySpace::kVmem); - Value vector_to_store = ConstantI32Vector(/*shape=*/{8, 2}, /*values=*/{1}); + Value vector_to_store = ConstantI32Vector(/*shape=*/{8}, /*values=*/{1}); Value mask = ConstantI1Vector( /*shape=*/{8}, /*values=*/{true, false, true, false, true, false, true, false}); @@ -452,6 +452,90 @@ TEST_F(TpuOpsVerificationTest, CompressStoreVregMismatchedMaskShape) { "Expected valueToStore dimension 0 to match mask dimension 0"))); } +TEST_F(TpuOpsVerificationTest, ExpandLoadVregValid) { + auto c0 = Create(0); + Value memref = AllocaI32({8}, MemorySpace::kVmem); + Value mask = ConstantI1Vector( + /*shape=*/{8}, + /*values=*/{true, false, true, false, true, false, true, false}); + auto el = Create( + /*result=*/VectorType::get({8}, i32()), + /*base=*/memref, + /*indices=*/ValueRange{c0}, + /*mask=*/mask); + + ASSERT_OK(VerifyOp(el)); +} + +// Packed types use a (lane_count, packing) vreg with a lane mask, so only +// dimension 0 has to match the mask. +TEST_F(TpuOpsVerificationTest, ExpandLoadVregValidPacked) { + auto c0 = Create(0); + Value memref = + Create( + GetMemRefType({16}, builder().getBF16Type(), MemorySpace::kVmem)) + .getMemref(); + Value mask = ConstantI1Vector( + /*shape=*/{8}, + /*values=*/{true, false, true, false, true, false, true, false}); + auto el = Create( + /*result=*/VectorType::get({8, 2}, builder().getBF16Type()), + /*base=*/memref, + /*indices=*/ValueRange{c0}, + /*mask=*/mask); + + ASSERT_OK(VerifyOp(el)); +} + +TEST_F(TpuOpsVerificationTest, ExpandLoadVregInvalidMemorySpace) { + auto c0 = Create(0); + Value memref = AllocaI32({8}, MemorySpace::kHbm); + Value mask = ConstantI1Vector( + /*shape=*/{8}, + /*values=*/{true, false, true, false, true, false, true, false}); + auto el = Create( + /*result=*/VectorType::get({8}, i32()), + /*base=*/memref, + /*indices=*/ValueRange{c0}, + /*mask=*/mask); + + ASSERT_THAT(VerifyOp(el), + StatusIs(_, HasSubstr("Expected base memref to be in VMEM."))); +} + +TEST_F(TpuOpsVerificationTest, ExpandLoadVregMismatchedMaskShape) { + auto c0 = Create(0); + Value memref = AllocaI32({8}, MemorySpace::kVmem); + Value mask = ConstantI1Vector(/*shape=*/{4}, + /*values=*/{true, true, true, true}); + auto el = Create( + /*result=*/VectorType::get({8}, i32()), + /*base=*/memref, + /*indices=*/ValueRange{c0}, + /*mask=*/mask); + + ASSERT_THAT( + VerifyOp(el), + StatusIs(_, + HasSubstr("Expected result dimension 0 to match mask dimension " + "0"))); +} + +TEST_F(TpuOpsVerificationTest, ExpandLoadVregMismatchedIndicesCount) { + auto c0 = Create(0); + Value memref = AllocaI32({8}, MemorySpace::kVmem); + Value mask = ConstantI1Vector( + /*shape=*/{8}, + /*values=*/{true, false, true, false, true, false, true, false}); + auto el = Create( + /*result=*/VectorType::get({8}, i32()), + /*base=*/memref, + /*indices=*/ValueRange{c0, c0}, + /*mask=*/mask); + + ASSERT_THAT(VerifyOp(el), StatusIs(_, HasSubstr("Expected 1 indices."))); +} + TEST_F(TpuOpsVerificationTest, VectorCompressStoreValid) { auto c0 = Create(0); Value memref = AllocaI32({4, 16}, MemorySpace::kVmem); diff --git a/third_party/xla/xla/pjrt/BUILD b/third_party/xla/xla/pjrt/BUILD index 832ba429b85b6c..bc2859c3567a6b 100644 --- a/third_party/xla/xla/pjrt/BUILD +++ b/third_party/xla/xla/pjrt/BUILD @@ -848,6 +848,7 @@ xla_cc_test( "@com_google_googletest//:gtest_main", "@llvm-project//mlir:IR", "@riegeli//riegeli/bytes:string_writer", + "@riegeli//riegeli/bytes:writer", "@stablehlo//:stablehlo_portable_api", ], ) diff --git a/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_client_test.cc b/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_client_test.cc index bd2f82bf978755..2feae674faa7b3 100644 --- a/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_client_test.cc +++ b/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_client_test.cc @@ -344,7 +344,8 @@ TEST(StreamExecutorGpuClientTest, PlatformVersionIsDerivedAtRuntime) { const absl::string_view version = client->platform_version(); EXPECT_NE(version, ""); EXPECT_TRUE(absl::StartsWith(version, "cuda ") || - absl::StartsWith(version, "rocm ")) + absl::StartsWith(version, "rocm ") || + absl::StartsWith(version, "oneapi")) << "unexpected platform version: " << version; } diff --git a/third_party/xla/xla/pjrt/mlir_to_hlo.cc b/third_party/xla/xla/pjrt/mlir_to_hlo.cc index 2fc7f207516248..6e43b8742321df 100644 --- a/third_party/xla/xla/pjrt/mlir_to_hlo.cc +++ b/third_party/xla/xla/pjrt/mlir_to_hlo.cc @@ -321,6 +321,11 @@ absl::Status SerializeToRiegeli(mlir::ModuleOp mlir_module, return absl::InvalidArgumentError("Writer cannot be null."); } + // If we are already in a failed state, no need to do anything else. + if (!writer->ok()) { + return writer->status(); + } + mlir::MLIRContext* context = mlir_module->getContext(); mlir::BaseScopedDiagnosticHandler diagnostic_handler(context); @@ -408,11 +413,23 @@ absl::Status SerializeToRiegeli(mlir::ModuleOp mlir_module, if (mlir::failed(mlir::stablehlo::serializePortableArtifact( mlir_module, target, os, /*allowOtherDialects=*/allow_mixed_serialization))) { + // If we have an I/O error, return it specifically before we return the + // generic serialization error. + if (!writer->ok()) { + return writer->status(); + } const absl::Status status = diagnostic_handler.ConsumeStatus(); return absl::InvalidArgumentError( absl::StrCat("Failed to serialize StableHLO to plugin version ", target, ";\n\nDetailed error from MLIR: ", status.message())); } + + // The raw_ostream APIs don't have support for reporting I/O errors, so we + // have to check them ourselves. + if (!writer->ok()) { + return writer->status(); + } + return absl::OkStatus(); } diff --git a/third_party/xla/xla/pjrt/mlir_to_hlo_test.cc b/third_party/xla/xla/pjrt/mlir_to_hlo_test.cc index 63d66f04e5611b..6a1a365c962846 100644 --- a/third_party/xla/xla/pjrt/mlir_to_hlo_test.cc +++ b/third_party/xla/xla/pjrt/mlir_to_hlo_test.cc @@ -15,6 +15,7 @@ limitations under the License. #include "xla/pjrt/mlir_to_hlo.h" +#include #include #include @@ -27,6 +28,7 @@ limitations under the License. #include "mlir/IR/MLIRContext.h" #include "mlir/IR/OwningOpRef.h" #include "riegeli/bytes/string_writer.h" +#include "riegeli/bytes/writer.h" #include "stablehlo/api/PortableApi.h" #include "xla/hlo/testlib/test.h" #include "xla/tsl/platform/statusor.h" @@ -284,5 +286,34 @@ TEST(MlirToHloTest, SerializeToRiegeliTest) { EXPECT_THAT(buffer, IsVhloArtifact("1.0.0")); } +class FailingWriter : public riegeli::Writer { + protected: + bool PushSlow(size_t, size_t) override { + Fail(absl::FailedPreconditionError("We were always going to fail.")); + return false; + } +}; + +TEST(MlirToHloTest, SerializeToRiegeliReportsIOErrors) { + constexpr absl::string_view kProgram = + R"( + func.func @add(%arg0: tensor<1x2xf32>) -> tensor<1x2xf32> { + %cst = stablehlo.constant dense<1.0> : tensor<1x2xf32> + %0 = stablehlo.add %arg0, %cst : tensor<1x2xf32> + return %0 : tensor<1x2xf32> + } + )"; + mlir::MLIRContext context; + ASSERT_OK_AND_ASSIGN(mlir::OwningOpRef module, + ParseMlirModuleString(kProgram, context)); + + FailingWriter writer; + EXPECT_THAT(SerializeToRiegeli(*module, /*requested_target=*/"1.0.0", + /*sdy_version=*/"0.0.1", /*inplace=*/false, + /*allow_mixed_serialization=*/true, &writer), + StatusIs(absl::StatusCode::kFailedPrecondition, + HasSubstr("We were always going to fail."))); +} + } // namespace } // namespace xla diff --git a/third_party/xla/xla/pjrt/transpose_kernels.h b/third_party/xla/xla/pjrt/transpose_kernels.h index da532efae24e83..90212a50fda7a4 100644 --- a/third_party/xla/xla/pjrt/transpose_kernels.h +++ b/third_party/xla/xla/pjrt/transpose_kernels.h @@ -200,6 +200,212 @@ HWY_INLINE void CombineRowsCPS(Cont&& cont, Vs... in) { } } +// Stream -> Pack adapter for loads +template +HWY_INLINE void LoadRowsCPS(D d, const char* a, int64_t lda, Cont&& cont, + Vs... vs) { + if constexpr (I == bs) { + std::forward(cont)(vs...); + } else { + auto v = LoadBytes(d, a + lda * I); + LoadRowsCPS(d, a, lda, std::forward(cont), + vs..., v); + } +} +// Pack -> Stream adapter for 128-bit stores +template +HWY_INLINE void StoreVecs128(D d, char* b, int64_t ldb, Vs... v) { + ForEachIndexed( + [&](auto idx_c, auto&& vec) { + constexpr size_t i = decltype(idx_c)::value; + StoreLanes(d, b, ldb, vec, i * stores_per_row, + std::make_index_sequence{}); + }, + v...); +} + +// 128-bit (single block) transpose +template +HWY_FLATTEN void TransposeMicroKernel128(const char* HWY_RESTRICT a, + int64_t lda, char* HWY_RESTRICT b, + int64_t ldb) { + using D = hn::Full128; + static constexpr size_t element_size = sizeof(T); + static constexpr size_t row_bytes = element_size * bs; + static_assert(row_bytes <= kBlockBytes); + const D d; + LoadRowsCPS<0, T, bs, D, row_bytes>(d, a, lda, [&](auto... loads) { + CombineRowsCPS( + [&](auto... vecs) { + constexpr size_t kNumVecs = sizeof...(vecs); + constexpr size_t kBytesInMatrix = element_size * bs * bs; + constexpr size_t kStoresPerRow = bs / kNumVecs; + constexpr size_t kElementsPerVec = kBytesInMatrix / (bs * kNumVecs); + InterleaveWithinBlocksCPS( + [&](auto... transposed) { + StoreVecs128(d, b, ldb, + transposed...); + }, + vecs...); + }, + loads...); + }); +} + +// 256-bit (two block) transposes +#if HWY_MIN_BYTES >= 32 +template +HWY_INLINE void StoreRows256(D d, char* b, int64_t ldb, Vs... v) { + ForEachIndexed( + [&](auto idx_c, auto&& vec) { + constexpr size_t i = decltype(idx_c)::value; + hn::StoreU(vec, d, reinterpret_cast*>(b + ldb * i)); + }, + v...); +} + +// Invokes a continuation with an explicit parameter pack, rearranging the +// arguments such that all even-indexed elements precede all odd-indexed +// elements. +template +HWY_INLINE void CallContWithSeparatedImpl(Cont&& cont, + std::index_sequence, + Vs&&... vs) { + std::forward(cont)(GetFromPack<2 * I>(std::forward(vs)...)..., + GetFromPack<2 * I + 1>(std::forward(vs)...)...); +} + +template +HWY_INLINE void CallContWithSeparated(Cont&& cont, Vs&&... vs) { + CallContWithSeparatedImpl(std::forward(cont), + std::make_index_sequence{}, + std::forward(vs)...); +} + +// Loads a square block of memory into 256-bit vectors, partitioning the +// data spatially. Invokes the continuation with the vectors corresponding +// to the left half of the block followed by the right half. +template +HWY_INLINE void BuildVecs256SquareCPS(D d, DHalf d_half, const char* a, + int64_t lda, Cont& cont, Vs... vs) { + if constexpr (I == bs / 2) { + CallContWithSeparated(cont, vs...); + } else { + auto* row0 = reinterpret_cast*>(a + lda * I); + auto row1 = + reinterpret_cast*>(a + lda * (I + bs / 2)); + auto row0_lo = hn::LoadU(d_half, row0); + auto row0_hi = hn::LoadU(d_half, row0 + kBlockBytes / sizeof(*row0)); + auto row1_lo = hn::LoadU(d_half, row1); + auto row1_hi = hn::LoadU(d_half, row1 + kBlockBytes / sizeof(*row0)); + auto v_lo = hn::Combine(d, row1_lo, row0_lo); + auto v_hi = hn::Combine(d, row1_hi, row0_hi); + BuildVecs256SquareCPS(d, d_half, a, lda, cont, vs..., v_lo, + v_hi); + } +} + +template +HWY_INLINE void BuildVecs256SquareCPS(D d, DHalf d_half, const char* a, + int64_t lda, Cont& cont) { + BuildVecs256SquareCPS<0, T, bs>(d, d_half, a, lda, cont); +} + +template +HWY_FLATTEN void TransposeMicroKernel256Square(const char* HWY_RESTRICT a, + int64_t lda, + char* HWY_RESTRICT b, + int64_t ldb) { + using D = hn::FixedTag; + using DHalf = hn::Half; + static constexpr size_t element_size = sizeof(T); + static constexpr size_t row_bytes = element_size * bs; + static_assert(row_bytes == 2 * kBlockBytes); + static_assert(bs % 2 == 0); + const D d; + const DHalf d_half; + auto cont = [&](auto... vecs) { + InterleaveWithinBlocksCPS( + [&](auto... transposed) { StoreRows256(d, b, ldb, transposed...); }, + vecs...); + }; + BuildVecs256SquareCPS(d, d_half, a, lda, cont); +} +// Permutes 64-bit lanes across 128-bit blocks in a 256-bit vector. +// Specifically, reorders the 4 64-bit lanes from [0, 1, 2, 3] to [0, 2, 1, 3]. +template +HWY_INLINE V ShuffleAcrossBlocks_3120(V v) { + const hn::DFromV d; + const hn::Repartition d64; + return hn::BitCast(d, + hn::Per4LaneBlockShuffle<3, 1, 2, 0>(hn::BitCast(d64, v))); +} + +template +HWY_INLINE void StoreRectPairs(DHalf d_half, char* b, int64_t ldb, Vs... v) { + ForEachIndexed( + [&](auto idx_c, auto&& vec) { + constexpr size_t i = decltype(idx_c)::value; + const auto lo = hn::LowerHalf(d_half, vec); + const auto hi = hn::UpperHalf(d_half, vec); + hn::StoreU(lo, d_half, + reinterpret_cast*>(b + ldb * (i * 2))); + hn::StoreU(hi, d_half, + reinterpret_cast*>(b + ldb * (i * 2 + 1))); + }, + v...); +} + +// Loads a rectangular block of memory into 256-bit vectors. Invokes the +// continuation with the constructed sequence of vectors. +template +HWY_INLINE void BuildVecs256RectCPS(D d, DHalf d_half, const char* a, + int64_t lda, Cont& cont, Vs... vs) { + if constexpr (I == bs / 2) { + cont(vs...); + } else { + auto lo = hn::LoadU( + d_half, reinterpret_cast*>(a + lda * I)); + auto hi = hn::LoadU(d_half, reinterpret_cast*>( + a + lda * (I + bs / 2))); + auto v = hn::Combine(d, hi, lo); + BuildVecs256RectCPS(d, d_half, a, lda, cont, vs..., v); + } +} + +template +HWY_INLINE void BuildVecs256RectCPS(D d, DHalf d_half, const char* a, + int64_t lda, Cont& cont) { + BuildVecs256RectCPS<0, T, bs>(d, d_half, a, lda, cont); +} + +template +HWY_FLATTEN void TransposeMicroKernel256Rect(const char* HWY_RESTRICT a, + int64_t lda, char* HWY_RESTRICT b, + int64_t ldb) { + using D = hn::FixedTag; + using DHalf = hn::Half; + static constexpr size_t element_size = sizeof(T); + static constexpr size_t row_bytes = element_size * bs; + static_assert(row_bytes == kBlockBytes); + static_assert(bs % 2 == 0); + const D d; + const DHalf d_half; + auto cont = [&](auto... vecs) { + InterleaveWithinBlocksCPS( + [&](auto... interleaved) { + StoreRectPairs(d_half, b, ldb, + ShuffleAcrossBlocks_3120(interleaved)...); + }, + vecs...); + }; + BuildVecs256RectCPS(d, d_half, a, lda, cont); +} +#endif // HWY_MIN_BYTES >= 32 + } // namespace HWY_NAMESPACE } // namespace xla HWY_AFTER_NAMESPACE(); diff --git a/third_party/xla/xla/python/ifrt/ir/BUILD b/third_party/xla/xla/python/ifrt/ir/BUILD index 4ced1e754822df..8a69d0b848b6d7 100644 --- a/third_party/xla/xla/python/ifrt/ir/BUILD +++ b/third_party/xla/xla/python/ifrt/ir/BUILD @@ -150,6 +150,18 @@ cc_library( ], ) +xla_cc_test( + name = "sharding_param_test", + srcs = ["sharding_param_test.cc"], + deps = [ + ":sharding_param", + "//xla/tsl/platform:errors", + "@com_google_absl//absl/status:status_matchers", + "@com_google_googletest//:gtest_main", + "@llvm-project//llvm:Support", + ], +) + tf_proto_library( name = "sharding_param_proto", srcs = ["sharding_param.proto"], diff --git a/third_party/xla/xla/python/ifrt/ir/sharding_param.cc b/third_party/xla/xla/python/ifrt/ir/sharding_param.cc index 4b7fb83df7389d..34b683ddf9d72c 100644 --- a/third_party/xla/xla/python/ifrt/ir/sharding_param.cc +++ b/third_party/xla/xla/python/ifrt/ir/sharding_param.cc @@ -151,6 +151,16 @@ absl::Status ShardingParam::MinorToMajor::verify() const { absl::StrJoin(axis_sizes, "x"))); } } + // The cumulative product of `axis_sizes` must fit in an `int`, since + // ToDeviceList() narrows to it. + int cum_size = 1; + for (const int size : axis_sizes) { + if (__builtin_mul_overflow(cum_size, size, &cum_size)) { + return absl::InvalidArgumentError( + absl::StrCat("`axis_sizes` product overflows: [", + absl::StrJoin(axis_sizes, "x"), "]")); + } + } return absl::OkStatus(); } @@ -248,7 +258,17 @@ absl::Status ShardingParam::verify() const { // distributed over the device mesh. int total_dim_shards_size = 1; for (const int dim_shard : dim_shards()) { - total_dim_shards_size *= dim_shard; + if (dim_shard <= 0) { + return absl::InvalidArgumentError( + absl::StrCat("`dim_shards` must be positive, got ", dim_shard, + " in [", absl::StrJoin(dim_shards(), "x"), "]")); + } + if (__builtin_mul_overflow(total_dim_shards_size, dim_shard, + &total_dim_shards_size)) { + return absl::InvalidArgumentError( + absl::StrCat("`dim_shards` product overflows: ", + absl::StrJoin(dim_shards(), "x"))); + } } if (NumDevices() % total_dim_shards_size != 0) { return absl::InvalidArgumentError(absl::StrCat( @@ -378,8 +398,10 @@ absl::StatusOr ShardingParam::FromProto( unreduced_axes = std::vector(proto.unreduced_axes().begin(), proto.unreduced_axes().end()); } - return ShardingParam(std::move(dim_shards), std::move(minor_to_major), - std::move(unreduced_axes)); + ShardingParam param(std::move(dim_shards), std::move(minor_to_major), + std::move(unreduced_axes)); + ABSL_RETURN_IF_ERROR(param.verify()); + return param; } absl::Status ShardingParam::ToProto(ShardingParamProto& proto, diff --git a/third_party/xla/xla/python/ifrt/ir/sharding_param_test.cc b/third_party/xla/xla/python/ifrt/ir/sharding_param_test.cc new file mode 100644 index 00000000000000..9b698d7aa5f75f --- /dev/null +++ b/third_party/xla/xla/python/ifrt/ir/sharding_param_test.cc @@ -0,0 +1,83 @@ +/* Copyright 2026 The OpenXLA Authors. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#include "xla/python/ifrt/ir/sharding_param.h" + +#include +#include + +#include +#include "absl/status/status_matchers.h" +#include "llvm/ADT/SmallVector.h" +#include "xla/tsl/platform/errors.h" + +namespace xla { +namespace ifrt { +namespace { + +using ::absl_testing::StatusIs; +using ::testing::HasSubstr; + +TEST(ShardingParamTest, AxisSizesProductOverflowIsRejected) { + // Two axis sizes whose product exceeds std::numeric_limits::max() + // (2 * 1,500,000,000 > INT_MAX), individually representable as plain ints. + ShardingParam::MinorToMajor minor_to_major; + minor_to_major.permutation = {0, 1}; + minor_to_major.axis_sizes = {2, 1500000000}; + + ShardingParam sharding(/*dim_shards=*/{1, 1}, minor_to_major); + + EXPECT_THAT(sharding.verify(), + StatusIs(tsl::error::INVALID_ARGUMENT, + HasSubstr("`axis_sizes` product overflows"))); +} + +TEST(ShardingParamTest, DimShardsProductOverflowIsRejected) { + ShardingParam::MinorToMajor minor_to_major; + minor_to_major.permutation = {0}; + minor_to_major.axis_sizes = {4}; + + ShardingParam sharding(/*dim_shards=*/{2, 1500000000}, minor_to_major); + + EXPECT_THAT(sharding.verify(), + StatusIs(tsl::error::INVALID_ARGUMENT, + HasSubstr("`dim_shards` product overflows"))); +} + +TEST(ShardingParamTest, NonPositiveDimShardIsRejected) { + ShardingParam::MinorToMajor minor_to_major; + minor_to_major.permutation = {0}; + minor_to_major.axis_sizes = {4}; + + ShardingParam sharding(/*dim_shards=*/{0}, minor_to_major); + + EXPECT_THAT(sharding.verify(), + StatusIs(tsl::error::INVALID_ARGUMENT, + HasSubstr("`dim_shards` must be positive"))); +} + +TEST(ShardingParamTest, ValidShardingIsAccepted) { + ShardingParam::MinorToMajor minor_to_major; + minor_to_major.permutation = {0, 1}; + minor_to_major.axis_sizes = {2, 4}; + + ShardingParam sharding(/*dim_shards=*/{2, 4}, minor_to_major); + + EXPECT_OK(sharding.verify()); +} + +} // namespace +} // namespace ifrt +} // namespace xla diff --git a/third_party/xla/xla/service/gpu/gb200_cross_compilation_test.cc b/third_party/xla/xla/service/gpu/gb200_cross_compilation_test.cc index 61d8fed3dc4275..5b32a62424680e 100644 --- a/third_party/xla/xla/service/gpu/gb200_cross_compilation_test.cc +++ b/third_party/xla/xla/service/gpu/gb200_cross_compilation_test.cc @@ -109,7 +109,9 @@ TEST_F(Gb200CrossCompilationTest, CrossCompilationToB200) { EXPECT_TRUE(!AutotunerCache::ResultCacheIsEmpty() || !entries.empty()); } -TEST_F(Gb200CrossCompilationTest, CrossCompilationToRTX6000PRO) { +// TODO(b/565341904): RTX6000PRO has compute capability 12, while gb200 +// only 10.a. Test used to pass as we didn't do any autotuning. +TEST_F(Gb200CrossCompilationTest, DISABLED_CrossCompilationToRTX6000PRO) { ASSERT_OK_AND_ASSIGN(std::unique_ptr result, CrossCompileTo(GpuModel::RTX6000PRO)); // Verify that the auto tuner ran. diff --git a/third_party/xla/xla/service/spmd/custom_call_handler.cc b/third_party/xla/xla/service/spmd/custom_call_handler.cc index f6d8472633c488..7e4931ce77b520 100644 --- a/third_party/xla/xla/service/spmd/custom_call_handler.cc +++ b/third_party/xla/xla/service/spmd/custom_call_handler.cc @@ -141,7 +141,7 @@ constexpr char kSPMDOpMultiRotate[] = "_SPMDInternalOp_MultiRotate"; constexpr char kSPMDOpMultiSlice[] = "_SPMDInternalOp_MultiSlice"; constexpr char kSPMDOpWrap[] = "_SPMDInternalOp_Wrap"; -bool IsTopKIndexOutputUsed(HloInstruction* hlo) { +bool IsTopKIndexOutputUsed(const HloInstruction* hlo) { if (hlo->IsRoot()) { return true; } @@ -162,7 +162,14 @@ bool IsTopKIndexOutputUsed(HloInstruction* hlo) { absl::Status SpmdPartitioningVisitor::HandleCustomCallTopK( HloInstruction* hlo) { - if (!hlo->operand(0)->has_sharding()) { + // TopkRewriter flattens rank >= 2 inputs to 2D and leaves rank 1 inputs as + // 1D; this handler assumes a 2D [batch, sort] input. + if (hlo->operand_count() != 1 || !hlo->operand(0)->has_sharding() || + !hlo->operand(0)->shape().IsArray() || + hlo->operand(0)->shape().dimensions().size() != 2 || + !hlo->shape().IsTuple() || hlo->shape().tuple_shapes().size() != 2 || + !hlo->shape().tuple_shapes(0).IsArray() || + hlo->shape().tuple_shapes(0).dimensions().size() != 2) { return DefaultAction(hlo); } @@ -195,19 +202,11 @@ absl::Status SpmdPartitioningVisitor::HandleCustomCallTopK( auto partitioned_input = GetPartitionedHlo(input).PadWithValue( CreateFirstWithType(element_type, &b_)); - auto partition_state = partitioned_input.state(); - auto replicated_sharding = HloSharding::Replicate(); - // If batch dimension is partitioned, partial replicated on sort dimension. - if (batch_dim_partition > 1) { - auto sharding_grouped = - hlo_sharding_util::GroupShardingOnDims(sharding, {batch_dim}); - partition_state = CreatePerGroupPartitioningState( - partitioned_input.state(), sharding_grouped.device_groups, - partitioned_input.state().b); - replicated_sharding = - hlo_sharding_util::PartiallyReplicateTiledShardingOnDims(sharding, - {sort_dim}); - } + const HloSharding replicated_sharding = + batch_dim_partition > 1 + ? hlo_sharding_util::PartiallyReplicateTiledShardingOnDims(sharding, + {sort_dim}) + : HloSharding::Replicate(); // Each partition needs to do TopK separately, thus the base shape // becomes [batch_size, k * shard_count]. @@ -243,7 +242,15 @@ absl::Status SpmdPartitioningVisitor::HandleCustomCallTopK( HloInstruction* slice_sort_value = nullptr; HloInstruction* slice_index_value = nullptr; + const int64_t partitioned_batch_size = + CeilOfRatio(batch_size, batch_dim_partition); if (IsTopKIndexOutputUsed(hlo)) { + const hlo_sharding_util::GroupedSharding sharding_grouped = + hlo_sharding_util::GroupShardingOnDims(sharding, {batch_dim}); + PartitionedHlo::PartitioningState partition_state = + CreatePerGroupPartitioningState(partitioned_input.state(), + sharding_grouped.device_groups, + partitioned_input.state().b); // Get index from TopK. HloInstruction* index_gte = b_.AddInstruction(HloInstruction::CreateGetTupleElement( @@ -280,21 +287,18 @@ absl::Status SpmdPartitioningVisitor::HandleCustomCallTopK( XlaBuilder b("Sort.Compare"); XlaComputation comparator = CreateScalarComparisonComputation( "compare-value-and-index", {input->shape().element_type(), S32}, - {Gt, Lt}, &b); + {GtTotalOrder, LtTotalOrder}, &b); ABSL_ASSIGN_OR_RETURN(HloComputation * compare_computation, XlaComputationToHloComputation(comparator, module_)); // Each partition needs to do TopK separately, thus the base shape for sort // becomes [ceil(batch_size / batch_dim_partition), k * shard_count]. const Shape sort_shape = ShapeUtil::MakeTupleShape( - {ShapeUtil::MakeShape( - hlo->operand(0)->shape().element_type(), - {CeilOfRatio(batch_size, batch_dim_partition), k * shard_count}), - ShapeUtil::MakeShape( - S32, - {CeilOfRatio(batch_size, batch_dim_partition), k * shard_count})}); - auto sort = b_.AddInstruction(HloInstruction::CreateSort( + {ShapeUtil::MakeShape(element_type, + {partitioned_batch_size, k * shard_count}), + ShapeUtil::MakeShape(S32, {partitioned_batch_size, k * shard_count})}); + HloInstruction* sort = b_.AddInstruction(HloInstruction::CreateSort( sort_shape, sort_dim, {replicated_value_gte, replicated_index_gte}, - compare_computation, true)); + compare_computation, /*is_stable=*/false)); sort->set_sharding( replicated_sharding.GetTupleSharding(sort->shape()).value()); PartitionedHlo replicated_sort(sort, replicated_shape, @@ -317,20 +321,19 @@ absl::Status SpmdPartitioningVisitor::HandleCustomCallTopK( // Sort only the replicated values. XlaBuilder b("Sort.Compare"); XlaComputation comparator = CreateScalarComparisonComputation( - "compare-value", {input->shape().element_type()}, {Gt}, &b); + "compare-value", {input->shape().element_type()}, {GtTotalOrder}, &b); ABSL_ASSIGN_OR_RETURN(HloComputation * compare_computation, XlaComputationToHloComputation(comparator, module_)); const Shape sort_shape = ShapeUtil::MakeShape( - hlo->operand(0)->shape().element_type(), - {CeilOfRatio(batch_size, batch_dim_partition), k * shard_count}); - auto sort = b_.AddInstruction( + element_type, {partitioned_batch_size, k * shard_count}); + HloInstruction* sort = b_.AddInstruction( HloInstruction::CreateSort(sort_shape, sort_dim, {replicated_value_gte}, - compare_computation, true)); + compare_computation, /*is_stable=*/false)); sort->set_sharding(replicated_sharding); - const Shape sort_replicated_shape = ShapeUtil::MakeShape( - hlo->operand(0)->shape().element_type(), {batch_size, k * shard_count}); + const Shape sort_replicated_shape = + ShapeUtil::MakeShape(element_type, {batch_size, k * shard_count}); PartitionedHlo replicated_sort(sort, sort_replicated_shape, MakePartitioningState()); @@ -342,7 +345,8 @@ absl::Status SpmdPartitioningVisitor::HandleCustomCallTopK( HloInstruction* dummy_index_constant = b_.AddInstruction( HloInstruction::CreateConstant(LiteralUtil::CreateR0(0))); slice_index_value = b_.AddInstruction(HloInstruction::CreateBroadcast( - ShapeUtil::MakeShape(S32, {batch_size, k}), dummy_index_constant, {})); + ShapeUtil::MakeShape(S32, {partitioned_batch_size, k}), + dummy_index_constant, {})); slice_index_value->set_sharding(replicated_sharding); } diff --git a/third_party/xla/xla/service/spmd/spmd_partitioner.cc b/third_party/xla/xla/service/spmd/spmd_partitioner.cc index 2b6feb4ca0b595..a866832ef944b1 100644 --- a/third_party/xla/xla/service/spmd/spmd_partitioner.cc +++ b/third_party/xla/xla/service/spmd/spmd_partitioner.cc @@ -7405,15 +7405,6 @@ absl::StatusOr SpmdPartitioner::RunImpl( *module, options_.report_instruction_count)); XLA_VLOG_LINES(1, logger.MakeReport()); - if (changed) { - HloPassPipeline pass("spmd-cleanup"); - pass.AddPass(/*remove_cross_partition_collective_ops=*/true); - pass.AddPass(); - pass.AddPass(/*remove_cross_partition_collective_ops=*/true); - pass.AddPass(/*is_layout_sensitive=*/false); - ABSL_RETURN_IF_ERROR(pass.Run(module, execution_threads).status()); - } - ABSL_RETURN_IF_ERROR(ClearShardingAttributes( module, num_replicas() * num_partitions(), execution_threads)); @@ -7424,6 +7415,15 @@ absl::StatusOr SpmdPartitioner::RunImpl( entry_root->set_sharding(std::move(final_sharding)); } + if (changed) { + HloPassPipeline pass("spmd-cleanup"); + pass.AddPass(/*remove_cross_partition_collective_ops=*/true); + pass.AddPass(); + pass.AddPass(/*remove_cross_partition_collective_ops=*/true); + pass.AddPass(/*is_layout_sensitive=*/false); + ABSL_RETURN_IF_ERROR(pass.Run(module, execution_threads).status()); + } + return changed; } diff --git a/third_party/xla/xla/service/spmd/spmd_partitioner_test.cc b/third_party/xla/xla/service/spmd/spmd_partitioner_test.cc index e7d679a29789b2..024a14599209d8 100644 --- a/third_party/xla/xla/service/spmd/spmd_partitioner_test.cc +++ b/third_party/xla/xla/service/spmd/spmd_partitioner_test.cc @@ -3749,9 +3749,13 @@ ENTRY %cluster_2013453984438090939__.47 ASSERT_OK_AND_ASSIGN(auto module, PartitionComputation(hlo_string, /*num_devices=*/2)); VLOG(1) << module->ToString(); - auto custom_call = FindInstruction(module.get(), "custom-call.1"); + const HloInstruction* custom_call = + FindInstruction(module.get(), "custom-call.1"); + ASSERT_NE(custom_call, nullptr); EXPECT_EQ(custom_call->operand(0)->shape().dimensions(1), 104832); - auto sort = FindInstruction(module.get(), "sort"); + const HloInstruction* sort = FindInstruction(module.get(), "sort"); + ASSERT_NE(sort, nullptr); + EXPECT_FALSE(Cast(sort)->is_stable()); EXPECT_EQ(sort->operand(0)->shape().dimensions(1), 4000); EXPECT_EQ(sort->operand(1)->shape().dimensions(1), 4000); } @@ -3779,9 +3783,13 @@ ENTRY entry { ASSERT_OK_AND_ASSIGN(auto module, PartitionComputation(hlo_string, /*num_devices=*/8)); - auto custom_call = FindInstruction(module.get(), "custom-call.1"); + const HloInstruction* custom_call = + FindInstruction(module.get(), "custom-call.1"); + ASSERT_NE(custom_call, nullptr); EXPECT_EQ(custom_call->operand(0)->shape().dimensions(1), 32128); - auto sort = FindInstruction(module.get(), "sort"); + const HloInstruction* sort = FindInstruction(module.get(), "sort"); + ASSERT_NE(sort, nullptr); + EXPECT_FALSE(Cast(sort)->is_stable()); EXPECT_EQ(sort->operand(0)->shape().dimensions(0), 1); EXPECT_EQ(sort->operand(0)->shape().dimensions(1), 2); EXPECT_EQ(sort->operand(1)->shape().dimensions(0), 1); @@ -3811,9 +3819,13 @@ ENTRY entry { ASSERT_OK_AND_ASSIGN(auto module, PartitionComputation(hlo_string, /*num_devices=*/8)); VLOG(1) << module->ToString(); - auto custom_call = FindInstruction(module.get(), "custom-call.1"); + const HloInstruction* custom_call = + FindInstruction(module.get(), "custom-call.1"); + ASSERT_NE(custom_call, nullptr); EXPECT_EQ(custom_call->operand(0)->shape().dimensions(1), 16064); - auto sort = FindInstruction(module.get(), "sort"); + const HloInstruction* sort = FindInstruction(module.get(), "sort"); + ASSERT_NE(sort, nullptr); + EXPECT_FALSE(Cast(sort)->is_stable()); EXPECT_EQ(sort->operand(0)->shape().dimensions(0), 2); EXPECT_EQ(sort->operand(0)->shape().dimensions(1), 4); EXPECT_EQ(sort->operand(1)->shape().dimensions(0), 2); @@ -17470,10 +17482,19 @@ ENTRY %entry { ASSERT_OK_AND_ASSIGN(auto module, PartitionComputation(hlo_string, /*num_devices=*/2)); VLOG(1) << module->ToString(); - auto sort_instruction = FindInstruction(module.get(), HloOpcode::kSort); + const HloInstruction* sort_instruction = + FindInstruction(module.get(), HloOpcode::kSort); + ASSERT_NE(sort_instruction, nullptr); EXPECT_THAT(sort_instruction, op::Shape("bf16[64,80]{1,0}")); - auto topk_instruction = FindInstruction(module.get(), HloOpcode::kCustomCall); - auto topk_operand = topk_instruction->operand(0); + EXPECT_FALSE(Cast(sort_instruction)->is_stable()); + EXPECT_EQ(Cast( + sort_instruction->to_apply()->root_instruction()) + ->order(), + Comparison::Order::kTotal); + const HloInstruction* topk_instruction = + FindInstruction(module.get(), HloOpcode::kCustomCall); + ASSERT_NE(topk_instruction, nullptr); + const HloInstruction* topk_operand = topk_instruction->operand(0); EXPECT_EQ(topk_instruction->custom_call_target(), "TopK"); EXPECT_THAT(topk_instruction, op::Shape("(bf16[64,40]{1,0}, s32[64,40]{1,0})")); @@ -17503,10 +17524,21 @@ ENTRY %entry { ASSERT_OK_AND_ASSIGN(auto module, PartitionComputation(hlo_string, /*num_devices=*/2)); VLOG(1) << module->ToString(); - auto sort_instruction = FindInstruction(module.get(), HloOpcode::kSort); - CHECK_NE(sort_instruction, nullptr); - auto topk_instruction = FindInstruction(module.get(), HloOpcode::kCustomCall); - auto topk_operand = topk_instruction->operand(0); + const HloInstruction* sort_instruction = + FindInstruction(module.get(), HloOpcode::kSort); + ASSERT_NE(sort_instruction, nullptr); + EXPECT_FALSE(Cast(sort_instruction)->is_stable()); + for (const HloInstruction* inst : + sort_instruction->to_apply()->instructions()) { + if (inst->opcode() == HloOpcode::kCompare) { + EXPECT_EQ(Cast(inst)->order(), + Comparison::Order::kTotal); + } + } + const HloInstruction* topk_instruction = + FindInstruction(module.get(), HloOpcode::kCustomCall); + ASSERT_NE(topk_instruction, nullptr); + const HloInstruction* topk_operand = topk_instruction->operand(0); EXPECT_EQ(topk_instruction->custom_call_target(), "TopK"); EXPECT_THAT(topk_instruction, op::Shape("(bf16[32,40]{1,0}, s32[32,40]{1,0})")); @@ -17584,10 +17616,19 @@ ENTRY %entry { ASSERT_OK_AND_ASSIGN(auto module, PartitionComputation(hlo_string, /*num_devices=*/2)); VLOG(1) << module->ToString(); - auto sort_instruction = FindInstruction(module.get(), HloOpcode::kSort); + const HloInstruction* sort_instruction = + FindInstruction(module.get(), HloOpcode::kSort); + ASSERT_NE(sort_instruction, nullptr); EXPECT_THAT(sort_instruction, op::Shape("bf16[32,40]{1,0}")); - auto topk_instruction = FindInstruction(module.get(), HloOpcode::kCustomCall); - auto topk_operand = topk_instruction->operand(0); + EXPECT_FALSE(Cast(sort_instruction)->is_stable()); + EXPECT_EQ(Cast( + sort_instruction->to_apply()->root_instruction()) + ->order(), + Comparison::Order::kTotal); + const HloInstruction* topk_instruction = + FindInstruction(module.get(), HloOpcode::kCustomCall); + ASSERT_NE(topk_instruction, nullptr); + const HloInstruction* topk_operand = topk_instruction->operand(0); EXPECT_EQ(topk_instruction->custom_call_target(), "TopK"); EXPECT_THAT(topk_instruction, op::Shape("(bf16[32,40]{1,0}, s32[32,40]{1,0})")); @@ -17616,16 +17657,89 @@ ENTRY %entry { ASSERT_OK_AND_ASSIGN(auto module, PartitionComputation(hlo_string, /*num_devices=*/2)); VLOG(1) << module->ToString(); - auto sort_instruction = FindInstruction(module.get(), HloOpcode::kSort); + const HloInstruction* sort_instruction = + FindInstruction(module.get(), HloOpcode::kSort); + ASSERT_NE(sort_instruction, nullptr); EXPECT_THAT(sort_instruction, op::Shape("bf16[64,80]{1,0}")); - auto topk_instruction = FindInstruction(module.get(), HloOpcode::kCustomCall); - auto topk_operand = topk_instruction->operand(0); + EXPECT_FALSE(Cast(sort_instruction)->is_stable()); + EXPECT_EQ(Cast( + sort_instruction->to_apply()->root_instruction()) + ->order(), + Comparison::Order::kTotal); + const HloInstruction* topk_instruction = + FindInstruction(module.get(), HloOpcode::kCustomCall); + ASSERT_NE(topk_instruction, nullptr); + const HloInstruction* topk_operand = topk_instruction->operand(0); EXPECT_EQ(topk_instruction->custom_call_target(), "TopK"); EXPECT_THAT(topk_instruction, op::Shape("(bf16[64,40]{1,0}, s32[64,40]{1,0})")); EXPECT_THAT(topk_operand, op::Shape("bf16[64,128000]{1,0}")); } +TEST_P(SpmdPartitioningTest, TopKCustomCallTwoDimsShardedValuesOnly) { + constexpr absl::string_view kHloString = R"( +HloModule module + +region_695.22546 { + Arg_2.22549 = s32[] parameter(2) + Arg_3.22550 = s32[] parameter(3) + Arg_0.22547 = bf16[] parameter(0) + Arg_1.22548 = bf16[] parameter(1) + ROOT compare.22551 = pred[] compare(Arg_0.22547, Arg_1.22548), direction=GT, order=TOTAL +} + +ENTRY %entry { + %multiply.43401 = bf16[64,256000]{1,0} parameter(0), sharding={devices=[2,2]<=[4]} + %custom-call = (bf16[64,40]{1,0}, s32[64,40]{1,0}) custom-call(bf16[64,256000]{1,0} %multiply.43401), custom_call_target="TopK", called_computations={%region_695.22546}, sharding={{devices=[2,2]<=[4]}, {devices=[2,2]<=[4]}} + ROOT %get-tuple-element.336 = bf16[64,40]{1,0} get-tuple-element((bf16[64,40]{1,0}, s32[64,40]{1,0}) %custom-call), index=0 +})"; + + ASSERT_OK_AND_ASSIGN(std::unique_ptr module, + PartitionComputation(kHloString, /*num_devices=*/4)); + const HloInstruction* sort_instruction = + FindInstruction(module.get(), HloOpcode::kSort); + ASSERT_NE(sort_instruction, nullptr); + EXPECT_THAT(sort_instruction, op::Shape("bf16[32,80]{1,0}")); + EXPECT_FALSE(Cast(sort_instruction)->is_stable()); + EXPECT_EQ(Cast( + sort_instruction->to_apply()->root_instruction()) + ->order(), + Comparison::Order::kTotal); + const HloInstruction* topk_instruction = + FindInstruction(module.get(), HloOpcode::kCustomCall); + ASSERT_NE(topk_instruction, nullptr); + EXPECT_THAT(topk_instruction, + op::Shape("(bf16[32,40]{1,0}, s32[32,40]{1,0})")); + EXPECT_THAT(topk_instruction->operand(0), op::Shape("bf16[32,128000]{1,0}")); +} + +TEST_P(SpmdPartitioningTest, TopKCustomCallRank1) { + constexpr absl::string_view kHloString = R"( +HloModule module + +region_695.22546 { + Arg_2.22549 = s32[] parameter(2) + Arg_3.22550 = s32[] parameter(3) + Arg_0.22547 = bf16[] parameter(0) + Arg_1.22548 = bf16[] parameter(1) + ROOT compare.22551 = pred[] compare(Arg_0.22547, Arg_1.22548), direction=GT, order=TOTAL +} + +ENTRY %entry { + %p0 = bf16[256000]{0} parameter(0), sharding={devices=[2]<=[2]} + %custom-call = (bf16[40]{0}, s32[40]{0}) custom-call(bf16[256000]{0} %p0), custom_call_target="TopK", called_computations={%region_695.22546}, sharding={{replicated}, {replicated}} + ROOT %get-tuple-element.336 = bf16[40]{0} get-tuple-element((bf16[40]{0}, s32[40]{0}) %custom-call), index=0 +})"; + + ASSERT_OK_AND_ASSIGN(std::unique_ptr module, + PartitionComputation(kHloString, /*num_devices=*/2)); + const HloInstruction* topk_instruction = + FindInstruction(module.get(), HloOpcode::kCustomCall); + ASSERT_NE(topk_instruction, nullptr); + EXPECT_THAT(topk_instruction, op::Shape("(bf16[40]{0}, s32[40]{0})")); + EXPECT_THAT(topk_instruction->operand(0), op::Shape("bf16[256000]{0}")); +} + TEST_P(SpmdPartitioningTest, TopKCustomCallManualSharding) { absl::string_view hlo_string = R"( HloModule module diff --git a/third_party/xla/xla/stream_executor/rocm/rocm_compute_capability.h b/third_party/xla/xla/stream_executor/rocm/rocm_compute_capability.h index 55c85b8d0595cb..1a6cb860c4a128 100644 --- a/third_party/xla/xla/stream_executor/rocm/rocm_compute_capability.h +++ b/third_party/xla/xla/stream_executor/rocm/rocm_compute_capability.h @@ -94,7 +94,9 @@ class RocmComputeCapability { "gfx1030", // RX68xx / RX69xx "gfx1100", // RX7900 "gfx1101", // RX7700 / RX7800 - "gfx1103", "gfx1150", "gfx1151", "gfx1200", "gfx1201", "gfx1250"}; + "gfx1103", "gfx1150", "gfx1151", "gfx1200", "gfx1201", + "gfx1250", // MI450 + }; bool is_supported_gfx_version() const { return IsThisGfxInAnyList(kSupportedGfxVersions); diff --git a/third_party/xla/xla/stream_executor/sycl/sycl_stream.cc b/third_party/xla/xla/stream_executor/sycl/sycl_stream.cc index f1c1220d911487..64195ed43ed918 100644 --- a/third_party/xla/xla/stream_executor/sycl/sycl_stream.cc +++ b/third_party/xla/xla/stream_executor/sycl/sycl_stream.cc @@ -234,7 +234,13 @@ absl::Status SyclStream::DoHostCallbackWithStatus( } absl::Status SyclStream::BlockHostUntilDone() { - stream_handle_->wait(); + // Called from ~SyclStream(), so exceptions must be caught, not propagated. + try { + stream_handle_->wait(); + } catch (const ::sycl::exception& e) { + return absl::InternalError(absl::StrCat( + "SYCL exception in SyclStream::BlockHostUntilDone: ", e.what())); + } return absl::OkStatus(); } @@ -272,7 +278,11 @@ absl::StatusOr> SyclStream::Create( SyclStream::~SyclStream() { // Wait for all pending operations to complete before destroying the stream. - BlockHostUntilDone().IgnoreError(); + absl::Status wait_status = BlockHostUntilDone(); + if (!wait_status.ok()) { + LOG(ERROR) << "BlockHostUntilDone failed during stream destruction: " + << wait_status; + } // Remove this stream from the executor's list of allocated streams. executor_->DeallocateStream(this);