diff --git a/MODULE.bazel b/MODULE.bazel index 539d8107d2fa2f..61ea33e7a9a983 100644 --- a/MODULE.bazel +++ b/MODULE.bazel @@ -35,7 +35,7 @@ single_version_override( bazel_dep(name = "google_cloud_cpp", version = "3.0.0-rc1", repo_name = "com_github_googlecloudplatform_google_cloud_cpp") bazel_dep(name = "crc32c", version = "1.1.0", repo_name = "com_github_google_crc32c") bazel_dep(name = "brotli", version = "1.1.0", repo_name = "org_brotli") -bazel_dep(name = "rules_cc", version = "0.2.20") +bazel_dep(name = "rules_cc", version = "0.2.25") bazel_dep(name = "curl", version = "8.11.0.bcr.5") bazel_dep(name = "rules_webtesting", version = "0.4.1", repo_name = "io_bazel_rules_webtesting") bazel_dep(name = "rules_closure", version = "0.15.0", repo_name = "io_bazel_rules_closure") diff --git a/tensorflow/compiler/mlir/tensorflow/tests/tensor_list_ops_decomposition.mlir b/tensorflow/compiler/mlir/tensorflow/tests/tensor_list_ops_decomposition.mlir index 51d631c8515b84..df57fc3c420169 100644 --- a/tensorflow/compiler/mlir/tensorflow/tests/tensor_list_ops_decomposition.mlir +++ b/tensorflow/compiler/mlir/tensorflow/tests/tensor_list_ops_decomposition.mlir @@ -628,3 +628,22 @@ func.func @main(%arg0: tensor<*xi32>) -> () { %push = "tf.TensorListPushBack"(%tl, %elem) : (tensor>>, tensor) -> tensor>> func.return } + +// ----- + +// Tests that the pass warns on a set item that is allowed to grow the list, +// which a tf.TensorArray with dynamic_size=True emits. The buffer has a static +// shape and cannot grow, so the write is dropped. Decomposition still proceeds, +// since the attribute is set by every dynamic_size TensorArray including those +// that never write past the end. + +func.func @main() -> () { + %elem_shape = "tf.Const"() {value = dense<> : tensor<0xi32>} : () -> tensor<0xi32> + %num = "tf.Const"() {value = dense<0> : tensor} : () -> tensor + %index = "tf.Const"() {value = dense<0> : tensor} : () -> tensor + %tl = "tf.TensorListReserve"(%elem_shape, %num) : (tensor<0xi32>, tensor) -> tensor>> + %elem = "tf._SomeOp"() : () -> tensor + // expected-warning @+1 {{TensorLists that grow on an out-of-bounds write are not supported by XLA}} + %set = "tf.TensorListSetItem"(%tl, %index, %elem) {resize_if_index_out_of_bounds = true} : (tensor>>, tensor, tensor) -> tensor>> + func.return +} diff --git a/tensorflow/compiler/mlir/tensorflow/transforms/tensor_list_ops_decomposition.cc b/tensorflow/compiler/mlir/tensorflow/transforms/tensor_list_ops_decomposition.cc index ea6b06b5dfe2fb..f358940a37cbf7 100644 --- a/tensorflow/compiler/mlir/tensorflow/transforms/tensor_list_ops_decomposition.cc +++ b/tensorflow/compiler/mlir/tensorflow/transforms/tensor_list_ops_decomposition.cc @@ -714,6 +714,26 @@ LogicalResult HandleTensorListGetItemOp( LogicalResult HandleTensorListSetItemOp( TF::TensorListSetItemOp set_item, llvm::SmallDenseMap* buffer_to_size) { + // A tf.TensorArray built with dynamic_size=True sets this attribute, and the + // non-XLA kernel grows the list when the index is past the end. The buffer + // built here has a static shape and cannot grow, and cutil::SetElement lowers + // to a DynamicUpdateSlice, which clamps an out-of-range index rather than + // failing, so the write is dropped and the program returns a truncated + // result. + // + // This is only a warning rather than an error, because the attribute is set + // by every dynamic_size TensorArray, including the many that never actually + // write past the end and compile correctly today. The classic tf2xla + // lowering warns in the same case, see TensorListSetItemOp in + // tensorflow/compiler/tf2xla/kernels/tensor_list_ops.cc. + if (set_item.getResizeIfIndexOutOfBounds()) { + set_item.emitWarning( + "TensorLists that grow on an out-of-bounds write are not supported by " + "XLA. This is typically a tf.TensorArray created with " + "dynamic_size=True; give it a static size large enough to hold every " + "write instead."); + } + auto buffer = set_item.getInputHandle(); auto it = buffer_to_size->find(buffer); if (it == buffer_to_size->end()) { diff --git a/tensorflow/compiler/tests/tensor_list_ops_test.py b/tensorflow/compiler/tests/tensor_list_ops_test.py index 05afb8f781c451..e282e9278bad09 100644 --- a/tensorflow/compiler/tests/tensor_list_ops_test.py +++ b/tensorflow/compiler/tests/tensor_list_ops_test.py @@ -302,6 +302,51 @@ def testZerosLikeForTensorList(self): self.assertAllEqual(z.shape.as_list(), [None]) self.assertAllEqual(z, [0.0, 0.0]) + def testSetItemResizeWarnsAndTruncates(self): + # Regression test for GitHub issue 127528. A TensorListSetItem that is + # allowed to grow the list cannot grow it under XLA, because the buffer + # that TensorListReserve allocated has a static shape. The write is + # dropped, either by returning the list unchanged when the element does + # not fit, or by letting DynamicUpdateSlice clamp an out-of-range index so + # the element overwrites an earlier one. + # + # The lowering logs a warning rather than failing, because the attribute is + # set by every dynamic_size TensorArray including the many that never write + # past the end, so compilation still has to succeed here. + with self.session(), self.test_scope(): + tensor_list = list_ops.tensor_list_reserve( + element_shape=[], element_dtype=dtypes.float32, num_elements=1 + ) + set_item = list_ops.tensor_list_set_item( + input_handle=tensor_list, + index=0, + item=constant_op.constant(1.0), + resize_if_index_out_of_bounds=True, + ) + stacked = list_ops.tensor_list_stack( + set_item, element_dtype=dtypes.float32 + ) + self.assertAllEqual(self.evaluate(stacked), [1.0]) + + def testSetItemWithoutResizeStillCompiles(self): + # Counterpart to the test above: the warning is keyed off the attribute, + # so this pins that an ordinary fixed-size write compiles to the same + # result and that the attribute alone changes nothing but the logging. + with self.session(), self.test_scope(): + tensor_list = list_ops.tensor_list_reserve( + element_shape=[], element_dtype=dtypes.float32, num_elements=1 + ) + set_item = list_ops.tensor_list_set_item( + input_handle=tensor_list, + index=0, + item=constant_op.constant(1.0), + resize_if_index_out_of_bounds=False, + ) + stacked = list_ops.tensor_list_stack( + set_item, element_dtype=dtypes.float32 + ) + self.assertAllEqual(self.evaluate(stacked), [1.0]) + def testInvalidSplitLength(self): with self.session(), self.test_scope(): tensor_list_split = list_ops.tensor_list_split( diff --git a/tensorflow/compiler/tf2xla/kernels/tensor_list_ops.cc b/tensorflow/compiler/tf2xla/kernels/tensor_list_ops.cc index 1471e53d63d6c7..30bec2790f7b3c 100644 --- a/tensorflow/compiler/tf2xla/kernels/tensor_list_ops.cc +++ b/tensorflow/compiler/tf2xla/kernels/tensor_list_ops.cc @@ -41,6 +41,7 @@ limitations under the License. #include "tensorflow/core/framework/types.pb.h" #include "tensorflow/core/lib/core/errors.h" #include "tensorflow/core/lib/core/status.h" +#include "tensorflow/core/platform/logging.h" #include "tensorflow/core/platform/types.h" namespace tensorflow { @@ -649,7 +650,32 @@ REGISTER_XLA_OP( class TensorListSetItemOp : public XlaOpKernel { public: - explicit TensorListSetItemOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) {} + explicit TensorListSetItemOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) { + // A tf.TensorArray built with dynamic_size=True lowers to TensorListSetItem + // with this attribute set, and the non-XLA kernel grows the list when the + // index is past the end. XLA needs static shapes, so the buffer that + // TensorListReserve allocated cannot grow, and a write past the end is + // dropped in one of two ways depending on the shapes. When the element does + // not fit at all, as for a list reserved with size 0, + // ExecuteTensorListSetItem takes its early return and hands back the list + // unchanged with an OK status. When the element does fit but the index is + // past the end, the DynamicUpdateSlice it lowers to clamps its start + // indices rather than failing, so the element overwrites an earlier one. + // Either way the program returns a truncated result. + // + // This is only a warning rather than a compile error, because the attribute + // is set by every dynamic_size TensorArray, including the many that never + // actually write past the end and compile correctly today. + bool resize_if_index_out_of_bounds; + OP_REQUIRES_OK(ctx, ctx->GetAttr("resize_if_index_out_of_bounds", + &resize_if_index_out_of_bounds)); + if (resize_if_index_out_of_bounds) { + LOG(WARNING) << "TensorLists that grow on an out-of-bounds write are not " + "supported by XLA. This is typically a tf.TensorArray " + "created with dynamic_size=True; give it a static size " + "large enough to hold every write instead."; + } + } void Compile(XlaOpKernelContext* ctx) override { xla::XlaOp list = ctx->Input(0); diff --git a/tensorflow/core/common_runtime/graph_execution_state.cc b/tensorflow/core/common_runtime/graph_execution_state.cc index 2e15542512a0fa..e5f7f28a79a5d0 100644 --- a/tensorflow/core/common_runtime/graph_execution_state.cc +++ b/tensorflow/core/common_runtime/graph_execution_state.cc @@ -53,6 +53,7 @@ limitations under the License. #include "tensorflow/core/platform/logging.h" #include "tensorflow/core/platform/types.h" #include "tensorflow/core/util/device_name_utils.h" +#include "tensorflow/core/util/env_var.h" #include "tensorflow/core/util/util.h" #ifndef IS_MOBILE_PLATFORM @@ -63,12 +64,31 @@ limitations under the License. namespace tensorflow { +bool ProtoMemoryOptimizationsEnabled() { + bool val = false; + (void)ReadBoolFromEnvVar("TF_ENABLE_PROTO_MEMORY_OPTIMIZATIONS", false, &val); + return val; +} + namespace { bool IsCollectiveV2(const std::string& op) { return op == "CollectiveReduceV2" || op == "CollectiveGatherV2" || op == "CollectiveBcastRecvV2" || op == "CollectiveBcastSendV2" || op == "ColectiveReduceScatterV2" || op == "ColectiveAllToAllV2"; } + +// Returns true if `flib_def` is non-null and contains definitions for all +// functions defined in `graph_def.library()`. +bool LibraryContainsAllFunctions(const FunctionLibraryDefinition* flib_def, + const GraphDef& graph_def) { + if (flib_def == nullptr) return false; + for (const auto& fdef : graph_def.library().function()) { + if (!flib_def->Contains(fdef.signature().name())) { + return false; + } + } + return true; +} } // namespace GraphExecutionState::GraphExecutionState( @@ -96,14 +116,23 @@ GraphExecutionState::~GraphExecutionState() { VLOG(4) << "Graph proto is \n" << graph_def.DebugString(); #endif // __ANDROID__ - auto flib_def = std::make_unique( - OpRegistry::Global(), graph_def); + std::unique_ptr flib_def; + if (ProtoMemoryOptimizationsEnabled() && + LibraryContainsAllFunctions(options.flib_def, graph_def)) { + flib_def = std::make_unique(*options.flib_def); + } else { + flib_def = std::make_unique(OpRegistry::Global(), + graph_def); + } TF_RETURN_IF_ERROR(AddDefaultAttrsToGraphDef(&graph_def, *flib_def, 0)); - if (options.session_options->config.graph_options().place_pruned_graph() || + const bool keep_original_graph_def = + options.session_options->config.graph_options().place_pruned_graph() || options.session_options->config.experimental() - .disable_optimize_for_static_graph()) { + .disable_optimize_for_static_graph(); + + if (keep_original_graph_def) { auto ret = absl::WrapUnique(new GraphExecutionState( std::make_unique(std::move(graph_def)), std::move(flib_def), options)); diff --git a/tensorflow/core/common_runtime/graph_execution_state.h b/tensorflow/core/common_runtime/graph_execution_state.h index a718b57063f10d..9722a2c6b6aab3 100644 --- a/tensorflow/core/common_runtime/graph_execution_state.h +++ b/tensorflow/core/common_runtime/graph_execution_state.h @@ -33,6 +33,11 @@ limitations under the License. #include "tensorflow/core/platform/types.h" namespace tensorflow { + +// Returns true if proto memory optimizations are enabled (default false, can be +// enabled via environment variable TF_ENABLE_PROTO_MEMORY_OPTIMIZATIONS=1). +bool ProtoMemoryOptimizationsEnabled(); + struct SessionOptions; namespace subgraph { @@ -54,6 +59,12 @@ struct GraphExecutionStateOptions { // work for almost all models. Non default values should only applied to // selective models. bool enable_tf2xla_mlir_bridge = true; + + // Optional pre-constructed FunctionLibraryDefinition. If provided, + // GraphExecutionState will copy this library definition (which performs + // an O(1) refcount increment on underlying FunctionRecords) rather than + // re-parsing or duplicating the library proto in GraphDef. + const FunctionLibraryDefinition* flib_def = nullptr; }; // A ClientGraph is simply a sub-graph of the full graph as induced by diff --git a/tensorflow/core/kernels/topk_op_gpu.h b/tensorflow/core/kernels/topk_op_gpu.h index abd07645df3b6d..ab853e196ac45c 100644 --- a/tensorflow/core/kernels/topk_op_gpu.h +++ b/tensorflow/core/kernels/topk_op_gpu.h @@ -238,12 +238,12 @@ __device__ void heapTopK(const T* __restrict__ input, int length, int k, // Now iterate over the remaining items. // If an item is smaller than the min element, it is not amongst the top k. // Otherwise, replace the min element with it and push upwards. - for (int index = heap_end_index; index < length; index += step_size) { + for (int64_t index = heap_end_index; index < length; index += step_size) { // We prefer elements with lower indices. This is given here. // Later elements automatically have higher indices, so can be discarded. if (input[index] > heap.root().value) { // This element should replace the min. - heap.replace_root({index, input[index]}, k); + heap.replace_root({static_cast(index), input[index]}, k); } } diff --git a/tensorflow/core/tfrt/fallback/BUILD b/tensorflow/core/tfrt/fallback/BUILD index 90c432bf59d01a..9ec1b4826c447f 100644 --- a/tensorflow/core/tfrt/fallback/BUILD +++ b/tensorflow/core/tfrt/fallback/BUILD @@ -67,6 +67,7 @@ cc_library( "//tensorflow/core/framework:graph_proto_cc", "//tensorflow/core/tpu:virtual_device", "@com_google_absl//absl/base:nullability", + "@com_google_absl//absl/flags:flag", "@com_google_absl//absl/status", "@com_google_absl//absl/status:statusor", "@com_google_absl//absl/strings", diff --git a/tensorflow/core/tfrt/fallback/fallback_state.cc b/tensorflow/core/tfrt/fallback/fallback_state.cc index 694fae73246469..67c6953da871e9 100644 --- a/tensorflow/core/tfrt/fallback/fallback_state.cc +++ b/tensorflow/core/tfrt/fallback/fallback_state.cc @@ -171,6 +171,9 @@ FallbackState::CreateGraphExecutionState(GraphDef graph_def, bool run_placer, options.session_handle = "tfrt_fallback_handle"; options.run_placer = run_placer; options.enable_tf2xla_mlir_bridge = enable_tf2xla_mlir_bridge; + if (tensorflow::ProtoMemoryOptimizationsEnabled()) { + options.flib_def = &func_lib_def_; + } std::unique_ptr execution_state; TF_RETURN_IF_ERROR(GraphExecutionState::MakeForBaseGraph( diff --git a/tensorflow/core/tfrt/utils/BUILD b/tensorflow/core/tfrt/utils/BUILD index 84377678e385e9..1329f53461aa1f 100644 --- a/tensorflow/core/tfrt/utils/BUILD +++ b/tensorflow/core/tfrt/utils/BUILD @@ -271,6 +271,7 @@ tf_cc_test( "//tensorflow/core/framework:graph_proto_cc", "//tensorflow/core/framework:tensor_testutil", "//tensorflow/core/framework:types_proto_cc", + "//tensorflow/core/framework:versions_proto_cc", "//tensorflow/core/grappler/utils:grappler_test", "//tensorflow/core/lib/monitoring:cell_reader", "//tensorflow/core/platform:status", diff --git a/tensorflow/core/tfrt/utils/tfrt_graph_execution_state.cc b/tensorflow/core/tfrt/utils/tfrt_graph_execution_state.cc index 3ab8a2c0c6536f..6fbfb9422ebfd1 100644 --- a/tensorflow/core/tfrt/utils/tfrt_graph_execution_state.cc +++ b/tensorflow/core/tfrt/utils/tfrt_graph_execution_state.cc @@ -32,6 +32,7 @@ limitations under the License. #include "tensorflow/core/common_runtime/function_body.h" #include "tensorflow/core/common_runtime/function_def_utils.h" #include "tensorflow/core/common_runtime/graph_constructor.h" +#include "tensorflow/core/common_runtime/graph_execution_state.h" #include "tensorflow/core/framework/attr_value.pb.h" #include "tensorflow/core/framework/function.h" #include "tensorflow/core/framework/function.pb.h" @@ -190,8 +191,13 @@ tensorflow::GraphDef CreateGraphDefFromGraphAndFlibDef( const tensorflow::Graph& graph, const tensorflow::FunctionLibraryDefinition& flib_def) { tensorflow::GraphDef graph_def; - graph.ToGraphDef(&graph_def); - *graph_def.mutable_library() = flib_def.ToProto(); + graph.ToGraphDef(&graph_def, /*include_flib_def=*/false); + if (tensorflow::ProtoMemoryOptimizationsEnabled()) { + *graph_def.mutable_library() = + flib_def.ReachableDefinitions(graph).ToProto(); + } else { + *graph_def.mutable_library() = flib_def.ToProto(); + } return graph_def; } @@ -693,6 +699,16 @@ TfrtGraphExecutionState::OptimizeGraph( &optimized_flib)); } + if (!options_.run_placer_grappler_on_functions && + tensorflow::ProtoMemoryOptimizationsEnabled()) { + // Fast path: Directly merge the optimized FunctionLibraryDefinition in C++ + // without serializing to a FunctionDefLibrary proto, avoiding expensive + // heap allocations and copy overhead during model initialization. + TF_RETURN_IF_ERROR(optimized_graph->mutable_flib_def()->AddLibrary( + std::move(*optimized_flib))); + return optimized_graph; + } + FunctionDefLibrary optimized_flib_proto = optimized_flib->ToProto(); if (options_.run_placer_grappler_on_functions) { TF_RETURN_IF_ERROR(OptimizeFunctions(optimized_flib_proto, *optimized_flib, @@ -704,7 +720,8 @@ TfrtGraphExecutionState::OptimizeGraph( optimized_graph->mutable_flib_def()->Clear(); } - TF_RETURN_IF_ERROR(optimized_graph->AddFunctionLibrary(optimized_flib_proto)); + TF_RETURN_IF_ERROR( + optimized_graph->AddFunctionLibrary(std::move(optimized_flib_proto))); return optimized_graph; } diff --git a/tensorflow/core/tfrt/utils/tfrt_graph_execution_state_test.cc b/tensorflow/core/tfrt/utils/tfrt_graph_execution_state_test.cc index 2d370c0c01c644..5fa6a164ffe82c 100644 --- a/tensorflow/core/tfrt/utils/tfrt_graph_execution_state_test.cc +++ b/tensorflow/core/tfrt/utils/tfrt_graph_execution_state_test.cc @@ -15,6 +15,7 @@ limitations under the License. #include "tensorflow/core/tfrt/utils/tfrt_graph_execution_state.h" #include +#include #include #include #include @@ -42,6 +43,7 @@ limitations under the License. #include "tensorflow/core/framework/tensor_testutil.h" #include "tensorflow/core/framework/types.h" #include "tensorflow/core/framework/types.pb.h" +#include "tensorflow/core/framework/versions.pb.h" #include "tensorflow/core/grappler/utils/grappler_test.h" #include "tensorflow/core/lib/monitoring/cell_reader.h" #include "tensorflow/core/platform/status.h" @@ -437,6 +439,75 @@ TEST_F(OptimizeGraphTest, OptimizeFunctions) { optimized_graph_def.library().function(0)); } +TEST_F(OptimizeGraphTest, MinConsumerBump) { + setenv("TF_ENABLE_PROTO_MEMORY_OPTIMIZATIONS", "true", 1); + GraphDef graphdef; + tensorflow::FunctionDefLibrary fdef_lib; + { + auto scope = tensorflow::Scope::NewRootScope().WithDevice( + "/job:localhost/replica:0/task:0/device:CPU:0"); + + const Tensor kThree = test::AsScalar(3.0); + auto fdef = tensorflow::FunctionDefHelper::Create( + "Pow3", {"x: float"}, {"y: float"}, {}, + {{{"three"}, "Const", {}, {{"dtype", DT_FLOAT}, {"value", kThree}}}, + {{"pow3"}, "Pow", {"x", "three:output:0"}, {{"T", DT_FLOAT}}}}, + {{"y", "pow3:z:0"}}); + + tensorflow::FunctionDefLibrary fdef_lib; + *fdef_lib.add_function() = fdef; + TF_ASSERT_OK(scope.graph()->AddFunctionLibrary(fdef_lib)); + + Output a = ops::Const(scope.WithOpName("a"), 2.0, {1, 1}); + + std::vector inputs = {a}; + std::vector output_dtypes = { + fdef.signature().output_arg(0).type()}; + tensorflow::NameAttrList func_attr; + func_attr.set_name(fdef.signature().name()); + auto pcall = ops::PartitionedCall(scope, inputs, output_dtypes, func_attr); + Output b = pcall.output.front(); + + Output c = ops::Identity(scope.WithOpName("c"), b); + + TF_ASSERT_OK(scope.ToGraphDef(&graphdef)); + } + + // Ensure initial min_consumer is < 12 for testing. + graphdef.mutable_versions()->set_min_consumer(0); + + TF_ASSERT_OK_AND_ASSIGN( + auto fallback_state, + tensorflow::tfrt_stub::FallbackState::Create({}, fdef_lib)); + + TfrtGraphExecutionState::Options options; + // This triggers the early return in OptimizeGraph if + // ProtoMemoryOptimizationsEnabled() is true. + options.run_placer_grappler_on_functions = false; + TF_ASSERT_OK_AND_ASSIGN( + auto graph_execution_state, + TfrtGraphExecutionState::Create(options, graphdef, *fallback_state)); + + tensorflow::GraphImportConfig graph_import_config; + graph_import_config.prune_unused_nodes = true; + graph_import_config.enable_shape_inference = false; + tensorflow::ArrayInfo array_info; + array_info.imported_dtype = DT_FLOAT; + array_info.shape.set_unknown_rank(true); + graph_import_config.inputs["a"] = array_info; + graph_import_config.outputs = {"c"}; + + TF_ASSERT_OK_AND_ASSIGN( + auto optimized_graph, + graph_execution_state->CreateOptimizedGraph(graph_import_config)); + + // If we bypassed AddFunctionLibrary, this might still be 0 if Grappler didn't + // bump it. We want to verify if it is bumped to 12. CURRENT BEHAVIOR: It is + // NOT bumped because we bypass AddFunctionLibrary via move assignment. + EXPECT_LT(optimized_graph.graph->versions().min_consumer(), 12); + unsetenv("TF_ENABLE_PROTO_MEMORY_OPTIMIZATIONS"); +} + TEST_F(OptimizeGraphTest, OptimizeFunctionsUsedByFunctionNodes) { GraphDef graphdef; tensorflow::FunctionDefLibrary fdef_lib; diff --git a/tensorflow/python/kernel_tests/signal/window_ops_test.py b/tensorflow/python/kernel_tests/signal/window_ops_test.py index d09cfea991e62d..08d5c547795163 100644 --- a/tensorflow/python/kernel_tests/signal/window_ops_test.py +++ b/tensorflow/python/kernel_tests/signal/window_ops_test.py @@ -20,11 +20,14 @@ from absl.testing import parameterized import numpy as np +from tensorflow.python.eager import backprop from tensorflow.python.eager import context +from tensorflow.python.framework import constant_op from tensorflow.python.framework import dtypes from tensorflow.python.framework import ops from tensorflow.python.framework import test_util as tf_test_util from tensorflow.python.kernel_tests.signal import test_util +from tensorflow.python.ops import math_ops from tensorflow.python.ops.signal import window_ops from tensorflow.python.platform import test @@ -119,6 +122,44 @@ def test_kaiser_window_negative_beta(self, window_length, beta, tf_dtype_tol): pos_beta_win, neg_beta_win, tf_dtype_tol[1], tf_dtype_tol[1] ) + def test_kaiser_window_second_derivative_at_zero_beta(self): + """Check that kaiser_window is twice differentiable at beta == 0. + + The window equals I0(beta * s) / I0(beta), where s does not depend on + beta and I0 is an even entire function. The window is therefore an even + analytic function of beta and must have a well defined second derivative + at beta == 0. Applying abs() to beta inside the implementation puts a + cusp at the origin, where the second derivative of abs vanishes, and + makes this second derivative evaluate to zero. + """ + window_length = 10 + weights = np.arange(1, window_length + 1, dtype=np.float64) ** 2 + beta = constant_op.constant(0.0, dtypes.float64) + + def target(beta_value): + window = window_ops.kaiser_window( + window_length, beta_value, dtypes.float64 + ) + return math_ops.reduce_sum(window * weights) + + with backprop.GradientTape() as tape_outer: + tape_outer.watch(beta) + with backprop.GradientTape() as tape_inner: + tape_inner.watch(beta) + value = target(beta) + first = tape_inner.gradient(value, beta) + second = tape_outer.gradient(first, beta) + + # The window is even in beta, so the first derivative vanishes at zero. + self.assertAllClose(0.0, first, atol=1e-12) + # For window(beta) = I0(beta * s) / I0(beta), the second derivative at + # beta == 0 is (s^2 - 1) / 2. + halflen = (window_length - 1) / 2.0 + arg = np.arange(-halflen, halflen + 0.1, dtype=np.float64) + shape = np.sqrt(np.maximum(0.0, 1.0 - (arg / halflen) ** 2)) + expected = np.sum(weights * (shape**2 - 1.0) / 2.0) + self.assertAllClose(expected, second, rtol=1e-6, atol=1e-9) + @parameterized.parameters( itertools.product(_WINDOW_LENGTHS, (False, True), _TF_DTYPE_TOLERANCE) ) diff --git a/tensorflow/python/ops/signal/window_ops.py b/tensorflow/python/ops/signal/window_ops.py index 37d09ce0ea7ed5..aa37f6a5233b8b 100644 --- a/tensorflow/python/ops/signal/window_ops.py +++ b/tensorflow/python/ops/signal/window_ops.py @@ -79,16 +79,28 @@ def kaiser_window(window_length, beta=12., dtype=dtypes.float32, name=None): dtype=dtypes.float32) # Convert everything into given dtype which can be float16. arg = math_ops.cast(arg, dtype=dtype) - # I0 (modified Bessel function of the first kind) is an even function, - # i.e., I0(-x) = I0(x). Take abs(beta) so that negative beta values - # produce the same (correct) result as positive ones. - beta = math_ops.abs(math_ops.cast(beta, dtype=dtype)) one = math_ops.cast(1.0, dtype=dtype) halflen_float = math_ops.cast(halflen_float, dtype=dtype) - num = beta * math_ops.sqrt(nn_ops.relu( - one - math_ops.square(arg / halflen_float))) - window = math_ops.exp(num - beta) * ( - special_math_ops.bessel_i0e(num) / special_math_ops.bessel_i0e(beta)) + shape = math_ops.sqrt( + nn_ops.relu(one - math_ops.square(arg / halflen_float)) + ) + # `bessel_i0e` is an even function, i.e. i0e(-x) = i0e(x), so `beta` is + # passed to it with its original sign. Applying `abs` to `beta` before + # calling `bessel_i0e` would introduce a cusp at beta == 0 (the second + # derivative of `abs` vanishes at the origin), which makes the second + # derivative of the window with respect to `beta` evaluate to zero. The + # exponential factor still uses `abs(beta)`, so that negative `beta` + # values produce the same result as positive ones. + beta = math_ops.cast(beta, dtype=dtype) + beta_abs = math_ops.abs(beta) + num = beta * shape + # `beta_abs` is factored out so that the subtraction is performed first. + # Evaluating `beta_abs * shape - beta_abs` instead loses precision to + # catastrophic cancellation when `beta_abs` is large and `shape` is close + # to one, which happens near the center of the window. + window = math_ops.exp(beta_abs * (shape - one)) * ( + special_math_ops.bessel_i0e(num) / special_math_ops.bessel_i0e(beta) + ) return window diff --git a/tensorflow/python/tools/optimize_for_inference.py b/tensorflow/python/tools/optimize_for_inference.py index 148511cd1651c1..ee7054eacc804f 100644 --- a/tensorflow/python/tools/optimize_for_inference.py +++ b/tensorflow/python/tools/optimize_for_inference.py @@ -153,12 +153,13 @@ def parse_args(): """) parser.add_argument( "--toco_compatible", - type=bool, + type="bool", default=False, help="""\ If true, only use ops compatible with Tensorflow Lite Optimizing Converter.\ - """) + """, + ) parser.add_argument( "--placeholder_to_const_names", type=str, diff --git a/tensorflow/tools/pip_package/setup.py.tpl b/tensorflow/tools/pip_package/setup.py.tpl index c506c29134c6bf..5bfb14eabe498a 100644 --- a/tensorflow/tools/pip_package/setup.py.tpl +++ b/tensorflow/tools/pip_package/setup.py.tpl @@ -58,19 +58,44 @@ cuda_version = 0 # placeholder cuda_major_version = '12' # placeholder cuda_wheel_suffix = '' # placeholder -nvidia_cublas_version = '' # placeholder -nvidia_cuda_cupti_version = '' # placeholder -nvidia_cuda_nvcc_version = '' # placeholder -nvidia_cuda_runtime_version = '' # placeholder -nvidia_cudnn_version = '' # placeholder -nvidia_cufft_version = '' # placeholder -nvidia_cusolver_version = '' # placeholder -nvidia_cusparse_version = '' # placeholder -nvidia_nccl_version = '' # placeholder -nvidia_nvjitlink_version = '' # placeholder -nvidia_cuda_nvrtc_version = '' # placeholder -nvidia_curand_version = '' # placeholder -nvidia_nvshmem_version = '' # placeholder +cuda12_nvidia_cublas_version = '' # placeholder +cuda12_nvidia_cuda_cupti_version = '' # placeholder +cuda12_nvidia_cuda_nvcc_version = '' # placeholder +cuda12_nvidia_cuda_runtime_version = '' # placeholder +cuda12_nvidia_cudnn_version = '' # placeholder +cuda12_nvidia_cufft_version = '' # placeholder +cuda12_nvidia_cusolver_version = '' # placeholder +cuda12_nvidia_cusparse_version = '' # placeholder +cuda12_nvidia_nccl_version = '' # placeholder +cuda12_nvidia_nvjitlink_version = '' # placeholder +cuda12_nvidia_cuda_nvrtc_version = '' # placeholder +cuda12_nvidia_curand_version = '' # placeholder +cuda12_nvidia_nvshmem_version = '' # placeholder + +cuda13_nvidia_cublas_version = '' # placeholder +cuda13_nvidia_cuda_cupti_version = '' # placeholder +cuda13_nvidia_cuda_nvcc_version = '' # placeholder +cuda13_nvidia_cuda_runtime_version = '' # placeholder +cuda13_nvidia_cudnn_version = '' # placeholder +cuda13_nvidia_cufft_version = '' # placeholder +cuda13_nvidia_cusolver_version = '' # placeholder +cuda13_nvidia_cusparse_version = '' # placeholder +cuda13_nvidia_nccl_version = '' # placeholder +cuda13_nvidia_nvjitlink_version = '' # placeholder +cuda13_nvidia_cuda_nvrtc_version = '' # placeholder +cuda13_nvidia_curand_version = '' # placeholder +cuda13_nvidia_nvshmem_version = '' # placeholder +cuda13_nvidia_nvvm_version = '' # placeholder +cuda13_nvidia_cuda_crt_version = '' # placeholder + +if str(cuda_major_version) == '13': + nvidia_cuda_nvcc_version = cuda13_nvidia_cuda_nvcc_version + nvidia_cuda_nvrtc_version = cuda13_nvidia_cuda_nvrtc_version + nvidia_cuda_runtime_version = cuda13_nvidia_cuda_runtime_version +else: + nvidia_cuda_nvcc_version = cuda12_nvidia_cuda_nvcc_version + nvidia_cuda_nvrtc_version = cuda12_nvidia_cuda_nvrtc_version + nvidia_cuda_runtime_version = cuda12_nvidia_cuda_runtime_version # We use the same setup.py for all tensorflow_* packages and for the nightly # equivalents (tf_nightly_*). The package is controlled from the argument line @@ -156,33 +181,67 @@ if collaborator_build: + ';platform_system=="Windows"', ] +NVIDIA_CUDA12_PACKAGES = [ + f'nvidia-cublas-cu12{cuda12_nvidia_cublas_version}', + f'nvidia-cuda-cupti-cu12{cuda12_nvidia_cuda_cupti_version}', + f'nvidia-cuda-nvcc-cu12{cuda12_nvidia_cuda_nvcc_version}', + f'nvidia-cuda-nvrtc-cu12{cuda12_nvidia_cuda_nvrtc_version}', + f'nvidia-cuda-runtime-cu12{cuda12_nvidia_cuda_runtime_version}', + f'nvidia-cudnn-cu12{cuda12_nvidia_cudnn_version}', + f'nvidia-cufft-cu12{cuda12_nvidia_cufft_version}', + f'nvidia-curand-cu12{cuda12_nvidia_curand_version}', + f'nvidia-cusolver-cu12{cuda12_nvidia_cusolver_version}', + f'nvidia-cusparse-cu12{cuda12_nvidia_cusparse_version}', + f'nvidia-nccl-cu12{cuda12_nvidia_nccl_version}', + f'nvidia-nvjitlink-cu12{cuda12_nvidia_nvjitlink_version}', + f'nvidia-nvshmem-cu12{cuda12_nvidia_nvshmem_version}', +] + +NVIDIA_CUDA13_PACKAGES = [ + f'nvidia-cublas{cuda13_nvidia_cublas_version}', + f'nvidia-cuda-cupti{cuda13_nvidia_cuda_cupti_version}', + f'nvidia-cuda-nvcc{cuda13_nvidia_cuda_nvcc_version}', + f'nvidia-cuda-nvrtc{cuda13_nvidia_cuda_nvrtc_version}', + f'nvidia-cuda-runtime{cuda13_nvidia_cuda_runtime_version}', + f'nvidia-cudnn-cu13{cuda13_nvidia_cudnn_version}', + f'nvidia-cufft{cuda13_nvidia_cufft_version}', + f'nvidia-curand{cuda13_nvidia_curand_version}', + f'nvidia-cusolver{cuda13_nvidia_cusolver_version}', + f'nvidia-cusparse{cuda13_nvidia_cusparse_version}', + f'nvidia-nccl-cu13{cuda13_nvidia_nccl_version}', + f'nvidia-nvjitlink{cuda13_nvidia_nvjitlink_version}', + f'nvidia-nvshmem-cu13{cuda13_nvidia_nvshmem_version}', + f'nvidia-nvvm{cuda13_nvidia_nvvm_version}', + f'nvidia-cuda-crt{cuda13_nvidia_cuda_crt_version}', +] + # Set up extra packages, which are optional sets of other Python package deps. -# E.g. "pip install tensorflow[and-cuda]" below installs the normal TF deps, -# plus the CUDA libraries listed. -EXTRA_PACKAGES = { - 'and-cuda': [ - # TODO(nluehr): set nvidia-* versions based on build components. - f'nvidia-cublas{cuda_wheel_suffix}{nvidia_cublas_version}', - f'nvidia-cuda-cupti{cuda_wheel_suffix}{nvidia_cuda_cupti_version}', - f'nvidia-cuda-nvcc{cuda_wheel_suffix}{nvidia_cuda_nvcc_version}', - f'nvidia-cuda-nvrtc{cuda_wheel_suffix}{nvidia_cuda_nvrtc_version}', - f'nvidia-cuda-runtime{cuda_wheel_suffix}{nvidia_cuda_runtime_version}', - f'nvidia-cudnn-cu13{nvidia_cudnn_version}' if cuda_major_version == '13' else f'nvidia-cudnn-cu12{nvidia_cudnn_version}', - f'nvidia-cufft{cuda_wheel_suffix}{nvidia_cufft_version}', - f'nvidia-curand{cuda_wheel_suffix}{nvidia_curand_version}', - f'nvidia-cusolver{cuda_wheel_suffix}{nvidia_cusolver_version}', - f'nvidia-cusparse{cuda_wheel_suffix}{nvidia_cusparse_version}', - f'nvidia-nccl-cu13{nvidia_nccl_version}' if cuda_major_version == '13' else f'nvidia-nccl-cu12{nvidia_nccl_version}', - f'nvidia-nvjitlink{cuda_wheel_suffix}{nvidia_nvjitlink_version}', - f'nvidia-nvshmem-cu13{nvidia_nvshmem_version}' if cuda_major_version == '13' else f'nvidia-nvshmem-cu12{nvidia_nvshmem_version}', - ], - 'gcs-filesystem': [ - ('tensorflow-io-gcs-filesystem>=0.23.1; ' - 'sys_platform!="win32" and python_version<"3.13"'), - ('tensorflow-io-gcs-filesystem>=0.23.1; ' - 'sys_platform=="win32" and python_version<"3.12"'), - ] -} +# E.g. "pip install tensorflow[and-cuda]" or "pip install tensorflow[cuda12]" +if str(cuda_major_version) == '13': + EXTRA_PACKAGES = { + 'and-cuda': NVIDIA_CUDA13_PACKAGES, + 'cuda': NVIDIA_CUDA13_PACKAGES, + 'cuda13': NVIDIA_CUDA13_PACKAGES, + 'cuda13-local': [], + } +else: + EXTRA_PACKAGES = { + 'and-cuda': NVIDIA_CUDA12_PACKAGES, + 'cuda': NVIDIA_CUDA12_PACKAGES, + 'cuda12': NVIDIA_CUDA12_PACKAGES, + 'cuda12-local': [], + } + +EXTRA_PACKAGES['gcs-filesystem'] = [ + ( + 'tensorflow-io-gcs-filesystem>=0.23.1; ' + 'sys_platform!="win32" and python_version<"3.13"' + ), + ( + 'tensorflow-io-gcs-filesystem>=0.23.1; ' + 'sys_platform=="win32" and python_version<"3.12"' + ), +] DOCLINES = __doc__.split('\n') diff --git a/tensorflow/workspace3.bzl b/tensorflow/workspace3.bzl index 06809b6207741a..bf557b54aa9769 100644 --- a/tensorflow/workspace3.bzl +++ b/tensorflow/workspace3.bzl @@ -60,19 +60,19 @@ def workspace(): tf_http_archive( name = "bazel_features", - sha256 = "4fd9922d464686820ffd8fcefa28ccffa147f7cdc6b6ac0d8b07fde565c65d66", - strip_prefix = "bazel_features-1.25.0", + sha256 = "094367e732ece23f334eaf84089a720b861d053beeba6a6a68356d3aee1dc32b", + strip_prefix = "bazel_features-1.50.0", urls = tf_mirror_urls( - "https://github.com/bazel-contrib/bazel_features/releases/download/v1.25.0/bazel_features-v1.25.0.tar.gz", + "https://github.com/bazel-contrib/bazel_features/releases/download/v1.50.0/bazel_features-v1.50.0.tar.gz", ), ) tf_http_archive( name = "rules_cc", - sha256 = "69e05df29f0010ba248ef8dafc1f084c8fd2f5c553da634422d8167f5c4b277b", - strip_prefix = "rules_cc-0.2.20", + sha256 = "bd7124a844d0403b4b353bcea34d6c8b2ba88dc26881c26c9ee668da89b71846", + strip_prefix = "rules_cc-0.2.25", urls = tf_mirror_urls( - "https://github.com/bazelbuild/rules_cc/releases/download/0.2.20/rules_cc-0.2.20.tar.gz", + "https://github.com/bazelbuild/rules_cc/releases/download/0.2.25/rules_cc-0.2.25.tar.gz", ), ) diff --git a/third_party/xla/.bant-macros b/third_party/xla/.bant-macros index 47774a7f8fe9e8..d3fb19952adabf 100644 --- a/third_party/xla/.bant-macros +++ b/third_party/xla/.bant-macros @@ -33,7 +33,6 @@ exhaustive_xla_test = bant_forward_args(cc_test()) # Binary rules xla_cc_binary = bant_forward_args(cc_binary()) -cc_binary = bant_forward_args(cc_binary()) # Library rules onednn_cc_library = bant_forward_args(cc_library()) diff --git a/third_party/xla/.github/workflows/scorecards-analysis.yml b/third_party/xla/.github/workflows/scorecards-analysis.yml index ec2c190fd2ef51..f1f84938ab300c 100644 --- a/third_party/xla/.github/workflows/scorecards-analysis.yml +++ b/third_party/xla/.github/workflows/scorecards-analysis.yml @@ -67,6 +67,6 @@ jobs: # Upload the results to GitHub's code scanning dashboard (optional). # Commenting out will disable upload of results to your repo's Code Scanning dashboard - name: "Upload to code-scanning" - uses: github/codeql-action/upload-sarif@cdf488f595d80d6e07e03d4674febd5ab45fa938 # v4.37.9 + uses: github/codeql-action/upload-sarif@b96794f015dfd88f77b49b1c93e0fa7110f94c63 # v4.38.0 with: sarif_file: results.sarif diff --git a/third_party/xla/BUILD.bazel b/third_party/xla/BUILD.bazel index 5a6d40abe0a80a..3a2613f00bbf71 100644 --- a/third_party/xla/BUILD.bazel +++ b/third_party/xla/BUILD.bazel @@ -28,6 +28,8 @@ license( license_kinds = ["@rules_license//licenses/spdx:Apache-2.0"], ) +exports_files([".bant-macros"]) + filegroup( name = "clang_tidy_config", srcs = [".clang-tidy"], diff --git a/third_party/xla/MODULE.bazel b/third_party/xla/MODULE.bazel index 432b20da54537f..9b994cf3db0c72 100644 --- a/third_party/xla/MODULE.bazel +++ b/third_party/xla/MODULE.bazel @@ -23,7 +23,7 @@ bazel_dep(name = "pybind11_abseil", version = "202402.0") bazel_dep(name = "pybind11_bazel", version = "3.0.0") bazel_dep(name = "pybind11_protobuf", version = "0.0.0-20250210-f02a2b7") bazel_dep(name = "re2", version = "2025-11-05.bcr.1", repo_name = "com_googlesource_code_re2") -bazel_dep(name = "rules_cc", version = "0.2.20") +bazel_dep(name = "rules_cc", version = "0.2.25") bazel_dep(name = "rules_java", version = "8.16.1") bazel_dep(name = "rules_license", version = "1.0.0") bazel_dep(name = "rules_python", version = "2.2.0") diff --git a/third_party/xla/build_tools/lint/BUILD b/third_party/xla/build_tools/lint/BUILD index bf17ce87f87517..f9ea2203214e5a 100644 --- a/third_party/xla/build_tools/lint/BUILD +++ b/third_party/xla/build_tools/lint/BUILD @@ -117,7 +117,9 @@ pytype_strict_library( py_test( name = "run_dwyu_test", srcs = ["run_dwyu_test.py"], + data = ["//:.bant-macros"], # copybara:comment_replace data = ["//third_party/xla:.bant-macros"], deps = [ + ":check_dwyu", ":run_dwyu", ], ) diff --git a/third_party/xla/build_tools/lint/run_dwyu_test.py b/third_party/xla/build_tools/lint/run_dwyu_test.py index 32f767327bde87..16d0ba4d0c6db7 100644 --- a/third_party/xla/build_tools/lint/run_dwyu_test.py +++ b/third_party/xla/build_tools/lint/run_dwyu_test.py @@ -22,9 +22,11 @@ import tempfile import unittest +from build_tools.lint import check_dwyu from build_tools.lint import run_dwyu _BANT = shutil.which(os.environ.get("BANT_BIN", "bant")) +_BANT_MACROS = pathlib.Path(run_dwyu.__file__).parents[2] / ".bant-macros" class BantWorkspaceTest(unittest.TestCase): @@ -42,7 +44,7 @@ def setUp(self): self.external.mkdir(parents=True) (self.external / "_main").symlink_to(self.root, target_is_directory=True) self.write_file("MODULE.bazel", 'module(name = "xla")\n') - self.write_file(".bant-macros", "") + self.write_file(".bant-macros", _BANT_MACROS.read_text()) self.write_file("third_party/tsl/WORKSPACE", "") def write_file(self, path, contents): @@ -138,7 +140,7 @@ def setUp(self): def check_targets(self, *targets): return subprocess.run( [ - sys.executable, + sys.executable or shutil.which("python3"), run_dwyu.__file__, "--bant", _BANT, @@ -229,10 +231,6 @@ def test_keeps_implementation_edge_but_reports_other_unused_dependency(self): ) def test_conditional_deps_macro_resolves_wrapped_dependency(self): - self.write_file( - ".bant-macros", - "if_cuda_is_configured = _arg_0\nif_static = _arg_0\n", - ) self.write_file( "xla/consumer/BUILD", """\ @@ -248,7 +246,25 @@ def test_conditional_deps_macro_resolves_wrapped_dependency(self): self.assertEqual(result.stdout, "") self.assertIn("Checked DWYU on 1 targets.", result.stderr) + def test_all_allowed_rules_supported_by_bant_macros(self): + rules = [] + labels = [] + for index, rule in enumerate(check_dwyu.DEFAULT_ALLOWED_RULES): + target = f"target_{index}_{rule}" + rules.append( + f'{rule}(name = "{target}", srcs = ["consumer.cc"],' + ' deps = ["//xla/platform:errors"])' + ) + labels.append(f"//xla/consumer:{target}") + self.write_file("xla/consumer/BUILD", "\n".join(rules) + "\n") + result = self.check_targets(*labels) + self.assertEqual(result.returncode, 0, result.stdout + result.stderr) + self.assertEqual(result.stdout, "") + self.assertIn( + f"Checked DWYU on {len(labels)} targets.", + result.stderr, + ) + if __name__ == "__main__": unittest.main() - diff --git a/third_party/xla/third_party/py/setup_py_nvidia_dependencies_util.py b/third_party/xla/third_party/py/setup_py_nvidia_dependencies_util.py index 904cdcf4ecae21..b0d135a34a7e12 100644 --- a/third_party/xla/third_party/py/setup_py_nvidia_dependencies_util.py +++ b/third_party/xla/third_party/py/setup_py_nvidia_dependencies_util.py @@ -48,35 +48,55 @@ def get_setup_py_content_with_nvidia_wheel_versions( ) -> str: nvidia_wheel_versions = {"12": {}, "13": {}} for line in nvidia_wheel_versions_data.splitlines(): + line = line.strip() + if not line or line.startswith("#"): + continue match = NVIDIA_WHEEL_VERSIONS_PATTERN.match(line) - if match: + if match and match.group(1).startswith("nvidia"): wheel_name = match.group(1).replace("-", "_") for suffix, version in {"_cu12": "12", "_cu13": "13", "": "13"}.items(): if not wheel_name.endswith(suffix): continue - wheel_name = wheel_name.replace(suffix, "") + "_version" + if suffix: + wheel_name = wheel_name[:-len(suffix)] + wheel_name = wheel_name + "_version" nvidia_wheel_versions[version][wheel_name] = match.group(2).strip() break - if not cuda_version: - cuda_version = "12" + cuda_version_str = str(cuda_version) if cuda_version else "12" setup_py_content = setup_py_content.replace( - "cuda_version = 0 # placeholder", f"cuda_version = {cuda_version}" + "cuda_version = 0 # placeholder", f"cuda_version = {cuda_version_str}" ) setup_py_content = setup_py_content.replace( "cuda_wheel_suffix = '' # placeholder", - "cuda_wheel_suffix = '-cu12'" if cuda_version == "12" else "cuda_wheel_suffix = ''", + ( + "cuda_wheel_suffix = '-cu12'" + if cuda_version_str == "12" + else "cuda_wheel_suffix = ''" + ), ) setup_py_content = setup_py_content.replace( "cuda_major_version = '12' # placeholder", - f"cuda_major_version = '{cuda_version}'", + f"cuda_major_version = '{cuda_version_str}'", ) - for version_name, version_value in nvidia_wheel_versions[ - str(cuda_version) - ].items(): + for version_name, version_value in nvidia_wheel_versions.get("12", {}).items(): setup_py_content = setup_py_content.replace( - f"{version_name} = '' # placeholder", + f"cuda12_{version_name} = '' # placeholder", + f"cuda12_{version_name} = '{version_value}'", + ) + for version_name, version_value in nvidia_wheel_versions.get("13", {}).items(): + setup_py_content = setup_py_content.replace( + f"cuda13_{version_name} = '' # placeholder", + f"cuda13_{version_name} = '{version_value}'", + ) + for version_name, version_value in nvidia_wheel_versions.get( + cuda_version_str, {} + ).items(): + setup_py_content = re.sub( + rf"^{re.escape(version_name)} = '' # placeholder$", f"{version_name} = '{version_value}'", + setup_py_content, + flags=re.MULTILINE, ) return setup_py_content diff --git a/third_party/xla/third_party/stablehlo/temporary.patch b/third_party/xla/third_party/stablehlo/temporary.patch index 7cc02539c261a5..b6c9e2d8cb1f8a 100644 --- a/third_party/xla/third_party/stablehlo/temporary.patch +++ b/third_party/xla/third_party/stablehlo/temporary.patch @@ -2017,6 +2017,104 @@ diff --ruN a/stablehlo/stablehlo/dialect/ChloEnums.td b/stablehlo/stablehlo/dial + #endif // STABLEHLO_DIALECT_CHLO_ENUMS +diff --ruN a/stablehlo/stablehlo/dialect/ChloOps.cpp b/stablehlo/stablehlo/dialect/ChloOps.cpp +--- stablehlo/stablehlo/dialect/ChloOps.cpp ++++ stablehlo/stablehlo/dialect/ChloOps.cpp +@@ -827,6 +827,10 @@ + if (adaptor.getScanDimSize().has_value()) { + dimSize = adaptor.getScanDimSize().value(); + } ++ int64_t dimBound = ShapedType::kDynamic; ++ // Prototype encoding attribute used to construct dialect type extensions ++ // if an output does not already have an encoding. ++ Attribute prototypeEncoding; + for (auto [i, type] : llvm::enumerate(adaptor.getInputs().getTypes())) { + auto inputType = dyn_cast(type); + if (!inputType) { +@@ -839,7 +843,26 @@ + if (dimSize == ShapedType::kDynamic) { + dimSize = inputType.getDimSize(dim); + } +- } ++ if (inputType.getEncoding()) { ++ prototypeEncoding = inputType.getEncoding(); ++ } ++ ArrayRef bounds = hlo::encodingToBounds(inputType.getEncoding()); ++ if (!bounds.empty() && dim < static_cast(bounds.size())) { ++ int64_t currentBound = bounds[dim]; ++ if (currentBound != ShapedType::kDynamic) { ++ if (dimBound == ShapedType::kDynamic) { ++ dimBound = currentBound; ++ } else { ++ dimBound = std::min(dimBound, currentBound); ++ } ++ } ++ } ++ } ++ ++ // If the scan dimension size is dynamic, propagate any bound inferred from ++ // the inputs; otherwise static dimensions do not need a bound. ++ const int64_t insertBound = ++ (dimSize == ShapedType::kDynamic) ? dimBound : ShapedType::kDynamic; + + for (auto [i, type] : llvm::enumerate(terminator->getOperands().getTypes())) { + auto resultType = dyn_cast(type); +@@ -849,10 +872,34 @@ + } + SmallVector shape(resultType.getShape().begin(), + resultType.getShape().end()); ++ Attribute encoding = resultType.getEncoding(); ++ // The first `numOutputs` terminator operands correspond to scan outputs, ++ // which have the scan dimension re-inserted. The remaining operands are ++ // carries whose shapes and encodings are preserved as-is. + if (i < numOutputs) { + shape.insert(std::next(shape.begin(), dim), dimSize); +- } +- inferredReturnShapes.emplace_back(shape, resultType.getElementType()); ++ SmallVector resultBounds = ++ llvm::to_vector(hlo::encodingToBounds(encoding)); ++ // Re-insert bounds along the scan dimension if the result has bounds or ++ // if the scan dimension has a bound. ++ if (!resultBounds.empty() || insertBound != ShapedType::kDynamic) { ++ // If the body result did not have bounds, initialize all dimensions to ++ // dynamic and borrow the prototype encoding to create dialect type ++ // extensions. ++ if (resultBounds.empty()) { ++ resultBounds.assign(resultType.getRank(), ShapedType::kDynamic); ++ if (!encoding) { ++ encoding = prototypeEncoding; ++ } ++ } ++ resultBounds.insert(std::next(resultBounds.begin(), dim), insertBound); ++ if (encoding) { ++ encoding = hlo::boundsToEncoding(encoding, resultBounds); ++ } ++ } ++ } ++ inferredReturnShapes.emplace_back(shape, resultType.getElementType(), ++ encoding); + } + + return success(); +@@ -902,7 +949,17 @@ + if (i < getInputs().size()) { + auto argShape = llvm::to_vector(argType.getShape()); + argShape.erase(std::next(argShape.begin(), dim)); +- argType = argType.clone(argShape); ++ Attribute encoding = argType.getEncoding(); ++ SmallVector argBounds = ++ llvm::to_vector(hlo::encodingToBounds(encoding)); ++ if (!argBounds.empty()) { ++ argBounds.erase(std::next(argBounds.begin(), dim)); ++ encoding = argBounds.empty() ++ ? Attribute() ++ : hlo::boundsToEncoding(encoding, argBounds); ++ } ++ argType = ++ RankedTensorType::get(argShape, argType.getElementType(), encoding); + } + if (!hlo::isCompatibleForHloTypeInference( + argType, bodyBlock.getArgument(i).getType())) { diff --ruN a/stablehlo/stablehlo/dialect/ReplicaGroupUtils.cpp b/stablehlo/stablehlo/dialect/ReplicaGroupUtils.cpp --- stablehlo/stablehlo/dialect/ReplicaGroupUtils.cpp +++ stablehlo/stablehlo/dialect/ReplicaGroupUtils.cpp @@ -2974,6 +3072,65 @@ diff --ruN a/stablehlo/stablehlo/tests/BUILD.bazel b/stablehlo/stablehlo/tests/B + "//stablehlo/tests:stablehlo_tests", + ], +) +diff --ruN a/stablehlo/stablehlo/tests/chlo/chlo_legalize_to_stablehlo.mlir b/stablehlo/stablehlo/tests/chlo/chlo_legalize_to_stablehlo.mlir +--- stablehlo/stablehlo/tests/chlo/chlo_legalize_to_stablehlo.mlir ++++ stablehlo/stablehlo/tests/chlo/chlo_legalize_to_stablehlo.mlir +@@ -5424,6 +5424,55 @@ + + // ----- + ++// CHECK-LABEL: func.func @scan_bounded( ++// CHECK-SAME: %[[ARG0:.*]]: tensor>, ++// CHECK-SAME: %[[ARG1:.*]]: tensor>) -> (tensor>, tensor>) { ++// CHECK: %[[GET_DIM_SIZE_0:.*]] = stablehlo.get_dimension_size %[[ARG0]], dim = 0 : (tensor>) -> tensor ++// CHECK: %[[CONVERT_0:.*]] = stablehlo.convert %[[GET_DIM_SIZE_0]] : (tensor) -> tensor ++// CHECK: %[[C0:.*]] = stablehlo.constant dense<0> : tensor ++// CHECK: %[[C0_I32:.*]] = stablehlo.constant dense<0> : tensor ++// CHECK: %[[SHAPE:.*]] = shape.shape_of %[[ARG0]] : tensor> -> tensor<2xindex> ++// CHECK: %[[INIT_OUT:.*]] = stablehlo.dynamic_broadcast_in_dim %[[C0_I32]], %[[SHAPE]], dims = [] : (tensor, tensor<2xindex>) -> tensor> ++// CHECK: %[[WHILE:.*]]:3 = stablehlo.while(%[[ITER_I:.*]] = %[[C0]], %[[ITER_CARRY:.*]] = %[[ARG1]], %[[ITER_OUT:.*]] = %[[INIT_OUT]]) : tensor, tensor>, tensor> ++// CHECK: cond { ++// CHECK: %[[CMP:.*]] = stablehlo.compare LT, %[[ITER_I]], %[[CONVERT_0]] : (tensor, tensor) -> tensor ++// CHECK: stablehlo.return %[[CMP]] : tensor ++// CHECK: } do { ++// CHECK: %[[C0_1:.*]] = stablehlo.constant dense<0> : tensor ++// CHECK: %[[I_RESHAPED:.*]] = stablehlo.reshape %[[ITER_I]] : (tensor) -> tensor<1xi64> ++// CHECK: %[[C0_RESHAPED:.*]] = stablehlo.reshape %[[C0_1]] : (tensor) -> tensor<1xi64> ++// CHECK: %[[START_INDICES:.*]] = stablehlo.concatenate %[[I_RESHAPED]], %[[C0_RESHAPED]], dim = 0 : (tensor<1xi64>, tensor<1xi64>) -> tensor<2xi64> ++// CHECK: %[[C1:.*]] = stablehlo.constant dense<1> : tensor ++// CHECK: %[[I_PLUS_1:.*]] = stablehlo.add %[[ITER_I]], %[[C1]] : tensor ++// CHECK: %[[GET_DIM_SIZE_1:.*]] = stablehlo.get_dimension_size %[[ARG0]], dim = 1 : (tensor>) -> tensor ++// CHECK: %[[CONVERT_1:.*]] = stablehlo.convert %[[GET_DIM_SIZE_1]] : (tensor) -> tensor ++// CHECK: %[[LIMIT_0:.*]] = stablehlo.reshape %[[I_PLUS_1]] : (tensor) -> tensor<1xi64> ++// CHECK: %[[LIMIT_1:.*]] = stablehlo.reshape %[[CONVERT_1]] : (tensor) -> tensor<1xi64> ++// CHECK: %[[LIMIT_INDICES:.*]] = stablehlo.concatenate %[[LIMIT_0]], %[[LIMIT_1]], dim = 0 : (tensor<1xi64>, tensor<1xi64>) -> tensor<2xi64> ++// CHECK: %[[STRIDES:.*]] = stablehlo.constant dense<1> : tensor<2xi64> ++// CHECK: %[[SLICE:.*]] = stablehlo.real_dynamic_slice %[[ARG0]], %[[START_INDICES]], %[[LIMIT_INDICES]], %[[STRIDES]] : (tensor>, tensor<2xi64>, tensor<2xi64>, tensor<2xi64>) -> tensor<1x?xi32, #stablehlo.bounds> ++// CHECK: %[[INPUT_ELEM:.*]] = stablehlo.reshape %[[SLICE]] : (tensor<1x?xi32, #stablehlo.bounds>) -> tensor> ++// CHECK: %[[ADD:.*]] = stablehlo.add %[[INPUT_ELEM]], %[[ITER_CARRY]] : tensor> ++// CHECK: %[[OUT_ELEM:.*]] = stablehlo.reshape %[[ADD]] : (tensor>) -> tensor<1x?xi32, #stablehlo.bounds> ++// CHECK: %[[C0_2:.*]] = stablehlo.constant dense<0> : tensor ++// CHECK: %[[UPDATED_OUT:.*]] = stablehlo.dynamic_update_slice %[[ITER_OUT]], %[[OUT_ELEM]], %[[ITER_I]], %[[C0_2]] : (tensor>, tensor<1x?xi32, #stablehlo.bounds>, tensor, tensor) -> tensor> ++// CHECK: %[[C1_2:.*]] = stablehlo.constant dense<1> : tensor ++// CHECK: %[[NEXT_I:.*]] = stablehlo.add %[[ITER_I]], %[[C1_2]] : tensor ++// CHECK: stablehlo.return %[[NEXT_I]], %[[ADD]], %[[UPDATED_OUT]] : tensor, tensor>, tensor> ++// CHECK: } ++// CHECK: return %[[WHILE]]#2, %[[WHILE]]#1 : tensor>, tensor> ++// CHECK: } ++func.func @scan_bounded(%arg0: tensor>, %arg1: tensor>) -> (tensor>, tensor>) { ++ %0:2 = chlo.scan(%arg0) inits(%arg1) dimension=0 { ++ ^bb0(%scan_arg0: tensor>, %scan_arg1: tensor>): ++ %1 = stablehlo.add %scan_arg0, %scan_arg1 : tensor> ++ stablehlo.return %1, %1 : tensor>, tensor> ++ } : (tensor>, tensor>) -> (tensor>, tensor>) ++ func.return %0#0, %0#1 : tensor>, tensor> ++} ++ ++// ----- ++ + // CHECK-LABEL: func.func @mulhi_i32( + // CHECK-SAME: %[[ARG0:.*]]: tensor<4xi32>, + // CHECK-SAME: %[[ARG1:.*]]: tensor<4xi32>) -> tensor<4xi32> { diff --ruN a/stablehlo/stablehlo/tests/interpret/all_gather.mlir b/stablehlo/stablehlo/tests/interpret/all_gather.mlir --- stablehlo/stablehlo/tests/interpret/all_gather.mlir +++ stablehlo/stablehlo/tests/interpret/all_gather.mlir @@ -3164,6 +3321,91 @@ diff --ruN a/stablehlo/stablehlo/tests/interpret/reduce_scatter.mlir b/stablehlo + func.return + } +} +diff --ruN a/stablehlo/stablehlo/tests/ops_chlo.mlir b/stablehlo/stablehlo/tests/ops_chlo.mlir +--- stablehlo/stablehlo/tests/ops_chlo.mlir ++++ stablehlo/stablehlo/tests/ops_chlo.mlir +@@ -485,6 +485,50 @@ + stablehlo.return %1, %2, %1, %2 : tensor<3xf32>, tensor<3xi32>, tensor<3xf32>, tensor<3xi32> + } : (tensor<2x3xf32>, tensor<2x3xi32>, tensor<3xf32>, tensor<3xi32>) -> (tensor<2x3xf32>, tensor<2x3xi32>, tensor<3xf32>, tensor<3xi32>) + func.return %0#0, %0#1 : tensor<2x3xf32>, tensor<2x3xi32> ++} ++ ++// ----- ++ ++// CHECK-LABEL: func @scan_bounded_dynamism ++func.func @scan_bounded_dynamism(%arg0: tensor>, %arg1: tensor>) -> (tensor>, tensor>) { ++ // CHECK: chlo.scan ++ // CHECK: } : (tensor>, tensor>) -> (tensor>, tensor>) ++ %0, %1 = chlo.scan (%arg0) inits (%arg1) dimension = 0 { ++ ^bb0(%input0: tensor>, %carry0: tensor>): ++ %2 = stablehlo.add %input0, %carry0 : tensor> ++ stablehlo.return %2, %2 : tensor>, tensor> ++ } : (tensor>, tensor>) -> (tensor>, tensor>) ++ func.return %0, %1 : tensor>, tensor> ++} ++ ++// ----- ++ ++// CHECK-LABEL: func @scan_bounded_dynamism_multiple_inputs ++func.func @scan_bounded_dynamism_multiple_inputs(%arg0: tensor>, %arg1: tensor>, %arg2: tensor>, %arg3: tensor>) -> (tensor>, tensor>, tensor>, tensor>) { ++ // CHECK: chlo.scan ++ // CHECK: } : (tensor>, tensor>, tensor>, tensor>) -> (tensor>, tensor>, tensor>, tensor>) ++ %0:4 = chlo.scan (%arg0, %arg2) inits (%arg1, %arg3) dimension = 0 { ++ ^bb0(%input0: tensor>, %input1: tensor>, %carry0: tensor>, %carry1: tensor>): ++ %1 = stablehlo.add %input0, %carry0 : tensor> ++ %2 = stablehlo.add %input1, %carry1 : tensor> ++ stablehlo.return %1, %2, %1, %2 : tensor>, tensor>, tensor>, tensor> ++ } : (tensor>, tensor>, tensor>, tensor>) -> (tensor>, tensor>, tensor>, tensor>) ++ func.return %0#0, %0#1, %0#2, %0#3 : tensor>, tensor>, tensor>, tensor> ++} ++ ++// ----- ++ ++// CHECK-LABEL: func @scan_bounded_dynamism_multiple_inputs_same_rank ++func.func @scan_bounded_dynamism_multiple_inputs_same_rank(%arg0: tensor>, %arg1: tensor>, %arg2: tensor>, %arg3: tensor>) -> (tensor>, tensor>, tensor>, tensor>) { ++ // CHECK: chlo.scan ++ // CHECK: } : (tensor>, tensor>, tensor>, tensor>) -> (tensor>, tensor>, tensor>, tensor>) ++ %0:4 = chlo.scan (%arg0, %arg2) inits (%arg1, %arg3) dimension = 0 { ++ ^bb0(%input0: tensor>, %input1: tensor>, %carry0: tensor>, %carry1: tensor>): ++ %1 = stablehlo.add %input0, %carry0 : tensor> ++ %2 = stablehlo.add %input1, %carry1 : tensor> ++ stablehlo.return %1, %2, %1, %2 : tensor>, tensor>, tensor>, tensor> ++ } : (tensor>, tensor>, tensor>, tensor>) -> (tensor>, tensor>, tensor>, tensor>) ++ func.return %0#0, %0#1, %0#2, %0#3 : tensor>, tensor>, tensor>, tensor> + } + + // ----- +@@ -561,6 +605,30 @@ + + // ----- + ++func.func @scan_bounded_dynamism_mismatch(%arg0: tensor>, %arg1: tensor>) -> (tensor>, tensor>) { ++ // expected-error @+1 {{'chlo.scan' op operand and body argument 0 are incompatible}} ++ %0, %1 = chlo.scan (%arg0) inits (%arg1) dimension = 0 { ++ ^bb0(%input0: tensor<4xf32>, %carry0: tensor>): ++ stablehlo.return %carry0, %carry0 : tensor>, tensor> ++ } : (tensor>, tensor>) -> (tensor>, tensor>) ++ func.return %0, %1 : tensor>, tensor> ++} ++ ++// ----- ++ ++func.func @scan_bounded_dynamism_result_rank_mismatch(%arg0: tensor>, %arg1: tensor>) -> (tensor, tensor>) { ++ // expected-error @+2 {{op failed to infer returned types}} ++ // expected-error @+1 {{'chlo.scan' op inferred type(s) 'tensor>', 'tensor>' are incompatible with return type(s) of operation 'tensor', 'tensor>'}} ++ %0, %1 = chlo.scan (%arg0) inits (%arg1) dimension = 0 { ++ ^bb0(%input0: tensor>, %carry0: tensor>): ++ %2 = stablehlo.add %input0, %carry0 : tensor> ++ stablehlo.return %2, %2 : tensor>, tensor> ++ } : (tensor>, tensor>) -> (tensor, tensor>) ++ func.return %0, %1 : tensor, tensor> ++} ++ ++// ----- ++ + // CHECK-LABEL: func @mulhi_i32 + func.func @mulhi_i32(%arg0: tensor<4xi32>, %arg1: tensor<4xi32>) -> tensor<4xi32> { + %0 = "chlo.mulhi"(%arg0, %arg1) : (tensor<4xi32>, tensor<4xi32>) -> tensor<4xi32> diff --ruN a/stablehlo/stablehlo/tests/transforms/mesh_axes_replica_group_compatibility.mlir b/stablehlo/stablehlo/tests/transforms/mesh_axes_replica_group_compatibility.mlir --- stablehlo/stablehlo/tests/transforms/mesh_axes_replica_group_compatibility.mlir +++ stablehlo/stablehlo/tests/transforms/mesh_axes_replica_group_compatibility.mlir @@ -3319,4 +3561,70 @@ diff --ruN a/stablehlo/stablehlo/tests/transforms/stablehlo_refine_shapes.mlir b %2 = stablehlo.dynamic_iota %1, dim = 0 : (tensor<1xi64>) -> tensor return %2 : tensor } +diff --ruN a/stablehlo/stablehlo/transforms/ChloLegalizeToStablehlo.cpp b/stablehlo/stablehlo/transforms/ChloLegalizeToStablehlo.cpp +--- stablehlo/stablehlo/transforms/ChloLegalizeToStablehlo.cpp ++++ stablehlo/stablehlo/transforms/ChloLegalizeToStablehlo.cpp +@@ -2930,8 +2930,19 @@ + SmallVector sliceShape(inputType.getShape().begin(), + inputType.getShape().end()); + sliceShape[dim] = 1; ++ Attribute sliceEncoding; ++ ArrayRef inputBounds = ++ hlo::encodingToBounds(inputType.getEncoding()); ++ if (!inputBounds.empty()) { ++ SmallVector sliceBounds(inputBounds.begin(), ++ inputBounds.end()); ++ sliceBounds[dim] = ShapedType::kDynamic; ++ sliceEncoding = ++ hlo::boundsToEncoding(inputType.getEncoding(), sliceBounds); ++ } + auto sliceType = +- RankedTensorType::get(sliceShape, inputType.getElementType()); ++ RankedTensorType::get(sliceShape, inputType.getElementType(), ++ sliceEncoding); + + Value slice; + if (inputType.hasStaticShape()) { +@@ -2980,9 +2991,18 @@ + for (int64_t d = 0; d < rank; ++d) { + if (d != dim) resultShape.push_back(sliceShape[d]); + } ++ Attribute resultEncoding; ++ if (!inputBounds.empty()) { ++ SmallVector resultBounds(inputBounds.begin(), ++ inputBounds.end()); ++ resultBounds.erase(std::next(resultBounds.begin(), dim)); ++ resultEncoding = ++ hlo::boundsToEncoding(inputType.getEncoding(), resultBounds); ++ } + Value reshaped = ReshapeOp::create( + rewriter, loc, +- RankedTensorType::get(resultShape, inputType.getElementType()), ++ RankedTensorType::get(resultShape, inputType.getElementType(), ++ resultEncoding), + slice); + slicedInputs.push_back(reshaped); + } +@@ -3004,9 +3024,20 @@ + SmallVector elementShape(outputType.getShape().begin(), + outputType.getShape().end()); + elementShape[dim] = 1; ++ Attribute elementEncoding; ++ ArrayRef outputBounds = ++ hlo::encodingToBounds(outputType.getEncoding()); ++ if (!outputBounds.empty()) { ++ SmallVector elementBounds(outputBounds.begin(), ++ outputBounds.end()); ++ elementBounds[dim] = ShapedType::kDynamic; ++ elementEncoding = ++ hlo::boundsToEncoding(outputType.getEncoding(), elementBounds); ++ } + Value reshapedElement = ReshapeOp::create( + rewriter, loc, +- RankedTensorType::get(elementShape, outputType.getElementType()), ++ RankedTensorType::get(elementShape, outputType.getElementType(), ++ elementEncoding), + newElement); + + int64_t outRank = outputType.getRank(); diff --git a/third_party/xla/workspace3.bzl b/third_party/xla/workspace3.bzl index 4ba3a8a3de1c9b..3f9a1f215739ac 100644 --- a/third_party/xla/workspace3.bzl +++ b/third_party/xla/workspace3.bzl @@ -56,17 +56,17 @@ def workspace(): tf_http_archive( name = "bazel_features", - sha256 = "4fd9922d464686820ffd8fcefa28ccffa147f7cdc6b6ac0d8b07fde565c65d66", - strip_prefix = "bazel_features-1.25.0", - urls = tf_mirror_urls("https://github.com/bazel-contrib/bazel_features/releases/download/v1.25.0/bazel_features-v1.25.0.tar.gz"), + sha256 = "094367e732ece23f334eaf84089a720b861d053beeba6a6a68356d3aee1dc32b", + strip_prefix = "bazel_features-1.50.0", + urls = tf_mirror_urls("https://github.com/bazel-contrib/bazel_features/releases/download/v1.50.0/bazel_features-v1.50.0.tar.gz"), ) tf_http_archive( name = "rules_cc", - sha256 = "69e05df29f0010ba248ef8dafc1f084c8fd2f5c553da634422d8167f5c4b277b", - strip_prefix = "rules_cc-0.2.20", + sha256 = "bd7124a844d0403b4b353bcea34d6c8b2ba88dc26881c26c9ee668da89b71846", + strip_prefix = "rules_cc-0.2.25", urls = tf_mirror_urls( - "https://github.com/bazelbuild/rules_cc/releases/download/0.2.20/rules_cc-0.2.20.tar.gz", + "https://github.com/bazelbuild/rules_cc/releases/download/0.2.25/rules_cc-0.2.25.tar.gz", ), ) diff --git a/third_party/xla/xla/backends/gpu/codegen/kernels/ptx_custom_kernel.cc b/third_party/xla/xla/backends/gpu/codegen/kernels/ptx_custom_kernel.cc index 9dd19b68021369..44883901699c6d 100644 --- a/third_party/xla/xla/backends/gpu/codegen/kernels/ptx_custom_kernel.cc +++ b/third_party/xla/xla/backends/gpu/codegen/kernels/ptx_custom_kernel.cc @@ -33,16 +33,6 @@ limitations under the License. namespace xla::gpu::kernel { namespace se = ::stream_executor; -namespace { -se::KernelArgsPackingSpec IdentityPackingSpec(int num_args) { - se::KernelArgsPackingSpec packing_spec; - for (int i = 0; i < num_args; i++) { - packing_spec.AddAddressArgument(i); - } - return packing_spec; -} -} // namespace - // Note: Make sure that the kernel_name matches the kernel name in the ptx, // otherwise you will get a "CUDA_ERROR_NOT_FOUND: named symbol not found.". // E.g. `.visible .entry AddI32(...)` would have a kernel name of "AddI32". @@ -54,7 +44,8 @@ absl::StatusOr GetPtxCustomKernel(std::string kernel_name, size_t shared_memory_bytes) { se::KernelLoaderSpec kernel_spec = se::KernelLoaderSpec::CreateCudaPtxInMemorySpec( - ptx, kernel_name, /*arity=*/num_args, IdentityPackingSpec(num_args)); + ptx, kernel_name, /*arity=*/num_args, + se::KernelArgsPackingSpec::Identity(num_args)); return CustomKernel(std::move(kernel_name), kernel_spec, block_dim, thread_dim, shared_memory_bytes); } @@ -65,7 +56,8 @@ absl::StatusOr GetPtxCustomKernel( se::ClusterDim cluster_dim, size_t shared_memory_bytes) { se::KernelLoaderSpec kernel_spec = se::KernelLoaderSpec::CreateCudaPtxInMemorySpec( - ptx, kernel_name, /*arity=*/num_args, IdentityPackingSpec(num_args)); + ptx, kernel_name, /*arity=*/num_args, + se::KernelArgsPackingSpec::Identity(num_args)); return CustomKernel(std::move(kernel_name), kernel_spec, block_dim, thread_dim, cluster_dim, shared_memory_bytes); } @@ -77,7 +69,7 @@ absl::StatusOr GetOwnedPtxCustomKernel( se::KernelLoaderSpec kernel_spec = se::KernelLoaderSpec::CreateOwningCudaPtxInMemorySpec( std::move(ptx), kernel_name, /*arity=*/num_args, - IdentityPackingSpec(num_args)); + se::KernelArgsPackingSpec::Identity(num_args)); return CustomKernel(std::move(kernel_name), kernel_spec, block_dim, thread_dim, shared_memory_bytes); } @@ -89,7 +81,7 @@ absl::StatusOr CreateOwnedCubinCustomKernel( se::KernelLoaderSpec kernel_spec = se::KernelLoaderSpec::CreateOwningCudaCubinInMemorySpec( std::move(cubin), kernel_name, /*arity=*/num_args, - IdentityPackingSpec(num_args)); + se::KernelArgsPackingSpec::Identity(num_args)); return CustomKernel(std::move(kernel_name), std::move(kernel_spec), block_dim, thread_dim, shared_memory_bytes); } diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/BUILD b/third_party/xla/xla/backends/gpu/codegen/triton/BUILD index 35be416a6d9edc..0449cd64363a0b 100644 --- a/third_party/xla/xla/backends/gpu/codegen/triton/BUILD +++ b/third_party/xla/xla/backends/gpu/codegen/triton/BUILD @@ -457,6 +457,7 @@ xla_test( "@com_google_absl//absl/status:statusor", "@com_google_absl//absl/strings", "@com_google_absl//absl/strings:str_format", + "@com_google_absl//absl/strings:string_view", "@com_google_absl//absl/time", "@com_google_absl//absl/types:span", "@com_google_googletest//:gtest", diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/dot_algorithms_test.cc b/third_party/xla/xla/backends/gpu/codegen/triton/dot_algorithms_test.cc index 26997ce8df3b22..81a06dc5487b9f 100644 --- a/third_party/xla/xla/backends/gpu/codegen/triton/dot_algorithms_test.cc +++ b/third_party/xla/xla/backends/gpu/codegen/triton/dot_algorithms_test.cc @@ -1142,6 +1142,11 @@ class NumericTestsForTriton : public TritonAlgorithmTest, }; TEST_P(NumericTestsForBlas, Infinity) { + // hipBLASLt emulates TF32 on gfx950 and returns NaN for inf*1. + if (GetParam() == PC::ALG_DOT_TF32_TF32_F32_X3 && GpuComputeComp().IsRocm() && + GpuComputeComp().rocm_compute_capability()->gfx9_mi350()) { + GTEST_SKIP() << "hipBLASLt FAST_TF32 inf*1 returns NaN on MI350 "; + } std::string hlo_text = HloText(); TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr module, GetOptimizedModule(hlo_text)); diff --git a/third_party/xla/xla/backends/gpu/runtime/ragged_all_to_all_thunk.cc b/third_party/xla/xla/backends/gpu/runtime/ragged_all_to_all_thunk.cc index 81855ce4aecfbd..2ddd40b68ff763 100644 --- a/third_party/xla/xla/backends/gpu/runtime/ragged_all_to_all_thunk.cc +++ b/third_party/xla/xla/backends/gpu/runtime/ragged_all_to_all_thunk.cc @@ -437,6 +437,7 @@ absl::Status RunNcclFallbackRaggedAllToAll( const int64_t* input_offsets = ragged_metadata_allocs[0]; const int64_t* send_sizes = ragged_metadata_allocs[1]; const int64_t* output_offsets = ragged_metadata_allocs[2]; + const int64_t* recv_sizes = ragged_metadata_allocs[3]; if (use_put_path) { XLA_VLOG_DEVICE(3, device_ordinal) @@ -445,12 +446,26 @@ absl::Status RunNcclFallbackRaggedAllToAll( int64_t element_byte_width = primitive_util::ByteWidth(element_type); se::DeviceAddressBase input_buffer = buffers[0].source_buffer; + // WaitSignal observes puts sent by `peer` to this rank, so its op_cnt must + // be the number of nonzero receives expected from that peer. + std::vector puts_from_peer(num_ranks, 0); + for (int peer = 0; peer < num_ranks; ++peer) { + for (int64_t i = 0; i < num_updates_per_replica; ++i) { + const int64_t idx = peer * num_updates_per_replica + i; + puts_from_peer[peer] += recv_sizes[idx] != 0; + } + } + Future<> future = gpu_comm->GroupExecute([&]() -> absl::Status { for (int peer = 0; peer < num_ranks; ++peer) { for (int64_t i = 0; i < num_updates_per_replica; ++i) { const int64_t idx = peer * num_updates_per_replica + i; const int64_t send_count = send_sizes[idx] * ragged_row_element_size; + // Skip zero-size sends: NCCL still pays ~12 us per LaunchPut even + // at count=0, causing O(n_devices) cost for sparse neighbor patterns. + if (send_count == 0) continue; + const se::DeviceAddressBase send_slice = GpuCollectives::Slice( input_buffer, element_type, input_offsets[idx] * ragged_row_element_size, send_count); @@ -472,8 +487,9 @@ absl::Status RunNcclFallbackRaggedAllToAll( GpuSignalDesc signal_desc(/*sig_idx=*/0, /*ctx=*/0); for (int peer = 0; peer < num_ranks; ++peer) { + if (puts_from_peer[peer] == 0) continue; ABSL_RETURN_IF_ERROR(comm.WaitSignal(RankId(peer), - /*op_cnt=*/num_updates_per_replica, + /*op_cnt=*/puts_from_peer[peer], signal_desc, GpuCollectives::On(stream)) .Await()); } @@ -483,7 +499,6 @@ absl::Status RunNcclFallbackRaggedAllToAll( XLA_VLOG_DEVICE(3, device_ordinal) << "RunFallbackNcclRaggedAllToAll: using Send/Recv path"; - const int64_t* recv_sizes = ragged_metadata_allocs[3]; Future<> future = gpu_comm->GroupExecute([&]() -> absl::Status { PrimitiveType element_type = buffers[0].element_type; @@ -494,23 +509,32 @@ absl::Status RunNcclFallbackRaggedAllToAll( for (int64_t i = 0; i < num_updates_per_replica; ++i) { for (int peer = 0; peer < num_ranks; ++peer) { int64_t idx = peer * num_updates_per_replica + i; - se::DeviceAddressBase send_slice = - GpuCollectives::Slice(input_buffer, element_type, - input_offsets[idx] * ragged_row_element_size, - send_sizes[idx] * ragged_row_element_size); - - se::DeviceAddressBase recv_slice = - GpuCollectives::Slice(output_buffer, element_type, - output_offsets[idx] * ragged_row_element_size, - recv_sizes[idx] * ragged_row_element_size); - - ABSL_RETURN_IF_ERROR(gpu_comm->LaunchSend( - send_slice, element_type, send_sizes[idx] * ragged_row_element_size, - RankId(peer), GpuCollectives::On(stream))); - - ABSL_RETURN_IF_ERROR(gpu_comm->LaunchRecv( - recv_slice, element_type, recv_sizes[idx] * ragged_row_element_size, - RankId(peer), GpuCollectives::On(stream))); + + // Send and receive refer to different slots. Skip each zero-size + // operation independently: NCCL pays ~12 us per LaunchSend/LaunchRecv + // even at count=0, causing O(n_devices) cost for sparse neighbor + // patterns (issue #46982). + if (send_sizes[idx] != 0) { + se::DeviceAddressBase send_slice = GpuCollectives::Slice( + input_buffer, element_type, + input_offsets[idx] * ragged_row_element_size, + send_sizes[idx] * ragged_row_element_size); + ABSL_RETURN_IF_ERROR( + gpu_comm->LaunchSend(send_slice, element_type, + send_sizes[idx] * ragged_row_element_size, + RankId(peer), GpuCollectives::On(stream))); + } + + if (recv_sizes[idx] != 0) { + se::DeviceAddressBase recv_slice = GpuCollectives::Slice( + output_buffer, element_type, + output_offsets[idx] * ragged_row_element_size, + recv_sizes[idx] * ragged_row_element_size); + ABSL_RETURN_IF_ERROR( + gpu_comm->LaunchRecv(recv_slice, element_type, + recv_sizes[idx] * ragged_row_element_size, + RankId(peer), GpuCollectives::On(stream))); + } } } diff --git a/third_party/xla/xla/backends/gpu/tests/ragged_all_to_all_e2e_test.cc b/third_party/xla/xla/backends/gpu/tests/ragged_all_to_all_e2e_test.cc index 60e09a30e872d1..1dec6a15f81264 100644 --- a/third_party/xla/xla/backends/gpu/tests/ragged_all_to_all_e2e_test.cc +++ b/third_party/xla/xla/backends/gpu/tests/ragged_all_to_all_e2e_test.cc @@ -891,6 +891,53 @@ TEST_P(RaggedAllToAllTest, RaggedAllToAll_2GPUs_PreservesInitialData) { EXPECT_TRUE(LiteralTestUtil::Equal(expected_outputs_[1], results[1])); } +// Regression test for issue #46982: zero-size peer slices must be skipped so +// that the collective's cost is O(actual_neighbors) not O(n_devices). +// +// Pattern: replica_0 sends 3 rows to replica_1; replica_1 sends 0 rows to +// replica_0. This creates one direction with all-zero sizes per peer, which +// previously caused ~12 us of unnecessary NCCL overhead for every zero-size +// slot in the Send/Recv or Put+Signal paths. +TEST_P(RaggedAllToAllTest, RaggedAllToAll_2GPUs_SparseNeighbors) { + absl::string_view kModuleReplicatedStr = R"( + HloModule module, num_partitions=1 + + ENTRY entry { + input = f32[4] parameter(0) + output = f32[4] parameter(1) + input_offsets = s32[2] parameter(2) + send_sizes = s32[2] parameter(3) + output_offsets = s32[2] parameter(4) + recv_sizes = s32[2] parameter(5) + ROOT ra2a = f32[4] ragged-all-to-all(input, output, input_offsets, + send_sizes, output_offsets, recv_sizes), replica_groups={{0,1}} + })"; + + const int64_t kNumReplicas = 2; + ASSERT_GE(device_count(), kNumReplicas) + << "Test requires at least " << kNumReplicas << " devices (" + << device_count() << " available)"; + + ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule( + kModuleReplicatedStr, kNumReplicas)); + + // Sparse pattern: replica_0 sends 3 rows to replica_1, and 1 row to itself. + // replica_1 sends 0 rows to replica_0 (zero-size slot), and 1 row to itself. + // The zero-size slot for replica_1->replica_0 must be skipped without hangs. + ASSERT_OK(CreateRandomTestData(module.get(), + /*input_sizes=*/{/*replica_0=*/{1, 3}, + /*replica_1=*/{0, 1}})); + + ASSERT_OK_AND_ASSIGN( + ExecutionResult execution_result, + ExecuteReplicated(std::move(module), GetInputLiteralPtrs())); + + const std::vector& results = execution_result.results; + ASSERT_EQ(results.size(), kNumReplicas); + EXPECT_TRUE(LiteralTestUtil::Equal(expected_outputs_[0], results[0])); + EXPECT_TRUE(LiteralTestUtil::Equal(expected_outputs_[1], results[1])); +} + TEST_P(RaggedAllToAllTest, RaggedAllToAll_8GPUs) { absl::string_view kModuleReplicatedStr = R"( HloModule module, num_partitions=1 diff --git a/third_party/xla/xla/backends/gpu/transforms/BUILD b/third_party/xla/xla/backends/gpu/transforms/BUILD index dc11f116a7cd84..f8107004477ca0 100644 --- a/third_party/xla/xla/backends/gpu/transforms/BUILD +++ b/third_party/xla/xla/backends/gpu/transforms/BUILD @@ -3018,6 +3018,7 @@ cc_library( "//xla:xla_data_proto_cc", "//xla/hlo/ir:hlo", "//xla/hlo/pass:hlo_pass", + "//xla/service:decision", "//xla/tsl/platform:errors", "//xla/tsl/platform:statusor", "@com_google_absl//absl/container:flat_hash_set", @@ -3025,6 +3026,7 @@ cc_library( "@com_google_absl//absl/status", "@com_google_absl//absl/status:status_macros", "@com_google_absl//absl/status:statusor", + "@com_google_absl//absl/strings", "@com_google_absl//absl/strings:string_view", "@llvm-project//mlir:AsmParser", "@llvm-project//mlir:IR", diff --git a/third_party/xla/xla/backends/gpu/transforms/composite_rewriter.cc b/third_party/xla/xla/backends/gpu/transforms/composite_rewriter.cc index aef59a56f37af0..855fe918a7285e 100644 --- a/third_party/xla/xla/backends/gpu/transforms/composite_rewriter.cc +++ b/third_party/xla/xla/backends/gpu/transforms/composite_rewriter.cc @@ -22,6 +22,7 @@ limitations under the License. #include "absl/status/status.h" #include "absl/status/status_macros.h" #include "absl/status/statusor.h" +#include "absl/strings/str_cat.h" #include "absl/strings/string_view.h" #include "mlir/AsmParser/AsmParser.h" #include "mlir/IR/Attributes.h" @@ -36,6 +37,7 @@ limitations under the License. #include "xla/hlo/ir/hlo_module.h" #include "xla/hlo/ir/hlo_opcode.h" #include "xla/literal.h" +#include "xla/service/decision.h" #include "xla/tsl/platform/errors.h" #include "xla/tsl/platform/statusor.h" #include "xla/util.h" @@ -99,151 +101,188 @@ absl::StatusOr ParseDimensionNumbers( return dnums; } +Decision IsScaledDotComposite(const HloInstruction* instruction) { + if (instruction->opcode() != HloOpcode::kCall) { + return Decision::Forbid("instruction is not a call"); + } + auto call = Cast(instruction); + if (!call->is_composite()) { + return Decision::Forbid( + absl::StrCat(instruction->name(), " is not composite")); + } + if (!call->has_frontend_attributes()) { + return Decision::Forbid("No frontend attributes"); + } + const auto& frontend_attrs = call->frontend_attributes().map(); + auto it = frontend_attrs.find("composite.name"); + if (it == frontend_attrs.end() || it->second != "xla.scaled_dot") { + return Decision::Forbid( + it != frontend_attrs.end() + ? absl::StrCat("composite.name is not xla.scaled_dot: ", it->second) + : "composite.name attribute is missing"); + } + return Decision::Allow(); +} + +Decision HasSupportedDimensionNumbers( + const DotDimensionNumbers& dot_dimension_numbers) { + if (dot_dimension_numbers.lhs_contracting_dimensions_size() != 1 || + dot_dimension_numbers.rhs_contracting_dimensions_size() != 1) { + return Decision::Forbid( + absl::StrCat("Expected 1 contracting dimension on LHS and RHS, got: ", + dot_dimension_numbers.DebugString())); + } + if (dot_dimension_numbers.lhs_batch_dimensions_size() > 1 || + dot_dimension_numbers.rhs_batch_dimensions_size() > 1) { + return Decision::Forbid( + absl::StrCat("Expected at most 1 batch dimension on LHS and RHS, got: ", + dot_dimension_numbers.DebugString())); + } + return Decision::Allow(); +} + +Decision IsValidBf16Scale(const HloInstruction* operand, + const HloInstruction* scale) { + if (scale->shape().dimensions().size() != + operand->shape().dimensions().size()) { + return Decision::Forbid("scale and operand rank mismatch for BF16"); + } + for (int64_t dim : scale->shape().dimensions()) { + if (dim != 1) { + return Decision::Forbid("scale dim != 1 for BF16"); + } + } + if (scale->opcode() != HloOpcode::kConstant) { + return Decision::Forbid("scale is not constant for BF16"); + } + if (!scale->literal().IsAllFloat(1.0)) { + return Decision::Forbid("scale is not 1.0 for BF16"); + } + return Decision::Allow(); +} + +Decision IsSupportedBf16(const HloInstruction* lhs, const HloInstruction* rhs, + const HloInstruction* lhs_scale, + const HloInstruction* rhs_scale) { + if (lhs_scale->shape().element_type() != BF16 || + rhs_scale->shape().element_type() != BF16) { + return Decision::Forbid("BF16 operands require BF16 scales"); + } + if (Decision d = IsValidBf16Scale(lhs, lhs_scale); !d) { + return d; + } + return IsValidBf16Scale(rhs, rhs_scale); +} + +Decision IsSupportedScaleAndOperand(const HloInstruction* operand, + const HloInstruction* scale, + int64_t contracting_dim) { + PrimitiveType op_type = operand->shape().element_type(); + PrimitiveType scale_type = scale->shape().element_type(); + + if (scale_type != F8E8M0FNU) { + return Decision::Forbid(absl::StrCat("Unsupported scale type: ", + PrimitiveType_Name(scale_type))); + } + + if (contracting_dim >= scale->shape().dimensions().size()) { + return Decision::Forbid("contracting_dim out of bounds for scale"); + } + int64_t operand_dim_size = operand->shape().dimensions(contracting_dim); + int64_t scale_dim_size = scale->shape().dimensions(contracting_dim); + + if (scale_dim_size == 0 || operand_dim_size % scale_dim_size != 0) { + return Decision::Forbid("operand_dim_size not divisible by scale_dim_size"); + } + int64_t scale_factor = operand_dim_size / scale_dim_size; + + if (op_type == F8E4M3FN || op_type == F8E5M2 || op_type == F4E2M1FN) { + if (scale_factor % 16 != 0) { + return Decision::Forbid( + absl::StrCat("scale_factor % 16 != 0: ", scale_factor)); + } + return Decision::Allow(); + } + + return Decision::Forbid( + absl::StrCat("Unsupported operand type: ", PrimitiveType_Name(op_type))); +} + +Decision IsSupportedScaledDot( + const HloCallInstruction* call, + const DotDimensionNumbers& dot_dimension_numbers) { + if (call->operand_count() < 4) { + return Decision::Forbid("Scaled dot call has fewer than 4 operands"); + } + const HloInstruction* lhs = call->operand(0); + const HloInstruction* rhs = call->operand(1); + const HloInstruction* lhs_scale = call->operand(2); + const HloInstruction* rhs_scale = call->operand(3); + + if (lhs->shape().element_type() == BF16 && + rhs->shape().element_type() == BF16) { + return IsSupportedBf16(lhs, rhs, lhs_scale, rhs_scale); + } + + int64_t lhs_contracting_dim = + dot_dimension_numbers.lhs_contracting_dimensions(0); + int64_t rhs_contracting_dim = + dot_dimension_numbers.rhs_contracting_dimensions(0); + + if (Decision d = + IsSupportedScaleAndOperand(lhs, lhs_scale, lhs_contracting_dim); + !d) { + return d; + } + return IsSupportedScaleAndOperand(rhs, rhs_scale, rhs_contracting_dim); +} + +absl::Status RewriteScaledDot( + HloComputation* computation, HloCallInstruction* call, + const DotDimensionNumbers& dot_dimension_numbers) { + PrecisionConfig precision{}; + precision.mutable_operand_precision()->Resize(2, PrecisionConfig::DEFAULT); + auto* scaled_dot = + computation->AddInstruction(HloInstruction::CreateScaledDot( + call->shape(), call->mutable_operand(0), call->mutable_operand(1), + call->mutable_operand(2), call->mutable_operand(3), + dot_dimension_numbers, precision)); + ABSL_RETURN_IF_ERROR(call->ReplaceAllUsesWith(scaled_dot)); + ABSL_RETURN_IF_ERROR(computation->RemoveInstruction(call)); + return absl::OkStatus(); +} + } // namespace absl::StatusOr CompositeRewriter::RewriteComputation( HloComputation* computation) { bool changed = false; for (HloInstruction* instruction : computation->MakeInstructionPostOrder()) { - if (instruction->opcode() != HloOpcode::kCall) { - continue; - } - VLOG(3) << "Found call instruction: " << instruction->name(); - auto call = Cast(instruction); - if (!call->is_composite()) { - VLOG(3) << instruction->name() << " is not composite"; - continue; - } - if (!call->has_frontend_attributes()) { - VLOG(3) << "No frontend attributes"; - continue; - } - auto frontend_attrs = call->frontend_attributes().map(); - auto key = "composite.name"; - if (!frontend_attrs.contains(key) || - frontend_attrs.at(key) != "xla.scaled_dot") { - VLOG(3) << key << " is not xla.scaled_dot: " << frontend_attrs.at(key); + if (Decision d = IsScaledDotComposite(instruction); !d) { + VLOG(3) << d.Explain(); continue; } - if (!frontend_attrs.contains("composite.attributes")) { + auto* call = Cast(instruction); + const auto& frontend_attrs = call->frontend_attributes().map(); + auto it = frontend_attrs.find("composite.attributes"); + if (it == frontend_attrs.end()) { return absl::InvalidArgumentError( "composite.attributes is not set for xla.scaled_dot"); } - ABSL_ASSIGN_OR_RETURN( - DotDimensionNumbers dot_dimension_numbers, - ParseDimensionNumbers(frontend_attrs.at("composite.attributes"))); - - if (dot_dimension_numbers.lhs_contracting_dimensions_size() != 1 || - dot_dimension_numbers.rhs_contracting_dimensions_size() != 1 || - dot_dimension_numbers.lhs_batch_dimensions_size() > 1 || - dot_dimension_numbers.rhs_batch_dimensions_size() > 1) { - LOG(ERROR) << "Unsupported dimension numbers: " - << dot_dimension_numbers.DebugString(); - continue; - } + ABSL_ASSIGN_OR_RETURN(DotDimensionNumbers dot_dimension_numbers, + ParseDimensionNumbers(it->second)); - const HloInstruction* lhs = call->operand(0); - const HloInstruction* rhs = call->operand(1); - const HloInstruction* lhs_scale = call->operand(2); - const HloInstruction* rhs_scale = call->operand(3); - - int64_t lhs_contracting_dim = - dot_dimension_numbers.lhs_contracting_dimensions(0); - int64_t rhs_contracting_dim = - dot_dimension_numbers.rhs_contracting_dimensions(0); - - auto is_supported = [&](const HloInstruction* operand, - const HloInstruction* scale, - int64_t contracting_dim) { - auto op_type = operand->shape().element_type(); - auto scale_type = scale->shape().element_type(); - auto is_fp8 = [](PrimitiveType type) { - return type == F8E4M3FN || type == F8E5M2; - }; - - if (op_type == BF16 && scale_type == BF16) { - if (scale->shape().dimensions().size() != - operand->shape().dimensions().size()) { - VLOG(3) << "scale and operand rank mismatch for BF16"; - return false; - } - for (int64_t dim : scale->shape().dimensions()) { - if (dim != 1) { - VLOG(3) << "scale dim != 1 for BF16"; - return false; - } - } - if (scale->opcode() != HloOpcode::kConstant) { - VLOG(3) << "scale is not constant for BF16"; - return false; - } - bool supported = scale->literal().IsAllFloat(1.0); - if (!supported) { - VLOG(3) << "scale is not 1.0 for BF16"; - } - return supported; - } - - auto is_supported_scale_type = [](PrimitiveType type) { - return type == F8E8M0FNU; - }; - - if (!is_supported_scale_type(scale_type)) { - VLOG(3) << "Unsupported scale type: " << PrimitiveType_Name(scale_type); - return false; - } - - if (contracting_dim >= scale->shape().dimensions().size()) { - VLOG(3) << "contracting_dim out of bounds for scale"; - return false; - } - int64_t operand_dim_size = operand->shape().dimensions(contracting_dim); - int64_t scale_dim_size = scale->shape().dimensions(contracting_dim); - - if (scale_dim_size == 0 || operand_dim_size % scale_dim_size != 0) { - VLOG(3) << "operand_dim_size not divisible by scale_dim_size"; - return false; - } - int64_t scale_factor = operand_dim_size / scale_dim_size; - - if (is_fp8(op_type)) { - bool supported = scale_factor % 16 == 0; - if (!supported) { - VLOG(3) << "scale_factor % 16 != 0: " << scale_factor; - } - return supported; - } - - if (op_type == F4E2M1FN) { - bool supported = scale_factor % 16 == 0; - if (!supported) { - VLOG(3) << "scale_factor % 16 != 0: " << scale_factor; - } - return supported; - } - - VLOG(3) << "Unsupported operand type: " << PrimitiveType_Name(op_type); - return false; - }; - - if (!is_supported(lhs, lhs_scale, lhs_contracting_dim)) { - VLOG(3) << "LHS not supported"; + if (Decision d = HasSupportedDimensionNumbers(dot_dimension_numbers); !d) { + LOG(ERROR) << d.Explain(); continue; } - if (!is_supported(rhs, rhs_scale, rhs_contracting_dim)) { - VLOG(3) << "RHS not supported"; + + if (Decision d = IsSupportedScaledDot(call, dot_dimension_numbers); !d) { + VLOG(3) << "Scaled dot composite operands not supported: " << d.Explain(); continue; } - PrecisionConfig precision{}; - precision.mutable_operand_precision()->Resize(2, PrecisionConfig::DEFAULT); - auto* scaled_dot = - computation->AddInstruction(HloInstruction::CreateScaledDot( - call->shape(), call->mutable_operand(0), call->mutable_operand(1), - call->mutable_operand(2), call->mutable_operand(3), - dot_dimension_numbers, precision)); - ABSL_RETURN_IF_ERROR(call->ReplaceAllUsesWith(scaled_dot)); - ABSL_RETURN_IF_ERROR(computation->RemoveInstruction(call)); + ABSL_RETURN_IF_ERROR(RewriteScaledDot(computation, call, dot_dimension_numbers)); changed = true; } return changed; diff --git a/third_party/xla/xla/codegen/emitters/BUILD b/third_party/xla/xla/codegen/emitters/BUILD index 50e9a0d16946bb..b3e3a248ebc56e 100644 --- a/third_party/xla/xla/codegen/emitters/BUILD +++ b/third_party/xla/xla/codegen/emitters/BUILD @@ -234,12 +234,13 @@ cc_library( "//xla/service:shaped_slice", "//xla/tsl/platform:errors", "//xla/tsl/platform:statusor", - "@com_google_absl//absl/algorithm:container", "@com_google_absl//absl/container:flat_hash_map", "@com_google_absl//absl/container:flat_hash_set", + "@com_google_absl//absl/functional:function_ref", "@com_google_absl//absl/status", "@com_google_absl//absl/status:status_macros", "@com_google_absl//absl/status:statusor", + "@com_google_absl//absl/strings", "@com_google_absl//absl/strings:string_view", "@com_google_absl//absl/types:span", "@llvm-project//llvm:Support", diff --git a/third_party/xla/xla/codegen/emitters/kernel_arguments.cc b/third_party/xla/xla/codegen/emitters/kernel_arguments.cc index 03ec8c97f81517..41d946c8b7b25e 100644 --- a/third_party/xla/xla/codegen/emitters/kernel_arguments.cc +++ b/third_party/xla/xla/codegen/emitters/kernel_arguments.cc @@ -16,17 +16,16 @@ limitations under the License. #include #include -#include #include #include #include -#include "absl/algorithm/container.h" #include "absl/container/flat_hash_map.h" #include "absl/container/flat_hash_set.h" #include "absl/status/status.h" #include "absl/status/status_macros.h" #include "absl/status/statusor.h" +#include "absl/strings/str_cat.h" #include "absl/strings/string_view.h" #include "absl/types/span.h" #include "llvm/ADT/STLExtras.h" @@ -125,13 +124,15 @@ void FillKernelArgumentAttributes( struct OutputArguments { std::vector output_arguments; + // Shape index of each entry of `output_arguments`, in the same order. + std::vector output_shape_indices; absl::flat_hash_set buffers_written; }; // Extract output arguments from an instruction's shape and return both // the arguments and the set of written buffer slices absl::StatusOr ExtractOutputArguments( - const BufferAssignment& buffer_assignment, + KernelArguments::SliceProvider slice_provider, const HloInstruction* hlo_instruction) { OutputArguments result; ABSL_RETURN_IF_ERROR(ShapeUtil::ForEachSubshapeWithStatus( @@ -139,56 +140,144 @@ absl::StatusOr ExtractOutputArguments( [&](const Shape& subshape, const ShapeIndex& index) { if (!subshape.IsArray()) return absl::OkStatus(); - ABSL_ASSIGN_OR_RETURN( - BufferAllocation::Slice slice, - buffer_assignment.GetUniqueSlice(hlo_instruction, index)); + ABSL_ASSIGN_OR_RETURN(BufferAllocation::Slice slice, + slice_provider(*hlo_instruction, index)); result.output_arguments.emplace_back(KernelArgument(subshape, slice)); + result.output_shape_indices.push_back(index); result.buffers_written.insert(slice); return absl::OkStatus(); })); return result; } -absl::StatusOr CreateKernelArguments( - const BufferAssignment& buffer_assignment, - const KernelArguments::BufferAlignment& buffer_alignment, + +// Validates `interleaved_output_indices` and returns the position that each +// result must take in the argument list, in result order. +// +// An empty `interleaved_output_indices` means "no interleaving": the results +// follow the operands. +absl::StatusOr> ResultPositions( + absl::Span interleaved_output_indices, size_t num_operands, + size_t num_results) { + std::vector result_positions; + result_positions.reserve(num_results); + + if (interleaved_output_indices.empty()) { + for (size_t i = 0; i < num_results; ++i) { + result_positions.push_back(num_operands + i); + } + return result_positions; + } + + if (interleaved_output_indices.size() != num_results) { + return absl::InvalidArgumentError(absl::StrCat( + "Expected an interleaving index for each of the ", num_results, + " results, got ", interleaved_output_indices.size(), " indices")); + } + + const size_t num_arguments = num_operands + num_results; + int64_t previous_index = -1; + for (const int32_t index : interleaved_output_indices) { + if (index < 0 || static_cast(index) >= num_arguments) { + return absl::InvalidArgumentError("Output index out of bounds"); + } + if (index <= previous_index) { + return absl::InvalidArgumentError( + absl::StrCat("Output indices must be strictly increasing, got ", + index, " after ", previous_index)); + } + previous_index = index; + result_positions.push_back(index); + } + return result_positions; +} + +} // namespace + +absl::StatusOr KernelArguments::CreateInternal( + SliceProvider slice_provider, const BufferAlignment& buffer_alignment, const HloInstruction* hlo_instruction, - absl::Span unmanaged_arguments) { - std::vector kernel_arguments; + absl::Span unmanaged_arguments, + absl::Span interleaved_output_indices) { absl::flat_hash_set no_invariant_operands = NonInvariantOperands(*hlo_instruction); + std::vector operand_arguments; + operand_arguments.reserve(hlo_instruction->operand_count()); for (auto [op_idx, operand] : llvm::enumerate(hlo_instruction->operands())) { ABSL_ASSIGN_OR_RETURN(BufferAllocation::Slice slice, - buffer_assignment.GetUniqueSlice(operand, {})); + slice_provider(*operand, {})); KernelArgument arg(operand->shape(), slice); if (no_invariant_operands.contains(op_idx)) { arg.set_invariant(false); } - kernel_arguments.emplace_back(std::move(arg)); + operand_arguments.push_back(std::move(arg)); } ABSL_ASSIGN_OR_RETURN(OutputArguments output_result, - ExtractOutputArguments(buffer_assignment, hlo_instruction)); + ExtractOutputArguments(slice_provider, hlo_instruction)); + + ABSL_ASSIGN_OR_RETURN( + std::vector result_positions, + ResultPositions(interleaved_output_indices, operand_arguments.size(), + output_result.output_arguments.size())); + + const size_t num_managed_arguments = + operand_arguments.size() + output_result.output_arguments.size(); + + std::vector kernel_arguments; + kernel_arguments.reserve(num_managed_arguments + unmanaged_arguments.size()); + ArgumentPositions positions; + positions.operands.reserve(operand_arguments.size()); + positions.results.reserve(output_result.output_arguments.size()); + + // `result_positions` is strictly increasing and holds exactly one position + // per result, so the remaining positions hold exactly the operands, in + // operand order. + size_t next_operand = 0; + size_t next_result = 0; + for (size_t position = 0; position < num_managed_arguments; ++position) { + if (next_result < result_positions.size() && + result_positions[next_result] == static_cast(position)) { + positions.results.emplace_back( + output_result.output_shape_indices[next_result], position); + kernel_arguments.push_back( + std::move(output_result.output_arguments[next_result])); + ++next_result; + continue; + } + positions.operands.push_back(position); + kernel_arguments.push_back(std::move(operand_arguments[next_operand])); + ++next_operand; + } - absl::c_move(output_result.output_arguments, - std::back_inserter(kernel_arguments)); for (const Shape& unmanaged_argument : unmanaged_arguments) { kernel_arguments.emplace_back(unmanaged_argument); } FillKernelArgumentAttributes(kernel_arguments, buffer_alignment, output_result.buffers_written); - return KernelArguments(std::move(kernel_arguments)); + return KernelArguments(std::move(kernel_arguments), std::move(positions)); } -} // namespace +absl::StatusOr KernelArguments::Create( + SliceProvider slice_provider, const BufferAlignment& buffer_alignment, + const HloInstruction* hlo_instruction, + absl::Span unmanaged_arguments) { + return CreateInternal(slice_provider, buffer_alignment, hlo_instruction, + unmanaged_arguments, + /*interleaved_output_indices=*/{}); +} absl::StatusOr KernelArguments::Create( const BufferAssignment& buffer_assignment, const BufferAlignment& buffer_alignment, const HloInstruction* hlo_instruction) { - return CreateKernelArguments(buffer_assignment, buffer_alignment, - hlo_instruction, {}); + auto slice_provider = [&buffer_assignment](const HloInstruction& instruction, + const ShapeIndex& index) { + return buffer_assignment.GetUniqueSlice(&instruction, index); + }; + return CreateInternal(slice_provider, buffer_alignment, hlo_instruction, {}, + /*interleaved_output_indices=*/{}); } absl::StatusOr KernelArguments::Create( @@ -196,8 +285,13 @@ absl::StatusOr KernelArguments::Create( const BufferAlignment& buffer_alignment, const HloInstruction* hlo_instruction, absl::Span unmanaged_arguments) { - return CreateKernelArguments(buffer_assignment, buffer_alignment, - hlo_instruction, unmanaged_arguments); + auto slice_provider = [&buffer_assignment](const HloInstruction& instruction, + const ShapeIndex& index) { + return buffer_assignment.GetUniqueSlice(&instruction, index); + }; + return CreateInternal(slice_provider, buffer_alignment, hlo_instruction, + unmanaged_arguments, + /*interleaved_output_indices=*/{}); } absl::StatusOr KernelArguments::Create( @@ -205,65 +299,43 @@ absl::StatusOr KernelArguments::Create( const BufferAlignment& buffer_alignment, const HloInstruction* hlo_instruction, absl::Span interleaved_output_indices) { - if (interleaved_output_indices.empty()) { - // Fall back to regular Create method when no interleaving is requested - return CreateKernelArguments(buffer_assignment, buffer_alignment, - hlo_instruction, {}); - } - - const auto& operands = hlo_instruction->operands(); - - ABSL_ASSIGN_OR_RETURN(OutputArguments output_result, - ExtractOutputArguments(buffer_assignment, hlo_instruction)); - auto& [output_arguments, buffers_written] = output_result; + auto slice_provider = [&buffer_assignment](const HloInstruction& instruction, + const ShapeIndex& index) { + return buffer_assignment.GetUniqueSlice(&instruction, index); + }; + return CreateInternal(slice_provider, buffer_alignment, hlo_instruction, {}, + interleaved_output_indices); +} - // Check bounds: all output indices must be valid positions - size_t total_positions = operands.size() + output_arguments.size(); - for (int32_t idx : interleaved_output_indices) { - if (idx < 0 || static_cast(idx) >= total_positions) { - return absl::InvalidArgumentError("Output index out of bounds"); - } +absl::StatusOr KernelArguments::OperandIndex( + int64_t operand_index) const { + if (!positions_.has_value()) { + return absl::FailedPreconditionError( + "KernelArguments were not created from an HLO instruction, so operand " + "positions are unknown"); } - - std::vector kernel_arguments; - kernel_arguments.reserve(total_positions); - - // Interleave the inputs and outputs according to the indices - size_t arg_idx = 0; - size_t output_pos = 0; - - for (size_t i = 0; i < total_positions; ++i) { - if (output_pos < interleaved_output_indices.size() && - interleaved_output_indices[output_pos] == static_cast(i)) { - // Place output at this position - if (output_pos >= output_arguments.size()) { - return absl::InvalidArgumentError("Invalid output position index"); - } - kernel_arguments.emplace_back(output_arguments[output_pos]); - ++output_pos; - } else { - // Place input at this position - if (arg_idx >= operands.size()) { - return absl::InvalidArgumentError( - "Not enough inputs for remaining positions"); - } - ABSL_ASSIGN_OR_RETURN(BufferAllocation::Slice slice, - buffer_assignment.GetUniqueSlice(operands[arg_idx], {})); - kernel_arguments.emplace_back( - KernelArgument(operands[arg_idx]->shape(), slice)); - ++arg_idx; - } + if (operand_index < 0 || + operand_index >= static_cast(positions_->operands.size())) { + return absl::OutOfRangeError( + absl::StrCat("No such operand: ", operand_index)); } + return positions_->operands[operand_index]; +} - // Verify we used all inputs and outputs - if (arg_idx != operands.size() || output_pos != output_arguments.size()) { - return absl::InvalidArgumentError("Did not use all inputs/outputs"); +absl::StatusOr KernelArguments::ResultIndex( + const ShapeIndex& shape_index) const { + if (!positions_.has_value()) { + return absl::FailedPreconditionError( + "KernelArguments were not created from an HLO instruction, so result " + "positions are unknown"); } - - FillKernelArgumentAttributes(kernel_arguments, buffer_alignment, - buffers_written); - - return KernelArguments(std::move(kernel_arguments)); + for (const auto& [index, position] : positions_->results) { + if (index == shape_index) { + return position; + } + } + return absl::OutOfRangeError( + absl::StrCat("No such result: ", shape_index.ToString())); } } // namespace xla::emitters diff --git a/third_party/xla/xla/codegen/emitters/kernel_arguments.h b/third_party/xla/xla/codegen/emitters/kernel_arguments.h index 9dc17aabeb2310..f8b131160854e2 100644 --- a/third_party/xla/xla/codegen/emitters/kernel_arguments.h +++ b/third_party/xla/xla/codegen/emitters/kernel_arguments.h @@ -16,15 +16,18 @@ limitations under the License. #define XLA_CODEGEN_EMITTERS_KERNEL_ARGUMENTS_H_ #include +#include #include #include +#include "absl/functional/function_ref.h" #include "absl/status/statusor.h" #include "absl/types/span.h" #include "xla/hlo/ir/hlo_instruction.h" #include "xla/service/buffer_assignment.h" #include "xla/service/shaped_slice.h" #include "xla/shape.h" +#include "xla/shape_util.h" namespace xla::emitters { @@ -101,6 +104,28 @@ class KernelArguments { int64_t constant_buffer_align_bytes; }; + // Resolves the buffer slice backing the subshape of `instruction` at + // `index`. + // + // Emitters that can redirect an instruction's buffers (for example + // `ThunkEmitter`, which lets specialized emitters override an allocation) + // must pass their own resolver instead of handing over a `BufferAssignment`, + // otherwise those overrides are silently ignored and the kernel would be + // launched on the wrong buffers. + using SliceProvider = + absl::FunctionRef( + const HloInstruction& instruction, const ShapeIndex& index)>; + + // Creates a KernelArguments object for the given HLO instruction, resolving + // buffer slices through `slice_provider`. Argument order is the same as for + // the `BufferAssignment` overloads below: operands in operand order, then + // the array leaves of the result shape in shape-index order, then + // `unmanaged_arguments`. + static absl::StatusOr Create( + SliceProvider slice_provider, const BufferAlignment& buffer_alignment, + const HloInstruction* hlo_instruction, + absl::Span unmanaged_arguments = {}); + // Creates a KernelArguments object for the given HLO instruction. // The unmanaged_arguments are added to the end of the list of input/output // arguments. @@ -122,6 +147,10 @@ class KernelArguments { // Example: If hlo_instruction->operands() has 3 elements and hlo_instruction // shape yields 2 output arguments, and interleaved_output_indices = {1, 4}: // - Final argument order will be: input0, output0, input1, input2, output1 + // + // `interleaved_output_indices` must be strictly increasing and must name a + // position for every output argument; an empty span requests no + // interleaving. Returns `InvalidArgument` otherwise. static absl::StatusOr Create( const BufferAssignment& buffer_assignment, const BufferAlignment& buffer_alignment, @@ -133,6 +162,25 @@ class KernelArguments { const std::vector& args() const { return args_; } + // Position of the argument that carries operand `operand_index` of the + // instruction these arguments were created for. + // + // Use this instead of hardcoding argument positions when building a + // `stream_executor::KernelArgsPackingSpec`, whose relocation indices refer + // to exactly this list. + // + // Returns `FailedPrecondition` if the arguments were not built by one of the + // `Create` factories, and `OutOfRange` if there is no such operand. + absl::StatusOr OperandIndex(int64_t operand_index) const; + + // Position of the argument that carries the result subshape at + // `shape_index`. For an instruction with a non-tuple result, that is + // `ResultIndex({})`; for a one-element tuple result, `ResultIndex({0})`. + // + // Returns `FailedPrecondition` if the arguments were not built by one of the + // `Create` factories, and `OutOfRange` if there is no such result. + absl::StatusOr ResultIndex(const ShapeIndex& shape_index) const; + std::vector GetArgumentShapedSlices() const { std::vector arg_slices; arg_slices.reserve(args_.size()); @@ -179,7 +227,29 @@ class KernelArguments { } private: + // Where each operand and each result of the originating instruction ended up + // in `args_`. Only populated by the `Create` factories. + struct ArgumentPositions { + // Position in `args_` of each operand, indexed by operand number. + std::vector operands; + // Position in `args_` of each result array leaf, paired with its shape + // index. Result shapes have few leaves, so a linear scan is fine. + std::vector> results; + }; + + KernelArguments(std::vector args, ArgumentPositions positions) + : args_(std::move(args)), positions_(std::move(positions)) {} + + // Shared implementation of all `Create` overloads. See the interleaving + // `Create` overload above for the meaning of `interleaved_output_indices`. + static absl::StatusOr CreateInternal( + SliceProvider slice_provider, const BufferAlignment& buffer_alignment, + const HloInstruction* hlo_instruction, + absl::Span unmanaged_arguments, + absl::Span interleaved_output_indices); + std::vector args_; + std::optional positions_; }; } // namespace xla::emitters diff --git a/third_party/xla/xla/codegen/emitters/kernel_arguments_test.cc b/third_party/xla/xla/codegen/emitters/kernel_arguments_test.cc index bc9ee1d0ae6a38..423fe75fd1bebc 100644 --- a/third_party/xla/xla/codegen/emitters/kernel_arguments_test.cc +++ b/third_party/xla/xla/codegen/emitters/kernel_arguments_test.cc @@ -43,6 +43,7 @@ limitations under the License. namespace xla::emitters { namespace { +using ::absl_testing::IsOkAndHolds; using ::absl_testing::StatusIs; using ::testing::ElementsAre; using ::testing::HasSubstr; @@ -417,5 +418,260 @@ TEST_F(KernelArgumentsTest, ElementsAre(false, true, false, true, true, true)); } +TEST_F(KernelArgumentsTest, + NoInvariantFrontendAnnotationPropagatesToInterleavedArguments) { + ASSERT_OK_AND_ASSIGN(std::unique_ptr module, + ParseAndReturnVerifiedModule(R"( + e { + a = s4[5] parameter(0) + b = s4[6] parameter(1) + c = s4[7] parameter(2) + t = tuple(c, b, a), + frontend_attributes={xla.no_invariant_operands="0,2"} + } + )")); + AliasInfo alias_info; + BufferAssigner::Options opts; + opts.allocate_buffers_for_constants = true; + ASSERT_OK_AND_ASSIGN( + std::unique_ptr assignment, + BufferAssigner::Run( + module.get(), std::make_unique(module.get()), + &BufferSizeBytes, &alias_info, [](LogicalBuffer::Color) { return 0; }, + std::move(opts))); + + // Argument order: operand0, result0, operand1, result1, operand2, result2. + const std::vector interleaved_output_indices = {1, 3, 5}; + ASSERT_OK_AND_ASSIGN( + KernelArguments kernel_arguments, + KernelArguments::Create(*assignment, gpu::GetDefaultBufferAlignment(), + module->entry_computation()->root_instruction(), + interleaved_output_indices)); + + std::vector invariant_flags; + invariant_flags.reserve(kernel_arguments.args().size()); + for (const KernelArgument& arg : kernel_arguments.args()) { + invariant_flags.push_back(arg.invariant()); + } + // Operands 0 and 2 are annotated as non-invariant, wherever they end up. + EXPECT_THAT(invariant_flags, + ElementsAre(false, true, true, true, false, true)); +} + +TEST_F(KernelArgumentsTest, + InterleavedOutputIndicesMustCoverAllOutputsInOrder) { + ASSERT_OK_AND_ASSIGN(std::unique_ptr module, + ParseAndReturnVerifiedModule(R"( + e { + a = s4[5] parameter(0) + b = s4[6] parameter(1) + c = s4[7] parameter(2) + t = tuple(c, b, a) + } + )")); + AliasInfo alias_info; + BufferAssigner::Options opts; + opts.allocate_buffers_for_constants = true; + ASSERT_OK_AND_ASSIGN( + std::unique_ptr assignment, + BufferAssigner::Run( + module.get(), std::make_unique(module.get()), + &BufferSizeBytes, &alias_info, [](LogicalBuffer::Color) { return 0; }, + std::move(opts))); + const HloInstruction* root = module->entry_computation()->root_instruction(); + + const std::vector too_few_indices = {1, 3}; + EXPECT_THAT( + KernelArguments::Create(*assignment, gpu::GetDefaultBufferAlignment(), + root, too_few_indices), + StatusIs(absl::StatusCode::kInvalidArgument, + HasSubstr("Expected an interleaving index for each of the 3 " + "results"))); + + const std::vector unordered_indices = {1, 5, 3}; + EXPECT_THAT( + KernelArguments::Create(*assignment, gpu::GetDefaultBufferAlignment(), + root, unordered_indices), + StatusIs(absl::StatusCode::kInvalidArgument, + HasSubstr("must be strictly increasing"))); +} + +TEST_F(KernelArgumentsTest, SliceProviderCanOverrideBufferAssignment) { + constexpr absl::string_view kHloString = R"( + HloModule module + + ENTRY entry { + param.0 = f32[1,2,3]{2,1,0} parameter(0) + param.1 = f32[1,2,3]{2,1,0} parameter(1) + ROOT add = f32[1,2,3]{2,1,0} add(param.0, param.1) + } + )"; + + ASSERT_OK_AND_ASSIGN(std::unique_ptr module, + ParseAndReturnVerifiedModule(kHloString)); + AliasInfo alias_info; + ASSERT_OK_AND_ASSIGN( + std::unique_ptr assignment, + BufferAssigner::Run( + module.get(), std::make_unique(module.get()), + &BufferSizeBytes, &alias_info, [](LogicalBuffer::Color) { return 0; }, + BufferAssigner::Options{/*allocate_buffers_for_constants=*/true})); + + const HloInstruction* root = module->entry_computation()->root_instruction(); + const HloInstruction* param0 = root->operand(0); + ASSERT_OK_AND_ASSIGN(BufferAllocation::Slice param1_slice, + assignment->GetUniqueSlice(root->operand(1), {})); + + // Redirect operand 0 onto operand 1's buffer, the way an emitter's + // allocation overrides do. A `BufferAssignment` alone cannot express this. + auto slice_provider = + [&](const HloInstruction& instruction, + const ShapeIndex& index) -> absl::StatusOr { + if (&instruction == param0) { + return param1_slice; + } + return assignment->GetUniqueSlice(&instruction, index); + }; + + ASSERT_OK_AND_ASSIGN( + KernelArguments kernel_arguments, + KernelArguments::Create(slice_provider, gpu::GetDefaultBufferAlignment(), + root)); + + ASSERT_THAT(kernel_arguments.args(), SizeIs(3)); + EXPECT_EQ(kernel_arguments.args()[0].slice(), param1_slice); + EXPECT_EQ(kernel_arguments.args()[1].slice(), param1_slice); +} + +TEST_F(KernelArgumentsTest, SliceProviderErrorsArePropagated) { + ASSERT_OK_AND_ASSIGN(std::unique_ptr module, + ParseAndReturnVerifiedModule(R"( + ENTRY entry { + p = f32[4] parameter(0) + ROOT add = f32[4] add(p, p) + } + )")); + + auto slice_provider = + [](const HloInstruction& instruction, + const ShapeIndex& index) -> absl::StatusOr { + return absl::NotFoundError("no slice for you"); + }; + + EXPECT_THAT( + KernelArguments::Create(slice_provider, gpu::GetDefaultBufferAlignment(), + module->entry_computation()->root_instruction()), + StatusIs(absl::StatusCode::kNotFound, HasSubstr("no slice for you"))); +} + +TEST_F(KernelArgumentsTest, OperandAndResultIndices) { + ASSERT_OK_AND_ASSIGN(std::unique_ptr module, + ParseAndReturnVerifiedModule(R"( + ENTRY main { + param0 = f32[10] parameter(0) + param1 = f32[20] parameter(1) + ROOT tuple_result = (f32[10], f32[20]) tuple(param0, param1) + } + )")); + HloInstruction* root = module->entry_computation()->root_instruction(); + + AliasInfo alias_info; + ASSERT_OK_AND_ASSIGN( + std::unique_ptr assignment, + BufferAssigner::Run( + module.get(), std::make_unique(module.get()), + &BufferSizeBytes, &alias_info, [](LogicalBuffer::Color) { return 0; }, + BufferAssigner::Options{/*allocate_buffers_for_constants=*/true})); + + ASSERT_OK_AND_ASSIGN( + KernelArguments kernel_arguments, + KernelArguments::Create(*assignment, gpu::GetDefaultBufferAlignment(), + root)); + + // Operands come first, then the array leaves of the result shape. + EXPECT_THAT(kernel_arguments.OperandIndex(0), IsOkAndHolds(0)); + EXPECT_THAT(kernel_arguments.OperandIndex(1), IsOkAndHolds(1)); + EXPECT_THAT(kernel_arguments.ResultIndex({0}), IsOkAndHolds(2)); + EXPECT_THAT(kernel_arguments.ResultIndex({1}), IsOkAndHolds(3)); + + EXPECT_THAT(kernel_arguments.OperandIndex(2), + StatusIs(absl::StatusCode::kOutOfRange)); + EXPECT_THAT(kernel_arguments.OperandIndex(-1), + StatusIs(absl::StatusCode::kOutOfRange)); + // The root is a tuple, so there is no array leaf at the top level. + EXPECT_THAT(kernel_arguments.ResultIndex({}), + StatusIs(absl::StatusCode::kOutOfRange)); +} + +TEST_F(KernelArgumentsTest, OperandAndResultIndicesFollowInterleaving) { + ASSERT_OK_AND_ASSIGN(std::unique_ptr module, + ParseAndReturnVerifiedModule(R"( + ENTRY main { + param0 = f32[10] parameter(0) + param1 = f32[20] parameter(1) + ROOT tuple_result = (f32[10], f32[20]) tuple(param0, param1) + } + )")); + HloInstruction* root = module->entry_computation()->root_instruction(); + + AliasInfo alias_info; + ASSERT_OK_AND_ASSIGN( + std::unique_ptr assignment, + BufferAssigner::Run( + module.get(), std::make_unique(module.get()), + &BufferSizeBytes, &alias_info, [](LogicalBuffer::Color) { return 0; }, + BufferAssigner::Options{/*allocate_buffers_for_constants=*/true})); + + // Argument order becomes: input0, output0, input1, output1. + std::vector interleaved_indices = {1, 3}; + ASSERT_OK_AND_ASSIGN( + KernelArguments kernel_arguments, + KernelArguments::Create(*assignment, gpu::GetDefaultBufferAlignment(), + root, interleaved_indices)); + + EXPECT_THAT(kernel_arguments.OperandIndex(0), IsOkAndHolds(0)); + EXPECT_THAT(kernel_arguments.ResultIndex({0}), IsOkAndHolds(1)); + EXPECT_THAT(kernel_arguments.OperandIndex(1), IsOkAndHolds(2)); + EXPECT_THAT(kernel_arguments.ResultIndex({1}), IsOkAndHolds(3)); +} + +TEST_F(KernelArgumentsTest, ResultIndexOfNonTupleResult) { + ASSERT_OK_AND_ASSIGN(std::unique_ptr module, + ParseAndReturnVerifiedModule(R"( + ENTRY main { + param0 = f32[10] parameter(0) + ROOT result = f32[10] add(param0, param0) + } + )")); + HloInstruction* root = module->entry_computation()->root_instruction(); + + AliasInfo alias_info; + ASSERT_OK_AND_ASSIGN( + std::unique_ptr assignment, + BufferAssigner::Run( + module.get(), std::make_unique(module.get()), + &BufferSizeBytes, &alias_info, [](LogicalBuffer::Color) { return 0; }, + BufferAssigner::Options{/*allocate_buffers_for_constants=*/true})); + + ASSERT_OK_AND_ASSIGN( + KernelArguments kernel_arguments, + KernelArguments::Create(*assignment, gpu::GetDefaultBufferAlignment(), + root)); + + EXPECT_THAT(kernel_arguments.ResultIndex({}), IsOkAndHolds(2)); + EXPECT_THAT(kernel_arguments.ResultIndex({0}), + StatusIs(absl::StatusCode::kOutOfRange)); +} + +TEST_F(KernelArgumentsTest, HandBuiltArgumentsHaveNoKnownPositions) { + KernelArguments kernel_arguments(std::vector{KernelArgument( + ShapeUtil::MakeShape(F32, {4}), BufferAllocation::Slice())}); + + EXPECT_THAT(kernel_arguments.OperandIndex(0), + StatusIs(absl::StatusCode::kFailedPrecondition)); + EXPECT_THAT(kernel_arguments.ResultIndex({}), + StatusIs(absl::StatusCode::kFailedPrecondition)); +} + } // namespace } // namespace xla::emitters diff --git a/third_party/xla/xla/hlo/separate_compilation/BUILD b/third_party/xla/xla/hlo/parallel/BUILD similarity index 93% rename from third_party/xla/xla/hlo/separate_compilation/BUILD rename to third_party/xla/xla/hlo/parallel/BUILD index 776e872a486520..c6f830de3fb2a3 100644 --- a/third_party/xla/xla/hlo/separate_compilation/BUILD +++ b/third_party/xla/xla/hlo/parallel/BUILD @@ -48,8 +48,6 @@ cc_library( "//xla/hlo/ir:hlo", "//xla/service:compilation_environments", "//xla/service:hlo_module_config", - "//xla/tsl/platform:errors", - "//xla/tsl/platform:statusor", "@com_google_absl//absl/container:flat_hash_map", "@com_google_absl//absl/container:flat_hash_set", "@com_google_absl//absl/hash", @@ -74,8 +72,6 @@ cc_library( "//xla/hlo/ir:hlo", "//xla/hlo/transforms/simplifiers:hlo_dce", "//xla/service:compilation_environments", - "//xla/tsl/platform:errors", - "//xla/tsl/platform:statusor", "@com_google_absl//absl/base:nullability", "@com_google_absl//absl/container:flat_hash_map", "@com_google_absl//absl/container:flat_hash_set", @@ -103,8 +99,6 @@ xla_cc_test( "//xla/service:cpu_plugin", # fixdeps: keep "//xla/service:platform_util", "//xla/stream_executor:platform", - "//xla/tsl/lib/core:status_test_util", - "//xla/tsl/platform:statusor", "@com_google_absl//absl/container:flat_hash_set", "@com_google_absl//absl/strings:string_view", "@com_google_googletest//:gtest", @@ -127,8 +121,6 @@ xla_cc_test( "//xla/service:hlo_verifier", "//xla/service:platform_util", "//xla/stream_executor:platform", - "//xla/tsl/lib/core:status_test_util", - "//xla/tsl/platform:statusor", "@com_google_absl//absl/container:flat_hash_set", "@com_google_absl//absl/log", "@com_google_absl//absl/strings", diff --git a/third_party/xla/xla/hlo/separate_compilation/hlo_linking_manifest.h b/third_party/xla/xla/hlo/parallel/hlo_linking_manifest.h similarity index 87% rename from third_party/xla/xla/hlo/separate_compilation/hlo_linking_manifest.h rename to third_party/xla/xla/hlo/parallel/hlo_linking_manifest.h index 1558a2f6c47788..408d2991badb60 100644 --- a/third_party/xla/xla/hlo/separate_compilation/hlo_linking_manifest.h +++ b/third_party/xla/xla/hlo/parallel/hlo_linking_manifest.h @@ -13,8 +13,8 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#ifndef XLA_HLO_SEPARATE_COMPILATION_HLO_LINKING_MANIFEST_H_ -#define XLA_HLO_SEPARATE_COMPILATION_HLO_LINKING_MANIFEST_H_ +#ifndef XLA_HLO_PARALLEL_HLO_LINKING_MANIFEST_H_ +#define XLA_HLO_PARALLEL_HLO_LINKING_MANIFEST_H_ #include @@ -24,7 +24,7 @@ limitations under the License. #include "xla/service/compilation_environments.h" #include "xla/service/hlo_module_config.h" -namespace xla::separate_compilation { +namespace xla::parallel { // Metadata to guide linking process of HLO modules. // @@ -46,5 +46,5 @@ struct HloLinkingManifest { std::unique_ptr compilation_environment; }; -} // namespace xla::separate_compilation -#endif // XLA_HLO_SEPARATE_COMPILATION_HLO_LINKING_MANIFEST_H_ +} // namespace xla::parallel +#endif // XLA_HLO_PARALLEL_HLO_LINKING_MANIFEST_H_ diff --git a/third_party/xla/xla/hlo/separate_compilation/hlo_module_linking.cc b/third_party/xla/xla/hlo/parallel/hlo_module_linking.cc similarity index 91% rename from third_party/xla/xla/hlo/separate_compilation/hlo_module_linking.cc rename to third_party/xla/xla/hlo/parallel/hlo_module_linking.cc index f0432ea964d55b..a0c9eb23db2917 100644 --- a/third_party/xla/xla/hlo/separate_compilation/hlo_module_linking.cc +++ b/third_party/xla/xla/hlo/parallel/hlo_module_linking.cc @@ -13,7 +13,7 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#include "xla/hlo/separate_compilation/hlo_module_linking.h" +#include "xla/hlo/parallel/hlo_module_linking.h" #include #include @@ -35,13 +35,13 @@ limitations under the License. #include "xla/hlo/ir/hlo_module.h" #include "xla/hlo/ir/hlo_opcode.h" #include "xla/hlo/ir/hlo_print_options.h" -#include "xla/hlo/separate_compilation/hlo_linking_manifest.h" +#include "xla/hlo/parallel/hlo_linking_manifest.h" #include "xla/hlo/transforms/simplifiers/hlo_dce.h" #include "xla/service/compilation_environments.h" -#include "xla/tsl/platform/errors.h" -#include "xla/tsl/platform/statusor.h" -namespace xla::separate_compilation { +// NOLINTBEGIN(clang-diagnostic-pre-c++20-compat) + +namespace xla::parallel { namespace { struct TraversalState { @@ -123,17 +123,17 @@ class HloLinker { } state.entered = true; - PushDependencies(state); + PushDependencies(state.principal); return absl::OkStatus(); } - // Pushes dependencies of `state.principal` onto stack_ if they are stubs + // Pushes dependencies of `principal` onto stack_ if they are stubs // that need to be linked. - void PushDependencies(const TraversalState& state) { - std::vector post_order_instrs = - state.principal->MakeInstructionPostOrder(); + void PushDependencies(const HloComputation* principal) { VLOG(6) << "Processing callees:"; - for (HloInstruction* caller : post_order_instrs) { + std::vector post_order_instrs = + principal->MakeInstructionPostOrder(); + for (const HloInstruction* caller : post_order_instrs) { if (caller->opcode() != HloOpcode::kCall) { continue; } @@ -146,8 +146,8 @@ class HloLinker { it != linking_manifest_.stub_links.end()) { // Only push if its principal hasn't been mapped yet. if (clone_context_.FindComputation(callee) == nullptr) { - const HloComputation* principal = it->second; - stack_.push({principal, callee, false}); + const HloComputation* callee_principal = it->second; + stack_.push({callee_principal, callee, false}); } } } @@ -184,7 +184,7 @@ class HloLinker { absl::flat_hash_set being_linked_; absl::flat_hash_map finished_principals_; - std::stack> stack_; + std::stack stack_; HloCloneContext clone_context_; int linked_count_ = 0; @@ -222,4 +222,6 @@ absl::StatusOr> LinkComputation( return linked_module; } -} // namespace xla::separate_compilation +} // namespace xla::parallel + +// NOLINTEND(clang-diagnostic-pre-c++20-compat) diff --git a/third_party/xla/xla/hlo/separate_compilation/hlo_module_linking.h b/third_party/xla/xla/hlo/parallel/hlo_module_linking.h similarity index 81% rename from third_party/xla/xla/hlo/separate_compilation/hlo_module_linking.h rename to third_party/xla/xla/hlo/parallel/hlo_module_linking.h index a6a68aae1d6e2c..4d55ac196fdc1f 100644 --- a/third_party/xla/xla/hlo/separate_compilation/hlo_module_linking.h +++ b/third_party/xla/xla/hlo/parallel/hlo_module_linking.h @@ -13,8 +13,8 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#ifndef XLA_HLO_SEPARATE_COMPILATION_HLO_MODULE_LINKING_H_ -#define XLA_HLO_SEPARATE_COMPILATION_HLO_MODULE_LINKING_H_ +#ifndef XLA_HLO_PARALLEL_HLO_MODULE_LINKING_H_ +#define XLA_HLO_PARALLEL_HLO_MODULE_LINKING_H_ #include @@ -22,9 +22,9 @@ limitations under the License. #include "absl/status/statusor.h" #include "xla/hlo/ir/hlo_computation.h" #include "xla/hlo/ir/hlo_module.h" -#include "xla/hlo/separate_compilation/hlo_linking_manifest.h" +#include "xla/hlo/parallel/hlo_linking_manifest.h" -namespace xla::separate_compilation { +namespace xla::parallel { // Link all the callees together with the `root_computation` in the provided // module. Returns the pointer to the linked version of `root_computation` @@ -38,6 +38,6 @@ absl::StatusOr> LinkComputation( const HloLinkingManifest& linking_manifest, const HloComputation* absl_nonnull root_computation); -} // namespace xla::separate_compilation +} // namespace xla::parallel -#endif // XLA_HLO_SEPARATE_COMPILATION_HLO_MODULE_LINKING_H_ +#endif // XLA_HLO_PARALLEL_HLO_MODULE_LINKING_H_ diff --git a/third_party/xla/xla/hlo/separate_compilation/hlo_module_linking_test.cc b/third_party/xla/xla/hlo/parallel/hlo_module_linking_test.cc similarity index 65% rename from third_party/xla/xla/hlo/separate_compilation/hlo_module_linking_test.cc rename to third_party/xla/xla/hlo/parallel/hlo_module_linking_test.cc index 62d611e9e73905..a5a3507af7e33b 100644 --- a/third_party/xla/xla/hlo/separate_compilation/hlo_module_linking_test.cc +++ b/third_party/xla/xla/hlo/parallel/hlo_module_linking_test.cc @@ -13,12 +13,13 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#include "xla/hlo/separate_compilation/hlo_module_linking.h" +#include "xla/hlo/parallel/hlo_module_linking.h" #include #include #include +#include #include #include "absl/container/flat_hash_set.h" #include "absl/log/log.h" @@ -27,17 +28,17 @@ limitations under the License. #include "xla/hlo/ir/hlo_computation.h" #include "xla/hlo/ir/hlo_instruction.h" #include "xla/hlo/ir/hlo_print_options.h" -#include "xla/hlo/separate_compilation/hlo_linking_manifest.h" -#include "xla/hlo/separate_compilation/hlo_module_splitting.h" +#include "xla/hlo/parallel/hlo_linking_manifest.h" +#include "xla/hlo/parallel/hlo_module_splitting.h" #include "xla/hlo/testlib/hlo_hardware_independent_test_base.h" #include "xla/service/compiler.h" #include "xla/service/hlo_verifier.h" #include "xla/service/platform_util.h" #include "xla/stream_executor/platform.h" -#include "xla/tsl/lib/core/status_test_util.h" -#include "xla/tsl/platform/statusor.h" -namespace xla::separate_compilation { +// NOLINTBEGIN(clang-diagnostic-pre-c++20-compat) + +namespace xla::parallel { namespace { // Function to normalize an HloModule by removing/replacing names. @@ -108,12 +109,13 @@ TEST_F(LinkingTest, SingleCallLinking) { ROOT %result = f32[] call(%p), to_apply=%comp } )"; - TF_ASSERT_OK_AND_ASSIGN(auto original_module, - ParseAndReturnVerifiedModule(module_text)); + ASSERT_OK_AND_ASSIGN(auto original_module, + ParseAndReturnVerifiedModule(module_text)); - TF_ASSERT_OK_AND_ASSIGN(auto module_split_group, - CreateHloModuleSplitGroup(*original_module)); + ASSERT_OK_AND_ASSIGN(auto module_split_group, + CreateHloModuleSplitGroup(*original_module)); // LOG the split group. + // NOLINTBEGIN(*-custom-deterministic-iteration-order) for (const auto& split : module_split_group->module_splits) { LOG(INFO) << "Split: " << split->submodule->name(); LOG(INFO) << "Split module: " << split->submodule->ToString(); @@ -143,32 +145,32 @@ TEST_F(LinkingTest, SingleCallLinking) { module_split_group->linking_manifest.stub_links) { LOG(INFO) << " " << stub->name() << " ==>> " << comp->name(); } + // NOLINTEND(*-custom-deterministic-iteration-order) const HloLinkingManifest& linking_manifest = module_split_group->linking_manifest; auto* original_root = FindComputation(original_module.get(), "main"); - TF_ASSERT_OK_AND_ASSIGN( - const HloComputation* split_group_root, - module_split_group->GetClonedComputation(original_root)); + ASSERT_OK_AND_ASSIGN(const HloComputation* split_group_root, + module_split_group->GetClonedComputation(original_root)); - TF_ASSERT_OK_AND_ASSIGN(auto linked_module, - LinkComputation(linking_manifest, split_group_root)); + ASSERT_OK_AND_ASSIGN(auto linked_module, + LinkComputation(linking_manifest, split_group_root)); HloVerifier verifier(HloVerifierOpts{}); - TF_ASSERT_OK(verifier.Run(linked_module.get())); + ASSERT_OK(verifier.Run(linked_module.get())); EXPECT_TRUE(AreHloModulesEquivalent(*original_module, *linked_module)); - TF_ASSERT_OK_AND_ASSIGN(stream_executor::Platform * platform, - PlatformUtil::GetPlatform("cpu")); - TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr compiler, - Compiler::GetForPlatform(platform->id())); - TF_EXPECT_OK(compiler->RunHloPasses(original_module->Clone(), - /*executor=*/nullptr, - Compiler::CompileOptions{})); + ASSERT_OK_AND_ASSIGN(stream_executor::Platform * platform, + PlatformUtil::GetPlatform("cpu")); + ASSERT_OK_AND_ASSIGN(std::unique_ptr compiler, + Compiler::GetForPlatform(platform->id())); + EXPECT_OK(compiler->RunHloPasses(original_module->Clone(), + /*executor=*/nullptr, + Compiler::CompileOptions{})); VLOG(6) << linked_module->ToString(); - TF_ASSERT_OK(compiler->RunHloPasses(std::move(linked_module), - /*executor=*/nullptr, - Compiler::CompileOptions{})); + ASSERT_OK(compiler->RunHloPasses(std::move(linked_module), + /*executor=*/nullptr, + Compiler::CompileOptions{})); } TEST_F(LinkingTest, ChainGraphLinking) { @@ -200,36 +202,35 @@ TEST_F(LinkingTest, ChainGraphLinking) { ROOT %result = f32[] add(%call_res, %p) } )"; - TF_ASSERT_OK_AND_ASSIGN(auto original_module, - ParseAndReturnVerifiedModule(module_text)); + ASSERT_OK_AND_ASSIGN(auto original_module, + ParseAndReturnVerifiedModule(module_text)); - TF_ASSERT_OK_AND_ASSIGN(auto module_split_group, - CreateHloModuleSplitGroup(*original_module)); + ASSERT_OK_AND_ASSIGN(auto module_split_group, + CreateHloModuleSplitGroup(*original_module)); const HloLinkingManifest& linking_manifest = module_split_group->linking_manifest; auto* original_root = FindComputation(original_module.get(), "main"); - TF_ASSERT_OK_AND_ASSIGN( - const HloComputation* split_group_root, - module_split_group->GetClonedComputation(original_root)); + ASSERT_OK_AND_ASSIGN(const HloComputation* split_group_root, + module_split_group->GetClonedComputation(original_root)); - TF_ASSERT_OK_AND_ASSIGN(auto linked_module, - LinkComputation(linking_manifest, split_group_root)); + ASSERT_OK_AND_ASSIGN(auto linked_module, + LinkComputation(linking_manifest, split_group_root)); HloVerifier verifier(HloVerifierOpts{}); - TF_ASSERT_OK(verifier.Run(linked_module.get())); + ASSERT_OK(verifier.Run(linked_module.get())); EXPECT_TRUE(AreHloModulesEquivalent(*original_module, *linked_module)); - TF_ASSERT_OK_AND_ASSIGN(stream_executor::Platform * platform, - PlatformUtil::GetPlatform("cpu")); - TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr compiler, - Compiler::GetForPlatform(platform->id())); + ASSERT_OK_AND_ASSIGN(stream_executor::Platform * platform, + PlatformUtil::GetPlatform("cpu")); + ASSERT_OK_AND_ASSIGN(std::unique_ptr compiler, + Compiler::GetForPlatform(platform->id())); VLOG(6) << linked_module->ToString(); - TF_EXPECT_OK(compiler->RunHloPasses(std::move(linked_module), - /*executor=*/nullptr, - Compiler::CompileOptions{})); - TF_EXPECT_OK(compiler->RunHloPasses(original_module->Clone(), - /*executor=*/nullptr, - Compiler::CompileOptions{})); + EXPECT_OK(compiler->RunHloPasses(std::move(linked_module), + /*executor=*/nullptr, + Compiler::CompileOptions{})); + EXPECT_OK(compiler->RunHloPasses(original_module->Clone(), + /*executor=*/nullptr, + Compiler::CompileOptions{})); } TEST_F(LinkingTest, DiamondGraphLinking) { @@ -273,38 +274,39 @@ TEST_F(LinkingTest, DiamondGraphLinking) { ROOT %result = f32[] add(%call_a, %call_b) } )"; - TF_ASSERT_OK_AND_ASSIGN(auto original_module, - ParseAndReturnVerifiedModule(module_text)); + ASSERT_OK_AND_ASSIGN(auto original_module, + ParseAndReturnVerifiedModule(module_text)); - TF_ASSERT_OK_AND_ASSIGN(auto module_split_group, - CreateHloModuleSplitGroup(*original_module)); + ASSERT_OK_AND_ASSIGN(auto module_split_group, + CreateHloModuleSplitGroup(*original_module)); const HloLinkingManifest& linking_manifest = module_split_group->linking_manifest; auto* original_root = FindComputation(original_module.get(), "main"); - TF_ASSERT_OK_AND_ASSIGN( - const HloComputation* split_group_root, - module_split_group->GetClonedComputation(original_root)); + ASSERT_OK_AND_ASSIGN(const HloComputation* split_group_root, + module_split_group->GetClonedComputation(original_root)); - TF_ASSERT_OK_AND_ASSIGN(auto linked_module, - LinkComputation(linking_manifest, split_group_root)); + ASSERT_OK_AND_ASSIGN(auto linked_module, + LinkComputation(linking_manifest, split_group_root)); HloVerifier verifier(HloVerifierOpts{}); - TF_ASSERT_OK(verifier.Run(linked_module.get())); + ASSERT_OK(verifier.Run(linked_module.get())); EXPECT_TRUE(AreHloModulesEquivalent(*original_module, *linked_module)); - TF_ASSERT_OK_AND_ASSIGN(stream_executor::Platform * platform, - PlatformUtil::GetPlatform("cpu")); - TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr compiler, - Compiler::GetForPlatform(platform->id())); - TF_EXPECT_OK(compiler->RunHloPasses(original_module->Clone(), - /*executor=*/nullptr, - Compiler::CompileOptions{})); + ASSERT_OK_AND_ASSIGN(stream_executor::Platform * platform, + PlatformUtil::GetPlatform("cpu")); + ASSERT_OK_AND_ASSIGN(std::unique_ptr compiler, + Compiler::GetForPlatform(platform->id())); + EXPECT_OK(compiler->RunHloPasses(original_module->Clone(), + /*executor=*/nullptr, + Compiler::CompileOptions{})); VLOG(6) << linked_module->ToString(); - TF_EXPECT_OK(compiler->RunHloPasses(std::move(linked_module), - /*executor=*/nullptr, - Compiler::CompileOptions{})); + EXPECT_OK(compiler->RunHloPasses(std::move(linked_module), + /*executor=*/nullptr, + Compiler::CompileOptions{})); } } // namespace -} // namespace xla::separate_compilation +} // namespace xla::parallel + +// NOLINTEND(clang-diagnostic-pre-c++20-compat) diff --git a/third_party/xla/xla/hlo/separate_compilation/hlo_module_splitting.cc b/third_party/xla/xla/hlo/parallel/hlo_module_splitting.cc similarity index 96% rename from third_party/xla/xla/hlo/separate_compilation/hlo_module_splitting.cc rename to third_party/xla/xla/hlo/parallel/hlo_module_splitting.cc index 02254a3b5fba2c..ae315bed1e0ec6 100644 --- a/third_party/xla/xla/hlo/separate_compilation/hlo_module_splitting.cc +++ b/third_party/xla/xla/hlo/parallel/hlo_module_splitting.cc @@ -13,7 +13,7 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#include "xla/hlo/separate_compilation/hlo_module_splitting.h" +#include "xla/hlo/parallel/hlo_module_splitting.h" #include #include @@ -39,14 +39,14 @@ limitations under the License. #include "xla/hlo/ir/hlo_instruction.h" #include "xla/hlo/ir/hlo_module.h" #include "xla/hlo/ir/hlo_opcode.h" -#include "xla/hlo/separate_compilation/hlo_linking_manifest.h" +#include "xla/hlo/parallel/hlo_linking_manifest.h" #include "xla/service/compilation_environments.h" #include "xla/service/hlo_module_config.h" #include "xla/status_macros.h" -#include "xla/tsl/platform/errors.h" -#include "xla/tsl/platform/statusor.h" -namespace xla::separate_compilation { +// NOLINTBEGIN(clang-diagnostic-pre-c++20-compat) + +namespace xla::parallel { namespace { constexpr absl::string_view kEntryName = "entry"; @@ -109,6 +109,7 @@ absl::StatusOr> ComposeMaps( const absl::flat_hash_map& first, const absl::flat_hash_map& second) { absl::flat_hash_map result; + // NOLINTNEXTLINE(*-custom-deterministic-iteration-order) for (const auto [k, kv] : first) { if (auto it = second.find(kv); it != second.end()) { result.insert({k, it->second}); @@ -123,6 +124,7 @@ template absl::Status MergeMapInto(absl::flat_hash_map& into, const absl::flat_hash_map& from, bool error_on_duplicate_key = true) { + // NOLINTNEXTLINE(*-custom-deterministic-iteration-order) for (const auto& [k, v] : from) { if (!into.insert({k, v}).second) { if (error_on_duplicate_key) { @@ -265,7 +267,7 @@ absl::StatusOr> CreateHloModuleSplit( if (replacement->parent() != submodule.get()) { replacement = clone_context.GetComputation(replacement); } - mapped_call_instruction->set_to_apply(callee_replacements[caller]); + mapped_call_instruction->set_to_apply(replacement); } entry_computation->SetAndSanitizeName(kEntryName); @@ -296,6 +298,7 @@ absl::StatusOr> CreateHloModuleSplitGroup( for (const auto& split : splits) { ABSL_ASSIGN_OR_RETURN(auto module_split, CreateHloModuleSplit(module, split)); module_splits.push_back(std::move(module_split)); + // NOLINTNEXTLINE(*-custom-deterministic-iteration-order) for (const auto* original_comp : split) { computation_address_book.insert( {original_comp, module_splits.back().get()}); @@ -311,6 +314,7 @@ absl::StatusOr> CreateHloModuleSplitGroup( for (const auto& split : module_splits) { VLOG(5) << "Split: " << split->submodule->name(); VLOG(5) << " Stub links:"; + // NOLINTNEXTLINE(*-custom-deterministic-iteration-order) for (const auto& [stub, comp] : split->stub_map) { VLOG(5) << " " << stub->name() << " ==>> " << comp->name() << "(" @@ -354,4 +358,6 @@ absl::StatusOr HloModuleSplitGroup::GetClonedComputation( return it2->second; } -} // namespace xla::separate_compilation +} // namespace xla::parallel + +// NOLINTEND(clang-diagnostic-pre-c++20-compat) diff --git a/third_party/xla/xla/hlo/separate_compilation/hlo_module_splitting.h b/third_party/xla/xla/hlo/parallel/hlo_module_splitting.h similarity index 93% rename from third_party/xla/xla/hlo/separate_compilation/hlo_module_splitting.h rename to third_party/xla/xla/hlo/parallel/hlo_module_splitting.h index 97a0031035511c..dbff1f44cefaec 100644 --- a/third_party/xla/xla/hlo/separate_compilation/hlo_module_splitting.h +++ b/third_party/xla/xla/hlo/parallel/hlo_module_splitting.h @@ -13,8 +13,8 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#ifndef XLA_HLO_SEPARATE_COMPILATION_HLO_MODULE_SPLITTING_H_ -#define XLA_HLO_SEPARATE_COMPILATION_HLO_MODULE_SPLITTING_H_ +#ifndef XLA_HLO_PARALLEL_HLO_MODULE_SPLITTING_H_ +#define XLA_HLO_PARALLEL_HLO_MODULE_SPLITTING_H_ #include #include @@ -26,9 +26,9 @@ limitations under the License. #include "xla/hlo/ir/hlo_computation.h" #include "xla/hlo/ir/hlo_instruction.h" #include "xla/hlo/ir/hlo_module.h" -#include "xla/hlo/separate_compilation/hlo_linking_manifest.h" +#include "xla/hlo/parallel/hlo_linking_manifest.h" -namespace xla::separate_compilation { +namespace xla::parallel { // Returns a list of sets of computations that can be split into separate // modules. Adjacent computations in the same set can be compiled together. @@ -108,5 +108,5 @@ struct HloModuleSplitGroup { absl::StatusOr> CreateHloModuleSplitGroup( const HloModule& module); -} // namespace xla::separate_compilation -#endif // XLA_HLO_SEPARATE_COMPILATION_HLO_MODULE_SPLITTING_H_ +} // namespace xla::parallel +#endif // XLA_HLO_PARALLEL_HLO_MODULE_SPLITTING_H_ diff --git a/third_party/xla/xla/hlo/separate_compilation/hlo_module_splitting_test.cc b/third_party/xla/xla/hlo/parallel/hlo_module_splitting_test.cc similarity index 73% rename from third_party/xla/xla/hlo/separate_compilation/hlo_module_splitting_test.cc rename to third_party/xla/xla/hlo/parallel/hlo_module_splitting_test.cc index 0eb36ac22118f9..373c3a4509b737 100644 --- a/third_party/xla/xla/hlo/separate_compilation/hlo_module_splitting_test.cc +++ b/third_party/xla/xla/hlo/parallel/hlo_module_splitting_test.cc @@ -13,7 +13,7 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#include "xla/hlo/separate_compilation/hlo_module_splitting.h" +#include "xla/hlo/parallel/hlo_module_splitting.h" #include @@ -26,10 +26,8 @@ limitations under the License. #include "xla/service/compiler.h" #include "xla/service/platform_util.h" #include "xla/stream_executor/platform.h" -#include "xla/tsl/lib/core/status_test_util.h" -#include "xla/tsl/platform/statusor.h" -namespace xla::separate_compilation { +namespace xla::parallel { namespace { using ::testing::UnorderedElementsAreArray; @@ -85,9 +83,8 @@ ENTRY %main() -> f32[] { ROOT %result = f32[] add(f32[] %res.01, f32[] %res.23) })"; - TF_ASSERT_OK_AND_ASSIGN(auto module, - ParseAndReturnVerifiedModule(module_text)); - TF_ASSERT_OK_AND_ASSIGN(auto splits, GroupComputationsForSplitting(*module)); + ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(module_text)); + ASSERT_OK_AND_ASSIGN(auto splits, GroupComputationsForSplitting(*module)); absl::flat_hash_set all_computations_in_splits; for (const auto& bucket : splits) { @@ -146,12 +143,11 @@ ENTRY %main() -> f32[] { ROOT %result = f32[] add(f32[] %res.01, f32[] %res.23) })"; - TF_ASSERT_OK_AND_ASSIGN(auto module, - ParseAndReturnVerifiedModule(module_text)); + ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(module_text)); auto* main = FindComputation(module.get(), "main"); - TF_ASSERT_OK_AND_ASSIGN(auto module_split, - CreateHloModuleSplit(*module, {main})); + ASSERT_OK_AND_ASSIGN(auto module_split, + CreateHloModuleSplit(*module, {main})); const int kMainModuleComputationCount = 5; // the main + 4 stubs // const int kMainModuleClonedComputations = 1; @@ -217,21 +213,20 @@ ENTRY %main() -> f32[] { ROOT %result = f32[] add(f32[] %res.01, f32[] %res.fu23) })"; - TF_ASSERT_OK_AND_ASSIGN(auto module, - ParseAndReturnVerifiedModule(module_text)); - TF_ASSERT_OK_AND_ASSIGN(auto module_split_group, - CreateHloModuleSplitGroup(*module)); - - TF_ASSERT_OK_AND_ASSIGN(stream_executor::Platform * platform, - PlatformUtil::GetPlatform("cpu")); - TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr compiler, - Compiler::GetForPlatform(platform->id())); - TF_ASSERT_OK(compiler->RunHloPasses(module->Clone(), /*executor=*/nullptr, - Compiler::CompileOptions{})); + ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(module_text)); + ASSERT_OK_AND_ASSIGN(auto module_split_group, + CreateHloModuleSplitGroup(*module)); + + ASSERT_OK_AND_ASSIGN(stream_executor::Platform * platform, + PlatformUtil::GetPlatform("cpu")); + ASSERT_OK_AND_ASSIGN(std::unique_ptr compiler, + Compiler::GetForPlatform(platform->id())); + ASSERT_OK(compiler->RunHloPasses(module->Clone(), /*executor=*/nullptr, + Compiler::CompileOptions{})); for (auto& split : module_split_group->module_splits) { - TF_EXPECT_OK(compiler->RunHloPasses(split->submodule->Clone(), - /*executor=*/nullptr, - Compiler::CompileOptions{})); + EXPECT_OK(compiler->RunHloPasses(split->submodule->Clone(), + /*executor=*/nullptr, + Compiler::CompileOptions{})); } } @@ -272,10 +267,9 @@ TEST_F(SplittingTest, SplitDiamondGraphModule) { } )"; - TF_ASSERT_OK_AND_ASSIGN(auto module, - ParseAndReturnVerifiedModule(module_text)); - TF_ASSERT_OK_AND_ASSIGN(auto module_split_group, - CreateHloModuleSplitGroup(*module)); + ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(module_text)); + ASSERT_OK_AND_ASSIGN(auto module_split_group, + CreateHloModuleSplitGroup(*module)); // Expect all computations to be assigned somewhere. for (auto original_comp : module->computations()) { @@ -285,16 +279,16 @@ TEST_F(SplittingTest, SplitDiamondGraphModule) { EXPECT_EQ(module_split_group->module_splits.size(), module->computation_count()); - TF_ASSERT_OK_AND_ASSIGN(stream_executor::Platform * platform, - PlatformUtil::GetPlatform("cpu")); - TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr compiler, - Compiler::GetForPlatform(platform->id())); - TF_ASSERT_OK(compiler->RunHloPasses(module->Clone(), /*executor=*/nullptr, - Compiler::CompileOptions{})); + ASSERT_OK_AND_ASSIGN(stream_executor::Platform * platform, + PlatformUtil::GetPlatform("cpu")); + ASSERT_OK_AND_ASSIGN(std::unique_ptr compiler, + Compiler::GetForPlatform(platform->id())); + ASSERT_OK(compiler->RunHloPasses(module->Clone(), /*executor=*/nullptr, + Compiler::CompileOptions{})); for (auto& split : module_split_group->module_splits) { - TF_EXPECT_OK(compiler->RunHloPasses(split->submodule->Clone(), - /*executor=*/nullptr, - Compiler::CompileOptions{})); + EXPECT_OK(compiler->RunHloPasses(split->submodule->Clone(), + /*executor=*/nullptr, + Compiler::CompileOptions{})); } } @@ -354,23 +348,22 @@ ENTRY %main() -> f32[] { ROOT %result = f32[] add(f32[] %res.01, f32[] %res.f) })"; - TF_ASSERT_OK_AND_ASSIGN(auto module, - ParseAndReturnVerifiedModule(module_text)); - TF_ASSERT_OK_AND_ASSIGN(auto module_split_group, - CreateHloModuleSplitGroup(*module)); - - TF_ASSERT_OK_AND_ASSIGN(stream_executor::Platform * platform, - PlatformUtil::GetPlatform("cpu")); - TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr compiler, - Compiler::GetForPlatform(platform->id())); - TF_ASSERT_OK(compiler->RunHloPasses(module->Clone(), /*executor=*/nullptr, - Compiler::CompileOptions{})); + ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(module_text)); + ASSERT_OK_AND_ASSIGN(auto module_split_group, + CreateHloModuleSplitGroup(*module)); + + ASSERT_OK_AND_ASSIGN(stream_executor::Platform * platform, + PlatformUtil::GetPlatform("cpu")); + ASSERT_OK_AND_ASSIGN(std::unique_ptr compiler, + Compiler::GetForPlatform(platform->id())); + ASSERT_OK(compiler->RunHloPasses(module->Clone(), /*executor=*/nullptr, + Compiler::CompileOptions{})); for (auto& split : module_split_group->module_splits) { - TF_EXPECT_OK(compiler->RunHloPasses(split->submodule->Clone(), - /*executor=*/nullptr, - Compiler::CompileOptions{})); + EXPECT_OK(compiler->RunHloPasses(split->submodule->Clone(), + /*executor=*/nullptr, + Compiler::CompileOptions{})); } } } // namespace -} // namespace xla::separate_compilation +} // namespace xla::parallel diff --git a/third_party/xla/xla/pjrt/c/pjrt_c_api_gpu_internal.cc b/third_party/xla/xla/pjrt/c/pjrt_c_api_gpu_internal.cc index a09787ab3e5790..39133603faafb7 100644 --- a/third_party/xla/xla/pjrt/c/pjrt_c_api_gpu_internal.cc +++ b/third_party/xla/xla/pjrt/c/pjrt_c_api_gpu_internal.cc @@ -84,7 +84,9 @@ namespace gpu_plugin { #if TENSORFLOW_USE_ROCM #define PJRT_GPU_PLUGIN_PLATFORM_NAME "ROCM" #elif TENSORFLOW_USE_SYCL -#define PJRT_GPU_PLUGIN_PLATFORM_NAME "ONEAPI" +// TODO(Intel-tf) this will be changed to ONEAPI +// when the SYCL backend has been renamed to ONEAPI. +#define PJRT_GPU_PLUGIN_PLATFORM_NAME "SYCL" #else #define PJRT_GPU_PLUGIN_PLATFORM_NAME "CUDA" #endif @@ -337,7 +339,7 @@ PJRT_Error* PJRT_GpuDeviceTopology_Create( if (plugin_platform == "ROCM") { platform_id = xla::RocmId(); platform_name = xla::RocmName(); - } else if (plugin_platform == "ONEAPI") { + } else if (plugin_platform == "SYCL") { platform_id = xla::OneapiId(); platform_name = xla::OneapiName(); } else { diff --git a/third_party/xla/xla/python/ifrt_proxy/common/array_util.cc b/third_party/xla/xla/python/ifrt_proxy/common/array_util.cc index 1bae9d8549fe88..b84b969d07a52a 100644 --- a/third_party/xla/xla/python/ifrt_proxy/common/array_util.cc +++ b/third_party/xla/xla/python/ifrt_proxy/common/array_util.cc @@ -131,10 +131,27 @@ absl::StatusOr ArrayMemRegion::FromZerothElementPointer( "byte_stride[", i, "] * (dim[", i, "] - 1) overflows: stride=", stride, " dim=", shape.dims()[i])); } - last_element_byte_offset += static_cast(product); + // Each product is independently in int64 range (checked above), but the + // uint64 accumulation of those products can still wrap. Example: + // shape=[2,2,2], strides=[INT64_MAX, INT64_MAX, 2] -> each mul succeeds + // and the running sum is 2^64 == 0. + const uint64_t term = static_cast(product); + if (__builtin_add_overflow(last_element_byte_offset, term, + &last_element_byte_offset)) { + return absl::InvalidArgumentError(absl::StrCat( + "byte-stride span overflows uint64 while accumulating byte_stride[", + i, "] * (dim[", i, "] - 1): stride=", stride, + " dim=", shape.dims()[i])); + } } } - return ArrayMemRegion(mem_region_start, last_element_byte_offset + byte_size); + uint64_t total_size; + if (__builtin_add_overflow(last_element_byte_offset, + static_cast(byte_size), &total_size)) { + return absl::InvalidArgumentError( + "byte-stride span overflows uint64: last-element offset + dtype size"); + } + return ArrayMemRegion(mem_region_start, total_size); } absl::StatusOr ArrayMemRegion::FromMinimalMemRegion( diff --git a/third_party/xla/xla/python/ifrt_proxy/common/array_util_test.cc b/third_party/xla/xla/python/ifrt_proxy/common/array_util_test.cc index 011991663a4e90..94ad8a658b95cc 100644 --- a/third_party/xla/xla/python/ifrt_proxy/common/array_util_test.cc +++ b/third_party/xla/xla/python/ifrt_proxy/common/array_util_test.cc @@ -41,6 +41,7 @@ using ::testing::TestWithParam; constexpr DType::Kind kF64 = DType::Kind::kF64; constexpr DType::Kind kS32 = DType::Kind::kS32; +constexpr DType::Kind kU8 = DType::Kind::kU8; constexpr DType::Kind kString = DType::Kind::kString; using Strides = std::vector; @@ -97,7 +98,10 @@ INSTANTIATE_TEST_SUITE_P( // Non-compact strides. TC{"NonCompactSingleDimension", kS32, {5}, Strides({16}), 68}, TC{"NonCompactDim0", kS32, {4, 3, 5}, Strides({120, 20, 4}), 420}, - TC{"PaddedElements", kS32, {4, 3, 5}, Strides({120, 40, 8}), 476}), + TC{"PaddedElements", kS32, {4, 3, 5}, Strides({120, 40, 8}), 476}, + // Products individually fit in int64; the accumulated span also fits. + TC{"U8ThreeDimsDefault", kU8, {2, 2, 2}, std::nullopt}, + TC{"U8ThreeDimsCompactStrides", kU8, {2, 2, 2}, Strides({4, 2, 1})}), testing::PrintToStringParamName()); TEST_P(ArrayMemRegionSuccess, TestCase) { const TC tc = GetParam(); @@ -148,7 +152,20 @@ INSTANTIATE_TEST_SUITE_P( TC{"ByteStrideOverflowMultiDim", kS32, {5, 3}, - Strides({INT64_C(4611686018427387904), 4})}), + Strides({INT64_C(4611686018427387904), 4})}, + // Each product stride*(dim-1) fits in int64 (the #46604 mul check + // succeeds) but the uint64 accumulation wraps: + // (INT64_MAX) + (INT64_MAX) + 2 == 2^64 == 0. + TC{"ByteStrideAccumOverflowUint64", + kU8, + {2, 2, 2}, + Strides({INT64_MAX, INT64_MAX, 2})}, + // Accumulation itself fits (INT64_MAX + INT64_MAX + 1 == UINT64_MAX); + // the final last-element-offset + dtype-size add is what overflows. + TC{"ByteStrideFinalSizeOverflowUint64", + kU8, + {2, 2, 2}, + Strides({INT64_MAX, INT64_MAX, 1})}), testing::PrintToStringParamName()); TEST_P(ArrayMemRegionFailure, TestCase) { const TC tc = GetParam(); diff --git a/third_party/xla/xla/service/buffer_assignment.cc b/third_party/xla/xla/service/buffer_assignment.cc index f9ee23be13925c..b8997931708ff2 100644 --- a/third_party/xla/xla/service/buffer_assignment.cc +++ b/third_party/xla/xla/service/buffer_assignment.cc @@ -680,7 +680,19 @@ BufferAllocationProto BufferAllocation::ToProto() const { } proto.set_is_constant(is_constant_); proto.set_maybe_live_out(maybe_live_out_); - for (const auto& buffer_offset_size : assigned_buffers_) { + using Entry = std::pair; + std::vector sorted; + sorted.reserve(assigned_buffers_.size()); + // NOLINTNEXTLINE(*-custom-deterministic-iteration-order) + for (const auto& kv : assigned_buffers_) { + sorted.push_back(&kv); + } + absl::c_sort(sorted, [](const Entry* a, const Entry* b) { + return a->first->id() < b->first->id(); + }); + proto.mutable_assigned()->Reserve(sorted.size()); + for (const Entry* ptr : sorted) { + const auto& buffer_offset_size = *ptr; BufferAllocationProto::Assigned* proto_assigned = proto.add_assigned(); proto_assigned->set_logical_buffer_id(buffer_offset_size.first->id()); proto_assigned->set_offset(buffer_offset_size.second.offset); @@ -688,12 +700,6 @@ BufferAllocationProto BufferAllocation::ToProto() const { proto_assigned->set_element_type( buffer_offset_size.first->shape().element_type()); } - absl::c_sort(*proto.mutable_assigned(), - [](const BufferAllocationProto::Assigned& assign1, - const BufferAllocationProto::Assigned& assign2) { - return assign1.logical_buffer_id() < - assign2.logical_buffer_id(); - }); return proto; } @@ -1592,11 +1598,21 @@ void BufferAssignment::ToProto(BufferAssignmentProto* proto) const { // because we need to do the HasAllocation check for each buffer. Otherwise // the buffer_size_ call might fail for some backends. const HloDataflowAnalysis& dataflow = this->dataflow_analysis(); + proto->mutable_logical_buffers()->Reserve(dataflow.values().size()); for (BufferValue::Id id = 0; id < dataflow.values().size(); id++) { auto& value = dataflow.values().at(id); if (HasAllocation(*value)) { - LogicalBufferProto proto_buffer = value->ToProto(buffer_size_); - proto->add_logical_buffers()->Swap(&proto_buffer); + LogicalBufferProto* pb = proto->add_logical_buffers(); + pb->set_id(value->id()); + pb->set_size(buffer_size_(*value)); + LogicalBufferProto::Location* loc = pb->mutable_defined_at(); + loc->set_instruction_id(value->instruction()->unique_id()); + for (int64_t i : value->index()) { + loc->add_shape_index(i); + } + if (value->has_color()) { + pb->set_color(value->color()); + } // Fill buffer aliases. for (const HloValue* alias : @@ -1607,13 +1623,17 @@ void BufferAssignment::ToProto(BufferAssignmentProto* proto) const { } BufferAssignmentProto::BufferAlias* proto_alias = proto->add_buffer_aliases(); - LogicalBufferProto::Location proto_alias_location = - BufferValue::ToLocationProto(*alias->instruction(), alias->index()); proto_alias->set_source_buffer_id(value->id()); - proto_alias->mutable_location()->Swap(&proto_alias_location); + LogicalBufferProto::Location* aloc = proto_alias->mutable_location(); + aloc->set_instruction_id(alias->instruction()->unique_id()); + for (int64_t i : alias->index()) { + aloc->add_shape_index(i); + } } } } + proto->mutable_buffer_allocations()->Reserve(Allocations().size()); + proto->mutable_peak_buffers()->Reserve(Allocations().size()); for (const BufferAllocation& allocation : Allocations()) { BufferAllocationProto proto_allocation = allocation.ToProto(); proto->add_buffer_allocations()->Swap(&proto_allocation); diff --git a/third_party/xla/xla/service/buffer_assignment_test.cc b/third_party/xla/xla/service/buffer_assignment_test.cc index 9a822813659134..68b1f8b58934c4 100644 --- a/third_party/xla/xla/service/buffer_assignment_test.cc +++ b/third_party/xla/xla/service/buffer_assignment_test.cc @@ -5689,5 +5689,46 @@ TEST_F(BufferAssignmentTest, FromProtoRejectsNegativeSize) { EXPECT_THAT(result.status().message(), ::testing::HasSubstr("negative")); } +TEST_F(BufferAssignmentTest, + ToProtoSortsAssignedBuffersAndPopulatesNestedTupleAliases) { + const char* const hlo_text = R"( + HloModule test + cond { + p = (f32[4]{0}, f32[8]{0}) parameter(0) + ROOT cond_res = pred[] constant(true) + } + body { + p = (f32[4]{0}, f32[8]{0}) parameter(0) + gte0 = f32[4]{0} get-tuple-element(p), index=0 + gte1 = f32[8]{0} get-tuple-element(p), index=1 + neg0 = f32[4]{0} negate(gte0) + ROOT out = (f32[4]{0}, f32[8]{0}) tuple(neg0, gte1) + } + ENTRY e { + p0 = (f32[4]{0}, f32[8]{0}) parameter(0) + ROOT loop = (f32[4]{0}, f32[8]{0}) while(p0), condition=cond, body=body + } + )"; + ASSERT_OK_AND_ASSIGN(std::unique_ptr module, + ParseAndReturnVerifiedModule(hlo_text)); + std::unique_ptr buffers = RunBufferAssignment(module.get()); + BufferAssignmentProto proto = buffers->ToProto(); + + EXPECT_FALSE(proto.logical_buffers().empty()); + EXPECT_FALSE(proto.buffer_aliases().empty()); + for (const BufferAllocationProto& alloc : proto.buffer_allocations()) { + for (int i = 1; i < alloc.assigned_size(); ++i) { + EXPECT_LT(alloc.assigned(i - 1).logical_buffer_id(), + alloc.assigned(i).logical_buffer_id()); + } + } + + ASSERT_OK_AND_ASSIGN( + std::unique_ptr roundtrip, + BufferAssignment::FromProto(proto, module.get(), &BufferSizeBytes, + &alias_info_)); + EXPECT_EQ(roundtrip->Allocations().size(), buffers->Allocations().size()); +} + } // namespace } // namespace xla diff --git a/third_party/xla/xla/stream_executor/BUILD b/third_party/xla/xla/stream_executor/BUILD index cb5d49d469c1ea..15ba54aaf48e28 100644 --- a/third_party/xla/xla/stream_executor/BUILD +++ b/third_party/xla/xla/stream_executor/BUILD @@ -370,6 +370,7 @@ cc_library( ":stream_executor_h", "//xla/service:device_assignment", "@com_google_absl//absl/base:core_headers", + "@com_google_absl//absl/container:btree", "@com_google_absl//absl/container:flat_hash_map", "@com_google_absl//absl/container:flat_hash_set", "@com_google_absl//absl/log", @@ -389,6 +390,7 @@ xla_cc_test( srcs = ["device_address_vmm_allocator_test.cc"], deps = [ ":device_address", + ":device_address_allocator", ":device_address_vmm_allocator", ":memory_allocation", ":memory_reservation", diff --git a/third_party/xla/xla/stream_executor/device_address_vmm_allocator.cc b/third_party/xla/xla/stream_executor/device_address_vmm_allocator.cc index 417518b1477768..9013f87dc5553c 100644 --- a/third_party/xla/xla/stream_executor/device_address_vmm_allocator.cc +++ b/third_party/xla/xla/stream_executor/device_address_vmm_allocator.cc @@ -132,6 +132,7 @@ void DeviceAddressVmmAllocator::AllocationRecord::MarkAllocatorStale( void DeviceAddressVmmAllocator::AllocationRecord::ReactivateAllocator( uint64_t new_size) { CHECK(allocator_stale()); + CHECK_EQ(allocator_deallocation_.pending.seqno, 0); CHECK(!allocator_address_mapping_.is_null()); allocator_address_ = DeviceAddressBase(allocator_address_.opaque(), new_size); allocator_stale_seqno_ = 0; @@ -154,14 +155,22 @@ void DeviceAddressVmmAllocator::AllocationRecord::MarkReservationStale( void DeviceAddressVmmAllocator::AllocationRecord::ReactivateReservation() { CHECK(reservation_stale()); + CHECK_EQ(reservation_alias_->deallocation.pending.seqno, 0); reservation_alias_->stale_seqno = 0; } void DeviceAddressVmmAllocator::AllocationRecord::CompleteStaleReservation() { CHECK(reservation_stale()); + CHECK_EQ(reservation_alias_->deallocation.pending.seqno, 0); reservation_alias_.reset(); } +DeviceAddressVmmAllocator::PendingDeallocationNode& +DeviceAddressVmmAllocator::AllocationRecord::reservation_deallocation() { + CHECK(has_reservation_alias()); + return reservation_alias_->deallocation; +} + // Interval between CPU polls of the GPU-written deallocation timeline while // waiting for deferred frees to become safe. The 50us value is a conservative // initial tradeoff: long enough to avoid busy-spinning a CPU core and short @@ -265,10 +274,8 @@ DeviceAddressVmmAllocator::~DeviceAddressVmmAllocator() { // no longer callable after the derived destructor has run, and the stream // is allowed to have been destroyed already. state->open_deallocation_batch_seqno = 0; - while (!state->pending_deallocations.empty()) { - PendingDeallocation pending = state->pending_deallocations.front(); - state->pending_deallocations.pop_front(); - CompletePendingDeallocation(*state, pending); + while (state->pending_head != nullptr) { + CompletePendingDeallocation(*state, *state->pending_head); } } @@ -319,16 +326,14 @@ absl::StatusOr DeviceAddressVmmAllocator::GetStream( absl::Status DeviceAddressVmmAllocator::DrainPendingDeallocations( PerDeviceState& state) { ABSL_RETURN_IF_ERROR(FlushOpenDeallocationBatch(state)); - if (state.pending_deallocations.empty()) { + if (state.pending_head == nullptr) { return absl::OkStatus(); } - uint64_t target_seqno = state.pending_deallocations.back().seqno; + uint64_t target_seqno = state.pending_tail->pending.seqno; ABSL_RETURN_IF_ERROR(WaitUntilSeqno(state, target_seqno)); - while (!state.pending_deallocations.empty() && - state.pending_deallocations.front().seqno <= target_seqno) { - PendingDeallocation pending = state.pending_deallocations.front(); - state.pending_deallocations.pop_front(); - CompletePendingDeallocation(state, pending); + while (state.pending_head != nullptr && + state.pending_head->pending.seqno <= target_seqno) { + CompletePendingDeallocation(state, *state.pending_head); } return absl::OkStatus(); } @@ -358,7 +363,8 @@ MemoryAllocation* DeviceAddressVmmAllocator::GetRawAllocation( // Allocator addresses are keyed directly by their VA. Stale records remain in // this map until deferred teardown completes, so expose their backing // allocation for diagnostics/reuse checks until the pending operation drains. - auto allocation_it = state->records_by_allocator_address.find(addr.opaque()); + auto allocation_it = + state->records_by_allocator_address.find(AddressStart(addr)); if (allocation_it != state->records_by_allocator_address.end() && allocation_it->second->allocator_matches(addr)) { return allocation_it->second->raw_allocation(); @@ -367,7 +373,7 @@ MemoryAllocation* DeviceAddressVmmAllocator::GetRawAllocation( // Reservation aliases created by Map() share one index. Only active aliases // are exposed; stale or already-unmapped aliases intentionally return // nullptr. - auto reservation_it = state->reservation_records.find(addr.opaque()); + auto reservation_it = state->reservation_records.find(AddressStart(addr)); if (reservation_it != state->reservation_records.end() && reservation_it->second->reservation_active() && reservation_it->second->reservation_matches(addr)) { @@ -385,7 +391,8 @@ MemoryReservation* DeviceAddressVmmAllocator::GetReservation( PerDeviceState* state = *state_or; absl::MutexLock lock(state->mu); - auto allocation_it = state->records_by_allocator_address.find(addr.opaque()); + auto allocation_it = + state->records_by_allocator_address.find(AddressStart(addr)); if (allocation_it != state->records_by_allocator_address.end() && allocation_it->second->allocator_active() && allocation_it->second->allocator_matches(addr)) { @@ -447,7 +454,7 @@ DeviceAddressVmmAllocator::TrackAllocatorAddressMappedAllocation( std::unique_ptr reservation, MemoryReservation::ScopedMapping mapping, int64_t memory_space, bool multi_device) { - void* va_ptr = allocator_address.opaque(); + uintptr_t va_ptr = AddressStart(allocator_address); CHECK(raw_allocation != nullptr); uint64_t physical_size = raw_allocation->address().size(); auto record = std::make_unique( @@ -465,7 +472,7 @@ std::optional DeviceAddressVmmAllocator::TryReuseMappedAllocation( PerDeviceState& state, const MappedAllocateRequest& request) { auto record_it = state.records_by_allocator_address.find( - request.reservation_address.opaque()); + AddressStart(request.reservation_address)); if (record_it == state.records_by_allocator_address.end()) { return std::nullopt; } @@ -479,20 +486,9 @@ DeviceAddressVmmAllocator::TryReuseMappedAllocation( return std::nullopt; } - auto pending_it = state.pending_deallocations.end(); - for (auto it = state.pending_deallocations.begin(); - it != state.pending_deallocations.end(); ++it) { - if (it->kind == PendingDeallocationKind::kAllocation && - it->addr.IsSameAs(record->allocator_address())) { - pending_it = it; - break; - } - } - CHECK(pending_it != state.pending_deallocations.end()); - DeviceAddressBase reused_mem(record->allocator_key(), request.size); - MoveAllocatorRecordToActive(state, *record, request.size); - state.pending_deallocations.erase(pending_it); + ErasePendingDeallocation(state, record->allocator_deallocation()); + record->ReactivateAllocator(request.size); return reused_mem; } @@ -520,16 +516,9 @@ DeviceAddressVmmAllocator::EnsureReservationAvailableForFreshMapping( if (stale_overlap.has_value()) { CHECK(!stale_overlap->is_active); AllocationRecord& record = *stale_overlap->record; - if (stale_overlap->is_allocator) { - pending_completion_key = PendingDeallocationKey{ - PendingDeallocationKind::kAllocation, - record.allocator_stale_seqno(), record.allocator_address()}; - } else { - CHECK(record.reservation_stale()); - pending_completion_key = PendingDeallocationKey{ - PendingDeallocationKind::kMap, record.reservation_stale_seqno(), - record.reservation_address()}; - } + pending_completion_key = stale_overlap->is_allocator + ? record.allocator_deallocation().key() + : record.reservation_deallocation().key(); } } if (!pending_completion_key.has_value()) { @@ -623,12 +612,14 @@ DeviceAddressVmmAllocator::TryWithPendingReclaim(PerDeviceState& state, // Target 1.1x the requested size to provide some headroom. uint64_t target_size = rounded_size + rounded_size / 10; - for (const PendingDeallocation& pending : state.pending_deallocations) { + for (const PendingDeallocationNode* node = state.pending_head; + node != nullptr; node = node->next) { + const PendingDeallocation& pending = node->pending; if (pending.kind == PendingDeallocationKind::kMap) { continue; } auto record_it = - state.records_by_allocator_address.find(pending.addr.opaque()); + state.records_by_allocator_address.find(AddressStart(pending.addr)); CHECK(record_it != state.records_by_allocator_address.end()); CHECK(record_it->second->allocator_stale()); if (record_it->second->memory_space() == reclaim_exempt_memory_space_) { @@ -639,8 +630,7 @@ DeviceAddressVmmAllocator::TryWithPendingReclaim(PerDeviceState& state, CHECK_GT(reclaimable_bytes, 0); accumulated_size += reclaimable_bytes; target_seqno = std::max(target_seqno, pending.seqno); - selected.push_back( - PendingDeallocationKey{pending.kind, pending.seqno, pending.addr}); + selected.push_back(node->key()); if (accumulated_size >= target_size) { break; } @@ -679,17 +669,18 @@ DeviceAddressVmmAllocator::Allocate(int device_ordinal, uint64_t size, auto try_reuse = [&]() -> std::optional { state->mu.AssertHeld(); uint64_t rounded_size = RoundUpToGranularity(*state, size); - for (auto it = state->pending_deallocations.begin(); - it != state->pending_deallocations.end(); ++it) { - if (it->kind != PendingDeallocationKind::kAllocation) { + for (PendingDeallocationNode* node = state->pending_head; node != nullptr; + node = node->next) { + const PendingDeallocation& pending = node->pending; + if (pending.kind != PendingDeallocationKind::kAllocation) { continue; } auto record_it = - state->records_by_allocator_address.find(it->addr.opaque()); + state->records_by_allocator_address.find(AddressStart(pending.addr)); CHECK(record_it != state->records_by_allocator_address.end()); AllocationRecord& record = *record_it->second; CHECK(record.allocator_stale()); - CHECK(record.allocator_matches(it->addr)); + CHECK(record.allocator_matches(pending.addr)); if (record.kind() != AllocationRecord::Kind::kAllocate) { continue; } @@ -706,8 +697,8 @@ DeviceAddressVmmAllocator::Allocate(int device_ordinal, uint64_t size, } DeviceAddressBase reused_mem(record.allocator_key(), size); - MoveAllocatorRecordToActive(*state, record, size); - state->pending_deallocations.erase(it); + ErasePendingDeallocation(*state, *node); + record.ReactivateAllocator(size); return reused_mem; } @@ -815,7 +806,7 @@ absl::Status DeviceAddressVmmAllocator::Deallocate(int device_ordinal, absl::MutexLock lock(state->mu); - auto record_it = state->records_by_allocator_address.find(mem.opaque()); + auto record_it = state->records_by_allocator_address.find(AddressStart(mem)); if (record_it == state->records_by_allocator_address.end() || !record_it->second->allocator_active() || !record_it->second->allocator_matches(mem)) { @@ -825,7 +816,7 @@ absl::Status DeviceAddressVmmAllocator::Deallocate(int device_ordinal, mem.opaque())); } AllocationRecord& record = *record_it->second; - auto reservation_it = state->reservation_records.find(mem.opaque()); + auto reservation_it = state->reservation_records.find(AddressStart(mem)); CHECK(reservation_it == state->reservation_records.end() || !reservation_it->second->reservation_active()); if (record.reservation_active()) { @@ -853,7 +844,8 @@ absl::Status DeviceAddressVmmAllocator::Deallocate(int device_ordinal, // Move the returned allocator address out of active ownership and keep its // mapping alive as stale state until the stream reaches `seqno`. record.MarkAllocatorStale(seqno); - state->pending_deallocations.push_back( + EnqueuePendingDeallocation( + *state, record.allocator_deallocation(), PendingDeallocation{PendingDeallocationKind::kAllocation, seqno, record.allocator_address(), reclaimable_bytes}); return absl::OkStatus(); @@ -902,46 +894,73 @@ DeviceAddressVmmAllocator::FindOverlappingRecord( } }; - auto check_record = [&](AllocationRecord* record, - DeviceAddressBase tracked_address, bool is_allocator, - bool is_active) -> std::optional { - if (matches(tracked_address)) { - return OverlappingRecord{record, tracked_address, is_allocator, - is_active}; + auto find_in_index = + [&](const auto& records, auto get_record, + bool is_allocator) -> std::optional { + auto check_record = + [&](const auto& entry) -> std::optional { + AllocationRecord* record = get_record(entry.second); + CHECK(is_allocator || record->has_reservation_alias()); + bool is_active = is_allocator ? record->allocator_active() + : record->reservation_active(); + if (!(is_active ? include_active : include_stale)) { + return std::nullopt; + } + DeviceAddressBase tracked_address = is_allocator + ? record->allocator_address() + : record->reservation_address(); + if (matches(tracked_address)) { + return OverlappingRecord{record, tracked_address, is_allocator, + is_active}; + } + return std::nullopt; + }; + + if (overlap_kind == OverlapKind::kExact) { + auto it = records.find(AddressStart(address)); + return it == records.end() ? std::nullopt : check_record(*it); + } + if (address.is_null() || address.size() == 0) { + return std::nullopt; } - return std::nullopt; - }; - if (include_allocator) { - for (const auto& [_, record_owner] : state.records_by_allocator_address) { - AllocationRecord* record = record_owner.get(); - bool include_record = (include_active && record->allocator_active()) || - (include_stale && record->allocator_stale()); - if (!include_record) { - continue; - } - if (auto overlap = - check_record(record, record->allocator_address(), - /*is_allocator=*/true, - /*is_active=*/record->allocator_active())) { + auto it = records.lower_bound(AddressStart(address)); + // Only the immediate predecessor can extend into the query: tracked ranges + // are disjoint, even while stale. Read the current range from the record + // because reusing an allocation can change its logical size. + if (it != records.begin()) { + auto previous = it; + --previous; + if (auto overlap = check_record(*previous)) { return overlap; } } - } - if (include_reservation) { - for (const auto& [_, record] : state.reservation_records) { - CHECK(record->has_reservation_alias()); - bool include_record = (include_active && record->reservation_active()) || - (include_stale && record->reservation_stale()); - if (!include_record) { - continue; - } - if (auto overlap = check_record( - record, record->reservation_address(), /*is_allocator=*/false, - /*is_active=*/record->reservation_active())) { + // With kBoth, the first successor either overlaps, exactly matches (so no + // other range overlaps), or starts beyond the query. Filtering by state can + // require checking additional successors, but never unrelated ranges. + for (; it != records.end() && it->first < AddressEnd(address); ++it) { + if (auto overlap = check_record(*it)) { return overlap; } } + return std::nullopt; + }; + + if (include_allocator) { + if (auto overlap = find_in_index( + state.records_by_allocator_address, + [](const std::unique_ptr& record) { + return record.get(); + }, + /*is_allocator=*/true)) { + return overlap; + } + } + if (include_reservation) { + return find_in_index( + state.reservation_records, + [](AllocationRecord* record) { return record; }, + /*is_allocator=*/false); } return std::nullopt; @@ -952,7 +971,7 @@ DeviceAddressVmmAllocator::ResolveMapSourceRecord( PerDeviceState& state, DeviceAddressBase source_address, uint64_t size) const { auto allocation_it = - state.records_by_allocator_address.find(source_address.opaque()); + state.records_by_allocator_address.find(AddressStart(source_address)); if (allocation_it == state.records_by_allocator_address.end() || !allocation_it->second->allocator_active() || !allocation_it->second->allocator_matches(source_address)) { @@ -1004,6 +1023,19 @@ absl::Status DeviceAddressVmmAllocator::ResolveAndMapAlias( ABSL_ASSIGN_OR_RETURN( source_record, ResolveMapSourceRecord(state, request.source_address, request.size)); + + // The source already identifies its one reservation alias. An exact + // stale alias still owns this entire VA range, so no other tracked range + // can overlap it. Reactivate it before any destination-index searches. + // Source validation above and all state changes remain under state.mu. + if (source_record->reservation_stale() && + source_record->reservation_matches(request.reservation_address)) { + ErasePendingDeallocation(state, + source_record->reservation_deallocation()); + source_record->ReactivateReservation(); + return absl::OkStatus(); + } + ABSL_RETURN_IF_ERROR( CheckNoPartialReservationOverlap(state, request.reservation_address)); @@ -1027,13 +1059,10 @@ absl::Status DeviceAddressVmmAllocator::ResolveAndMapAlias( } if (source_record->reservation_stale()) { - CHECK(source_record->has_reservation_alias()); - if (!source_record->reservation_matches(request.reservation_address)) { - pending_completion_key = - PendingDeallocationKey{PendingDeallocationKind::kMap, - source_record->reservation_stale_seqno(), - source_record->reservation_address()}; - } + // An exact alias was reused above; any remaining source alias must be + // retired before mapping the source at a different destination. + pending_completion_key = + source_record->reservation_deallocation().key(); } if (!pending_completion_key.has_value()) { @@ -1043,22 +1072,12 @@ absl::Status DeviceAddressVmmAllocator::ResolveAndMapAlias( if (stale_reservation_overlap.has_value()) { AllocationRecord& stale_record = *stale_reservation_overlap->record; - // A deferred UnMap() leaves the old mapping valid. Reuse it when it - // aliases the requested physical allocation; otherwise wait before - // overwriting it. + // Matching source aliases were reactivated above. This stale + // destination belongs to another source; wait before overwriting it. CHECK(stale_record.has_reservation_alias()); CHECK(stale_record.reservation_matches(request.reservation_address)); - if (stale_record.raw_allocation() == - source_record->raw_allocation()) { - stale_record.ReactivateReservation(); - ErasePendingDeallocation(state, PendingDeallocationKind::kMap, - request.reservation_address); - return absl::OkStatus(); - } pending_completion_key = - PendingDeallocationKey{PendingDeallocationKind::kMap, - stale_record.reservation_stale_seqno(), - stale_record.reservation_address()}; + stale_record.reservation_deallocation().key(); } } @@ -1067,11 +1086,8 @@ absl::Status DeviceAddressVmmAllocator::ResolveAndMapAlias( state, request.reservation_address, AddressRole::kAllocator, RecordState::kStale, OverlapKind::kExact); if (stale_allocator_overlap.has_value()) { - AllocationRecord& stale_record = *stale_allocator_overlap->record; pending_completion_key = - PendingDeallocationKey{PendingDeallocationKind::kAllocation, - stale_record.allocator_stale_seqno(), - stale_record.allocator_address()}; + stale_allocator_overlap->record->allocator_deallocation().key(); } } @@ -1090,8 +1106,8 @@ absl::Status DeviceAddressVmmAllocator::ResolveAndMapAlias( << ", actual=" << mapped.opaque(); source_record->AddActiveReservationAlias(std::move(mapping)); - auto mapping_insert_result = - state.reservation_records.emplace(mapped.opaque(), source_record); + auto mapping_insert_result = state.reservation_records.emplace( + AddressStart(mapped), source_record); CHECK(mapping_insert_result.second); return absl::OkStatus(); } @@ -1144,17 +1160,18 @@ DeviceAddressVmmAllocator::OpenBatchSize(const PerDeviceState& state) const { // Entries are appended in non-decreasing seqno order, so the open batch is // the trailing run carrying the open seqno. Walking it costs at most // kMaxOpenDeallocationBatchEntries steps. - for (auto it = state.pending_deallocations.rbegin(); - it != state.pending_deallocations.rend(); ++it) { - if (it->seqno != state.open_deallocation_batch_seqno) { + for (const PendingDeallocationNode* node = state.pending_tail; + node != nullptr; node = node->previous) { + const PendingDeallocation& pending = node->pending; + if (pending.seqno != state.open_deallocation_batch_seqno) { break; } ++size.entries; if (std::numeric_limits::max() - size.bytes < - it->reclaimable_bytes) { + pending.reclaimable_bytes) { size.bytes = std::numeric_limits::max(); } else { - size.bytes += it->reclaimable_bytes; + size.bytes += pending.reclaimable_bytes; } } return size; @@ -1216,25 +1233,38 @@ absl::Status DeviceAddressVmmAllocator::FlushOpenDeallocationBatch( return absl::OkStatus(); } -void DeviceAddressVmmAllocator::ErasePendingDeallocation( - PerDeviceState& state, PendingDeallocationKind kind, - DeviceAddressBase addr) { - for (auto it = state.pending_deallocations.begin(); - it != state.pending_deallocations.end(); ++it) { - if (it->kind == kind && it->addr.IsSameAs(addr)) { - state.pending_deallocations.erase(it); - return; - } +void DeviceAddressVmmAllocator::EnqueuePendingDeallocation( + PerDeviceState& state, PendingDeallocationNode& node, + PendingDeallocation pending) { + CHECK_EQ(node.pending.seqno, 0); + CHECK_GT(pending.seqno, 0); + node.pending = pending; + node.previous = state.pending_tail; + node.next = nullptr; + if (state.pending_tail != nullptr) { + state.pending_tail->next = &node; + } else { + state.pending_head = &node; } + state.pending_tail = &node; } -void DeviceAddressVmmAllocator::MoveAllocatorRecordToActive( - PerDeviceState& state, AllocationRecord& record, uint64_t new_size) { - void* allocator_va = record.allocator_key(); - auto record_it = state.records_by_allocator_address.find(allocator_va); - CHECK(record_it != state.records_by_allocator_address.end()); - CHECK_EQ(record_it->second.get(), &record); - record.ReactivateAllocator(new_size); +void DeviceAddressVmmAllocator::ErasePendingDeallocation( + PerDeviceState& state, PendingDeallocationNode& node) { + CHECK_GT(node.pending.seqno, 0); + if (node.previous != nullptr) { + node.previous->next = node.next; + } else { + CHECK_EQ(state.pending_head, &node); + state.pending_head = node.next; + } + if (node.next != nullptr) { + node.next->previous = node.previous; + } else { + CHECK_EQ(state.pending_tail, &node); + state.pending_tail = node.previous; + } + node = {}; } absl::Status DeviceAddressVmmAllocator::WaitUntilSeqno(PerDeviceState& state, @@ -1259,19 +1289,20 @@ absl::Status DeviceAddressVmmAllocator::WaitUntilSeqno(PerDeviceState& state, void DeviceAddressVmmAllocator::CompleteReadyAllocatorDeallocationsForReclaim( PerDeviceState& state, uint64_t completed_seqno) { std::vector selected; - for (const PendingDeallocation& pending : state.pending_deallocations) { + for (const PendingDeallocationNode* node = state.pending_head; + node != nullptr; node = node->next) { + const PendingDeallocation& pending = node->pending; if (pending.seqno > completed_seqno || pending.kind == PendingDeallocationKind::kMap) { continue; } auto record_it = - state.records_by_allocator_address.find(pending.addr.opaque()); + state.records_by_allocator_address.find(AddressStart(pending.addr)); if (record_it != state.records_by_allocator_address.end() && record_it->second->memory_space() == reclaim_exempt_memory_space_) { continue; } - selected.push_back( - PendingDeallocationKey{pending.kind, pending.seqno, pending.addr}); + selected.push_back(node->key()); } for (const PendingDeallocationKey& key : selected) { CompletePendingDeallocationByKey(state, key); @@ -1280,22 +1311,22 @@ void DeviceAddressVmmAllocator::CompleteReadyAllocatorDeallocationsForReclaim( void DeviceAddressVmmAllocator::CompletePendingDeallocationByKey( PerDeviceState& state, const PendingDeallocationKey& key) { - for (auto it = state.pending_deallocations.begin(); - it != state.pending_deallocations.end(); ++it) { - if (it->kind == key.kind && it->seqno == key.seqno && - it->addr.IsSameAs(key.addr)) { - PendingDeallocation pending = *it; - state.pending_deallocations.erase(it); - CompletePendingDeallocation(state, pending); + for (PendingDeallocationNode* node = state.pending_head; node != nullptr; + node = node->next) { + if (node->pending.kind == key.kind && node->pending.seqno == key.seqno && + node->pending.addr.IsSameAs(key.addr)) { + CompletePendingDeallocation(state, *node); return; } } } void DeviceAddressVmmAllocator::CompletePendingDeallocation( - PerDeviceState& state, const PendingDeallocation& pending) { + PerDeviceState& state, PendingDeallocationNode& node) { + const PendingDeallocation pending = node.pending; + ErasePendingDeallocation(state, node); if (pending.kind == PendingDeallocationKind::kMap) { - auto record_it = state.reservation_records.find(pending.addr.opaque()); + auto record_it = state.reservation_records.find(AddressStart(pending.addr)); CHECK(record_it != state.reservation_records.end()); AllocationRecord& record = *record_it->second; CHECK(record.reservation_stale()); @@ -1308,7 +1339,7 @@ void DeviceAddressVmmAllocator::CompletePendingDeallocation( } auto record_it = - state.records_by_allocator_address.find(pending.addr.opaque()); + state.records_by_allocator_address.find(AddressStart(pending.addr)); CHECK(record_it != state.records_by_allocator_address.end()); CHECK_EQ(pending.kind, PendingDeallocationKind::kAllocation); CHECK(record_it->second->allocator_stale()); @@ -1329,10 +1360,8 @@ void DeviceAddressVmmAllocator::CompletePendingDeallocation( // sequence numbers are assigned in that same order. Without this invariant // the alias could sit in a still-open batch with no stream marker at all. CHECK_LE(record.reservation_stale_seqno(), pending.seqno); - CompletePendingDeallocationByKey( - state, PendingDeallocationKey{PendingDeallocationKind::kMap, - record.reservation_stale_seqno(), - record.reservation_address()}); + CompletePendingDeallocationByKey(state, + record.reservation_deallocation().key()); CHECK(!record.has_reservation_alias()); } uint64_t physical_size = record.raw_allocation()->address().size(); @@ -1362,7 +1391,7 @@ absl::Status DeviceAddressVmmAllocator::UnMap(int device_ordinal, // UnMap() only accepts the exact active reservation range previously created // by Map(). Allocator addresses and subranges are not valid UnMap() inputs. auto reservation_it = - state->reservation_records.find(reservation_address.opaque()); + state->reservation_records.find(AddressStart(reservation_address)); if (reservation_it == state->reservation_records.end()) { return absl::NotFoundError(absl::StrFormat( "UnMap() requires an exact active reservation range created by Map(): " @@ -1397,7 +1426,8 @@ absl::Status DeviceAddressVmmAllocator::UnMap(int device_ordinal, // stream marker will be enqueued for the whole batch when it is flushed. uint64_t seqno = GetOrCreateOpenDeallocationBatchSeqno(*state); record->MarkReservationStale(seqno); - state->pending_deallocations.push_back( + EnqueuePendingDeallocation( + *state, record->reservation_deallocation(), PendingDeallocation{PendingDeallocationKind::kMap, seqno, reservation_address, /*reclaimable_bytes=*/0}); return absl::OkStatus(); diff --git a/third_party/xla/xla/stream_executor/device_address_vmm_allocator.h b/third_party/xla/xla/stream_executor/device_address_vmm_allocator.h index 2b1891aa304c08..2e4df779996dd7 100644 --- a/third_party/xla/xla/stream_executor/device_address_vmm_allocator.h +++ b/third_party/xla/xla/stream_executor/device_address_vmm_allocator.h @@ -17,13 +17,13 @@ limitations under the License. #define XLA_STREAM_EXECUTOR_DEVICE_ADDRESS_VMM_ALLOCATOR_H_ #include -#include #include #include #include #include #include "absl/base/thread_annotations.h" +#include "absl/container/btree_map.h" #include "absl/container/flat_hash_map.h" #include "absl/status/status.h" #include "absl/status/statusor.h" @@ -292,6 +292,43 @@ class DeviceAddressVmmAllocator : public DeviceAddressAllocator { kMap, }; + // Stable identity for one pending operation. A batch shares one sequence + // number, so the operation kind and address are also needed to select an + // entry after WaitUntilSeqno() temporarily releases state.mu. + struct PendingDeallocationKey { + PendingDeallocationKind kind; + uint64_t seqno; + DeviceAddressBase addr; + }; + + // Snapshot of a stream-ordered deferred operation. Copy before unlinking its + // node: completing the operation can destroy the record that owns the node. + struct PendingDeallocation { + PendingDeallocationKind kind; + uint64_t seqno; + // Allocator address for allocation deallocations; reservation address for + // kMap. + DeviceAddressBase addr; + // Physical bytes released when this entry completes; 0 for kMap. Store the + // bytes charged when enqueuing so open-batch accounting remains consistent. + uint64_t reclaimable_bytes; + }; + + // Intrusive queue node owned by an allocation or its reservation alias. + // Records stay at stable addresses while queued, allowing cancellation + // without searching the queue or allocating separate queue storage. + struct PendingDeallocationNode { + // A zero seqno means this node is not queued. + PendingDeallocation pending{PendingDeallocationKind::kAllocation, 0, + DeviceAddressBase(), 0}; + PendingDeallocationNode* previous = nullptr; + PendingDeallocationNode* next = nullptr; + + PendingDeallocationKey key() const { + return {pending.kind, pending.seqno, pending.addr}; + } + }; + // Lifetime record for one raw physical allocation. // // The record is owned by records_by_allocator_address while either the @@ -347,6 +384,11 @@ class DeviceAddressVmmAllocator : public DeviceAddressAllocator { } uint64_t reservation_stale_seqno() const; + PendingDeallocationNode& allocator_deallocation() { + return allocator_deallocation_; + } + PendingDeallocationNode& reservation_deallocation(); + void MarkAllocatorStale(uint64_t seqno); void ReactivateAllocator(uint64_t new_size); @@ -361,6 +403,7 @@ class DeviceAddressVmmAllocator : public DeviceAddressAllocator { MemoryReservation::ScopedMapping mapping; // Zero while active; the deferred UnMap() sequence number while stale. uint64_t stale_seqno = 0; + PendingDeallocationNode deallocation = {}; }; Kind kind_; @@ -379,33 +422,7 @@ class DeviceAddressVmmAllocator : public DeviceAddressAllocator { std::optional reservation_alias_; uint64_t allocator_stale_seqno_ = 0; - }; - - // Queue entry for a stream-ordered deferred operation. The heavy resources - // live in AllocationRecord; this entry only says which stale address becomes - // safe to complete when the GPU timeline reaches `seqno`. - struct PendingDeallocation { - PendingDeallocationKind kind; - // GPU stream sequence number recorded at deallocation time. When the - // pinned_timeline value reaches this seqno, the memory is safe to free. - uint64_t seqno; - // Allocator address for allocation deallocations; reservation address for - // kMap. - DeviceAddressBase addr; - // Physical bytes released when this entry completes; 0 for kMap, which only - // drops a reservation alias. Stored here rather than re-derived from the - // AllocationRecord so that open-batch accounting reads exactly the value - // that was accounted when the entry was queued. - uint64_t reclaimable_bytes; - }; - - // Stable identity for one pending operation. A batch shares one sequence - // number, so the operation kind and address are also needed to select an - // entry after WaitUntilSeqno() temporarily releases state.mu. - struct PendingDeallocationKey { - PendingDeallocationKind kind; - uint64_t seqno; - DeviceAddressBase addr; + PendingDeallocationNode allocator_deallocation_; }; struct PerDeviceState { @@ -455,25 +472,30 @@ class DeviceAddressVmmAllocator : public DeviceAddressAllocator { // The next deferred operation opens a new batch with seqno 2. // // Only the sequence number is stored. The batch's entry and byte totals are - // the trailing run of `pending_deallocations` carrying this seqno (entries + // the trailing run of the pending queue carrying this seqno (entries // are appended in non-decreasing seqno order, so that run is contiguous), // and are computed on demand by OpenDeallocationBatchSize(). Keeping them // derived means cancelling an entry -- which happens whenever a stale // record is reused -- needs no batch bookkeeping at all. uint64_t open_deallocation_batch_seqno ABSL_GUARDED_BY(mu) = 0; - std::deque pending_deallocations ABSL_GUARDED_BY(mu); - // Owns AllocationRecord objects. Key is the allocator address pointer - // (`AllocationRecord::allocator_address().opaque()`), including the - // reservation-derived allocator address returned by the mapped Allocate() - // overload. Allocator-address active/stale state is stored in + PendingDeallocationNode* pending_head ABSL_GUARDED_BY(mu) = nullptr; + PendingDeallocationNode* pending_tail ABSL_GUARDED_BY(mu) = nullptr; + // Owns AllocationRecord objects. Key is the allocator address as uintptr_t, + // including the reservation-derived address returned by the mapped + // Allocate() overload. Allocator-address active/stale state is stored in // AllocationRecord::allocator_active()/allocator_stale(). - absl::flat_hash_map> + // Both address indexes are ordered by range start. Ranges within each + // index are disjoint, including stale ranges retained until deferred + // teardown. This permits overlap checks starting at lower_bound() instead + // of scanning every allocation. The pointed-to records stay valid when + // insertion or erasure invalidates btree iterators. + absl::btree_map> records_by_allocator_address ABSL_GUARDED_BY(mu); - // Reservation-address index. Keys are reservation alias pointers - // (`AllocationRecord::reservation_address().opaque()`) created by Map(). + // Reservation-address index. Keys are reservation alias addresses as + // uintptr_t, created by Map(). // Active/stale state is stored in the record. - absl::flat_hash_map reservation_records + absl::btree_map reservation_records ABSL_GUARDED_BY(mu); }; @@ -677,14 +699,15 @@ class DeviceAddressVmmAllocator : public DeviceAddressAllocator { absl::Status DrainPendingDeallocations(PerDeviceState& state) ABSL_EXCLUSIVE_LOCKS_REQUIRED(state.mu); - // Removes the matching pending entry when a stale record is reused. - void ErasePendingDeallocation(PerDeviceState& state, - PendingDeallocationKind kind, - DeviceAddressBase addr) + // Appends a record-owned node in stream sequence order. + void EnqueuePendingDeallocation(PerDeviceState& state, + PendingDeallocationNode& node, + PendingDeallocation pending) ABSL_EXCLUSIVE_LOCKS_REQUIRED(state.mu); - void MoveAllocatorRecordToActive(PerDeviceState& state, - AllocationRecord& record, uint64_t new_size) + // Unlinks a queued node in constant time, before reusing or completing it. + void ErasePendingDeallocation(PerDeviceState& state, + PendingDeallocationNode& node) ABSL_EXCLUSIVE_LOCKS_REQUIRED(state.mu); // Waits for the device timeline to reach `target_seqno`. Temporarily releases @@ -699,12 +722,12 @@ class DeviceAddressVmmAllocator : public DeviceAddressAllocator { uint64_t completed_seqno) ABSL_EXCLUSIVE_LOCKS_REQUIRED(state.mu); - // Completes a pending operation whose stream sequence has passed by dropping - // its ScopedMappings, allocator-owned reservation, and raw allocation - // reference. This is where VA unmap, reservation release, and PA budget - // accounting happen. + // Unlinks and completes a pending operation whose stream sequence has passed. + // Copies its metadata before releasing the mappings and owning record, which + // can destroy the node. This is where VA unmap, reservation release, and PA + // budget accounting happen. void CompletePendingDeallocation(PerDeviceState& state, - const PendingDeallocation& pending) + PendingDeallocationNode& node) ABSL_EXCLUSIVE_LOCKS_REQUIRED(state.mu); // Finds, erases, and completes the selected pending entry if it is still diff --git a/third_party/xla/xla/stream_executor/device_address_vmm_allocator_test.cc b/third_party/xla/xla/stream_executor/device_address_vmm_allocator_test.cc index a4417c2886db08..9dac5e156ebb8b 100644 --- a/third_party/xla/xla/stream_executor/device_address_vmm_allocator_test.cc +++ b/third_party/xla/xla/stream_executor/device_address_vmm_allocator_test.cc @@ -18,7 +18,10 @@ limitations under the License. #include #include #include +#include #include +#include +#include #include #include @@ -27,6 +30,7 @@ limitations under the License. #include "absl/status/statusor.h" #include "absl/types/span.h" #include "xla/stream_executor/device_address.h" +#include "xla/stream_executor/device_address_allocator.h" #include "xla/stream_executor/memory_allocation.h" #include "xla/stream_executor/memory_reservation.h" #include "xla/stream_executor/mock_platform.h" @@ -72,6 +76,7 @@ class TestMemoryReservation final : public MemoryReservation { } int active_mapping_count() const { return active_mapping_count_; } + int mapping_count() const { return mapping_count_; } private: absl::Status Map(size_t reservation_offset, size_t allocation_offset, @@ -82,6 +87,7 @@ class TestMemoryReservation final : public MemoryReservation { return absl::InvalidArgumentError("mapping range is out of bounds"); } ++active_mapping_count_; + ++mapping_count_; return absl::OkStatus(); } @@ -101,6 +107,7 @@ class TestMemoryReservation final : public MemoryReservation { std::unique_ptr storage_; uint64_t size_; int active_mapping_count_ = 0; + int mapping_count_ = 0; }; class TestDeviceAddressVmmAllocator final : public DeviceAddressVmmAllocator { @@ -120,6 +127,7 @@ class TestDeviceAddressVmmAllocator final : public DeviceAddressVmmAllocator { } int allocation_count() const { return allocation_count_; } + int timeline_write_count() const { return timeline_write_count_; } protected: absl::Status InitializeDeviceState(PerDeviceState& state) override { @@ -152,6 +160,7 @@ class TestDeviceAddressVmmAllocator final : public DeviceAddressVmmAllocator { absl::Status EnqueueDeferredDeallocation(PerDeviceState& state, uint64_t seqno) override { + ++timeline_write_count_; __atomic_store_n(state.pinned_timeline, seqno, __ATOMIC_RELEASE); return absl::OkStatus(); } @@ -167,6 +176,7 @@ class TestDeviceAddressVmmAllocator final : public DeviceAddressVmmAllocator { uint64_t physical_size_padding_; std::function on_device_destroy_; int allocation_count_ = 0; + int timeline_write_count_ = 0; }; class DeviceAddressVmmAllocatorTest : public ::testing::Test { @@ -404,5 +414,260 @@ TEST_F(DeviceAddressVmmAllocatorTest, EXPECT_EQ(allocator->allocation_count(), 2); } +TEST_F(DeviceAddressVmmAllocatorTest, + OverlapLookupPreservesRangeBoundariesForActiveAndStaleAddresses) { + for (bool alias : {false, true}) { + for (bool stale : {false, true}) { + SCOPED_TRACE(testing::Message() + << "alias=" << alias << " stale=" << stale); + TestMemoryReservation reservation(16 * kGranularity); + ASSERT_OK_AND_ASSIGN(auto allocator, + TestDeviceAddressVmmAllocator::Create( + &platform_, {Config(UINT64_MAX)})); + ASSERT_OK_AND_ASSIGN( + auto tracked, + alias ? allocator->Allocate(/*device_ordinal=*/0, 4 * kGranularity) + : allocator->Allocate( + /*device_ordinal=*/0, 4 * kGranularity, + /*retry_on_failure=*/false, /*memory_space=*/0, + &reservation, 4 * kGranularity, 4 * kGranularity)); + if (alias) { + ASSERT_THAT(allocator->Map(0, tracked.cref(), &reservation, + 4 * kGranularity, 4 * kGranularity), + absl_testing::IsOk()); + } + if (stale) { + if (alias) { + ASSERT_THAT(allocator->UnMap(0, &reservation, 4 * kGranularity, + 4 * kGranularity), + absl_testing::IsOk()); + } else { + ASSERT_THAT(allocator->Deallocate(0, tracked.Release()), + absl_testing::IsOk()); + } + } + ASSERT_OK_AND_ASSIGN(auto source, + allocator->Allocate(0, 16 * kGranularity)); + struct Range { + uint64_t offset; + uint64_t size; + }; + // Intersect the successor, lie inside the predecessor, intersect its + // end, share its start with a different size, or contain the whole range. + const Range partial_overlaps[] = {{2, 3}, {5, 1}, {7, 2}, + {4, 2}, {4, 5}, {2, 8}}; + for (Range range : partial_overlaps) { + SCOPED_TRACE(testing::Message() + << "offset=" << range.offset << " size=" << range.size); + EXPECT_THAT(allocator->Map(0, source.cref(), &reservation, + range.offset * kGranularity, + range.size * kGranularity), + StatusIs(absl::StatusCode::kFailedPrecondition)); + EXPECT_THAT(allocator->Allocate( + 0, range.size * kGranularity, false, 0, &reservation, + range.offset * kGranularity, range.size * kGranularity), + StatusIs(absl::StatusCode::kFailedPrecondition)); + // Rejecting an overlap must not drain the stale mapping. + EXPECT_EQ(reservation.active_mapping_count(), 1); + } + + // Exactly adjacent ranges are disjoint on either side of the record. + ASSERT_THAT( + allocator->Map(0, source.cref(), &reservation, 0, 4 * kGranularity), + absl_testing::IsOk()); + ASSERT_OK_AND_ASSIGN( + auto right, + allocator->Allocate(0, 4 * kGranularity, false, 0, &reservation, + 8 * kGranularity, 4 * kGranularity)); + auto exact = + allocator->Allocate(0, 4 * kGranularity, false, 0, &reservation, + 4 * kGranularity, 4 * kGranularity); + if (stale) { + ASSERT_THAT(exact, absl_testing::IsOk()); + } else { + EXPECT_THAT(exact, StatusIs(absl::StatusCode::kAlreadyExists)); + } + ASSERT_THAT(allocator->UnMap(0, &reservation, 0, 4 * kGranularity), + absl_testing::IsOk()); + if (alias && !stale) { + ASSERT_THAT(allocator->UnMap(0, &reservation, 4 * kGranularity, + 4 * kGranularity), + absl_testing::IsOk()); + } + } + } +} + +TEST_F(DeviceAddressVmmAllocatorTest, + ManyReservationAliasesReuseAndRemoveOrderedRecords) { + constexpr int kCount = 512; + TestMemoryReservation reservation((2 * kCount + 1) * kGranularity); + ASSERT_OK_AND_ASSIGN(auto allocator, TestDeviceAddressVmmAllocator::Create( + &platform_, {Config(UINT64_MAX)})); + std::vector> sources; + for (int i = 0; i < kCount; ++i) { + ASSERT_OK_AND_ASSIGN(auto source, allocator->Allocate(0, kGranularity)); + sources.push_back(std::move(source)); + } + // Insert in an order that interleaves low and high addresses, leaving gaps + // between aliases. Records must remain valid across index rebalancing. + for (int n = 0; n < kCount; ++n) { + int i = (n % 2 == 0) ? n / 2 : kCount - 1 - n / 2; + ASSERT_THAT(allocator->Map(0, sources[i].cref(), &reservation, + (2 * i + 1) * kGranularity, kGranularity), + absl_testing::IsOk()); + } + for (int i = 0; i < kCount; ++i) { + ASSERT_THAT(allocator->UnMap(0, &reservation, (2 * i + 1) * kGranularity, + kGranularity), + absl_testing::IsOk()); + } + for (int i = kCount - 1; i >= 0; --i) { + ASSERT_THAT(allocator->Map(0, sources[i].cref(), &reservation, + (2 * i + 1) * kGranularity, kGranularity), + absl_testing::IsOk()); + } + EXPECT_EQ(reservation.active_mapping_count(), kCount); + EXPECT_EQ(reservation.mapping_count(), kCount); + EXPECT_EQ(allocator->allocation_count(), kCount); + + // A query starting in a gap must still find the next occupied range. + ASSERT_OK_AND_ASSIGN(auto source, allocator->Allocate(0, 4 * kGranularity)); + EXPECT_THAT(allocator->Map(0, source.cref(), &reservation, + kCount * kGranularity, 4 * kGranularity), + StatusIs(absl::StatusCode::kFailedPrecondition)); + + for (int i = 0; i < kCount; ++i) { + ASSERT_THAT(allocator->UnMap(0, &reservation, (2 * i + 1) * kGranularity, + kGranularity), + absl_testing::IsOk()); + } + ASSERT_THAT(allocator->SynchronizePendingOperations(0), absl_testing::IsOk()); + EXPECT_EQ(reservation.active_mapping_count(), 0); + // The removed alias ranges can now be covered by a larger mapping. + ASSERT_THAT(allocator->Map(0, source.cref(), &reservation, + kCount * kGranularity, 4 * kGranularity), + absl_testing::IsOk()); + ASSERT_THAT(allocator->UnMap(0, &reservation, kCount * kGranularity, + 4 * kGranularity), + absl_testing::IsOk()); +} + +TEST_F(DeviceAddressVmmAllocatorTest, + StaleAliasReuseValidatesSourceAndFullReservationRange) { + TestMemoryReservation reservation(2 * kGranularity); + ASSERT_OK_AND_ASSIGN(auto allocator, TestDeviceAddressVmmAllocator::Create( + &platform_, {Config(UINT64_MAX)})); + ASSERT_OK_AND_ASSIGN(auto source, allocator->Allocate(0, 2 * kGranularity)); + ASSERT_THAT(allocator->Map(0, source.cref(), &reservation, 0, kGranularity), + absl_testing::IsOk()); + ASSERT_THAT(allocator->UnMap(0, &reservation, 0, kGranularity), + absl_testing::IsOk()); + + // Matching the start is insufficient: a larger range must still be rejected, + // without retiring the stale alias or its pending operation. + EXPECT_THAT( + allocator->Map(0, source.cref(), &reservation, 0, 2 * kGranularity), + StatusIs(absl::StatusCode::kFailedPrecondition)); + EXPECT_EQ(reservation.active_mapping_count(), 1); + EXPECT_EQ(allocator->timeline_write_count(), 0); + + ASSERT_THAT(allocator->Map(0, source.cref(), &reservation, 0, kGranularity), + absl_testing::IsOk()); + EXPECT_THAT(allocator->Map(0, source.cref(), &reservation, 0, kGranularity), + StatusIs(absl::StatusCode::kAlreadyExists)); + ASSERT_THAT(allocator->SynchronizePendingOperations(0), absl_testing::IsOk()); + EXPECT_EQ(reservation.mapping_count(), 1); + EXPECT_EQ(reservation.active_mapping_count(), 1); + EXPECT_EQ(allocator->timeline_write_count(), 0); + + ASSERT_THAT(allocator->UnMap(0, &reservation, 0, kGranularity), + absl_testing::IsOk()); + DeviceAddressBase stale_source = source.Release(); + ASSERT_THAT(allocator->Deallocate(0, stale_source), absl_testing::IsOk()); + EXPECT_THAT(allocator->Map(0, stale_source, &reservation, 0, kGranularity), + StatusIs(absl::StatusCode::kNotFound)); + ASSERT_THAT(allocator->SynchronizePendingOperations(0), absl_testing::IsOk()); + EXPECT_EQ(reservation.active_mapping_count(), 0); +} + +TEST_F(DeviceAddressVmmAllocatorTest, + MixedPendingOperationsCanBeCancelledAndRequeuedInAnyOrder) { + constexpr int kCount = 4; + TestMemoryReservation reservation(2 * kCount * kGranularity); + ASSERT_OK_AND_ASSIGN(auto allocator, TestDeviceAddressVmmAllocator::Create( + &platform_, {Config(UINT64_MAX)})); + std::vector> temps; + std::vector> sources; + for (int i = 0; i < kCount; ++i) { + ASSERT_OK_AND_ASSIGN( + auto temp, allocator->Allocate(0, kGranularity, false, 0, &reservation, + i * kGranularity, kGranularity)); + temps.push_back(std::move(temp)); + ASSERT_OK_AND_ASSIGN(auto source, allocator->Allocate(0, kGranularity)); + ASSERT_THAT(allocator->Map(0, source.cref(), &reservation, + (kCount + i) * kGranularity, kGranularity), + absl_testing::IsOk()); + sources.push_back(std::move(source)); + } + + for (int step = 0; step < 3; ++step) { + SCOPED_TRACE(step); + // Interleave allocator deallocations and alias unmaps in the same batch. + for (int i = 0; i < kCount; ++i) { + ASSERT_THAT(allocator->Deallocate(0, temps[i].Release()), + absl_testing::IsOk()); + ASSERT_THAT(allocator->UnMap(0, &reservation, (kCount + i) * kGranularity, + kGranularity), + absl_testing::IsOk()); + } + // Cancel from the middle, head, tail, and finally the remaining entries. + // The following iteration requeues the same record-owned nodes. + for (int i : {1, 0, 3, 2}) { + ASSERT_OK_AND_ASSIGN( + temps[i], allocator->Allocate(0, kGranularity, false, 0, &reservation, + i * kGranularity, kGranularity)); + ASSERT_THAT(allocator->Map(0, sources[i].cref(), &reservation, + (kCount + i) * kGranularity, kGranularity), + absl_testing::IsOk()); + } + ASSERT_THAT(allocator->SynchronizePendingOperations(0), + absl_testing::IsOk()); + EXPECT_EQ(allocator->timeline_write_count(), 0); + EXPECT_EQ(allocator->allocation_count(), 2 * kCount); + EXPECT_EQ(reservation.mapping_count(), 2 * kCount); + EXPECT_EQ(reservation.active_mapping_count(), 2 * kCount); + } + + // Leave one alias pending between cancelled entries. Draining must release + // only that alias, then permit a fresh mapping at the same address. + for (int i = 0; i < kCount; ++i) { + ASSERT_THAT(allocator->UnMap(0, &reservation, (kCount + i) * kGranularity, + kGranularity), + absl_testing::IsOk()); + } + for (int i : {3, 0, 2}) { + ASSERT_THAT(allocator->Map(0, sources[i].cref(), &reservation, + (kCount + i) * kGranularity, kGranularity), + absl_testing::IsOk()); + } + ASSERT_THAT(allocator->SynchronizePendingOperations(0), absl_testing::IsOk()); + EXPECT_EQ(reservation.active_mapping_count(), 2 * kCount - 1); + EXPECT_EQ(allocator->timeline_write_count(), 1); + ASSERT_THAT(allocator->Map(0, sources[1].cref(), &reservation, + (kCount + 1) * kGranularity, kGranularity), + absl_testing::IsOk()); + + for (int i = 0; i < kCount; ++i) { + ASSERT_THAT(allocator->UnMap(0, &reservation, (kCount + i) * kGranularity, + kGranularity), + absl_testing::IsOk()); + } + temps.clear(); + ASSERT_THAT(allocator->SynchronizePendingOperations(0), absl_testing::IsOk()); + EXPECT_EQ(reservation.active_mapping_count(), 0); + EXPECT_EQ(allocator->timeline_write_count(), 2); +} + } // namespace } // namespace stream_executor diff --git a/third_party/xla/xla/stream_executor/kernel_args_packing_spec.cc b/third_party/xla/xla/stream_executor/kernel_args_packing_spec.cc index d174e3ac179384..d52f92f94957d0 100644 --- a/third_party/xla/xla/stream_executor/kernel_args_packing_spec.cc +++ b/third_party/xla/xla/stream_executor/kernel_args_packing_spec.cc @@ -90,6 +90,39 @@ KernelArgPackingSpec KernelArgPackingSpec::BuildArgRelocation( KernelArgPackingRelocation::Kind::kBits64Absolute, argument_index)}); } +KernelArgsPackingSpec KernelArgsPackingSpec::Identity(int num_args) { + CHECK_GE(num_args, 0); + std::vector kernel_arguments; + kernel_arguments.reserve(num_args); + for (int i = 0; i < num_args; ++i) { + kernel_arguments.push_back(KernelArgPackingSpec::BuildArgRelocation(i)); + } + return KernelArgsPackingSpec(std::move(kernel_arguments)); +} + +std::optional KernelArgsPackingSpec::MaxArgumentIndex() const { + std::optional max_index; + for (const KernelArgPackingSpec& kernel_argument : kernel_arguments_) { + std::optional index = kernel_argument.relocation_argument_index(); + if (index.has_value() && (!max_index.has_value() || *index > *max_index)) { + max_index = index; + } + } + return max_index; +} + +absl::Status KernelArgsPackingSpec::Validate(size_t num_available_args) const { + std::optional max_index = MaxArgumentIndex(); + if (max_index.has_value() && + static_cast(*max_index) >= num_available_args) { + return absl::InvalidArgumentError(absl::StrFormat( + "Kernel arguments packing spec refers to device buffer %d, but only " + "%zu buffers are available", + *max_index, num_available_args)); + } + return absl::OkStatus(); +} + absl::StatusOr> KernelArgsPackingSpec::BuildArguments( absl::Span> args, diff --git a/third_party/xla/xla/stream_executor/kernel_args_packing_spec.h b/third_party/xla/xla/stream_executor/kernel_args_packing_spec.h index 586fcfc963eeca..ceec408db33ec3 100644 --- a/third_party/xla/xla/stream_executor/kernel_args_packing_spec.h +++ b/third_party/xla/xla/stream_executor/kernel_args_packing_spec.h @@ -24,6 +24,7 @@ limitations under the License. #include #include +#include "absl/status/status.h" #include "absl/status/statusor.h" #include "absl/types/span.h" #include "xla/stream_executor/kernel_args.h" @@ -66,6 +67,14 @@ class KernelArgPackingSpec { // Builds KernelArgPackingSpec that refers to the given arg number. static KernelArgPackingSpec BuildArgRelocation(int argument_index); + // Returns the index of the argument this spec relocates, or `std::nullopt` + // if this spec holds a constant. + std::optional relocation_argument_index() const { + return relocation_.has_value() + ? std::make_optional(relocation_->argument_index()) + : std::nullopt; + } + // Builds KernelArgPackingSpec that refers to a constant. The value must be // trivially copyable. template @@ -146,6 +155,11 @@ class KernelArgsPackingSpec { std::vector kernel_arguments) : kernel_arguments_(std::move(kernel_arguments)) {} + // Builds the packing spec that passes the first `num_args` device buffer + // pointers to the kernel unchanged, in order. This is the right spec for + // kernels whose signature is exactly "all inputs, then all outputs". + static KernelArgsPackingSpec Identity(int num_args); + // Adds a single argument packing spec to the kernel arguments packing spec. void AddArgument(KernelArgPackingSpec spec) { kernel_arguments_.push_back(std::move(spec)); @@ -163,6 +177,20 @@ class KernelArgsPackingSpec { kernel_arguments_.push_back(KernelArgPackingSpec::BuildFor(value)); } + // Number of arguments the kernel will be launched with. + size_t size() const { return kernel_arguments_.size(); } + + // Returns the largest device buffer index referenced by any relocation, or + // `std::nullopt` if the spec contains no relocations. + std::optional MaxArgumentIndex() const; + + // Returns an error if any relocation refers to a device buffer index that is + // out of range for `num_available_args` buffers. `BuildArguments` performs + // the same check, but only once the kernel is actually launched; calling + // `Validate` while emitting lets a mismatch between the packing spec and the + // kernel arguments surface at compile time instead. + absl::Status Validate(size_t num_available_args) const; + // Materializes the argument buffers for this packing spec. The `args` span // must contain at least the number of arguments referenced in the packing // spec, otherwise an error will be returned. diff --git a/third_party/xla/xla/stream_executor/kernel_args_packing_spec_test.cc b/third_party/xla/xla/stream_executor/kernel_args_packing_spec_test.cc index 281d455fae6359..36dac56a368bb0 100644 --- a/third_party/xla/xla/stream_executor/kernel_args_packing_spec_test.cc +++ b/third_party/xla/xla/stream_executor/kernel_args_packing_spec_test.cc @@ -17,6 +17,7 @@ limitations under the License. #include #include +#include #include #include @@ -232,5 +233,60 @@ TEST(KernelArgsPackingSpecTest, FromProto) { ElementsAre(0x34, 0x12, 0x00, 0x00)); } +TEST(KernelArgsPackingSpecTest, Identity) { + KernelArgsPackingSpec spec = KernelArgsPackingSpec::Identity(3); + + EXPECT_EQ(spec.size(), 3); + EXPECT_THAT(spec.ToProto(), IsOkAndHolds(EqualsProto(R"pb( + kernel_arguments { + relocations { kind: KIND_BITS64_ABSOLUTE argument_index: 0 } + } + kernel_arguments { + relocations { kind: KIND_BITS64_ABSOLUTE argument_index: 1 } + } + kernel_arguments { + relocations { kind: KIND_BITS64_ABSOLUTE argument_index: 2 } + } + )pb"))); +} + +TEST(KernelArgsPackingSpecTest, IdentityWithZeroArguments) { + KernelArgsPackingSpec spec = KernelArgsPackingSpec::Identity(0); + + EXPECT_EQ(spec.size(), 0); + EXPECT_EQ(spec.MaxArgumentIndex(), std::nullopt); +} + +TEST(KernelArgsPackingSpecTest, MaxArgumentIndex) { + KernelArgsPackingSpec spec; + EXPECT_EQ(spec.MaxArgumentIndex(), std::nullopt); + + spec.AddConstantArgument(0x1234); + EXPECT_EQ(spec.MaxArgumentIndex(), std::nullopt); + + spec.AddAddressArgument(/*argument_index=*/5); + spec.AddAddressArgument(/*argument_index=*/2); + EXPECT_EQ(spec.MaxArgumentIndex(), 5); +} + +TEST(KernelArgsPackingSpecTest, Validate) { + KernelArgsPackingSpec spec; + spec.AddAddressArgument(/*argument_index=*/0); + spec.AddConstantArgument(0x1234); + spec.AddAddressArgument(/*argument_index=*/2); + + EXPECT_OK(spec.Validate(/*num_available_args=*/3)); + EXPECT_OK(spec.Validate(/*num_available_args=*/4)); + EXPECT_THAT(spec.Validate(/*num_available_args=*/2), + StatusIs(absl::StatusCode::kInvalidArgument)); +} + +TEST(KernelArgsPackingSpecTest, ValidateConstantsOnlySpecNeedsNoArguments) { + KernelArgsPackingSpec spec; + spec.AddConstantArgument(0x1234); + + EXPECT_OK(spec.Validate(/*num_available_args=*/0)); +} + } // namespace } // namespace stream_executor diff --git a/third_party/xla/xla/stream_executor/rocm/hip_blas_lt.cc b/third_party/xla/xla/stream_executor/rocm/hip_blas_lt.cc index cf063977e54d87..9c2543e98d5fcb 100644 --- a/third_party/xla/xla/stream_executor/rocm/hip_blas_lt.cc +++ b/third_party/xla/xla/stream_executor/rocm/hip_blas_lt.cc @@ -91,6 +91,11 @@ void GroupGemmUpdateArgs( int8_t bias_type, bool has_matrix_bias); namespace { +// Upper bound for the heuristic query only, never allocated. Generous on +// purpose: a low ceiling would hide complex kernels that need scratch and +// divert the matmul to rocBLAS. +constexpr size_t kComplexProbeMaxWorkspaceBytes = 64 * 1024 * 1024; + template absl::Status SetAttr(hipblasLtMatrixLayout_t handle, hipblasLtMatrixLayoutAttribute_t attr, T value) { @@ -311,21 +316,8 @@ auto BlasLt::RegularMatmulPlan::GetAlgorithms(size_t max_algorithm_count, return algorithms; } -absl::StatusOr BlasLt::GetMatmulPlan( +absl::StatusOr BlasLt::GetHipBlasLtMatmulPlan( const gpu::GemmConfig& cfg, Epilogue epilogue) const { - // hipBLASLt has no complex GEMM kernels; redirect C64/C128 to rocBLAS while - // still consuming the cublasLt matmul custom call emitted by GemmRewriter. - // TODO(magaonka-amd): Once we get a hipBLASLt that supports complex GEMMs, - // clean up this routing code. - if (cfg.output_layout.dtype == xla::C64 || - cfg.output_layout.dtype == xla::C128) { - TF_RET_CHECK(epilogue == Epilogue::kDefault) - << "rocBLAS complex fallback does not support epilogues (bias, GELU, " - "etc.); got epilogue=" - << static_cast(epilogue); - return std::make_unique(*this, cfg); - } - auto lhs_layout = cfg.lhs_layout, rhs_layout = cfg.rhs_layout, output_layout = cfg.output_layout, c_layout = cfg.c_layout; @@ -566,6 +558,36 @@ absl::StatusOr BlasLt::GetMatmulPlan( return plan; } +absl::StatusOr BlasLt::GetMatmulPlan( + const gpu::GemmConfig& cfg, Epilogue epilogue) const { + if (cfg.output_layout.dtype != xla::C64 && + cfg.output_layout.dtype != xla::C128) { + return GetHipBlasLtMatmulPlan(cfg, epilogue); + } + + // Older hipBLASLt rejects the complex matmul descriptor outright, and a newer + // one can still have no kernels for this device, so test both rather than the + // ROCm version XLA was built against. + auto plan = GetHipBlasLtMatmulPlan(cfg, epilogue); + if (plan.ok()) { + auto algorithms = (*plan)->GetAlgorithms(/*max_algorithm_count=*/1, + kComplexProbeMaxWorkspaceBytes); + if (algorithms.ok() && !algorithms->empty()) { + return plan; + } + } + + TF_RET_CHECK(epilogue == Epilogue::kDefault) + << "rocBLAS complex fallback does not support epilogues (bias, GELU, " + "etc.); got epilogue=" + << static_cast(epilogue); + VLOG(1) << "hipBLASLt offers no complex GEMM algorithm here; routing " + << xla::primitive_util::LowercasePrimitiveTypeName( + cfg.output_layout.dtype) + << " matmul to rocBLAS"; + return std::make_unique(*this, cfg); +} + absl::Status BlasLt::RocBlasGemmPlan::ExecuteOnStream( Stream* stream, const gpu::BlasLt::MemoryArgs& args, blas::ProfileResult* profile_result) const { diff --git a/third_party/xla/xla/stream_executor/rocm/hip_blas_lt.h b/third_party/xla/xla/stream_executor/rocm/hip_blas_lt.h index b71f97e6341bed..8a08f42b855408 100644 --- a/third_party/xla/xla/stream_executor/rocm/hip_blas_lt.h +++ b/third_party/xla/xla/stream_executor/rocm/hip_blas_lt.h @@ -193,8 +193,8 @@ class BlasLt : public gpu::BlasLt { int8_t bias_type_ = 0; }; // class GroupedMatmulPlan - // Executes complex (C64/C128) matmuls via rocBLAS (rocblas_cgemm/zgemm), - // since hipBLASLt has no complex GEMM kernels in current ROCm releases. + // Executes complex (C64/C128) matmuls via rocBLAS (rocblas_cgemm/zgemm) when + // the hipBLASLt loaded at runtime has no algorithm for them. class RocBlasGemmPlan : public gpu::BlasLt::MatmulPlan { public: friend class BlasLt; @@ -239,6 +239,10 @@ class BlasLt : public gpu::BlasLt { ~BlasLt() override = default; private: + // Fails if hipBLASLt does not handle the requested types. + absl::StatusOr GetHipBlasLtMatmulPlan( + const gpu::GemmConfig& cfg, Epilogue epilogue) const; + StreamExecutor* executor_; mutable absl::Mutex mu_; Owned handle_ ABSL_GUARDED_BY(mu_); diff --git a/third_party/xla/xla/tests/BUILD b/third_party/xla/xla/tests/BUILD index 05ada6ce5eb56f..d0e406542294a3 100644 --- a/third_party/xla/xla/tests/BUILD +++ b/third_party/xla/xla/tests/BUILD @@ -1522,10 +1522,7 @@ xla_test( # Turn on logging so that VLOG statements don't appear uncovered to zapfhahn. args = if_google(["--vmodule=convolution_emitter=7"]), shard_count = 40, - tags = [ - "cuda-only", - "optonly", - ], + tags = ["optonly"], deps = [ ":client_library_test_runner_mixin", ":hlo_pjrt_interpreter_reference_mixin", @@ -1649,7 +1646,6 @@ xla_test( backend_args = {"gpu": ["--xla_backend_extra_options=xla_gpu_experimental_conv_disable_layout_heuristic"]}, backends = ["gpu"], shard_count = 25, - tags = ["cuda-only"], deps = [ ":client_library_test_runner_mixin", ":hlo_pjrt_interpreter_reference_mixin", @@ -3013,10 +3009,7 @@ xla_test( "interpreter", ], shard_count = 30, - tags = [ - "cuda-only", - "optonly", - ], + tags = ["optonly"], use_legacy_runtime = True, deps = [ ":literal_test_util", diff --git a/third_party/xla/xla/tools/cost_model/BUILD b/third_party/xla/xla/tools/cost_model/BUILD index 3c8026b2be1cea..6517af28e5af9c 100644 --- a/third_party/xla/xla/tools/cost_model/BUILD +++ b/third_party/xla/xla/tools/cost_model/BUILD @@ -1,8 +1,6 @@ +load("@local_config_rocm//rocm:build_defs.bzl", "if_rocm_is_configured") load("@rules_cc//cc:cc_library.bzl", "cc_library") -load( - "//xla:xla.default.bzl", - "xla_cc_test", -) +load("//xla/tests:build_defs.bzl", "xla_test") load( "//xla/tsl/platform/default:cuda_build_defs.bzl", "if_cuda_is_configured", @@ -30,15 +28,17 @@ cc_library( "@com_google_absl//absl/types:span", ] + if_cuda_is_configured([ "//xla/stream_executor:cuda_platform", + ]) + if_rocm_is_configured([ + "//xla/stream_executor:rocm_platform", ]), ) -xla_cc_test( +xla_test( name = "gpu_bandwidth_benchmark_test", srcs = ["gpu_bandwidth_benchmark_test.cc"], - tags = [ - "cuda-only", - "requires-gpu-sm90", + backends = [ + "h100", + "amdgpu_any", ], deps = [ ":gpu_bandwidth_benchmark", diff --git a/third_party/xla/xla/tsl/framework/bfc_allocator.cc b/third_party/xla/xla/tsl/framework/bfc_allocator.cc index 51de93e13aa239..4ebdb5d0e42681 100644 --- a/third_party/xla/xla/tsl/framework/bfc_allocator.cc +++ b/third_party/xla/xla/tsl/framework/bfc_allocator.cc @@ -22,6 +22,7 @@ limitations under the License. #include #include #include +#include #include #include #include @@ -526,6 +527,11 @@ void* BFCAllocator::AllocateRawInternal(size_t alignment, size_t num_bytes, VLOG(2) << "tried to allocate 0 bytes"; return nullptr; } + if (ABSL_PREDICT_FALSE(num_bytes > std::numeric_limits::max() - + (kMinAllocationSize - 1))) { + VLOG(2) << "allocation size cannot be safely rounded: " << num_bytes; + return nullptr; + } // First, always allocate memory of at least kMinAllocationSize // bytes, and always allocate multiples of kMinAllocationSize bytes // so all memory addresses are nicely byte aligned. diff --git a/third_party/xla/xla/tsl/framework/bfc_allocator_test.cc b/third_party/xla/xla/tsl/framework/bfc_allocator_test.cc index 51284bb6e38055..9e202ab1909308 100644 --- a/third_party/xla/xla/tsl/framework/bfc_allocator_test.cc +++ b/third_party/xla/xla/tsl/framework/bfc_allocator_test.cc @@ -153,6 +153,17 @@ TEST(BFCAllocatorTest, DefaultAlignment) { alloc.DeallocateRaw(ptr); } +TEST(BFCAllocatorTest, RejectsAllocationWhoseRoundedSizeWouldOverflow) { + BFCAllocator::Options opts; + opts.allow_retry_on_failure = false; + BFCAllocator alloc(std::make_unique(), + /*total_memory=*/1 << 20, /*name=*/"test", opts); + + EXPECT_EQ( + alloc.AllocateRaw(kAlignment, std::numeric_limits::max() - 1), + nullptr); +} + TEST(BFCAllocatorTest, OomLogsAllocationAnnotations) { BFCAllocator::Options opts; opts.allow_growth = false; diff --git a/third_party/xla/xla/tsl/platform/default/dso_loader.cc b/third_party/xla/xla/tsl/platform/default/dso_loader.cc index e79365fd562c67..ed9ecb415e5322 100644 --- a/third_party/xla/xla/tsl/platform/default/dso_loader.cc +++ b/third_party/xla/xla/tsl/platform/default/dso_loader.cc @@ -73,6 +73,36 @@ absl::StatusOr GetDsoHandle(const std::string& name, VLOG(1) << message; return absl::Status(absl::StatusCode::kFailedPrecondition, message); } +absl::StatusOr GetCudaDsoHandle(const std::string& name, + absl::string_view version) { + auto handle = GetDsoHandle(name, version); + if (handle.ok()) return handle; + auto primary_status = handle.status(); + + // Determine alternative major version across CUDA 12 and CUDA 13 ecosystems. + // Note: cuDNN does not support cross-major loading (cuDNN 8 lacks v9 entry + // points). + absl::string_view alt_version = ""; + if (name == "cufft" || name == "cusolver") { + alt_version = (version == "12") ? "11" : (version == "11" ? "12" : ""); + } else if (name == "cublas" || name == "cublasLt" || name == "cudart" || + name == "nvrtc" || name == "cusparse" || name == "cupti") { + alt_version = (version == "13") ? "12" : (version == "12" ? "13" : ""); + } + + if (!alt_version.empty()) { + handle = GetDsoHandle(name, alt_version); + if (handle.ok()) return handle; + } + + // Only CUPTI supports unversioned library fallback. + if (name == "cupti") { + handle = GetDsoHandle(name, ""); + if (handle.ok()) return handle; + } + + return primary_status; +} } // namespace namespace DsoLoader { @@ -95,51 +125,47 @@ absl::StatusOr GetNvmlDsoHandle() { } absl::StatusOr GetNvrtcDsoHandle() { - return GetDsoHandle("nvrtc", GetCudaRtVersion()); + return GetCudaDsoHandle("nvrtc", GetCudaRtVersion()); } absl::StatusOr GetCudaRuntimeDsoHandle() { - return GetDsoHandle("cudart", GetCudaRtVersion()); + return GetCudaDsoHandle("cudart", GetCudaRtVersion()); } absl::StatusOr GetCublasDsoHandle() { - return GetDsoHandle("cublas", GetCublasVersion()); + return GetCudaDsoHandle("cublas", GetCublasVersion()); } absl::StatusOr GetCublasLtDsoHandle() { - return GetDsoHandle("cublasLt", GetCublasVersion()); + return GetCudaDsoHandle("cublasLt", GetCublasVersion()); } absl::StatusOr GetCufftDsoHandle() { - return GetDsoHandle("cufft", GetCufftVersion()); + return GetCudaDsoHandle("cufft", GetCufftVersion()); } absl::StatusOr GetCusolverDsoHandle() { - return GetDsoHandle("cusolver", GetCusolverVersion()); + return GetCudaDsoHandle("cusolver", GetCusolverVersion()); } absl::StatusOr GetCusparseDsoHandle() { - return GetDsoHandle("cusparse", GetCusparseVersion()); + return GetCudaDsoHandle("cusparse", GetCusparseVersion()); } absl::StatusOr GetCuptiDsoHandle() { - // Load specific version of CUPTI this is built. - auto status_or_handle = GetDsoHandle("cupti", GetCuptiVersion()); - if (status_or_handle.ok()) return status_or_handle; - // Load whatever libcupti.so user specified. - return GetDsoHandle("cupti", ""); + return GetCudaDsoHandle("cupti", GetCuptiVersion()); } absl::StatusOr GetCudnnDsoHandle() { - return GetDsoHandle("cudnn", GetCudnnVersion()); + return GetCudaDsoHandle("cudnn", GetCudnnVersion()); } absl::StatusOr GetNcclDsoHandle() { - return GetDsoHandle("nccl", GetNcclVersion()); + return GetCudaDsoHandle("nccl", GetNcclVersion()); } absl::StatusOr GetNvshmemDsoHandle() { - return GetDsoHandle("nvshmem_host", GetNvshmemVersion()); + return GetCudaDsoHandle("nvshmem_host", GetNvshmemVersion()); } absl::StatusOr GetNvInferDsoHandle() {