diff --git a/tensorflow/compiler/jit/BUILD b/tensorflow/compiler/jit/BUILD index a93dc10be415fd..98333a3c525ece 100644 --- a/tensorflow/compiler/jit/BUILD +++ b/tensorflow/compiler/jit/BUILD @@ -588,10 +588,9 @@ cc_library( tf_cuda_cc_test( name = "pjrt_compile_util_test", srcs = ["pjrt_compile_util_test.cc"], - tags = [ + tags = tf_cuda_tests_tags() + [ "config-cuda-only", "no_oss", # This test only runs with GPU. - "requires-gpu-nvidia", "xla", ], deps = [ @@ -1976,10 +1975,9 @@ tf_cc_test( tf_cuda_cc_test( name = "device_compiler_test", srcs = ["device_compiler_test.cc"], - tags = [ + tags = tf_cuda_tests_tags() + [ "config-cuda-only", "no_oss", # This test only runs with GPU. - "requires-gpu-nvidia", "xla", ], deps = [ diff --git a/tensorflow/compiler/tests/BUILD b/tensorflow/compiler/tests/BUILD index 56b3fc99746d7d..2ac8634e76f779 100644 --- a/tensorflow/compiler/tests/BUILD +++ b/tensorflow/compiler/tests/BUILD @@ -2935,7 +2935,6 @@ tf_cuda_cc_test( "no_pip", # TODO(b/149738646): fix pip install so these tests run on kokoro pip "no_rocm", # ROCmSoftwarePlatform #958 "noasan", # TODO(b/201651800) - "requires-gpu-nvidia", ] + tf_cuda_tests_tags(), deps = [":randomized_tests_library"], ) @@ -2956,7 +2955,6 @@ tf_cuda_cc_test( "no_pip", # TODO(b/149738646): fix pip install so these tests run on kokoro pip "no_rocm", # ROCmSoftwarePlatform #958 "noasan", # TODO(b/201651800) - "requires-gpu-nvidia", ] + tf_cuda_tests_tags(), deps = [":randomized_tests_library"], ) diff --git a/tensorflow/core/common_runtime/eager/BUILD b/tensorflow/core/common_runtime/eager/BUILD index 8935b2fc2a4b84..db40793f7e7373 100644 --- a/tensorflow/core/common_runtime/eager/BUILD +++ b/tensorflow/core/common_runtime/eager/BUILD @@ -326,11 +326,15 @@ tf_cc_test( "//tensorflow/core/kernels:cast_op", "//tensorflow/core/kernels:logging_ops", "//tensorflow/core/kernels:math", + "//tensorflow/core/platform:env", "//tensorflow/core/platform:refcount", "//tensorflow/core/platform:status", "@com_google_absl//absl/status", + "@com_google_absl//absl/synchronization", "@com_google_absl//absl/types:span", "@com_google_googletest//:gtest_main", + "@tsl//tsl/platform:refcount", + "@xla//xla/tsl/platform:env", "@xla//xla/tsl/platform:status", ], ) diff --git a/tensorflow/core/common_runtime/eager/context.cc b/tensorflow/core/common_runtime/eager/context.cc index ff5ae93053c3c2..14bdac9ee10358 100644 --- a/tensorflow/core/common_runtime/eager/context.cc +++ b/tensorflow/core/common_runtime/eager/context.cc @@ -18,6 +18,7 @@ limitations under the License. #include #include #include +#include #include #include #include @@ -27,6 +28,7 @@ limitations under the License. // clang-format off // Required for IS_MOBILE_PLATFORM #include "absl/status/status.h" +#include "absl/strings/match.h" #include "absl/strings/str_cat.h" #include "tensorflow/c/eager/immediate_execution_context.h" #include "tensorflow/core/common_runtime/function.h" @@ -618,6 +620,28 @@ void EagerContext::WaitForAndCloseRemoteContexts() { #endif // !IS_MOBILE_PLATFORM } +namespace internal { +bool IsTensorFlowExecutorThread() { + std::string name; + return tsl::Env::Default()->GetCurrentThreadName(&name) && + absl::StartsWith(name, "tf_"); +} +} // namespace internal + +void EagerContext::Release() { + // If we are on a TF executor thread, we must not execute the destructor + // synchronously to avoid a deadlock when it waits on the same thread pool. + // We unconditionally detach a thread rather than checking RefCountIsOne() + // to prevent TOCTOU races where another thread drops the last reference + // concurrently. + if (internal::IsTensorFlowExecutorThread()) { + env_->StartDetachedThread(ThreadOptions(), "eager_context_cleanup", + [this]() { this->Unref(); }); + } else { + Unref(); + } +} + EagerContext::~EagerContext() { // TODO(iga): Add a separate API method to shutdown EagerContext so that we // don't send RPCs and block in destructor. diff --git a/tensorflow/core/common_runtime/eager/context.h b/tensorflow/core/common_runtime/eager/context.h index 0225869ba4eac6..bd02317bd3d172 100644 --- a/tensorflow/core/common_runtime/eager/context.h +++ b/tensorflow/core/common_runtime/eager/context.h @@ -79,6 +79,11 @@ namespace eager { class RemoteMgr; } // namespace eager +namespace internal { +// Returns true if the calling thread's name begins with "tf_". +bool IsTensorFlowExecutorThread(); +} // namespace internal + // Check the value of the environment variable, // `TF_REMOTE_HANDLE_SKIP_WAIT_FOR_READY` from its cached copy in memory and if // not cached, reads from the environment variable. @@ -105,7 +110,7 @@ class EagerContext : public ImmediateExecutionContext, public core::RefCounted { CollectiveExecutorMgrInterface* collective_executor_mgr = nullptr, bool run_eager_op_as_function = false, bool jit_compile_rewrite = false); - void Release() override { Unref(); } + void Release() override; AbstractTensorInterface* CreateInt64Scalar(int64_t value) override; AbstractTensorInterface* CreateUint64Scalar(uint64_t value) override; diff --git a/tensorflow/core/common_runtime/eager/context_test.cc b/tensorflow/core/common_runtime/eager/context_test.cc index 590abf83871f67..9aef88f09d38c7 100644 --- a/tensorflow/core/common_runtime/eager/context_test.cc +++ b/tensorflow/core/common_runtime/eager/context_test.cc @@ -15,18 +15,22 @@ limitations under the License. #include "tensorflow/core/common_runtime/eager/context.h" +#include #include #include +#include #include #include #include "absl/status/status.h" +#include "absl/synchronization/notification.h" #include "absl/types/span.h" #include "tensorflow/c/eager/abstract_tensor_handle.h" #include "tensorflow/c/eager/immediate_execution_context.h" #include "tensorflow/c/eager/immediate_execution_operation.h" #include "tensorflow/c/eager/immediate_execution_tensor_handle.h" #include "xla/tsl/lib/core/status_test_util.h" +#include "xla/tsl/platform/env.h" #include "xla/tsl/platform/status.h" #include "tensorflow/core/common_runtime/composite_device.h" #include "tensorflow/core/common_runtime/device_mgr.h" @@ -36,11 +40,13 @@ limitations under the License. #include "tensorflow/core/framework/tensor.h" #include "tensorflow/core/framework/tensor_testutil.h" #include "tensorflow/core/framework/types.h" +#include "tensorflow/core/platform/env.h" #include "tensorflow/core/platform/refcount.h" #include "tensorflow/core/platform/status.h" #include "tensorflow/core/platform/test.h" #include "tensorflow/core/platform/types.h" #include "tensorflow/core/public/session_options.h" +#include "tsl/platform/refcount.h" namespace tensorflow { namespace { @@ -69,12 +75,15 @@ class EagerContextTest : public ::testing::Test { EagerContext* context() { return context_.get(); } void InitContext(const SessionOptions& opts, - ContextDevicePlacementPolicy policy, bool async = false) { + ContextDevicePlacementPolicy policy, bool async = false, + tsl::core::RefCountPtr rendezvous = nullptr) { ASSERT_EQ(context_, nullptr); - InitDeviceManager(); + if (!device_manager_) { + InitDeviceManager(); + } context_ = core::RefCountPtr(new EagerContext( opts, policy, async, device_manager_.get(), - /*device_mgr_owned=*/false, /*rendezvous=*/nullptr, + /*device_mgr_owned=*/false, std::move(rendezvous), /*cluster_flr=*/nullptr, /*collective_executor_mgr=*/nullptr, /*run_eager_op_as_function=*/true)); } @@ -414,5 +423,70 @@ TEST_F(EagerContextTest, ReuseGlobalRendezvous) { TestGlobalRendezvous(context(), true); } +TEST_F(EagerContextTest, IsTensorFlowExecutorThreadTest) { + { + std::unique_ptr t(tsl::Env::Default()->StartThread( + ThreadOptions(), "tf_test", + []() { EXPECT_TRUE(internal::IsTensorFlowExecutorThread()); })); + } + { + std::unique_ptr t(tsl::Env::Default()->StartThread( + ThreadOptions(), "other_name", + []() { EXPECT_FALSE(internal::IsTensorFlowExecutorThread()); })); + } +} + +class NotificationRendezvous : public Rendezvous { + public: + explicit NotificationRendezvous(absl::Notification* n, int64_t* thread_id) + : n_(n), destructor_thread_id_(thread_id) {} + ~NotificationRendezvous() override { + if (destructor_thread_id_) { + *destructor_thread_id_ = tsl::Env::Default()->GetCurrentThreadId(); + } + n_->Notify(); + } + + absl::Status Send(const ParsedKey& key, const Args& args, const Tensor& val, + const bool is_dead) override { + return absl::OkStatus(); + } + void RecvAsync(const ParsedKey& key, const Args& args, + DoneCallback done) override {} + void StartAbort(const absl::Status& status) override {} + + private: + absl::Notification* n_; + int64_t* destructor_thread_id_ = nullptr; +}; + +TEST_F(EagerContextTest, ReleaseOnTfThread) { + absl::Notification n; + int64_t rendezvous_destructor_thread_id = 0; + int64_t tf_thread_id = 0; + auto rendezvous = core::RefCountPtr( + new NotificationRendezvous(&n, &rendezvous_destructor_thread_id)); + InitDeviceManager(); + InitContext(SessionOptions(), DEVICE_PLACEMENT_EXPLICIT, false, + std::move(rendezvous)); + + { + std::unique_ptr t( + tsl::Env::Default()->StartThread(ThreadOptions(), "tf_test", [&]() { + tf_thread_id = tsl::Env::Default()->GetCurrentThreadId(); + // Drop the last reference to context_. Because the thread is named + // "tf_test", EagerContext::Release() will detach a thread to execute + // Unref(). + context_.release()->Release(); + })); + } + + // Wait deterministically for the detached thread to finish the deletion + // by waiting for the EagerContext's rendezvous member to be destroyed. + n.WaitForNotification(); + + EXPECT_NE(rendezvous_destructor_thread_id, tf_thread_id); +} + } // namespace } // namespace tensorflow diff --git a/tensorflow/dtensor/python/tests/test_util.py b/tensorflow/dtensor/python/tests/test_util.py index d6e875a5b5fe15..d8ebcb6b199e4e 100644 --- a/tensorflow/dtensor/python/tests/test_util.py +++ b/tensorflow/dtensor/python/tests/test_util.py @@ -54,7 +54,9 @@ DEFAULT_TOL = 1e-5 -_DEFAULT_GPU_MEMORY_LIMIT = 1024 # 1G +# Limit each virtual GPU to 512MB so 8 virtual devices plus CUDA/cuSolver +# runtime overhead fit within smaller GPU memory partitions. +_DEFAULT_GPU_MEMORY_LIMIT = 512 def get_use_xla_spmd(device_type): diff --git a/tensorflow/lite/c/CMakeLists.txt b/tensorflow/lite/c/CMakeLists.txt index 931f4372c5f2e2..ebe9ec3114b605 100644 --- a/tensorflow/lite/c/CMakeLists.txt +++ b/tensorflow/lite/c/CMakeLists.txt @@ -18,17 +18,10 @@ project(tensorflow-lite-c C CXX) option(TFLITE_C_BUILD_SHARED_LIBS "Build shared libraries" ON) -set(TF_SOURCE_DIR "" CACHE PATH - "Directory that contains the TensorFlow project" +get_filename_component(TFLITE_SOURCE_DIR + "${CMAKE_CURRENT_LIST_DIR}/.." + ABSOLUTE ) -if (NOT TF_SOURCE_DIR) - get_filename_component(TF_SOURCE_DIR - "${CMAKE_CURRENT_LIST_DIR}/../../../" - ABSOLUTE - ) -endif() - -set(TFLITE_SOURCE_DIR "${TF_SOURCE_DIR}/tensorflow/lite") add_subdirectory( "${TFLITE_SOURCE_DIR}" @@ -67,7 +60,7 @@ add_library(tensorflowlite_c ${TFLITE_C_LIBTYPE} ${TFLITE_SOURCE_DIR}/core/c/c_api_experimental.cc ${TFLITE_SOURCE_DIR}/core/c/common.cc ${TFLITE_SOURCE_DIR}/core/c/operator.cc - ${TF_SOURCE_DIR}/tensorflow/compiler/mlir/lite/core/c/tflite_types.h + ${TFLITE_SOURCE_DIR}/converter/core/c/tflite_types.h builtin_op_data.h c_api.h c_api_experimental.h diff --git a/tensorflow/lite/core/c/common.cc b/tensorflow/lite/core/c/common.cc index 5924be1c1fe3e0..3af6a0ffc7a0b0 100644 --- a/tensorflow/lite/core/c/common.cc +++ b/tensorflow/lite/core/c/common.cc @@ -265,6 +265,20 @@ TfLiteQuantization TfLiteQuantizationClone(const TfLiteQuantization& src) { dst_params->quantized_dimension = src_params->quantized_dimension; break; } + case kTfLiteBlockwiseQuantizationV2: { + dst.params = calloc(1, sizeof(TfLiteBlockwiseQuantizationV2)); + if (!dst.params) return TfLiteQuantization(); + const TfLiteBlockwiseQuantizationV2* const src_params = + reinterpret_cast(src.params); + TfLiteBlockwiseQuantizationV2* const dst_params = + reinterpret_cast(dst.params); + dst_params->blocksize = src_params->blocksize; + dst_params->scale = src_params->scale; + dst_params->zero_point = src_params->zero_point; + dst_params->quantized_dimension = src_params->quantized_dimension; + dst_params->block_shape = TfLiteIntArrayCopy(src_params->block_shape); + break; + } case kTfLiteMultiAxisQuantization: { dst.params = calloc(1, sizeof(TfLiteMultiAxisQuantization)); if (!dst.params) return TfLiteQuantization(); @@ -386,6 +400,17 @@ void TfLiteQuantizationFree(TfLiteQuantization* quantization) { reinterpret_cast(quantization->params); free(q_params); } + if (quantization->type == kTfLiteBlockwiseQuantizationV2) { + TfLiteBlockwiseQuantizationV2* q_params = + reinterpret_cast(quantization->params); + if (q_params) { + if (q_params->block_shape) { + TfLiteIntArrayFree(q_params->block_shape); + q_params->block_shape = nullptr; + } + free(q_params); + } + } if (quantization->type == kTfLiteMultiAxisQuantization) { TfLiteMultiAxisQuantization* q_params = reinterpret_cast(quantization->params); diff --git a/tensorflow/lite/core/c/common.h b/tensorflow/lite/core/c/common.h index 06271c1f848c08..66c8bf6fad6402 100644 --- a/tensorflow/lite/core/c/common.h +++ b/tensorflow/lite/core/c/common.h @@ -358,6 +358,9 @@ typedef enum TfLiteQuantizationType { kTfLiteBlockwiseQuantization = 2, /// Multi-axis quantization. kTfLiteMultiAxisQuantization = 3, + /// N-D blockwise quantization. + /// Corresponds to TfLiteBlockwiseQuantizationV2. + kTfLiteBlockwiseQuantizationV2 = 4, } TfLiteQuantizationType; /// Structure specifying the quantization used by the tensor, if-any. @@ -397,6 +400,24 @@ typedef struct TfLiteBlockwiseQuantization { int32_t quantized_dimension; } TfLiteBlockwiseQuantization; +/// Parameters for N-D blockwise quantization. +/// `block_shape->data[d]` is the extent of a single block along dimension `d`, +/// so the scale and zero point tensors have shape +/// `ceil(tensor_shape / block_shape)`. `block_shape` is owned by this struct. +typedef struct TfLiteBlockwiseQuantizationV2 { + // Index of the tensor containing the scales. + int32_t scale; + // Index of the tensor containing the zero points. + int32_t zero_point; + // Quantization blocksize. Only expresses blocking along the last dimension. + // Ignored when `block_shape` is non-null. + int32_t blocksize; + int32_t quantized_dimension; + // N-D block shape, with one entry per dimension of the quantized tensor. + // Owned by this struct. + TfLiteIntArray* block_shape; +} TfLiteBlockwiseQuantizationV2; + /// Parameters for multi-axis quantization. The scales and zero_points fields /// refer to tensor indices in the same subgraph. If zero_points is -1, zero /// point is assumed to be 0. For per-channel quantization, blocksize is 0. diff --git a/tensorflow/lite/core/c/common_test.cc b/tensorflow/lite/core/c/common_test.cc index e334812735ee13..7620ddb3da2634 100644 --- a/tensorflow/lite/core/c/common_test.cc +++ b/tensorflow/lite/core/c/common_test.cc @@ -180,6 +180,37 @@ TEST(Quantization, TestMultiAxisQuantizationCloneAndFree) { TfLiteQuantizationFree(&quantization); } +TEST(Quantization, TestBlockwiseQuantizationV2CloneAndFree) { + auto* params = reinterpret_cast( + malloc(sizeof(TfLiteBlockwiseQuantizationV2))); + params->scale = 1; + params->zero_point = 2; + params->blocksize = 0; + params->quantized_dimension = 0; + params->block_shape = BuildTfLiteArray({1, 1, 32}).release(); + TfLiteQuantization quantization = { + /*type=*/kTfLiteBlockwiseQuantizationV2, + /*params=*/params, + }; + + TfLiteQuantization clone = TfLiteQuantizationClone(quantization); + + EXPECT_THAT(clone.type, Eq(kTfLiteBlockwiseQuantizationV2)); + ASSERT_THAT(clone.params, Not(AnyOf(nullptr, quantization.params))); + auto* clone_params = + reinterpret_cast(clone.params); + EXPECT_THAT(clone_params->scale, Eq(params->scale)); + EXPECT_THAT(clone_params->zero_point, Eq(params->zero_point)); + EXPECT_THAT(clone_params->blocksize, Eq(params->blocksize)); + EXPECT_THAT(clone_params->quantized_dimension, + Eq(params->quantized_dimension)); + EXPECT_THAT(clone_params->block_shape, TfLiteArrayIs(params->block_shape)); + EXPECT_NE(clone_params->block_shape, params->block_shape); + + TfLiteQuantizationFree(&clone); + TfLiteQuantizationFree(&quantization); +} + TEST(Sparsity, TestSparsityFree) { TfLiteTensor t = {}; // Set these values, otherwise TfLiteTensorFree has uninitialized values. diff --git a/tensorflow/lite/core/interpreter_builder.cc b/tensorflow/lite/core/interpreter_builder.cc index c09d55d3d72bb8..bea6d38df1f19e 100644 --- a/tensorflow/lite/core/interpreter_builder.cc +++ b/tensorflow/lite/core/interpreter_builder.cc @@ -506,16 +506,67 @@ TfLiteStatus InterpreterBuilder::ParseQuantization( "Blockwise quantization details are missing."); return kTfLiteError; } - quantization->type = kTfLiteBlockwiseQuantization; - auto* blockwise_quantization = + // `block_shape` supersedes `block_size`. It is only meaningful if it has + // exactly one entry per dimension of the tensor, so reject anything else + // rather than silently falling back to `block_size`, which would produce + // wrong values. + const auto* block_shape = src_quant->block_shape(); + if (block_shape && !block_shape->empty()) { + if (!dims.empty() && + static_cast(block_shape->size()) != dims.size()) { + TF_LITE_REPORT_ERROR( + error_reporter_, + "BlockwiseQuantization block_shape must have one entry per " + "dimension. Expected %d entries, got %d.", + static_cast(dims.size()), block_shape->size()); + return kTfLiteError; + } + for (int i = 0; i < block_shape->size(); ++i) { + if (block_shape->Get(i) <= 0) { + TF_LITE_REPORT_ERROR(error_reporter_, + "BlockwiseQuantization block_shape entries must " + "be positive. Value %d at index %d is invalid.", + block_shape->Get(i), i); + return kTfLiteError; + } + } + auto* blockwise_quantization_v2 = + reinterpret_cast( + malloc(sizeof(TfLiteBlockwiseQuantizationV2))); + if (!blockwise_quantization_v2) { + return kTfLiteError; + } + blockwise_quantization_v2->scale = src_quant->scales(); + blockwise_quantization_v2->zero_point = src_quant->zero_points(); + blockwise_quantization_v2->quantized_dimension = + src_quantization->quantized_dimension(); + blockwise_quantization_v2->blocksize = src_quant->block_size(); + blockwise_quantization_v2->block_shape = + TfLiteIntArrayCreate(static_cast(block_shape->size())); + if (!blockwise_quantization_v2->block_shape) { + free(blockwise_quantization_v2); + return kTfLiteError; + } + for (int i = 0; i < block_shape->size(); ++i) { + blockwise_quantization_v2->block_shape->data[i] = block_shape->Get(i); + } + quantization->type = kTfLiteBlockwiseQuantizationV2; + quantization->params = reinterpret_cast(blockwise_quantization_v2); + return kTfLiteOk; + } + auto* blockwise_quantization_v1 = reinterpret_cast( malloc(sizeof(TfLiteBlockwiseQuantization))); - blockwise_quantization->scale = src_quant->scales(); - blockwise_quantization->zero_point = src_quant->zero_points(); - blockwise_quantization->quantized_dimension = + if (!blockwise_quantization_v1) { + return kTfLiteError; + } + blockwise_quantization_v1->scale = src_quant->scales(); + blockwise_quantization_v1->zero_point = src_quant->zero_points(); + blockwise_quantization_v1->quantized_dimension = src_quantization->quantized_dimension(); - blockwise_quantization->blocksize = src_quant->block_size(); - quantization->params = reinterpret_cast(blockwise_quantization); + blockwise_quantization_v1->blocksize = src_quant->block_size(); + quantization->type = kTfLiteBlockwiseQuantization; + quantization->params = reinterpret_cast(blockwise_quantization_v1); return kTfLiteOk; } if (!src_quantization->scale() || src_quantization->scale()->empty()) { diff --git a/tensorflow/lite/core/model_test.cc b/tensorflow/lite/core/model_test.cc index 124a66904956fb..81fadef874c314 100644 --- a/tensorflow/lite/core/model_test.cc +++ b/tensorflow/lite/core/model_test.cc @@ -124,6 +124,45 @@ std::string BuildMultiAxisQuantizationModelBuffer( builder.GetSize()); } +std::string BuildBlockwiseQuantizationModelBuffer( + const std::vector& block_shape, int32_t block_size = 32) { + flatbuffers::FlatBufferBuilder builder; + + std::vector scale_shape = {2, 4, 1}; + auto scale_tensor = CreateTensorDirect(builder, &scale_shape, + TensorType_FLOAT32, 0, "scales"); + std::vector zero_point_shape = {2, 4, 1}; + auto zero_point_tensor = CreateTensorDirect( + builder, &zero_point_shape, TensorType_INT32, 0, "zero_points"); + auto blockwise_quantization = CreateBlockwiseQuantizationDirect( + builder, /*scales=*/0, /*zero_points=*/1, block_size, &block_shape); + auto quantization = CreateQuantizationParameters( + builder, /*min=*/0, /*max=*/0, /*scale=*/0, /*zero_point=*/0, + QuantizationDetails_BlockwiseQuantization, blockwise_quantization.Union(), + /*quantized_dimension=*/0); + std::vector weight_shape = {2, 4, 3}; + auto weight_tensor = CreateTensorDirect( + builder, &weight_shape, TensorType_INT8, 0, "moe_weight", quantization); + + std::vector> tensors = { + scale_tensor, zero_point_tensor, weight_tensor}; + std::vector inputs = {2}; + std::vector outputs = {2}; + std::vector> operators; + auto subgraph = + CreateSubGraphDirect(builder, &tensors, &inputs, &outputs, &operators); + std::vector> subgraphs = {subgraph}; + std::vector> buffers = { + CreateBufferDirect(builder)}; + std::vector> operator_codes; + auto model = + CreateModelDirect(builder, TFLITE_SCHEMA_VERSION, &operator_codes, + &subgraphs, "blockwise quantization test", &buffers); + FinishModelBuffer(builder, model); + return std::string(reinterpret_cast(builder.GetBufferPointer()), + builder.GetSize()); +} + } // namespace TEST(BasicFlatBufferModel, TestNonExistentFiles) { @@ -223,6 +262,90 @@ TEST(BasicFlatBufferModel, EXPECT_EQ(interpreter, nullptr); } +TEST(BasicFlatBufferModel, TestBlockwiseQuantizationBlockShapeInInterpreter) { + // Weight is {2, 4, 3}, so a block_shape of {1, 1, 3} means one scale per + // (dim0, dim1) pair, i.e. the scale tensor shape {2, 4, 1}. + std::string model_buffer = + BuildBlockwiseQuantizationModelBuffer(/*block_shape=*/{1, 1, 3}); + auto model = FlatBufferModel::BuildFromBuffer(model_buffer.data(), + model_buffer.size()); + ASSERT_TRUE(model); + + std::unique_ptr interpreter; + ASSERT_EQ(InterpreterBuilder(*model, TrivialResolver())(&interpreter), + kTfLiteOk); + ASSERT_NE(interpreter, nullptr); + + TfLiteTensor* tensor = interpreter->tensor(2); + ASSERT_NE(tensor, nullptr); + EXPECT_EQ(tensor->quantization.type, kTfLiteBlockwiseQuantizationV2); + ASSERT_NE(tensor->quantization.params, nullptr); + auto* quantization = reinterpret_cast( + tensor->quantization.params); + EXPECT_EQ(quantization->scale, 0); + EXPECT_EQ(quantization->zero_point, 1); + ASSERT_NE(quantization->block_shape, nullptr); + TfLiteIntArray* expected_block_shape = TfLiteIntArrayCreate(3); + expected_block_shape->data[0] = 1; + expected_block_shape->data[1] = 1; + expected_block_shape->data[2] = 3; + EXPECT_TRUE( + TfLiteIntArrayEqual(quantization->block_shape, expected_block_shape)); + TfLiteIntArrayFree(expected_block_shape); +} + +// Models written before block_shape existed must keep working, and must be +// distinguishable from models that set it. +TEST(BasicFlatBufferModel, TestBlockwiseQuantizationWithoutBlockShape) { + std::string model_buffer = + BuildBlockwiseQuantizationModelBuffer(/*block_shape=*/{}); + auto model = FlatBufferModel::BuildFromBuffer(model_buffer.data(), + model_buffer.size()); + ASSERT_TRUE(model); + + std::unique_ptr interpreter; + ASSERT_EQ(InterpreterBuilder(*model, TrivialResolver())(&interpreter), + kTfLiteOk); + ASSERT_NE(interpreter, nullptr); + + TfLiteTensor* tensor = interpreter->tensor(2); + ASSERT_NE(tensor, nullptr); + EXPECT_EQ(tensor->quantization.type, kTfLiteBlockwiseQuantization); + auto* quantization = reinterpret_cast( + tensor->quantization.params); + ASSERT_NE(quantization, nullptr); + EXPECT_EQ(quantization->blocksize, 32); +} + +TEST(BasicFlatBufferModel, TestRejectsBlockwiseQuantizationBadRankBlockShape) { + // The weight is rank 3, so a rank 2 block_shape is ambiguous and must be + // rejected rather than silently falling back to block_size. + std::string model_buffer = + BuildBlockwiseQuantizationModelBuffer(/*block_shape=*/{1, 3}); + auto model = FlatBufferModel::BuildFromBuffer(model_buffer.data(), + model_buffer.size()); + ASSERT_TRUE(model); + + std::unique_ptr interpreter; + EXPECT_NE(InterpreterBuilder(*model, TrivialResolver())(&interpreter), + kTfLiteOk); + EXPECT_EQ(interpreter, nullptr); +} + +TEST(BasicFlatBufferModel, + TestRejectsBlockwiseQuantizationNonPositiveBlockShape) { + std::string model_buffer = + BuildBlockwiseQuantizationModelBuffer(/*block_shape=*/{1, 0, 3}); + auto model = FlatBufferModel::BuildFromBuffer(model_buffer.data(), + model_buffer.size()); + ASSERT_TRUE(model); + + std::unique_ptr interpreter; + EXPECT_NE(InterpreterBuilder(*model, TrivialResolver())(&interpreter), + kTfLiteOk); + EXPECT_EQ(interpreter, nullptr); +} + TEST(BasicFlatBufferModel, TestNullDestination) { auto model = FlatBufferModel::BuildFromFile( "tensorflow/lite/testdata/empty_model.bin"); diff --git a/tensorflow/lite/tools/cmake/modules/Findtensorflow_headers.cmake b/tensorflow/lite/tools/cmake/modules/Findtensorflow_headers.cmake index 7fef6f8c81c55e..1320e23af68123 100644 --- a/tensorflow/lite/tools/cmake/modules/Findtensorflow_headers.cmake +++ b/tensorflow/lite/tools/cmake/modules/Findtensorflow_headers.cmake @@ -25,7 +25,7 @@ if(TENSORFLOW_HEADERS_INCLUDE_DIR) endif() set(_tensorflow_headers_base_url - "https://raw.githubusercontent.com/tensorflow/tensorflow/v2.21.0-rc0" + "https://raw.githubusercontent.com/tensorflow/tensorflow/v2.21.0" ) set(_tensorflow_headers_dir "${CMAKE_BINARY_DIR}/tensorflow_headers") # || diff --git a/tensorflow/python/eager/backprop_test.py b/tensorflow/python/eager/backprop_test.py index 2f3a92bdebafb3..42a69cda8dda82 100644 --- a/tensorflow/python/eager/backprop_test.py +++ b/tensorflow/python/eager/backprop_test.py @@ -1864,6 +1864,9 @@ def _func(x): array_ops.reshape(def_function.function(f)(x), [-1]), rtol=1e-3) + @test_util.run_without_tensor_float_32( + 'Avoid TF32 conv2d in finite-difference Jacobian test' + ) def test_grad_jacobian_conv(self): def _inner(x): kernel = array_ops.ones([3, 3, 1, 9]) diff --git a/tensorflow/python/kernel_tests/array_ops/scatter_nd_ops_test.py b/tensorflow/python/kernel_tests/array_ops/scatter_nd_ops_test.py index 4ba544a40c8f31..7fcf51ab63af1a 100644 --- a/tensorflow/python/kernel_tests/array_ops/scatter_nd_ops_test.py +++ b/tensorflow/python/kernel_tests/array_ops/scatter_nd_ops_test.py @@ -153,7 +153,11 @@ def _VariableRankTest(self, self.evaluate(tf_scatter(ref_var, indices, updates)) # Compare - tol = 1e-6 if vtype != dtypes.bfloat16.as_numpy_dtype else 1e-2 + tol = ( + 1e-2 + if vtype in (np.float16, dtypes.bfloat16.as_numpy_dtype) + else 1e-6 + ) self.assertAllClose(new, self.evaluate(ref_var), rtol=tol, atol=tol) def _VariableRankTests(self, np_scatter, tf_scatter): diff --git a/third_party/xla/third_party/mkl_dnn/mkldnn_v1_async_aarch64.BUILD b/third_party/xla/third_party/mkl_dnn/mkldnn_v1_async_aarch64.BUILD new file mode 100644 index 00000000000000..3f46656b6bf194 --- /dev/null +++ b/third_party/xla/third_party/mkl_dnn/mkldnn_v1_async_aarch64.BUILD @@ -0,0 +1,196 @@ +# Copyright 2026 The OpenXLA Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== + +load("@bazel_skylib//rules:expand_template.bzl", "expand_template") + +exports_files(["LICENSE"]) + +_DNNL_COPTS_THREADPOOL = [ + "-fopenmp-simd", + "-fexceptions", + "-DDNNL_ENABLE_MAX_CPU_ISA", +] + +_CMAKE_COMMON_LIST = { + "#cmakedefine DNNL_USE_RT_OBJECTS_IN_PRIMITIVE_CACHE": "#undef DNNL_USE_RT_OBJECTS_IN_PRIMITIVE_CACHE", + "#cmakedefine DNNL_WITH_LEVEL_ZERO": "#undef DNNL_WITH_LEVEL_ZERO", + "#cmakedefine DNNL_GPU_RUNTIME DNNL_RUNTIME_${DNNL_GPU_RUNTIME}": "#define DNNL_GPU_RUNTIME DNNL_RUNTIME_NONE", + "#cmakedefine DNNL_GPU_VENDOR DNNL_VENDOR_${DNNL_GPU_VENDOR}": "#define DNNL_GPU_VENDOR DNNL_VENDOR_NONE", + "#cmakedefine DNNL_WITH_SYCL": "#undef DNNL_WITH_SYCL", + "#cmakedefine DNNL_ENABLE_STACK_CHECKER": "#undef DNNL_ENABLE_STACK_CHECKER", + "#cmakedefine ONEDNN_BUILD_GRAPH": "#undef ONEDNN_BUILD_GRAPH", + "#cmakedefine DNNL_EXPERIMENTAL_SPARSE": "#undef DNNL_EXPERIMENTAL_SPARSE", + "#cmakedefine DNNL_EXPERIMENTAL": "#undef DNNL_EXPERIMENTAL", + "#cmakedefine01 DNNL_EXPERIMENTAL_GROUPED_MEMORY": "#define DNNL_EXPERIMENTAL_GROUPED_MEMORY 0", + "#cmakedefine DNNL_SAFE_RBP": "#undef DNNL_SAFE_RBP", + "#cmakedefine01 BUILD_TRAINING": "#define BUILD_TRAINING 1", + "#cmakedefine01 BUILD_INFERENCE": "#define BUILD_INFERENCE 0", + "#cmakedefine01 BUILD_PRIMITIVE_ALL": "#define BUILD_PRIMITIVE_ALL 1", + "#cmakedefine01 BUILD_BATCH_NORMALIZATION": "#define BUILD_BATCH_NORMALIZATION 0", + "#cmakedefine01 BUILD_BINARY": "#define BUILD_BINARY 0", + "#cmakedefine01 BUILD_CONCAT": "#define BUILD_CONCAT 0", + "#cmakedefine01 BUILD_CONVOLUTION": "#define BUILD_CONVOLUTION 0", + "#cmakedefine01 BUILD_DECONVOLUTION": "#define BUILD_DECONVOLUTION 0", + "#cmakedefine01 BUILD_ELTWISE": "#define BUILD_ELTWISE 0", + "#cmakedefine01 BUILD_GATED_MLP": "#define BUILD_GATED_MLP 0", + "#cmakedefine01 BUILD_GEMM_KERNELS_ALL": "#define BUILD_GEMM_KERNELS_ALL 1", + "#cmakedefine01 BUILD_GEMM_KERNELS_NONE": "#define BUILD_GEMM_KERNELS_NONE 0", + "#cmakedefine01 BUILD_GEMM_SSE41": "#define BUILD_GEMM_SSE41 0", + "#cmakedefine01 BUILD_GEMM_AVX2": "#define BUILD_GEMM_AVX2 0", + "#cmakedefine01 BUILD_GEMM_AVX512": "#define BUILD_GEMM_AVX512 0", + "#cmakedefine01 BUILD_GROUP_NORMALIZATION": "#define BUILD_GROUP_NORMALIZATION 0", + "#cmakedefine01 BUILD_INNER_PRODUCT": "#define BUILD_INNER_PRODUCT 0", + "#cmakedefine01 BUILD_LAYER_NORMALIZATION": "#define BUILD_LAYER_NORMALIZATION 0", + "#cmakedefine01 BUILD_LRN": "#define BUILD_LRN 0", + "#cmakedefine01 BUILD_MATMUL": "#define BUILD_MATMUL 0", + "#cmakedefine01 BUILD_POOLING": "#define BUILD_POOLING 0", + "#cmakedefine01 BUILD_PRELU": "#define BUILD_PRELU 0", + "#cmakedefine01 BUILD_REDUCTION": "#define BUILD_REDUCTION 0", + "#cmakedefine01 BUILD_REORDER": "#define BUILD_REORDER 0", + "#cmakedefine01 BUILD_RESAMPLING": "#define BUILD_RESAMPLING 0", + "#cmakedefine01 BUILD_RNN": "#define BUILD_RNN 0", + "#cmakedefine01 BUILD_SHUFFLE": "#define BUILD_SHUFFLE 0", + "#cmakedefine01 BUILD_SOFTMAX": "#define BUILD_SOFTMAX 0", + "#cmakedefine01 BUILD_SUM": "#define BUILD_SUM 0", + "#cmakedefine01 BUILD_PRIMITIVE_CPU_ISA_ALL": "#define BUILD_PRIMITIVE_CPU_ISA_ALL 0", + "#cmakedefine01 BUILD_SSE41": "#define BUILD_SSE41 0", + "#cmakedefine01 BUILD_AVX2": "#define BUILD_AVX2 0", + "#cmakedefine01 BUILD_AVX512": "#define BUILD_AVX512 0", + "#cmakedefine01 BUILD_AMX": "#define BUILD_AMX 0", + "#cmakedefine01 BUILD_GEN9": "#define BUILD_GEN9 0", + "#cmakedefine01 BUILD_GEN11": "#define BUILD_GEN11 0", + "#cmakedefine01 BUILD_SDPA": "#define BUILD_SDPA 0", + "#cmakedefine01 BUILD_XE2": "#define BUILD_XE2 0", + "#cmakedefine01 BUILD_XE3P": "#define BUILD_XE3P 0", + "#cmakedefine01 BUILD_XE3": "#define BUILD_XE3 0", + "#cmakedefine01 BUILD_XELP": "#define BUILD_XELP 0", + "#cmakedefine01 BUILD_XEHPG": "#define BUILD_XEHPG 0", + "#cmakedefine01 BUILD_XEHPC": "#define BUILD_XEHPC 0", + "#cmakedefine01 BUILD_XEHP": "#define BUILD_XEHP 0", + # SYCL specific settings + "#cmakedefine DNNL_SYCL_CUDA": "#undef DNNL_SYCL_CUDA", + "#cmakedefine DNNL_SYCL_HIP": "#undef DNNL_SYCL_HIP", + "#cmakedefine DNNL_SYCL_GENERIC": "#undef DNNL_SYCL_GENERIC", + "#cmakedefine DNNL_DISABLE_GPU_REF_KERNELS": "#undef DNNL_DISABLE_GPU_REF_KERNELS", + "#cmakedefine DNNL_EXPERIMENTAL_SYCL_KERNEL_COMPILER": "#undef DNNL_EXPERIMENTAL_SYCL_KERNEL_COMPILER", + "#cmakedefine01 BUILD_PRIMITIVE_GPU_ISA_ALL": "#define BUILD_PRIMITIVE_GPU_ISA_ALL 0", +} + +_DNNL_RUNTIME_THREADPOOL = { + "#cmakedefine DNNL_CPU_THREADING_RUNTIME DNNL_RUNTIME_${DNNL_CPU_THREADING_RUNTIME}": "#define DNNL_CPU_THREADING_RUNTIME DNNL_RUNTIME_THREADPOOL", + "#cmakedefine DNNL_CPU_RUNTIME DNNL_RUNTIME_${DNNL_CPU_RUNTIME}": "#define DNNL_CPU_RUNTIME DNNL_RUNTIME_THREADPOOL", +} + +_DNNL_RUNTIME_THREADPOOL.update(_CMAKE_COMMON_LIST) + +expand_template( + name = "dnnl_config_h", + out = "include/oneapi/dnnl/dnnl_config.h", + substitutions = select({ + "//conditions:default": _DNNL_RUNTIME_THREADPOOL, + }), + template = "include/oneapi/dnnl/dnnl_config.h.in", +) + +expand_template( + name = "dnnl_version_h", + out = "include/oneapi/dnnl/dnnl_version.h", + substitutions = { + "@DNNL_VERSION_MAJOR@": "3", + "@DNNL_VERSION_MINOR@": "14", + "@DNNL_VERSION_PATCH@": "0", + "@DNNL_VERSION_HASH@": "N/A", + }, + template = "include/oneapi/dnnl/dnnl_version.h.in", +) + +expand_template( + name = "dnnl_version_hash_h", + out = "include/oneapi/dnnl/dnnl_version_hash.h", + substitutions = { + "@DNNL_VERSION_HASH@": "e0d8b940aba7ece3d4374f80e508bbcdb8f141db", + }, + template = "include/oneapi/dnnl/dnnl_version_hash.h.in", +) + +_INCLUDES_LIST = [ + "include", + "src", + "src/common", + "src/cpu", + "src/cpu/gemm", + "third_party", + "third_party/ittnotify", + "third_party/spdlog", + "third_party/xbyak_aarch64", + "third_party/xbyak_aarch64/src", + "third_party/xbyak_aarch64/xbyak_aarch64", +] + +_TEXTUAL_HDRS_LIST = glob([ + "include/**/*", + "include/*", + "src/common/*.hpp", + "src/common/**/*.h", + "src/cpu/*.hpp", + "src/cpu/**/*.hpp", + "src/cpu/jit_utils/**/*.hpp", + "third_party/**/*.h", + "third_party/ittnotify/**/*.h", + "third_party/xbyak_aarch64/**/*.h", +]) + [ + ":dnnl_config_h", + ":dnnl_version_h", + ":dnnl_version_hash_h", +] + +alias( + name = "mkl_dnn", + actual = ":onednn_cpu", + visibility = ["//visibility:public"], +) + +cc_library( + name = "onednn_cpu", + srcs = glob( + [ + "src/common/*.cpp", + "src/cpu/*.cpp", + "src/cpu/**/*.cpp", + "src/cpu/jit_utils/**/*.cpp", + "third_party/ittnotify/*.c", + "third_party/ittnotify/ittnotify/*.c", + "third_party/xbyak_aarch64/src/*.cpp", + ], + exclude = [ + "src/cpu/aarch64/acl_*.cpp", + "src/cpu/aarch64/matmul/acl_*.cpp", + "src/cpu/aarch64/reorder/acl_*.cpp", + "src/cpu/ppc64/**", + "src/cpu/rv64/**", + "src/cpu/s390x/**", + "src/cpu/x64/**", + "src/cpu/sycl/**", + "src/graph/**", + "src/xpu/**", + ], + ), + copts = select({ + "//conditions:default": _DNNL_COPTS_THREADPOOL, + }), + includes = _INCLUDES_LIST, + textual_hdrs = _TEXTUAL_HDRS_LIST, + visibility = ["//visibility:public"], +) diff --git a/third_party/xla/third_party/mkl_dnn/onednn_async_aarch64_pr_5790.patch b/third_party/xla/third_party/mkl_dnn/onednn_async_aarch64_pr_5790.patch new file mode 100644 index 00000000000000..8b9f1992be01b1 --- /dev/null +++ b/third_party/xla/third_party/mkl_dnn/onednn_async_aarch64_pr_5790.patch @@ -0,0 +1,339 @@ +From ca13a085c0594ebab0fa75fcb2fa47582e27c925 Mon Sep 17 00:00:00 2001 +From: Siddhartha Menon +Date: Wed, 16 Sep 2026 14:59:13 +0100 +Subject: [PATCH 1/2] cpu: aarch64: brgemm: utils: add async-safe scale helpers + +The DEFINE_ARG_SCALES_BUFFER macros rely on stack allocated variables +that make them unsuitable for asynchronous runtimes. + +The long-term fix would be to move towards in-kernel scale processing as +done in f744f73c. However, as that is a more complex fix, this change +serves as an intermediate fix to get async-compatibility for the +high-priority implementations while we work on the more permanent +solution. + +Signed-off-by: Siddhartha Menon +--- + src/cpu/aarch64/brgemm/brgemm_utils.hpp | 49 ++++++++++++++++++++++++- + 1 file changed, 47 insertions(+), 2 deletions(-) + +diff --git a/src/cpu/aarch64/brgemm/brgemm_utils.hpp b/src/cpu/aarch64/brgemm/brgemm_utils.hpp +index b7ab1d28ed9..89760e871bc 100644 +--- a/src/cpu/aarch64/brgemm/brgemm_utils.hpp ++++ b/src/cpu/aarch64/brgemm/brgemm_utils.hpp +@@ -1,7 +1,7 @@ + /******************************************************************************* + * Copyright 2022 Intel Corporation + * Copyright 2024 FUJITSU LIMITED +-* Copyright 2024-2025 Arm Ltd. and affiliates ++* Copyright 2024-2026 Arm Ltd. and affiliates + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. +@@ -20,8 +20,8 @@ + #define CPU_AARCH64_BRGEMM_BRGEMM_UTILS_HPP + + #include "cpu/aarch64/brgemm/brgemm.hpp" +- + #include "cpu/aarch64/cpu_isa_traits.hpp" ++#include "cpu/ref_io_helper.hpp" + + #include "common/c_types_map.hpp" + +@@ -30,6 +30,51 @@ namespace impl { + namespace cpu { + namespace aarch64 { + ++namespace scale_utils { ++inline size_t scales_simd_w() { ++ return simd_elems(data_type::f32, get_max_cpu_isa()); ++} ++ ++inline float load_scale(const primitive_attr_t *attr, int arg, ++ const void *scales, dim_t idx = 0) { ++ if (attr->scales_.has_default_values(arg)) return 1.f; ++ ++ const auto scales_dt = attr->scales_.get_data_type(arg); ++ return cpu::io::load_float_value(scales_dt, scales, idx); ++} ++ ++inline void precompute_oscales(float *oscales, const primitive_attr_t *attr, ++ const void *src_scales, const void *wei_scales, dim_t oc, ++ float scale_adjust_factor) { ++ if (oscales == nullptr) return; ++ ++ const float src_scale = load_scale(attr, DNNL_ARG_SRC, src_scales); ++ const auto &wei_scales_attr = attr->scales_.get(DNNL_ARG_WEIGHTS); ++ const bool is_oc_scale = wei_scales_attr.get_mask() > 0; ++ ++ if (is_oc_scale) { ++ for (dim_t c = 0; c < oc; c++) { ++ oscales[c] = src_scale ++ * load_scale(attr, DNNL_ARG_WEIGHTS, wei_scales, c) ++ * scale_adjust_factor; ++ } ++ } else { ++ const float scale = src_scale ++ * load_scale(attr, DNNL_ARG_WEIGHTS, wei_scales) ++ * scale_adjust_factor; ++ utils::array_set(oscales, scale, scales_simd_w()); ++ } ++} ++ ++inline void precompute_inv_dst_scales(float *dst_scales_precomputed, ++ const primitive_attr_t *attr, const void *dst_scales) { ++ if (dst_scales_precomputed == nullptr) return; ++ ++ const float dst_scale = 1.f / load_scale(attr, DNNL_ARG_DST, dst_scales); ++ utils::array_set(dst_scales_precomputed, dst_scale, scales_simd_w()); ++} ++} // namespace scale_utils ++ + status_t init_kernel_datatype( + brgemm_desc_t *brg, data_type_t dt_a, data_type_t dt_b); + + +From 1a7e69fa15b7c504c83ce3e9bc7a20d1c0c0d214 Mon Sep 17 00:00:00 2001 +From: Siddhartha Menon +Date: Tue, 18 Aug 2026 11:16:34 +0100 +Subject: [PATCH 2/2] cpu: aarch64: brgemm_matmul: make async compatible + +Signed-off-by: Siddhartha Menon +--- + src/cpu/aarch64/matmul/brgemm_matmul.cpp | 98 ++++++++++++++----- + src/cpu/aarch64/matmul/brgemm_matmul.hpp | 10 +- + .../aarch64/matmul/brgemm_matmul_reorders.cpp | 7 +- + 3 files changed, 82 insertions(+), 33 deletions(-) + +diff --git a/src/cpu/aarch64/matmul/brgemm_matmul.cpp b/src/cpu/aarch64/matmul/brgemm_matmul.cpp +index bcdf31e2e33..ac72eecb5ef 100644 +--- a/src/cpu/aarch64/matmul/brgemm_matmul.cpp ++++ b/src/cpu/aarch64/matmul/brgemm_matmul.cpp +@@ -16,16 +16,20 @@ + * limitations under the License. + *******************************************************************************/ + ++#include ++ + #include "common/c_types_map.hpp" ++#include "common/compiler_workarounds.hpp" + #include "common/dnnl_thread.hpp" + #include "common/memory_tracking.hpp" ++#include "common/nstl.hpp" ++#include "common/primitive.hpp" + #include "common/type_helpers.hpp" + #include "common/utils.hpp" + +-#include "cpu/cpu_primitive.hpp" + #include "cpu/matmul/matmul_utils.hpp" +-#include "cpu/scale_utils.hpp" + ++#include "cpu/aarch64/brgemm/brgemm_utils.hpp" + #include "cpu/aarch64/injectors/jit_uni_binary_injector.hpp" + #include "cpu/aarch64/matmul/brgemm_matmul.hpp" + +@@ -228,7 +232,22 @@ status_t brgemm_matmul_t::pd_t::init(const engine_t *engine) { + + auto scratchpad = scratchpad_registry().registrar(); + init_scratchpad(scratchpad, bgmmc_); +- book_precomputed_scales(scratchpad, attr()->scales_, N()); ++ ++ constexpr size_t page_alignment_bytes = 4096; ++ if (bgmmc_.with_scales) { ++ const size_t scales_size = bgmmc_.is_oscale_per_n ++ ? nstl::max(static_cast(N()), ++ scale_utils::scales_simd_w()) ++ : scale_utils::scales_simd_w(); ++ ++ scratchpad.template book( ++ key_precomputed_scales, scales_size, page_alignment_bytes); ++ } ++ ++ if (bgmmc_.with_dst_scales) { ++ scratchpad.template book(key_matmul_dst_scales, ++ scale_utils::scales_simd_w(), page_alignment_bytes); ++ } + + const bool is_B_transposed = one_of(bgmmc_.wei_tag, abdc, ba, acb, adbc, + abced, abcdfe, abcdegf, abcdefhg, abcdefgih, abcdefghji, +@@ -284,32 +303,59 @@ status_t brgemm_matmul_t::init(engine_t *engine) { + + template + status_t brgemm_matmul_t::execute_body(const exec_ctx_t &ctx) const { +- DEFINE_ARG_SCALES_BUFFER(src_scales, DNNL_ARG_SRC); +- DEFINE_ARG_SCALES_BUFFER(wei_scales, DNNL_ARG_WEIGHTS); +- DEFINE_ARG_SCALES_BUFFER(dst_scales, DNNL_ARG_DST); +- + const auto src_d = ctx.memory_mdw(DNNL_ARG_SRC, pd()->src_md()); + const auto weights_d = ctx.memory_mdw(DNNL_ARG_WEIGHTS, pd()->weights_md()); + const auto dst_d = ctx.memory_mdw(DNNL_ARG_DST, pd()->dst_md()); + matmul_helper_t helper(src_d, weights_d, dst_d); + +- const auto &scratchpad = ctx.get_scratchpad_grantor(); +- const float *oscales = precompute_scales( +- scratchpad, src_scales, wei_scales, pd()->N(), pd()->attr()); ++ const void *src_scales ++ = CTX_IN_MEM(const void *, DNNL_ARG_ATTR_SCALES | DNNL_ARG_SRC); ++ const void *wei_scales ++ = CTX_IN_MEM(const void *, DNNL_ARG_ATTR_SCALES | DNNL_ARG_WEIGHTS); ++ const void *dst_scales ++ = CTX_IN_MEM(const void *, DNNL_ARG_ATTR_SCALES | DNNL_ARG_DST); + +- brg_matmul_exec_ctx_t brgmm_ctx(ctx, pd(), oscales, dst_scales, helper); ++ const auto &scratchpad = ctx.get_scratchpad_grantor(); + + const auto &bgmmc = pd()->get_brgemm_matmul_conf(); +- const bool use_buffer_a +- = bgmmc.use_buffer_a || bgmmc.use_buffer_a_tail_only; +- const int num_threads = brgmm_ctx.get_num_threads_for_parallelization(); +- +- const int M_chunks = brgmm_ctx.get_M_chunks(); +- const int M_chunk_size = brgmm_ctx.get_M_chunk_size(); +- const int M_chunk_tail = brgmm_ctx.get_M_chunk_tail(); +- parallel(num_threads, [&](const int ithr, const int nthr) { ++ float *oscales = bgmmc.with_scales ++ ? scratchpad.template get(key_precomputed_scales) ++ : nullptr; ++ ++ float *inv_dst_scales = bgmmc.with_dst_scales ++ ? scratchpad.template get(key_matmul_dst_scales) ++ : nullptr; ++ ++ if (bgmmc.with_scales || bgmmc.with_dst_scales) { ++ parallel(1, [= COMPAT_THIS_CAPTURE](int, int) { ++ scale_utils::precompute_oscales(oscales, pd()->attr(), src_scales, ++ wei_scales, pd()->N(), 1.f); ++ ++ scale_utils::precompute_inv_dst_scales( ++ inv_dst_scales, pd()->attr(), dst_scales); ++ }); ++ } ++ ++ const auto &brgmm_ctx_ptr = std::make_shared( ++ ctx, pd(), oscales, inv_dst_scales, helper); ++ ++ const int num_threads ++ = brgmm_ctx_ptr->get_num_threads_for_parallelization(); ++ ++ parallel(num_threads, ++ [= COMPAT_THIS_CAPTURE](const int ithr, const int nthr) { ++ const auto &brgmm_ctx = *brgmm_ctx_ptr; ++ + const int ithr_bmn = brgmm_ctx.get_thread_idx_for_bmn(ithr); + const int ithr_k = brgmm_ctx.get_thread_idx_for_k(ithr); ++ const int M_chunks = brgmm_ctx.get_M_chunks(); ++ const int M_chunk_size = brgmm_ctx.get_M_chunk_size(); ++ const int M_chunk_tail = brgmm_ctx.get_M_chunk_tail(); ++ ++ const auto &bgmmc = pd()->get_brgemm_matmul_conf(); ++ const bool use_buffer_a ++ = bgmmc.use_buffer_a || bgmmc.use_buffer_a_tail_only; ++ + if (ithr_bmn < 0 || ithr_k < 0) return; + int start {0}, end {0}; + balance211(brgmm_ctx.get_parallel_work_amount(), +@@ -363,7 +409,7 @@ status_t brgemm_matmul_t::execute_body(const exec_ctx_t &ctx) const { + } + }); + +- maybe_reduce_partial_results_and_apply_postops(brgmm_ctx); ++ maybe_reduce_partial_results_and_apply_postops(brgmm_ctx_ptr); + + return status::success; + } +@@ -506,13 +552,17 @@ void brgemm_matmul_t::compute_kernel( + + template + void brgemm_matmul_t::maybe_reduce_partial_results_and_apply_postops( +- const brg_matmul_exec_ctx_t &brgmm_ctx) const { +- if (!brgmm_ctx.parallel_reduction_is_used()) return; ++ const std::shared_ptr &brgmm_ctx_ptr) const { ++ if (!brgmm_ctx_ptr->parallel_reduction_is_used()) return; + + const auto &bgmmc = pd()->get_brgemm_matmul_conf(); +- const int num_threads = brgmm_ctx.get_num_threads_for_parallelization(); ++ const int num_threads ++ = brgmm_ctx_ptr->get_num_threads_for_parallelization(); ++ ++ parallel(num_threads, ++ [= COMPAT_THIS_CAPTURE](const int ithr, const int nthr) { ++ const auto &brgmm_ctx = *brgmm_ctx_ptr; + +- parallel(num_threads, [&](const int ithr, const int nthr) { + const int nthr_k = brgmm_ctx.get_num_threads_for_k(); + const int ithr_bmn = brgmm_ctx.get_thread_idx_for_bmn(ithr); + const int ithr_k = brgmm_ctx.get_thread_idx_for_k(ithr); +diff --git a/src/cpu/aarch64/matmul/brgemm_matmul.hpp b/src/cpu/aarch64/matmul/brgemm_matmul.hpp +index 558dad80c19..a3c15d5cf54 100644 +--- a/src/cpu/aarch64/matmul/brgemm_matmul.hpp ++++ b/src/cpu/aarch64/matmul/brgemm_matmul.hpp +@@ -1,7 +1,7 @@ + /******************************************************************************* + * Copyright 2021 Intel Corporation + * Copyright 2024 FUJITSU LIMITED +-* Copyright 2025 Arm Ltd. and affiliates ++* Copyright 2025-2026 Arm Ltd. and affiliates + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at +@@ -18,15 +18,13 @@ + #ifndef CPU_AARCH64_MATMUL_BRGEMM_MATMUL_HPP + #define CPU_AARCH64_MATMUL_BRGEMM_MATMUL_HPP + ++#include ++ + #include "common/c_types_map.hpp" + #include "common/primitive.hpp" +-#include "common/type_helpers.hpp" + + #include "cpu/matmul/cpu_matmul_pd.hpp" + +-#include "cpu/aarch64/brgemm/brgemm.hpp" +-#include "cpu/aarch64/brgemm/brgemm_containers.hpp" +-#include "cpu/aarch64/brgemm/brgemm_utils.hpp" + #include "cpu/aarch64/cpu_reducer.hpp" + #include "cpu/aarch64/matmul/brgemm_matmul_copy_utils.hpp" + #include "cpu/aarch64/matmul/brgemm_matmul_utils.hpp" +@@ -88,7 +86,7 @@ struct brgemm_matmul_t : public primitive_t { + void copy_b_chunk_in_buffer(const brg_matmul_exec_ctx_t &brgmm_ctx, + int ithr, int b_idx, int n_blk_idx, int k_blk_idx) const; + void maybe_reduce_partial_results_and_apply_postops( +- const brg_matmul_exec_ctx_t &brgmm_ctx) const; ++ const std::shared_ptr &brgmm_ctx_ptr) const; + void accumulate( + char *result_ptr, const char *reduce_ptr, size_t size) const; + +diff --git a/src/cpu/aarch64/matmul/brgemm_matmul_reorders.cpp b/src/cpu/aarch64/matmul/brgemm_matmul_reorders.cpp +index 2b1c7683839..da8792aca19 100644 +--- a/src/cpu/aarch64/matmul/brgemm_matmul_reorders.cpp ++++ b/src/cpu/aarch64/matmul/brgemm_matmul_reorders.cpp +@@ -1,7 +1,7 @@ + /******************************************************************************* + * Copyright 2022 Intel Corporation + * Copyright 2024 FUJITSU LIMITED +-* Copyright 2025 Arm Ltd. and affiliates ++* Copyright 2025-2026 Arm Ltd. and affiliates + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. +@@ -16,9 +16,10 @@ + * limitations under the License. + *******************************************************************************/ + ++#include "common/compiler_workarounds.hpp" + #include "common/dnnl_thread.hpp" +-#include "cpu/aarch64/cpu_isa_traits.hpp" + ++#include "cpu/aarch64/cpu_isa_traits.hpp" + #include "cpu/aarch64/matmul/brgemm_matmul_reorders.hpp" + + namespace dnnl { +@@ -182,7 +183,7 @@ status_t brgemm_matmul_copy_reorder_t::execute_body( + : (dt_sz) * (md).blk_off((d0), (d1))) + + parallel_nd(kernel_conf.batch, div_up(kernel_conf.N, kernel_conf.N_blk), +- [&](dim_t batch, dim_t n_blk_idx) { ++ [= COMPAT_THIS_CAPTURE](dim_t batch, dim_t n_blk_idx) { + const auto n = n_blk_idx * kernel_conf.N_blk; + const bool is_N_tail = (kernel_conf.N - n < kernel_conf.N_blk); + auto ker_exec_ctx = matmul::jit_brgemm_matmul_copy_b_t::ctx_t(); diff --git a/third_party/xla/third_party/mkl_dnn/onednn_async_aarch64_pr_5828.patch b/third_party/xla/third_party/mkl_dnn/onednn_async_aarch64_pr_5828.patch new file mode 100644 index 00000000000000..636ade65ea368f --- /dev/null +++ b/third_party/xla/third_party/mkl_dnn/onednn_async_aarch64_pr_5828.patch @@ -0,0 +1,257 @@ +From 0da69cb9a47b3dab8c5d2e24021477e19ad5f3be Mon Sep 17 00:00:00 2001 +From: Siddhartha Menon +Date: Tue, 11 Aug 2026 14:50:53 +0100 +Subject: [PATCH 1/2] cpu: aarch64: make brgconv async compatible + +Signed-off-by: Siddhartha Menon +--- + src/cpu/aarch64/jit_brgemm_conv.cpp | 70 ++++++++++++++++++++++------- + 1 file changed, 53 insertions(+), 17 deletions(-) + +diff --git a/src/cpu/aarch64/jit_brgemm_conv.cpp b/src/cpu/aarch64/jit_brgemm_conv.cpp +index 42f5c53b9c7..97affba10e2 100644 +--- a/src/cpu/aarch64/jit_brgemm_conv.cpp ++++ b/src/cpu/aarch64/jit_brgemm_conv.cpp +@@ -19,15 +19,16 @@ + #include + + #include "common/c_types_map.hpp" ++#include "common/compiler_workarounds.hpp" + #include "common/dnnl_thread.hpp" + #include "common/type_helpers.hpp" + #include "common/utils.hpp" ++ ++#include "cpu/aarch64/brgemm/brgemm_utils.hpp" + #include "cpu/aarch64/cpu_isa_traits.hpp" + #include "cpu/aarch64/jit_brgemm_conv.hpp" + #include "cpu/aarch64/jit_brgemm_conv_comp_pad_kernel.hpp" + #include "cpu/aarch64/jit_brgemm_conv_utils.hpp" +-#include "cpu/cpu_primitive.hpp" +-#include "cpu/scale_utils.hpp" + + namespace dnnl { + namespace impl { +@@ -548,9 +549,17 @@ status_t brgemm_convolution_fwd_t::pd_t::init(const engine_t *engine) { + + auto scratchpad = scratchpad_registry().registrar(); + brgemm_convolution_utils::init_scratchpad(scratchpad, jcp_); +- if (jcp_.with_scales) +- book_precomputed_scales(scratchpad, attr()->scales_, OC(), +- jcp_.scale_adjust_factor != 1.0f); ++ if (need_postwork || jcp_.with_scales) { ++ const size_t scales_size = jcp_.is_oc_scale ++ ? nstl::max(static_cast(OC()), ++ scale_utils::scales_simd_w()) ++ : scale_utils::scales_simd_w(); ++ scratchpad.template book(key_precomputed_scales, scales_size, ++ brgemm_convolution_utils::P4K); ++ } ++ if (!attr()->scales_.has_default_values(DNNL_ARG_DST)) ++ scratchpad.template book(key_conv_dst_scales, ++ scale_utils::scales_simd_w(), brgemm_convolution_utils::P4K); + + return status::success; + } +@@ -1061,7 +1070,7 @@ status_t brgemm_convolution_fwd_t::init(engine_t *engine) { + } + template + struct brgemm_convolution_fwd_t::brgemm_thread_ctx_t { +- brgemm_thread_ctx_t(brgemm_exec_ctx_t &brgemm_ctx_, int ithr_, ++ brgemm_thread_ctx_t(const brgemm_exec_ctx_t &brgemm_ctx_, int ithr_, + brgemm_batch_element_t *__restrict brg_batch_, char *c_buffer_, + char *wsp_tile_) + : brgemm_ctx(brgemm_ctx_) +@@ -1070,7 +1079,7 @@ struct brgemm_convolution_fwd_t::brgemm_thread_ctx_t { + , c_buffer(c_buffer_) + , wsp_tile(wsp_tile_) {} + +- brgemm_exec_ctx_t &brgemm_ctx; ++ const brgemm_exec_ctx_t &brgemm_ctx; + int ithr {0}; + brgemm_batch_element_t *__restrict brg_batch {nullptr}; + char *c_buffer {nullptr}; +@@ -1097,13 +1106,22 @@ status_t brgemm_convolution_fwd_t::execute(const exec_ctx_t &ctx) const { + const int32_t *dst_zero_points = CTX_IN_MEM( + const int32_t *, DNNL_ARG_ATTR_ZERO_POINTS | DNNL_ARG_DST); + +- DEFINE_ARG_SCALES_BUFFER(src_scales, DNNL_ARG_SRC); +- DEFINE_ARG_SCALES_BUFFER(wei_scales, DNNL_ARG_WEIGHTS); +- DEFINE_ARG_SCALES_BUFFER(dst_scales, DNNL_ARG_DST); +- +- const float *oscales = precompute_scales(ctx.get_scratchpad_grantor(), +- src_scales, wei_scales, _pd->OC(), _pd->attr(), +- jcp.scale_adjust_factor); ++ const void *src_scales ++ = CTX_IN_MEM(const void *, DNNL_ARG_ATTR_SCALES | DNNL_ARG_SRC); ++ const void *wei_scales ++ = CTX_IN_MEM(const void *, DNNL_ARG_ATTR_SCALES | DNNL_ARG_WEIGHTS); ++ const void *dst_scales ++ = CTX_IN_MEM(const void *, DNNL_ARG_ATTR_SCALES | DNNL_ARG_DST); ++ ++ if (!_pd->attr()->scales_.has_default_values(DNNL_ARG_SRC) ++ && src_scales == nullptr) ++ return status::invalid_arguments; ++ if (!_pd->attr()->scales_.has_default_values(DNNL_ARG_WEIGHTS) ++ && wei_scales == nullptr) ++ return status::invalid_arguments; ++ if (!_pd->attr()->scales_.has_default_values(DNNL_ARG_DST) ++ && dst_scales == nullptr) ++ return status::invalid_arguments; + + brgemm_exec_ctx_t brgemm_ctx(ctx, _pd); + +@@ -1126,6 +1144,15 @@ status_t brgemm_convolution_fwd_t::execute(const exec_ctx_t &ctx) const { + : nullptr; + + const auto &scratchpad = ctx.get_scratchpad_grantor(); ++ const bool needs_oscales = _pd->need_postwork || jcp.with_scales; ++ float *const oscales = needs_oscales ++ ? scratchpad.template get(key_precomputed_scales) ++ : nullptr; ++ float *const inv_dst_scales_precomputed ++ = !pd()->attr()->scales_.has_default_values(DNNL_ARG_DST) ++ ? scratchpad.template get(key_conv_dst_scales) ++ : nullptr; ++ + brgemm_batch_element_t *const __restrict brg_batch_global + = brgemm_convolution_utils::uses_batch_elements( + jcp.brg_type, jcp.exec_type) +@@ -1155,6 +1182,15 @@ status_t brgemm_convolution_fwd_t::execute(const exec_ctx_t &ctx) const { + + cal_compensation(wei, src_zp_comp_base, s8s8_comp_base); + ++ if (needs_oscales || inv_dst_scales_precomputed) { ++ parallel(1, [= COMPAT_THIS_CAPTURE](const int, const int) { ++ scale_utils::precompute_oscales(oscales, pd()->attr(), src_scales, ++ wei_scales, pd()->OC(), pd()->jcp_.scale_adjust_factor); ++ scale_utils::precompute_inv_dst_scales( ++ inv_dst_scales_precomputed, pd()->attr(), dst_scales); ++ }); ++ } ++ + // --------------- Parallel section ------------------------------ + const dim_t work_amount = static_cast(jcp.mb) * jcp.ngroups + * jcp.nb_oc * jcp.nb_od * jcp.nb_oh * jcp.nb_ow; +@@ -1163,7 +1199,7 @@ status_t brgemm_convolution_fwd_t::execute(const exec_ctx_t &ctx) const { + // or made ic_chunks = 1 if use_buffer + // or (looks more general) increase buffer size to store several rows + +- parallel(jcp.nthr, [&](const int ithr, const int nthr) { ++ parallel(jcp.nthr, [= COMPAT_THIS_CAPTURE](const int ithr, const int nthr) { + if (ithr >= work_amount) return; + + brgemm_batch_element_t *const __restrict brg_batch = brg_batch_global +@@ -1213,7 +1249,7 @@ status_t brgemm_convolution_fwd_t::execute(const exec_ctx_t &ctx) const { + = jcp.src_zero_point ? src_zp_comp_base : nullptr; + btc.s8s8_comp_ptr + = jcp.s8s8_compensation_required ? s8s8_comp_base : nullptr; +- btc.dst_scales = dst_scales; ++ btc.dst_scales = inv_dst_scales_precomputed; + + if (jcp.exec_type == exec_trans && (last_n != n || last_g != g)) { + if (!jcp.copy_block_only) +@@ -1292,7 +1328,7 @@ status_t brgemm_convolution_fwd_t::cal_compensation( + <= platform::get_per_core_cache_size(1)); + const int nthr = is_small_shape ? 1 : jcp.nthr; + +- parallel(nthr, [&](const int ithr, const int nthr) { ++ parallel(nthr, [= COMPAT_THIS_CAPTURE](const int ithr, const int nthr) { + if (ithr >= work_amount) return; + + dim_t start {0}, end {0}; + +From b368dc607fe128ad8f7d0c0b2a7064c8a42f952e Mon Sep 17 00:00:00 2001 +From: Siddhartha Menon +Date: Wed, 16 Sep 2026 16:28:32 +0100 +Subject: [PATCH 2/2] cpu: aarch64: conv: disable threadpool for unsupported + convs + +Signed-off-by: Siddhartha Menon +--- + src/cpu/aarch64/jit_brdgmm_dw_conv.cpp | 4 +++- + src/cpu/aarch64/jit_brgemm_1x1_conv.cpp | 3 +++ + src/cpu/aarch64/jit_sve_1x1_convolution.hpp | 3 +++ + src/cpu/aarch64/jit_sve_convolution.hpp | 4 ++++ + src/cpu/aarch64/jit_uni_dw_convolution.hpp | 4 ++++ + 5 files changed, 17 insertions(+), 1 deletion(-) + +diff --git a/src/cpu/aarch64/jit_brdgmm_dw_conv.cpp b/src/cpu/aarch64/jit_brdgmm_dw_conv.cpp +index 7e8dd5c6f15..9844ee0eff7 100644 +--- a/src/cpu/aarch64/jit_brdgmm_dw_conv.cpp ++++ b/src/cpu/aarch64/jit_brdgmm_dw_conv.cpp +@@ -1,7 +1,7 @@ + /******************************************************************************* + * Copyright 2021 Intel Corporation + * Copyright 2024-2026 FUJITSU LIMITED +-* Copyright 2025 Arm Ltd. and affiliates ++* Copyright 2025-2026 Arm Ltd. and affiliates + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. +@@ -110,6 +110,8 @@ status_t brdgmm_dw_convolution_fwd_t::pd_t::init(const engine_t *engine) { + auto skip_mask = skip_mask_t::post_ops; + if (is_int8) skip_mask |= skip_mask_t::scales; + ++ VDISPATCH_CONV(DNNL_CPU_THREADING_RUNTIME != DNNL_RUNTIME_THREADPOOL, ++ VERBOSE_UNSUPPORTED_THREADPOOL_RUNTIME); + VDISPATCH_CONV(is_fwd(), VERBOSE_BAD_PROPKIND); + VDISPATCH_CONV(set_default_alg_kind(alg_kind::convolution_direct), + VERBOSE_BAD_ALGORITHM); +diff --git a/src/cpu/aarch64/jit_brgemm_1x1_conv.cpp b/src/cpu/aarch64/jit_brgemm_1x1_conv.cpp +index 81411a68233..6eb5e8e0caa 100644 +--- a/src/cpu/aarch64/jit_brgemm_1x1_conv.cpp ++++ b/src/cpu/aarch64/jit_brgemm_1x1_conv.cpp +@@ -58,6 +58,9 @@ status_t brgemm_1x1_convolution_fwd_t::pd_t::init(const engine_t *engine) { + | skip_mask_t::zero_points; + if (one_of(src_type, u8, s8)) skip_mask |= skip_mask_t::scales; + ++ VDISPATCH_CONV(DNNL_CPU_THREADING_RUNTIME != DNNL_RUNTIME_THREADPOOL, ++ VERBOSE_UNSUPPORTED_THREADPOOL_RUNTIME); ++ + bool ok = is_fwd() && set_default_alg_kind(alg_kind::convolution_direct) + && expect_data_types(src_type, wei_type, data_type::undef, dst_type, + data_type::undef) +diff --git a/src/cpu/aarch64/jit_sve_1x1_convolution.hpp b/src/cpu/aarch64/jit_sve_1x1_convolution.hpp +index 4ba4823b452..b8c69eee912 100644 +--- a/src/cpu/aarch64/jit_sve_1x1_convolution.hpp ++++ b/src/cpu/aarch64/jit_sve_1x1_convolution.hpp +@@ -62,6 +62,9 @@ struct jit_sve_1x1_convolution_fwd_t : public primitive_t { + return status::unimplemented; + } + #endif ++ VDISPATCH_CONV( ++ DNNL_CPU_THREADING_RUNTIME != DNNL_RUNTIME_THREADPOOL, ++ VERBOSE_UNSUPPORTED_THREADPOOL_RUNTIME); + + bool ok = true && is_fwd() + && set_default_alg_kind(alg_kind::convolution_direct) +diff --git a/src/cpu/aarch64/jit_sve_convolution.hpp b/src/cpu/aarch64/jit_sve_convolution.hpp +index de5dd26ff7c..932b7ed9d86 100644 +--- a/src/cpu/aarch64/jit_sve_convolution.hpp ++++ b/src/cpu/aarch64/jit_sve_convolution.hpp +@@ -53,6 +53,10 @@ struct jit_sve_convolution_fwd_t : public primitive_t { + } + #endif + ++ VDISPATCH_CONV( ++ DNNL_CPU_THREADING_RUNTIME != DNNL_RUNTIME_THREADPOOL, ++ VERBOSE_UNSUPPORTED_THREADPOOL_RUNTIME); ++ + bool ok = true && mayiuse(isa) && is_fwd() + && set_default_alg_kind(alg_kind::convolution_direct) + && expect_data_types(src_type, wei_type, dst_type, dst_type, +diff --git a/src/cpu/aarch64/jit_uni_dw_convolution.hpp b/src/cpu/aarch64/jit_uni_dw_convolution.hpp +index f0991572183..2180f6209c7 100644 +--- a/src/cpu/aarch64/jit_uni_dw_convolution.hpp ++++ b/src/cpu/aarch64/jit_uni_dw_convolution.hpp +@@ -47,6 +47,10 @@ struct jit_uni_dw_convolution_fwd_t : public primitive_t { + jit_uni_dw_convolution_fwd_t); + + status_t init(const engine_t *engine) { ++ VDISPATCH_CONV( ++ DNNL_CPU_THREADING_RUNTIME != DNNL_RUNTIME_THREADPOOL, ++ VERBOSE_UNSUPPORTED_THREADPOOL_RUNTIME); ++ + bool ok = true && is_fwd() + && set_default_alg_kind(alg_kind::convolution_direct) + && expect_data_types(src_type, src_type, data_type::undef, diff --git a/third_party/xla/third_party/mkl_dnn/onednn_async_aarch64_pr_6033.patch b/third_party/xla/third_party/mkl_dnn/onednn_async_aarch64_pr_6033.patch new file mode 100644 index 00000000000000..88d8145f321527 --- /dev/null +++ b/third_party/xla/third_party/mkl_dnn/onednn_async_aarch64_pr_6033.patch @@ -0,0 +1,206 @@ +From 5a76c778c08f1fa730686e347500eecf9c05e6d5 Mon Sep 17 00:00:00 2001 +From: Gassan +Date: Mon, 14 Sep 2026 11:22:28 +0000 +Subject: [PATCH] cpu: aarch64: tune SVE128 f32 BRGEMM matmul + +Prefer N16 blocking and cap K panels at 256 for eligible SVE128 FP32 matmuls. A full K256-by-N16 B panel contains 16 KiB of FP32 data. Batch up to 16 K panels per call to amortize kernel entry and keep C accumulators live across panels. Search larger M blocks when batch and N provide enough parallel work, reusing B across more rows. + +Alternate two B registers to load one column ahead in the existing indexed loop; retain the existing loop-selection policy. +--- + src/cpu/aarch64/brgemm/jit_brgemm_kernel.cpp | 33 ++++++---- + .../aarch64/matmul/brgemm_matmul_utils.cpp | 66 +++++++++++-------- + 2 files changed, 57 insertions(+), 42 deletions(-) + +diff --git a/src/cpu/aarch64/brgemm/jit_brgemm_kernel.cpp b/src/cpu/aarch64/brgemm/jit_brgemm_kernel.cpp +index 8a251273e7d..ed799e48e96 100644 +--- a/src/cpu/aarch64/brgemm/jit_brgemm_kernel.cpp ++++ b/src/cpu/aarch64/brgemm/jit_brgemm_kernel.cpp +@@ -243,7 +243,8 @@ struct jit_brgemm_kernel_t : public jit_generator_t { + + ZReg load(int ld = 0) const { + if (n_bcast_1_load) { +- return ZReg(brg.bd_block); ++ // Two B registers between A and C enable one-ahead loads. ++ return ZReg(brg.bd_block + ld % 2); + } else { + // Starts off from lowest accm register, and continues descending + int idx = max_effective_vregs - 1 - (brg.ld_block2 * brg.bd_block) +@@ -1640,6 +1641,7 @@ void jit_brgemm_kernel_t::gemm_microkernel(int bd_block2, bool is_bdb_tail, + auto A_stride_bytes = static_cast(brg.typesize_A * brg.LDA); + + if (n_bcast_1_load) { ++ const bool load_b_early = ld_block2 > 1; + // Use a tmp to store the pointer to the first element of A in the tile + // We can't use reg_aux_A directly because we increment this by a quadword when needed + XReg reg_A_ptr = X_TMP_4; +@@ -1660,6 +1662,20 @@ void jit_brgemm_kernel_t::gemm_microkernel(int bd_block2, bool is_bdb_tail, + // Bump by quadword + add(reg_A_ptr, reg_A_ptr, 16); + } ++ auto load_b = [&](const int ld) { ++ const auto mask = is_ld_tail ? ld_tail_mask : P_ALL_ONE; ++ const int b_offset = B_offset(ld, rd); ++ if (!use_mul_vl(b_offset - b_base_offset, 4, cpu_sveLen)) { ++ add_vl_or_imm(reg_tmp_, x_addr, b_offset - b_base_offset, ++ X_TMP_0); ++ b_base_offset = b_offset; ++ x_addr = reg_tmp_; ++ } ++ const auto b_mul_vl = compute_off_mul_vl( ++ b_offset - b_base_offset, 4, cpu_sveLen); ++ ld1w(load(ld).s, mask / T_z, ptr(x_addr, b_mul_vl, MUL_VL)); ++ }; ++ if (load_b_early) load_b(0); + if (quadword_index == 0) { + // If loading a quadword would take us past the end of rd_loop, we need to use a mask + bool quadword_is_too_much = (rd + 4 * brg.rd_step) > rd_loop; +@@ -1675,19 +1691,10 @@ void jit_brgemm_kernel_t::gemm_microkernel(int bd_block2, bool is_bdb_tail, + } + } + for (int ld = 0; ld < ld_block2; ld++) { +- const auto mask = is_ld_tail ? ld_tail_mask : P_ALL_ONE; +- const int b_offset = B_offset(ld, rd); +- if (!use_mul_vl(b_offset - b_base_offset, 4, cpu_sveLen)) { +- add_vl_or_imm(reg_tmp_, x_addr, b_offset - b_base_offset, +- X_TMP_0); +- b_base_offset = b_offset; +- x_addr = reg_tmp_; +- } +- auto b_mul_vl = compute_off_mul_vl( +- b_offset - b_base_offset, 4, cpu_sveLen); +- ld1w(load(ld).s, mask / T_z, ptr(x_addr, b_mul_vl, MUL_VL)); ++ const int next_ld = ld + load_b_early; ++ if (next_ld < ld_block2) load_b(next_ld); + for (int bd = bd_b; bd < bd_e; bd++) { +- dot_product(accm(ld_block2, bd, ld), load(), bcst(bd), ++ dot_product(accm(ld_block2, bd, ld), load(ld), bcst(bd), + quadword_index); + } + } +diff --git a/src/cpu/aarch64/matmul/brgemm_matmul_utils.cpp b/src/cpu/aarch64/matmul/brgemm_matmul_utils.cpp +index 6c48b613538..d8ec2119337 100644 +--- a/src/cpu/aarch64/matmul/brgemm_matmul_utils.cpp ++++ b/src/cpu/aarch64/matmul/brgemm_matmul_utils.cpp +@@ -79,21 +79,20 @@ int get_default_n_block( + case BA16a16b2a: + case BA16a16b4a: return 16; + default: { +- if (bgmmc.N == 16 || bgmmc.N == 32 || bgmmc.N == 64) return bgmmc.N; +- if (!mayiuse(sve_512)) { +- if (bgmmc.N <= 16) +- return 16; +- else { +- // It is observed that for M,K>512, N block of 64 works better provided that thread distribution is not hindered. +- if (bgmmc.N / 64 >= bgmmc.nthr && bgmmc.K > 512 +- && bgmmc.M > 512) +- return 64; +- else +- return 32; +- } +- +- } else ++ if (one_of(bgmmc.N, 16, 32, 64)) return bgmmc.N; ++ // Prefer N16 for SVE128 FP32 before the generic N64 rule below. ++ // A full K256-by-N16 FP32 panel contains 16 KiB. ++ if (bgmmc.isa == sve_128 ++ && everyone_is(f32, bgmmc.src_dt, bgmmc.wei_dt) ++ && bgmmc.M >= 128 && bgmmc.N > 64 && bgmmc.K > 128) ++ return 16; ++ if (mayiuse(sve_512)) return 64; ++ if (bgmmc.N <= 16) return 16; ++ // For remaining large M/K shapes, use N64 only when N provides ++ // at least one block per thread. ++ if (bgmmc.N / 64 >= bgmmc.nthr && bgmmc.K > 512 && bgmmc.M > 512) + return 64; ++ return 32; + } + } + } +@@ -701,8 +700,8 @@ float compute_blocking_heuristic_sve_256(brgemm_matmul_conf_t &bgmmc, + return best_imbalance; + } + +-// This is a direct copy of compute_blocking_heuristic_sve_256 +-// with basic parameters halved; TODO: develop a good sve_128 heuristic ++// Derived from compute_blocking_heuristic_sve_256 with baseline parameters ++// halved; SVE128 FP32 overrides are applied below. + float compute_blocking_heuristic_sve_128(brgemm_matmul_conf_t &bgmmc, + const brgemm_matmul_conf_utils_t &bm_conf_utils, + const matmul_brgemm_blocking_params_t::matmul_params_t &matmul, +@@ -710,22 +709,34 @@ float compute_blocking_heuristic_sve_128(brgemm_matmul_conf_t &bgmmc, + + const int nthr = bgmmc.nthr; + +- const int max_m_blk = nstl::min(128, matmul.M); +- int min_m = (matmul.batch > 1) ? 16 : 64; +- int min_m_blk = nstl::min(min_m, matmul.M); // max_m_blk +- ++ const bool is_f32 = bm_conf_utils.is_f32(); + int n_blk = bgmmc.N_blk; + const int n_chunks = div_up(matmul.N, n_blk); + const int max_n_chunks = bgmmc.use_buffer_a ? 4 : 1; + const int n_chunks_start = nstl::min(max_n_chunks, n_chunks); ++ const size_t max_parallel = matmul.batch * n_chunks; ++ ++ // Search larger M blocks when batch and N already provide enough ++ // parallel work, reusing each B panel across more rows. ++ const bool use_large_m_block = is_f32 && matmul.K > 128 && matmul.K <= 512 ++ && n_chunks > 1 && max_parallel >= static_cast(nthr); ++ const int max_m_blk = nstl::min(use_large_m_block ? 1024 : 128, matmul.M); ++ const int min_m = matmul.batch > 1 ? 16 : 64; ++ int min_m_blk = nstl::min(min_m, matmul.M); + + int default_k_blk = (matmul.M >= 256) ? 512 : 64; + int k_blk = nstl::min(matmul.K, default_k_blk); ++ int batch_size = 1; ++ // Limit K panels to 256 elements. Batch up to 16 panels per call to ++ // amortize kernel entry and keep C accumulators live across panels. ++ if (is_f32 && matmul.M >= 128) { ++ k_blk = nstl::min(matmul.K, 256); ++ batch_size = nstl::min(16, div_up(matmul.K, k_blk)); ++ } + int start_nthr_k = 1; + + // for cases with low parallel work, reduce 'min_m_blk' to + // increase potential parallelization balance. +- const size_t max_parallel = matmul.batch * n_chunks; + const bool low_parallel_work = static_cast(nthr) > max_parallel; + if (low_parallel_work) { + +@@ -765,7 +776,7 @@ float compute_blocking_heuristic_sve_128(brgemm_matmul_conf_t &bgmmc, + + matmul_brgemm_blocking_params_t cur_params(matmul, nthr); + cur_params.update_params( +- 1, m_blk, n_chunk_size, n_blk, 1, k_blk, nthr_k); ++ 1, m_blk, n_chunk_size, n_blk, batch_size, k_blk, nthr_k); + + float cur_imbalance = cur_params.get_imbalance(); + if (cur_imbalance < best_imbalance) { +@@ -807,9 +818,8 @@ status_t compute_blocking_heuristic(brgemm_matmul_conf_t &bgmmc, + // choose a larger 'M_blk'. + // + // N_blk: +- // - ideally 64 (from 'get_default_n_block()'). +- // - can be reduced to 32 to improve performance for some shapes, as +- // well as increasing parallelization search space. ++ // - selected by the weights layout or 'get_default_n_block()'. ++ // - smaller blocks can improve cache locality and parallelization. + // + // N_Chunks: + // - No different as long as thread/work balance is the same. +@@ -818,13 +828,11 @@ status_t compute_blocking_heuristic(brgemm_matmul_conf_t &bgmmc, + // reuse. + // + // K_blk: +- // - block size variation '512 <= K_blk < 1024' has negligible +- // performance difference. However, Some cases benefit from higher +- // block size. ++ // - SVE128 f32 caps K_blk at 256 for M >= 128. + // - can parallelize if not enough work; notice: requires reduction! + // + // Batch_Size: +- // - unused. ++ // - number of K panels per kernel call, up to 16 for SVE128 f32. + case sme: + case sve_512: + best_imbalance = compute_blocking_heuristic_sve_512( diff --git a/third_party/xla/third_party/mkl_dnn/workspace.bzl b/third_party/xla/third_party/mkl_dnn/workspace.bzl index 679d4a4dd4187d..76568d72997469 100644 --- a/third_party/xla/third_party/mkl_dnn/workspace.bzl +++ b/third_party/xla/third_party/mkl_dnn/workspace.bzl @@ -41,6 +41,21 @@ def repo(): urls = tf_mirror_urls("https://github.com/oneapi-src/oneDNN/archive/refs/tags/v3.12.3.tar.gz"), ) + tf_http_archive( + name = "onednn_async_aarch64", + build_file = "//third_party/mkl_dnn:mkldnn_v1_async_aarch64.BUILD", + # Rename the gtest module in oneDNN's third_party to avoid + # conflict with Google's gtest module + patch_file = [ + "//third_party/mkl_dnn:onednn_async_aarch64_pr_5790.patch", + "//third_party/mkl_dnn:onednn_async_aarch64_pr_5828.patch", + "//third_party/mkl_dnn:onednn_async_aarch64_pr_6033.patch", + ], + sha256 = "cceec79d37474df43589eb951602f1a820eedbbb61e58021dce2faaff1da827b", + strip_prefix = "oneDNN-80fb15f600887e22c0c7989ce6201723e6c11536", + urls = tf_mirror_urls("https://github.com/uxlfoundation/oneDNN/archive/80fb15f600887e22c0c7989ce6201723e6c11536.tar.gz"), + ) + tf_http_archive( name = "mkl_dnn_acl_compatible", build_file = "//third_party/mkl_dnn:mkldnn_acl.BUILD", diff --git a/third_party/xla/third_party/stablehlo/temporary.patch b/third_party/xla/third_party/stablehlo/temporary.patch index c0fd93daf55cd0..48889076fb4a16 100644 --- a/third_party/xla/third_party/stablehlo/temporary.patch +++ b/third_party/xla/third_party/stablehlo/temporary.patch @@ -123,6 +123,109 @@ diff --ruN a/stablehlo/stablehlo/conversions/linalg/transforms/LegalizeToLinalgU /// Extracts integer values from the attribute |elements|. SmallVector extract1DVector(DenseIntElementsAttr elements); +diff --ruN a/stablehlo/stablehlo/conversions/linalg/transforms/MapStablehloToScalarOp.h b/stablehlo/stablehlo/conversions/linalg/transforms/MapStablehloToScalarOp.h +--- stablehlo/stablehlo/conversions/linalg/transforms/MapStablehloToScalarOp.h ++++ stablehlo/stablehlo/conversions/linalg/transforms/MapStablehloToScalarOp.h +@@ -204,7 +204,8 @@ + Value operator()(Location loc, ArrayRef resultTypes, + ArrayRef /*argTypes*/, ValueRange args, OpBuilder* b) { + return StdScalarOp::create(*b, loc, resultTypes, args, +- ArrayRef()); ++ typename StdScalarOp::Properties{}, ++ /*discardableAttributes=*/{}); + } + }; + +@@ -215,7 +216,8 @@ + Type elementType = getElementTypeOrSelf(argTypes.front()); + if (SupportedType{}(elementType)) { + return StdScalarOp::create(*b, loc, resultTypes, args, +- ArrayRef()); ++ typename StdScalarOp::Properties{}, ++ /*discardableAttributes=*/{}); + } + return MapStablehloOpToScalarOpImpl{}(loc, resultTypes, argTypes, + args, b); +@@ -688,32 +690,39 @@ + // A boolean value is considered to be unsigned when converting to + // floating-point. Otherwise, it will become `-1`. + if (IsUnsignedIntegerType{}(sourceType) && +- mlir::arith::UIToFPOp::areCastCompatible(convertedSourceType, +- targetType)) { +- return mlir::arith::UIToFPOp::create(*b, loc, resultTypes, args, +- ArrayRef()); ++ arith::UIToFPOp::areCastCompatible(convertedSourceType, targetType)) { ++ return arith::UIToFPOp::create(*b, loc, resultTypes, args, ++ typename arith::UIToFPOp::Properties{}, ++ /*discardableAttributes=*/{}); + } +- if (mlir::arith::SIToFPOp::areCastCompatible(sourceType, targetType)) { +- return mlir::arith::SIToFPOp::create(*b, loc, resultTypes, args, +- ArrayRef()); ++ if (arith::SIToFPOp::areCastCompatible(sourceType, targetType)) { ++ return arith::SIToFPOp::create(*b, loc, resultTypes, args, ++ typename arith::SIToFPOp::Properties{}, ++ /*discardableAttributes=*/{}); + } + if (isa(sourceType) && isa(targetType)) { + auto src = cast(sourceType); + auto res = cast(targetType); + if (src.getWidth() > res.getWidth()) { +- return mlir::arith::TruncFOp::create(*b, loc, resultTypes, args, +- ArrayRef()); ++ return arith::TruncFOp::create(*b, loc, resultTypes, args, ++ typename arith::TruncFOp::Properties{}, ++ /*discardableAttributes=*/{}); + } + if (src.getWidth() < res.getWidth()) { +- return mlir::arith::ExtFOp::create(*b, loc, resultTypes, args, +- ArrayRef()); ++ return arith::ExtFOp::create(*b, loc, resultTypes, args, ++ typename arith::ExtFOp::Properties{}, ++ /*discardableAttributes=*/{}); + } + // There's no direct conversion between different 16 bit floating point + // types, so go through 32 bit float. + if (sourceType != targetType) { + assert(sourceType.isBF16() || targetType.isBF16()); +- Value ext = arith::ExtFOp::create(*b, loc, b->getF32Type(), args); +- return arith::TruncFOp::create(*b, loc, resultTypes, ext); ++ Value ext = arith::ExtFOp::create(*b, loc, b->getF32Type(), args, ++ typename arith::ExtFOp::Properties{}, ++ /*discardableAttributes=*/{}); ++ return arith::TruncFOp::create(*b, loc, resultTypes, ext, ++ typename arith::TruncFOp::Properties{}, ++ /*discardableAttributes=*/{}); + } + // No conversion is needed for identical float types. + return args.front(); +@@ -738,17 +747,20 @@ + auto src = cast(sourceType); + auto res = cast(targetType); + if (src.getWidth() > res.getWidth()) { +- return mlir::arith::TruncIOp::create(*b, loc, resultTypes, args, +- ArrayRef()); ++ return arith::TruncIOp::create(*b, loc, resultTypes, args, ++ typename arith::TruncIOp::Properties{}, ++ /*discardableAttributes=*/{}); + } + if (src.getWidth() < res.getWidth()) { + // Special case boolean values, so they get casted to `1` instead of `-1`. + if (IsUnsignedIntegerType{}(src)) { +- return mlir::arith::ExtUIOp::create(*b, loc, resultTypes, args, +- ArrayRef()); ++ return arith::ExtUIOp::create(*b, loc, resultTypes, args, ++ typename arith::ExtUIOp::Properties{}, ++ /*discardableAttributes=*/{}); + } +- return mlir::arith::ExtSIOp::create(*b, loc, resultTypes, args, +- ArrayRef()); ++ return arith::ExtSIOp::create(*b, loc, resultTypes, args, ++ typename arith::ExtSIOp::Properties{}, ++ /*discardableAttributes=*/{}); + } + // No conversion is needed for the same width integers + return args.front(); diff --ruN a/stablehlo/stablehlo/conversions/tosa/tests/binary.mlir b/stablehlo/stablehlo/conversions/tosa/tests/binary.mlir --- stablehlo/stablehlo/conversions/tosa/tests/binary.mlir +++ stablehlo/stablehlo/conversions/tosa/tests/binary.mlir @@ -834,6 +937,21 @@ diff --ruN a/stablehlo/stablehlo/dialect/StablehloEnums.td b/stablehlo/stablehlo +} #endif // STABLEHLO_DIALECT_STABLEHLO_ENUMS +diff --ruN a/stablehlo/stablehlo/dialect/StablehloOps.td b/stablehlo/stablehlo/dialect/StablehloOps.td +--- stablehlo/stablehlo/dialect/StablehloOps.td ++++ stablehlo/stablehlo/dialect/StablehloOps.td +@@ -2525,8 +2525,9 @@ + "::mlir::TypeRange":$resultTypes, "::mlir::ValueRange":$inputs, + "::mlir::ArrayRef<::mlir::NamedAttribute>":$attributes + ), [{ +- build($_builder, $_state, resultTypes, inputs, attributes, +- /*regionsCount=*/0); ++ $_state.addOperands(inputs); ++ buildPropertiesAndDiscardableAttributes($_state, attributes); ++ $_state.addTypes(resultTypes); + }]>, + OpBuilder<(ins + "::mlir::TypeRange":$resultTypes, "::mlir::ValueRange":$inputs, diff --ruN a/stablehlo/stablehlo/dialect/VhloBytecode.h b/stablehlo/stablehlo/dialect/VhloBytecode.h --- stablehlo/stablehlo/dialect/VhloBytecode.h +++ stablehlo/stablehlo/dialect/VhloBytecode.h @@ -1115,7 +1233,7 @@ diff --ruN a/stablehlo/stablehlo/reference/Api.cpp b/stablehlo/stablehlo/referen public: DefaultInterpreterFallback(const InterpreterConfiguration& config) - : config(config){}; -+ : config(config) {}; ++ : config(config) {} virtual llvm::Error operator()(Operation& op, Scope& scope, Process* process) final { @@ -1856,6 +1974,20 @@ diff --ruN a/stablehlo/stablehlo/tests/CheckOps.h b/stablehlo/stablehlo/tests/Ch uint64_t min_ulp_difference, uint64_t max_ulp_difference); +diff --ruN a/stablehlo/stablehlo/tests/TestUtils.cpp b/stablehlo/stablehlo/tests/TestUtils.cpp +--- stablehlo/stablehlo/tests/TestUtils.cpp ++++ stablehlo/stablehlo/tests/TestUtils.cpp +@@ -60,7 +60,9 @@ + + // Replace with custom call to avoid pattern reapplication + auto customCall = stablehlo::CustomCallOp::create( +- rewriter, op->getLoc(), op->getResultTypes(), *broadcastedOperands); ++ rewriter, op->getLoc(), op->getResultTypes(), *broadcastedOperands, ++ typename stablehlo::CustomCallOp::Properties{}, ++ /*discardableAttributes=*/{}); + customCall.setCallTargetName("numpy_broadcasted"); + customCall.setHasSideEffect(true); + rewriter.replaceOp(op, customCall); 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 @@ -2231,6 +2363,24 @@ diff --ruN a/stablehlo/stablehlo/transforms/ChloDecompositionUtils.h b/stablehlo diff --ruN a/stablehlo/stablehlo/transforms/ChloLegalizeToStablehlo.cpp b/stablehlo/stablehlo/transforms/ChloLegalizeToStablehlo.cpp --- stablehlo/stablehlo/transforms/ChloLegalizeToStablehlo.cpp +++ stablehlo/stablehlo/transforms/ChloLegalizeToStablehlo.cpp +@@ -2520,7 +2520,7 @@ + } + Value dotGeneral = mlir::stablehlo::DotGeneralOp::create( + rewriter, op.getLoc(), resultType, ValueRange{lhsMasked, rhsReshaped}, +- attributes); ++ typename DotGeneralOp::Properties{}, attributes); + + if (i == 0) { + out = dotGeneral; +@@ -2672,7 +2672,7 @@ + } + Value dotGeneral = mlir::stablehlo::DotGeneralOp::create( + rewriter, op.getLoc(), RankedTensorType::get(resultShape, outDType), +- {lhsMasked, rhs}, attributes); ++ {lhsMasked, rhs}, typename DotGeneralOp::Properties{}, attributes); + + results.push_back(dotGeneral); + start = limit; @@ -2930,8 +2930,18 @@ SmallVector sliceShape(inputType.getShape().begin(), inputType.getShape().end()); @@ -2409,6 +2559,113 @@ diff --ruN a/stablehlo/stablehlo/transforms/StablehloComplexMathExpander.cpp b/s populateWithGenerated(*patterns); } +diff --ruN a/stablehlo/stablehlo/transforms/StablehloLegalizeQuantToMath.cpp b/stablehlo/stablehlo/transforms/StablehloLegalizeQuantToMath.cpp +--- stablehlo/stablehlo/transforms/StablehloLegalizeQuantToMath.cpp ++++ stablehlo/stablehlo/transforms/StablehloLegalizeQuantToMath.cpp +@@ -865,8 +865,9 @@ + Type resultType, Value& lhs, Value& rhs, + ArrayRef attrs, + const DotLikeDimensionNumbers& dims) { +- return stablehlo::DotGeneralOp::create(builder, loc, resultType, +- ArrayRef{lhs, rhs}, attrs); ++ return stablehlo::DotGeneralOp::create( ++ builder, loc, resultType, ArrayRef{lhs, rhs}, ++ typename DotGeneralOp::Properties{}, attrs); + } + + // Template specialization for Convolution op. +@@ -915,8 +916,9 @@ + } + } + } +- return stablehlo::ConvolutionOp::create(builder, loc, resultType, +- ArrayRef{lhs, rhs}, newAttrs); ++ return stablehlo::ConvolutionOp::create( ++ builder, loc, resultType, ArrayRef{lhs, rhs}, ++ typename ConvolutionOp::Properties{}, newAttrs); + } + + template +diff --ruN a/stablehlo/stablehlo/transforms/StablehloLegalizeQuantizedOpToQDQ.cpp b/stablehlo/stablehlo/transforms/StablehloLegalizeQuantizedOpToQDQ.cpp +--- stablehlo/stablehlo/transforms/StablehloLegalizeQuantizedOpToQDQ.cpp ++++ stablehlo/stablehlo/transforms/StablehloLegalizeQuantizedOpToQDQ.cpp +@@ -87,8 +87,9 @@ + llvm::map_to_vector(origOp->getResultTypes(), + [](Type t) { return getQuantExpressedType(t); }); + auto origAttrs = origOp->getAttrs(); +- auto newOp = StablehloOpType::create(rewriter, op.getLoc(), newResultTypes, +- dequantizedOperands, origAttrs) ++ auto newOp = StablehloOpType::create( ++ rewriter, op.getLoc(), newResultTypes, dequantizedOperands, ++ typename StablehloOpType::Properties{}, origAttrs) + .getOperation(); + + SmallVector quantizedResults; +diff --ruN a/stablehlo/stablehlo/transforms/StablehloLegalizeToVhlo.cpp b/stablehlo/stablehlo/transforms/StablehloLegalizeToVhlo.cpp +--- stablehlo/stablehlo/transforms/StablehloLegalizeToVhlo.cpp ++++ stablehlo/stablehlo/transforms/StablehloLegalizeToVhlo.cpp +@@ -1131,15 +1131,17 @@ + // This can almost be done in a generic fashion, except for + // stablehlo.case that uses a variadic number of regions which means an + // additional argument for the generic builder. +- StablehloToVhloOp vhloOp; ++ using VhloOpTy = StablehloToVhloOp; ++ VhloOpTy vhloOp; + if constexpr (StablehloOpTy::template hasTrait< + OpTrait::VariadicRegions>()) { +- vhloOp = StablehloToVhloOp::create( +- rewriter, stablehloOp.getLoc(), vhloTypes, vhloOperands, vhloAttrs, +- stablehloOp.getNumRegions()); ++ vhloOp = VhloOpTy::create(rewriter, stablehloOp.getLoc(), vhloTypes, ++ vhloOperands, typename VhloOpTy::Properties{}, ++ vhloAttrs, stablehloOp.getNumRegions()); + } else { +- vhloOp = StablehloToVhloOp::create( +- rewriter, stablehloOp.getLoc(), vhloTypes, vhloOperands, vhloAttrs); ++ vhloOp = VhloOpTy::create(rewriter, stablehloOp.getLoc(), vhloTypes, ++ vhloOperands, typename VhloOpTy::Properties{}, ++ vhloAttrs); + } + + for (auto [stablehloRegion, vhloRegion] : +diff --ruN a/stablehlo/stablehlo/transforms/StablehloRefineArguments.cpp b/stablehlo/stablehlo/transforms/StablehloRefineArguments.cpp +--- stablehlo/stablehlo/transforms/StablehloRefineArguments.cpp ++++ stablehlo/stablehlo/transforms/StablehloRefineArguments.cpp +@@ -57,7 +57,7 @@ + builder.getI64TensorAttr(refinedType.getShape())); + return stablehlo::CustomCallOp::create( + builder, operand.getLoc(), operand.getType(), +- ValueRange{operand, constant}, ++ ValueRange{operand, constant}, typename CustomCallOp::Properties{}, + llvm::SmallVector{ + builder.getNamedAttr( + "call_target_name", +diff --ruN a/stablehlo/stablehlo/transforms/VhloLegalizeToStablehlo.cpp b/stablehlo/stablehlo/transforms/VhloLegalizeToStablehlo.cpp +--- stablehlo/stablehlo/transforms/VhloLegalizeToStablehlo.cpp ++++ stablehlo/stablehlo/transforms/VhloLegalizeToStablehlo.cpp +@@ -1121,15 +1121,17 @@ + // Convert the VHLO operation to a StableHLO equivalent. This can almost be + // done in a generic fashion, except for ops with a variadic number of + // regions which means an additional argument for the generic builder. +- VhloToStablehloOp stablehloOp; ++ using StablehloOpTy = VhloToStablehloOp; ++ StablehloOpTy stablehloOp; + if constexpr (VhloOpTy::template hasTrait()) { +- stablehloOp = VhloToStablehloOp::create( ++ stablehloOp = StablehloOpTy::create( + rewriter, vhloOp.getLoc(), stablehloTypes, stablehloOperands, +- stablehloAttrs, vhloOp.getNumRegions()); ++ typename StablehloOpTy::Properties{}, stablehloAttrs, ++ vhloOp.getNumRegions()); + } else { +- stablehloOp = VhloToStablehloOp::create( ++ stablehloOp = StablehloOpTy::create( + rewriter, vhloOp.getLoc(), stablehloTypes, stablehloOperands, +- stablehloAttrs); ++ typename StablehloOpTy::Properties{}, stablehloAttrs); + } + + for (auto [vhloRegion, stablehloRegion] : diff --ruN a/stablehlo/stablehlo/transforms/optimization/Passes.h b/stablehlo/stablehlo/transforms/optimization/Passes.h --- stablehlo/stablehlo/transforms/optimization/Passes.h +++ stablehlo/stablehlo/transforms/optimization/Passes.h diff --git a/third_party/xla/third_party/tsl/tsl/platform/BUILD b/third_party/xla/third_party/tsl/tsl/platform/BUILD index c4570495a91101..9482149821cd3d 100644 --- a/third_party/xla/third_party/tsl/tsl/platform/BUILD +++ b/third_party/xla/third_party/tsl/tsl/platform/BUILD @@ -1214,6 +1214,7 @@ tsl_cc_test( ":blocking_counter", ":random", ":unbounded_work_queue", + "@com_google_absl//absl/memory", "@com_google_absl//absl/synchronization", "@com_google_googletest//:gtest_main", "@xla//xla/tsl/platform:env", diff --git a/third_party/xla/third_party/tsl/tsl/platform/unbounded_work_queue_test.cc b/third_party/xla/third_party/tsl/tsl/platform/unbounded_work_queue_test.cc index 32e5e32419ff49..ea24a27f3f5f5e 100644 --- a/third_party/xla/third_party/tsl/tsl/platform/unbounded_work_queue_test.cc +++ b/third_party/xla/third_party/tsl/tsl/platform/unbounded_work_queue_test.cc @@ -18,6 +18,7 @@ limitations under the License. #include #include +#include "absl/memory/memory.h" #include "absl/synchronization/mutex.h" #include "xla/tsl/platform/env.h" #include "xla/tsl/platform/test.h" @@ -104,26 +105,5 @@ TEST_F(UnboundedWorkQueueTest, RacyDestructor) { EXPECT_LE(NumClosuresExecuted(), num_closures); } -TEST(UnboundedWorkQueueOptionsTest, CustomStackSize) { - ThreadOptions thread_options; - thread_options.stack_size = 2 * 1024 * 1024; // 2 MB - UnboundedWorkQueue queue(Env::Default(), "custom_stack_test", thread_options); - BlockingCounter counter(10); - for (int i = 0; i < 10; ++i) { - queue.Schedule([&counter]() { counter.DecrementCount(); }); - } - counter.Wait(); -} - -TEST(UnboundedWorkQueueOptionsTest, DefaultThreadOptions) { - UnboundedWorkQueue queue(Env::Default(), "default_options_test", - ThreadOptions{}); - BlockingCounter counter(10); - for (int i = 0; i < 10; ++i) { - queue.Schedule([&counter]() { counter.DecrementCount(); }); - } - counter.Wait(); -} - } // namespace } // namespace tsl diff --git a/third_party/xla/xla/backends/cpu/nanort/ifrt_client.cc b/third_party/xla/xla/backends/cpu/nanort/ifrt_client.cc index 41c84a692fe175..18723e997480f0 100644 --- a/third_party/xla/xla/backends/cpu/nanort/ifrt_client.cc +++ b/third_party/xla/xla/backends/cpu/nanort/ifrt_client.cc @@ -1029,8 +1029,6 @@ class NanoExecutable final ifrt::UserContextRef user_context() const override { return user_context_; } - int num_devices() const override { return 1; } - int64_t SizeOfGeneratedCodeInBytes() const override { return 0; } absl::StatusOr GetCompiledMemoryStats() const override { @@ -1087,10 +1085,6 @@ class NanoExecutable final return absl::UnimplementedError("GetCostAnalysis is not implemented."); } - absl::Span addressable_devices() const override { - return client_->addressable_devices(); - } - std::optional devices() const override { return devices_; } diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/collective_emitter.cc b/third_party/xla/xla/backends/gpu/codegen/triton/collective_emitter.cc index f44abb7fc7ccf7..a40bcbd0aed1b3 100644 --- a/third_party/xla/xla/backends/gpu/codegen/triton/collective_emitter.cc +++ b/third_party/xla/xla/backends/gpu/codegen/triton/collective_emitter.cc @@ -681,8 +681,9 @@ class AllReduceEmitter { mlir::triton::gpu::BarrierOp::create(builder_, mlir::triton::gpu::AddrSpace::Local); mtx::BlockBarrierOp::create(builder_, signal_buffers_, device_rank_, - signal_value, - builder_.getI32IntegerAttr(ctx_.world_size)); + signal_value, /*signal_slot=*/nullptr, + builder_.getI32IntegerAttr(ctx_.world_size), + /*signal_stride=*/nullptr); return mlir::success(); } @@ -1317,8 +1318,9 @@ absl::Status EmitCollectiveEntryBarrier(mlir::ModuleOp module, // Inter-block barrier via signal flags. This blocks until all // remote ranks have also signaled. mtx::BlockBarrierOp::create(builder, signal_buffers_arg, rank_arg, - signal_value, - builder.getI32IntegerAttr(world_size)); + signal_value, /*signal_slot=*/nullptr, + builder.getI32IntegerAttr(world_size), + /*signal_stride=*/nullptr); return absl::OkStatus(); } diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/ir/tests/invalid.mlir b/third_party/xla/xla/backends/gpu/codegen/triton/ir/tests/invalid.mlir index a6bb30ae570008..bf680cada9aaeb 100644 --- a/third_party/xla/xla/backends/gpu/codegen/triton/ir/tests/invalid.mlir +++ b/third_party/xla/xla/backends/gpu/codegen/triton/ir/tests/invalid.mlir @@ -75,3 +75,20 @@ func.func @insert_wrong_shape(%arg0: tensor<8xbf16>, %arg1: !tt.ptr) { triton_xla.insert %arg0 into %arg1 as memref<16xbf16, #xtile.layout<[0]>> [0][16][1] : tensor<8xbf16> return } + +// ----- + +func.func @block_barrier_stride_without_slot(%ptr: !tt.ptr, %rank: i32, %sig: i32) { + // expected-error @+1 {{'triton_xla.block_barrier' op signal_stride requires signal_slot to be specified}} + triton_xla.block_barrier %ptr, %rank, %sig, {world_size = 2 : i32, signal_stride = 4 : i32} : (!tt.ptr, i32, i32) -> () + return +} + +// ----- + +func.func @block_barrier_slot_without_stride(%ptr: !tt.ptr, %rank: i32, %sig: i32, %slot: i32) { + // expected-error @+1 {{'triton_xla.block_barrier' op signal_slot requires a positive signal_stride}} + triton_xla.block_barrier %ptr, %rank, %sig, %slot, {world_size = 2 : i32} : (!tt.ptr, i32, i32, i32) -> () + return +} + diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/ir/triton_xla_ops.cc b/third_party/xla/xla/backends/gpu/codegen/triton/ir/triton_xla_ops.cc index 7f25bf5bcbd5e8..98ae71a0bb3c3f 100644 --- a/third_party/xla/xla/backends/gpu/codegen/triton/ir/triton_xla_ops.cc +++ b/third_party/xla/xla/backends/gpu/codegen/triton/ir/triton_xla_ops.cc @@ -323,6 +323,19 @@ LogicalResult PtrToMemrefOp::verify() { return success(); } +LogicalResult BlockBarrierOp::verify() { + if (getSignalStride() < 0) { + return emitOpError("signal_stride must be non-negative"); + } + if (getSignalStride() > 0 && !getSignalSlot()) { + return emitOpError("signal_stride requires signal_slot to be specified"); + } + if (getSignalSlot() && getSignalStride() == 0) { + return emitOpError("signal_slot requires a positive signal_stride"); + } + return success(); +} + } // namespace mlir::triton::xla #define GET_OP_CLASSES diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/ir/triton_xla_ops.td b/third_party/xla/xla/backends/gpu/codegen/triton/ir/triton_xla_ops.td index a65644be2066a3..27bbca78c5aad7 100644 --- a/third_party/xla/xla/backends/gpu/codegen/triton/ir/triton_xla_ops.td +++ b/third_party/xla/xla/backends/gpu/codegen/triton/ir/triton_xla_ops.td @@ -331,8 +331,40 @@ def TTXLA_BlockBarrierOp : TTXLA_Op<"block_barrier", []> { It takes a pointer of pointers to signal buffers (one from each participating GPU), the rank of the current device, and a signal value to write/wait for. The signal values on consequent calls to this operation must - increase monotonically for the same signal buffers pointer. A block on one - GPU synchronizes with blocks on other GPUs with the same block id. + increase monotonically for the same signal buffers pointer. + + Each GPU `r`'s signal buffer (`$signal_buffers[r]`) is logically a 2D array + of `i32` flags with shape `[num_blocks][world_size]`, flattened as + `block_idx * world_size + slot_idx`. + + This barrier supports two synchronization modes: + + ### 1. Symmetric Mode (`$signal_stride == 0`, `$signal_slot` omitted) + Used when block `block_idx` on every GPU produces and consumes data for the + same `block_idx`: + - **Signal (Write)**: Block `block_idx` on GPU `r` atomically writes + `$signal_value` to index `block_idx * world_size + r` in every peer GPU's + signal buffer (`$signal_buffers[0..world_size)`). + - **Wait (Read)**: Block `block_idx` on GPU `r` spin-waits on its + own buffer `$signal_buffers[r]` until all `world_size` slots at + `block_idx * world_size + [0..world_size)` reach `>= $signal_value`. + + ### 2. Asymmetric Mode (`$signal_stride > 0`, `$signal_slot` provided). + Used when `world_size` blocks on the same GPU (spaced by `$signal_stride` in + block ID) each produce `$signal_slot` (`0 <= $signal_slot < world_size`) + of a tile that will be consumed by a single target block on every peer GPU: + - Let `base_block = block_idx - $signal_slot * $signal_stride` be the base + block index within the stride, and + `target_block = base_block + r * $signal_stride` be the + consumer block ID on every GPU that reads the tile produced on + `r`. + - **Signal (Write)**: Block `block_idx` on GPU `r` atomically writes + `$signal_value` to index `target_block * world_size + $signal_slot` in + every peer GPU's signal buffer (`$signal_buffers[0..world_size)`). + - **Wait (Read)**: Block `block_idx` on GPU `r` spin-waits on its + own buffer `$signal_buffers[r]` until all `world_size` producer + slots at `block_idx * world_size + [0..world_size)` reach + `>= $signal_value`. *NOTE* Block orchestration order is not guaranteed across GPUs, so there is a chance of a deadlock should the number of blocks be greater @@ -347,17 +379,6 @@ def TTXLA_BlockBarrierOp : TTXLA_Op<"block_barrier", []> { of participating GPUs would usually be much smaller than the number of threads per block. - The synchronization logic is as follows: - 1. Each GPU (identified by `device_rank`) writes the `signal` value to an - index `blockId * world_size + device_rank` in *every* buffer provided in - `$signal_buffers`. - 2. After writing, GPU `r` polls its own buffer (`$signal_buffers[r]`). - 3. It waits until the values at locations corresponding to all other ranks - (i.e., `blockId * world_size + [0..world_size)` are greater than or equal - to the `signal` value. - 4. When the wait is over, all GPUs working on the same `blockId` have - successfully synchronized. - This operation ensures that GPUs in the same synchronization group do not get more than one "wave" (or chunk) of data processing ahead of each other. For clarity, a wave here is a set of work items that are processed together. @@ -376,19 +397,25 @@ def TTXLA_BlockBarrierOp : TTXLA_Op<"block_barrier", []> { I32:$device_rank, // The signal value to write/wait for. I32:$signal_value, + // Optional producer slot index in [0, world_size) for asymmetric barriers. + Optional:$signal_slot, // The number of devices participating in the barrier. // The length of the first dimension of the signal buffers should be the // same as this value. - I32Attr:$world_size + I32Attr:$world_size, + // Optional block-ID stride between consumer blocks for asymmetric barriers. + DefaultValuedOptionalAttr:$signal_stride ); // This is a barrier op and produces no results. let results = (outs); let assemblyFormat = [{ - $signal_buffers `,` $device_rank `,` $signal_value `,` attr-dict `:` + $signal_buffers `,` $device_rank `,` $signal_value `,` ($signal_slot^ `,`)? attr-dict `:` functional-type(operands, results) }]; + + let hasVerifier = 1; } def TTXLA_GetRankOp : TTXLA_Op<"get_rank", [Pure]> { diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/reduce/reduce_concat_diamond.hlo b/third_party/xla/xla/backends/gpu/codegen/triton/tests/reduce/reduce_concat_diamond.hlo new file mode 100644 index 00000000000000..5046db172c47dd --- /dev/null +++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/reduce/reduce_concat_diamond.hlo @@ -0,0 +1,46 @@ +// RUN: fusion_to_triton --xla_gpu_experimental_enable_tiling_propagation %s | FileCheck %s +// RUN: triton_test_correctness --xla_gpu_experimental_enable_tiling_propagation --abs_error_bound=1e-4 --rel_error_bound=1e-4 %s + +add_region { + p0 = f32[] parameter(0) + p1 = f32[] parameter(1) + ROOT add = f32[] add(p0, p1) +} + +fusion { + p = f32[16,32] parameter(0) + c_zero = f32[] constant(0) + red = f32[16] reduce(p, c_zero), dimensions={1}, to_apply=add_region, backend_config={sizes:[8]} + bcast = f32[16,32] broadcast(red), dimensions={0} + x = f32[16,32] add(p, bcast) + slice0 = f32[16,16] slice(x), slice={[0:16], [0:16]} + slice1 = f32[16,16] slice(x), slice={[0:16], [16:32]} + ROOT concat = f32[16,32] concatenate(slice0, slice1), dimensions={1} +} + +ENTRY main { + p = f32[16,32] parameter(0) + ROOT r = f32[16,32] fusion(p), kind=kCustom, calls=fusion, + backend_config={ + "fusion_backend_config":{ + "kind":"__triton", + "block_level_fusion_config":{ + "output_tiles":[{"sizes":["4", "16"]}], + "num_warps":"1", + "num_ctas":"1", + "num_stages":"1"}}} +} + +// `red` feeds both concat operands, so tiling duplicates it into both concat +// regions and its reduction loop is emitted once per `scf.if` branch. + +// CHECK-LABEL: xtile.entry_func @triton_fn +// CHECK: %[[IF:.*]] = scf.if +// CHECK: %[[THEN_ACC:.*]] = scf.for +// CHECK: arith.addf +// CHECK: "tt.reduce"(%[[THEN_ACC]]) +// CHECK: } else { +// CHECK: %[[ELSE_ACC:.*]] = scf.for +// CHECK: arith.addf +// CHECK: "tt.reduce"(%[[ELSE_ACC]]) +// CHECK: xtile.insert %[[IF]] into %arg1 diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/transforms/tests/triton_xla_block_barrier.mlir b/third_party/xla/xla/backends/gpu/codegen/triton/transforms/tests/triton_xla_block_barrier.mlir index a14aa686b49e5d..7fcbfbf52b4d08 100644 --- a/third_party/xla/xla/backends/gpu/codegen/triton/transforms/tests/triton_xla_block_barrier.mlir +++ b/third_party/xla/xla/backends/gpu/codegen/triton/transforms/tests/triton_xla_block_barrier.mlir @@ -48,4 +48,47 @@ tt.func @block_barrier_kernel( triton_xla.block_barrier %ptr, %rank, %signal_value, { world_size = 8 : i32 } : (!tt.ptr, i32, i32) -> () tt.return +} + +// CHECK-LABEL: block_barrier_asymmetric_kernel +// CHECK-SAME: %[[PTR:[^:]+]]: !tt.ptr, +// CHECK-SAME: %[[SIGNAL_VALUE:[^:]+]]: i32, +// CHECK-SAME: %[[RANK:[^:]+]]: i32, +// CHECK-SAME: %[[SIGNAL_SLOT:[^:]+]]: i32 +tt.func @block_barrier_asymmetric_kernel( + %ptr : !tt.ptr, %signal_value : i32, %rank: i32, %signal_slot: i32) { + // CHECK-NEXT: %[[STRIDE:.+]] = arith.constant 4 : i32 + // CHECK-NEXT: %[[WORLD_SIZE:.+]] = arith.constant 2 : i32 + // CHECK-NEXT: %[[TID:.+]] = triton_xla.get_tid + // CHECK-NEXT: %[[PROGRAM_ID:.+]] = tt.get_program_id x + // CHECK-NEXT: %[[COND:.+]] = arith.cmpi ult, %[[TID]], %[[WORLD_SIZE]] + // CHECK-NEXT: scf.if %[[COND]] { + // CHECK-NEXT: %[[SPLAT_PTR:.+]] = tt.splat %[[PTR]] + // CHECK-NEXT: %[[RANGE:.+]] = tt.make_range {end = 2 : i32, start = 0 : i32} + // CHECK-NEXT: %[[ADD_PTR:.+]] = tt.addptr %[[SPLAT_PTR]], %[[RANGE]] + // CHECK-NEXT: %[[LOAD:.+]] = tt.load %[[ADD_PTR]] + // CHECK-NEXT: %[[INT_TO_PTR:.+]] = tt.int_to_ptr %[[LOAD]] + // CHECK-NEXT: %[[MULI:.+]] = arith.muli %[[PROGRAM_ID]], %[[WORLD_SIZE]] + // CHECK-NEXT: %[[S_OFFSET:.+]] = arith.muli %[[SIGNAL_SLOT]], %[[STRIDE]] + // CHECK-NEXT: %[[BASE_T:.+]] = arith.subi %[[PROGRAM_ID]], %[[S_OFFSET]] + // CHECK-NEXT: %[[RANK_OFFSET:.+]] = arith.muli %[[RANK]], %[[STRIDE]] + // CHECK-NEXT: %[[WRITE_BLOCK:.+]] = arith.addi %[[BASE_T]], %[[RANK_OFFSET]] + // CHECK-NEXT: %[[WRITE_BLOCK_OFFSET:.+]] = arith.muli %[[WRITE_BLOCK]], %[[WORLD_SIZE]] + // CHECK-NEXT: %[[WRITE_OFFSET:.+]] = arith.addi %[[WRITE_BLOCK_OFFSET]], %[[SIGNAL_SLOT]] + // CHECK-NEXT: %[[SPLAT_WRITE_OFFSET:.+]] = tt.splat %[[WRITE_OFFSET]] + // CHECK-NEXT: %[[ADD_PTR_2:.+]] = tt.addptr %[[INT_TO_PTR]], %[[SPLAT_WRITE_OFFSET]] + // CHECK-NEXT: triton_xla.atomic_write sys, release, %[[ADD_PTR_2]], %[[SIGNAL_VALUE]] + // CHECK-NEXT: %[[ADD_PTR_3:.+]] = tt.addptr %[[PTR]], %[[RANK]] + // CHECK-NEXT: %[[LOAD_2:.+]] = tt.load %[[ADD_PTR_3]] + // CHECK-NEXT: %[[INT_TO_PTR_2:.+]] = tt.int_to_ptr %[[LOAD_2]] + // CHECK-NEXT: %[[ADD_PTR_4:.+]] = tt.addptr %[[INT_TO_PTR_2]], %[[MULI]] + // CHECK-NEXT: %[[SPLAT_ADD_PTR_4:.+]] = tt.splat %[[ADD_PTR_4]] + // CHECK-NEXT: %[[ADD_PTR_5:.+]] = tt.addptr %[[SPLAT_ADD_PTR_4]], %[[RANGE]] + // CHECK-NEXT: triton_xla.atomic_spin_wait sys, acquire, %[[ADD_PTR_5]], greater_than_or_equal_to, %[[SIGNAL_VALUE]] + // CHECK-NEXT: } + // CHECK-NEXT: ttg.barrier local + // CHECK-NEXT: tt.return + triton_xla.block_barrier %ptr, %rank, %signal_value, %signal_slot, { world_size = 2 : i32, signal_stride = 4 : i32 } : + (!tt.ptr, i32, i32, i32) -> () + tt.return } \ No newline at end of file diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/transforms/triton_xla_lower_block_barrier_pass.cc b/third_party/xla/xla/backends/gpu/codegen/triton/transforms/triton_xla_lower_block_barrier_pass.cc index 13bded4da392c4..41f48e48010301 100644 --- a/third_party/xla/xla/backends/gpu/codegen/triton/transforms/triton_xla_lower_block_barrier_pass.cc +++ b/third_party/xla/xla/backends/gpu/codegen/triton/transforms/triton_xla_lower_block_barrier_pass.cc @@ -51,6 +51,7 @@ LogicalResult LowerBlockBarrierOp(BlockBarrierOp block_barrier, const mlir::TypedValue signal_value = block_barrier.getSignalValue(); const int32_t world_size = block_barrier.getWorldSize(); + const int32_t signal_stride = block_barrier.getSignalStride(); // Triton magic constant. constexpr auto kGlobalAddressSpace = mlir::triton::PtrAddrSpace::Global; @@ -126,11 +127,51 @@ LogicalResult LowerBlockBarrierOp(BlockBarrierOp block_barrier, builder, tensor_of_ptr_to_i32_type, signal_buffer_i64); auto block_offset = mlir::arith::MulIOp::create( builder, i32_type, block_id, world_size_op); - auto block_offset_plus_rank = - mlir::arith::AddIOp::create(builder, i32_type, block_offset, rank); + mlir::Value write_offset; + if (signal_stride > 0) { + // Suppose we have 4 blocks (0, 1, 2, 3), 2 ranks (world_size = 2), + // and signal_stride = 2. + // + // Blocks 1 and 3 work on the same scratch tile and are + // `signal_stride` blocks apart: + // - Block 1 has signal_slot = 0 (first block in the pair). + // - Block 3 has signal_slot = 1 (second block in the pair). + // + // 1. Find `base_block`, the first block in the pair (block 1): + // `slot_offset = signal_slot * signal_stride` is how many blocks + // this block is ahead of `base_block`. + // Subtracting it gives `base_block = block_id - slot_offset`: + // - Block 1: base_block = 1 - (0 * 2) = 1 + // - Block 3: base_block = 3 - (1 * 2) = 1 + // + // 2. Find `target_block`, the block on peer GPUs that will read this + // rank's data (`base_block + rank * signal_stride`). + // + // 3. Write to `signal_slot` inside `target_block`'s slots + // (`target_block * world_size + signal_slot`), so blocks 1 and 3 + // land in adjacent slots 0 and 1 for `target_block`. + mlir::Value signal_slot = block_barrier.getSignalSlot(); + auto stride_op = mlir::arith::ConstantOp::create( + builder, builder.getI32IntegerAttr(signal_stride)); + auto slot_offset = mlir::arith::MulIOp::create( + builder, i32_type, signal_slot, stride_op); + auto base_block = mlir::arith::SubIOp::create(builder, i32_type, + block_id, slot_offset); + auto rank_offset = + mlir::arith::MulIOp::create(builder, i32_type, rank, stride_op); + auto target_block = mlir::arith::AddIOp::create( + builder, i32_type, base_block, rank_offset); + auto target_block_offset = mlir::arith::MulIOp::create( + builder, i32_type, target_block, world_size_op); + write_offset = mlir::arith::AddIOp::create( + builder, i32_type, target_block_offset, signal_slot); + } else { + write_offset = mlir::arith::AddIOp::create(builder, i32_type, + block_offset, rank); + } // -> tensor auto block_offset_plus_rank_tensor = mlir::triton::SplatOp::create( - builder, i32_tensor_type, block_offset_plus_rank); + builder, i32_tensor_type, write_offset); // SignalBuffers[0..WorldSize][block_id][rank] // -> tensor> auto signal_addresses = mlir::triton::AddPtrOp::create( diff --git a/third_party/xla/xla/backends/gpu/runtime/BUILD b/third_party/xla/xla/backends/gpu/runtime/BUILD index 0d3efec73c29e9..8fcdd36e3b9583 100644 --- a/third_party/xla/xla/backends/gpu/runtime/BUILD +++ b/third_party/xla/xla/backends/gpu/runtime/BUILD @@ -1955,6 +1955,8 @@ cc_library( ":collective_params", ":collective_thunk", ":collective_thunk_proto_cc", + ":command", + ":command_state", ":thunk", ":thunk_proto_cc", ":traced_command", @@ -1974,6 +1976,7 @@ cc_library( "//xla/service/gpu:gpu_constants", "//xla/service/gpu:launch_dimensions", "//xla/service/gpu:stream_executor_util", + "//xla/stream_executor:command_buffer", "//xla/stream_executor:device_address", "//xla/stream_executor:device_address_handle", "//xla/stream_executor:kernel", @@ -1981,6 +1984,7 @@ cc_library( "//xla/stream_executor:memory_space", "//xla/stream_executor:stream", "//xla/stream_executor:stream_executor_h", + "//xla/stream_executor:trace_command_buffer_factory", "//xla/stream_executor/gpu:all_reduce_kernel", "//xla/stream_executor/gpu:collective_kernel_metadata", "//xla/tsl/util:safe_reinterpret_cast", @@ -1988,6 +1992,7 @@ cc_library( "@com_google_absl//absl/algorithm:container", "@com_google_absl//absl/base:core_headers", "@com_google_absl//absl/container:flat_hash_map", + "@com_google_absl//absl/functional:overload", "@com_google_absl//absl/log", "@com_google_absl//absl/log:vlog_is_on", "@com_google_absl//absl/status", diff --git a/third_party/xla/xla/backends/gpu/runtime/collective_kernel_thunk.cc b/third_party/xla/xla/backends/gpu/runtime/collective_kernel_thunk.cc index 223d7479a5f77a..3a7060359a1f1c 100644 --- a/third_party/xla/xla/backends/gpu/runtime/collective_kernel_thunk.cc +++ b/third_party/xla/xla/backends/gpu/runtime/collective_kernel_thunk.cc @@ -20,10 +20,12 @@ limitations under the License.*/ #include #include #include +#include #include #include "absl/algorithm/container.h" #include "absl/container/flat_hash_map.h" +#include "absl/functional/overload.h" #include "absl/log/log.h" #include "absl/log/vlog_is_on.h" #include "absl/status/status.h" @@ -45,6 +47,7 @@ limitations under the License.*/ #include "xla/backends/gpu/runtime/collective_params.h" #include "xla/backends/gpu/runtime/collective_thunk.h" #include "xla/backends/gpu/runtime/collective_thunk.pb.h" +#include "xla/backends/gpu/runtime/command_state.h" #include "xla/backends/gpu/runtime/thunk.h" #include "xla/backends/gpu/runtime/thunk.pb.h" #include "xla/core/collectives/rank_id.h" @@ -58,6 +61,7 @@ limitations under the License.*/ #include "xla/service/gpu/stream_executor_util.h" #include "xla/shape_util.h" #include "xla/status_macros.h" +#include "xla/stream_executor/command_buffer.h" #include "xla/stream_executor/device_address.h" #include "xla/stream_executor/device_address_handle.h" #include "xla/stream_executor/gpu/collective_kernel_metadata.h" @@ -66,6 +70,7 @@ limitations under the License.*/ #include "xla/stream_executor/memory_space.h" #include "xla/stream_executor/stream.h" #include "xla/stream_executor/stream_executor.h" +#include "xla/stream_executor/trace_command_buffer_factory.h" #include "xla/tsl/util/safe_reinterpret_cast.h" #include "xla/tsl/util/tied_ref.h" #include "xla/util.h" @@ -288,6 +293,20 @@ absl::StatusOr> CollectParamToPeers( return param_to_peers_ptrs; } +// Recorded input -> symmetric scratch copy that precedes a kernel launch node. +struct RecordedCopy { + const se::CommandBuffer::Command* command = nullptr; + void* dst = nullptr; + void* src = nullptr; +}; + +// Per command buffer state: maps the recorded kernel launch node (the command +// returned from Record) to the copy command that precedes it, so both can be +// updated in place. +struct CollectiveKernelRecordState : public CommandState { + absl::flat_hash_map + launch_to_copy; +}; } // namespace absl::Status CollectiveKernelThunk::IsSupported( @@ -629,8 +648,8 @@ absl::Status CollectiveKernelThunk::Initialize(const InitializeParams& params) { return absl::OkStatus(); } -absl::Status CollectiveKernelThunk::ExecuteOnStream( - const ExecuteParams& params) { +absl::StatusOr +CollectiveKernelThunk::PrepareLaunch(const ExecuteParams& params) { se::Stream* stream = params.stream; TF_RET_CHECK(stream != nullptr); ABSL_ASSIGN_OR_RETURN( @@ -668,8 +687,10 @@ absl::Status CollectiveKernelThunk::ExecuteOnStream( has_multimem) + kernel_spec_.scratch_buffers.size(); + LaunchPlan plan; + plan.kernel = state->kernel.get(); ABSL_ASSIGN_OR_RETURN( - std::vector kernel_args, + plan.kernel_args, BuildKernelArguments(kernel_spec_, buffers_, params, state->rank, state->invocation_count, state->metadata, clique_key, num_parameters)); @@ -687,21 +708,104 @@ absl::Status CollectiveKernelThunk::ExecuteOnStream( // rank's input slice. scratch_allocations[0] is the signal buffer. TF_RET_CHECK(memory_state->scratch_allocations.size() >= 2) << "Expected at least 2 scratch allocations for D2D copy."; - se::DeviceAddressBase input_addr = - params.buffer_allocations->GetDeviceAddress( - buffers_[0].source_buffer.slice); - const int64_t copy_size = GetInputSizeBytes(); + plan.copy_dst = memory_state->scratch_allocations[1].address_ptr(); + plan.copy_src = params.buffer_allocations->GetDeviceAddress( + buffers_[0].source_buffer.slice); + plan.copy_size = GetInputSizeBytes(); + } + return plan; +} + +absl::Status CollectiveKernelThunk::ExecuteOnStream( + const ExecuteParams& params) { + ABSL_ASSIGN_OR_RETURN(LaunchPlan plan, PrepareLaunch(params)); + if (plan.copy_dst != nullptr) { ABSL_RETURN_IF_ERROR( - stream->Memcpy(memory_state->scratch_allocations[1].address_ptr(), - input_addr, copy_size)); - VLOG(3) << "D2D copy: src=" << input_addr.opaque() << " dst=" - << memory_state->scratch_allocations[1].address().opaque() - << " size=" << copy_size - << " invocation=" << state->invocation_count; + params.stream->Memcpy(plan.copy_dst, plan.copy_src, plan.copy_size)); } + return ExecuteKernelOnStream(*plan.kernel, plan.kernel_args, + launch_dimensions_, + /*cluster_dim=*/std::nullopt, params.stream); +} - return ExecuteKernelOnStream(*state->kernel, kernel_args, launch_dimensions_, - /*cluster_dim=*/std::nullopt, stream); +absl::StatusOr CollectiveKernelThunk::Record( + const Thunk::ExecuteParams& execute_params, + const RecordParams& record_params, RecordAction record_action, + se::CommandBuffer* command_buffer) { + // The kernel is recorded as an explicit graph kernel node. For kernels with a + // host-side sync count PrepareLaunch bumps the per-stream invocation_count, + // so every create/update carries the up-to-date barrier signal value in the + // kernel node parameters. Kernels with a device-side sync count get constant + // arguments and only need updates when buffer addresses change. + ABSL_ASSIGN_OR_RETURN(LaunchPlan plan, PrepareLaunch(execute_params)); + + ABSL_ASSIGN_OR_RETURN(std::unique_ptr packed_args, + se::PackKernelArgs(absl::MakeConstSpan(plan.kernel_args), + plan.kernel->metadata())); + auto* state = record_params.state.GetOrCreate( + this, command_buffer); + + // The input -> symmetric scratch copy is recorded as a (tiny) stream-captured + // child graph. An explicit memcpy node targeting symmetric (collective) + // memory crashes CUPTI during graph instantiation when profiling is enabled. + // Its parameters only depend on buffer addresses, so it is re-captured only + // when those change. + auto trace_copy = + [&]() -> absl::StatusOr> { + return se::TraceCommandBufferFactory::Create( + execute_params.stream->parent(), + execute_params.command_buffer_trace_stream, [&](se::Stream* stream) { + return stream->Memcpy(plan.copy_dst, plan.copy_src, plan.copy_size); + }); + }; + + auto create = [&](const RecordCreate& action) + -> absl::StatusOr { + absl::Span deps = + action.dependencies; + const se::CommandBuffer::Command* copy_cmd = nullptr; + if (plan.copy_dst != nullptr) { + ABSL_ASSIGN_OR_RETURN(std::unique_ptr nested, trace_copy()); + ABSL_ASSIGN_OR_RETURN(copy_cmd, + command_buffer->CreateChildCommand(*nested, deps)); + deps = absl::MakeConstSpan(©_cmd, 1); + } + ABSL_ASSIGN_OR_RETURN( + const se::CommandBuffer::Command* launch_cmd, + command_buffer->CreateLaunch( + launch_dimensions_.thread_counts_per_block(), + launch_dimensions_.block_counts(), /*cluster_dims=*/std::nullopt, + *plan.kernel, *packed_args, deps, priority())); + if (copy_cmd != nullptr) { + state->launch_to_copy[launch_cmd] = RecordedCopy{ + copy_cmd, plan.copy_dst->opaque(), plan.copy_src.opaque()}; + } + return launch_cmd; + }; + auto update = [&](const RecordUpdate& action) + -> absl::StatusOr { + if (plan.copy_dst != nullptr) { + auto it = state->launch_to_copy.find(action.command); + TF_RET_CHECK(it != state->launch_to_copy.end()) + << "Missing recorded D2D copy command for collective kernel update"; + RecordedCopy& copy = it->second; + if (copy.dst != plan.copy_dst->opaque() || + copy.src != plan.copy_src.opaque()) { + ABSL_ASSIGN_OR_RETURN(std::unique_ptr nested, + trace_copy()); + ABSL_RETURN_IF_ERROR( + command_buffer->UpdateChildCommand(copy.command, *nested)); + copy.dst = plan.copy_dst->opaque(); + copy.src = plan.copy_src.opaque(); + } + } + ABSL_RETURN_IF_ERROR(command_buffer->UpdateLaunch( + action.command, launch_dimensions_.thread_counts_per_block(), + launch_dimensions_.block_counts(), /*cluster_dims=*/std::nullopt, + *plan.kernel, *packed_args)); + return action.command; + }; + return std::visit(absl::Overload(create, update), record_action); } Thunk::BufferUses CollectiveKernelThunk::buffer_uses() const { diff --git a/third_party/xla/xla/backends/gpu/runtime/collective_kernel_thunk.h b/third_party/xla/xla/backends/gpu/runtime/collective_kernel_thunk.h index 760b3643accba4..3ff46900d050a8 100644 --- a/third_party/xla/xla/backends/gpu/runtime/collective_kernel_thunk.h +++ b/third_party/xla/xla/backends/gpu/runtime/collective_kernel_thunk.h @@ -33,6 +33,7 @@ limitations under the License.*/ #include "xla/backends/gpu/runtime/collective_kernel_thunk.pb.h" #include "xla/backends/gpu/runtime/collective_params.h" #include "xla/backends/gpu/runtime/collective_thunk.h" +#include "xla/backends/gpu/runtime/command.h" #include "xla/backends/gpu/runtime/thunk.h" #include "xla/backends/gpu/runtime/thunk.pb.h" #include "xla/backends/gpu/runtime/traced_command.h" @@ -40,10 +41,12 @@ limitations under the License.*/ #include "xla/core/collectives/symmetric_memory.h" #include "xla/service/buffer_assignment.h" #include "xla/service/gpu/launch_dimensions.h" +#include "xla/stream_executor/command_buffer.h" #include "xla/stream_executor/device_address.h" #include "xla/stream_executor/device_address_handle.h" #include "xla/stream_executor/gpu/all_reduce_kernel.h" #include "xla/stream_executor/kernel.h" +#include "xla/stream_executor/kernel_args.h" #include "xla/stream_executor/stream.h" #include "xla/tsl/util/tied_ref.h" @@ -118,6 +121,25 @@ class CollectiveKernelThunk : public TracedCommand { // Execute the kernel on all devices. absl::Status ExecuteOnStream(const ExecuteParams& params) final; + // Records the collective kernel into a command buffer as an explicit kernel + // node, preceded by a (captured) input-to-scratch copy. On update the nodes + // are updated in place. + absl::StatusOr Record( + const Thunk::ExecuteParams& execute_params, + const RecordParams& record_params, RecordAction record_action, + se::CommandBuffer* command_buffer) override; + + // Kernels that take the barrier signal value from the host-side + // `invocation_count` argument must be updated on every execution. Kernels + // with a device-side sync count keep the counter in device memory, so their + // recorded command only changes when buffer addresses change. + bool requires_update_on_execute() const override { + return !kernel_spec_.codegen_config.device_sync_count; + } + + // Recorded with explicit graph construction APIs, not stream tracing. + bool IsTracedCommand() const override { return false; } + BufferUses buffer_uses() const override; static absl::StatusOr> FromProto( @@ -161,6 +183,20 @@ class CollectiveKernelThunk : public TracedCommand { // Returns the input size in bytes for the collective. int64_t GetInputSizeBytes() const; + // Everything needed to launch the collective kernel once, either eagerly on + // a stream or as command buffer nodes. + struct LaunchPlan { + se::Kernel* kernel = nullptr; + std::vector kernel_args; + // Optional input -> symmetric scratch copy performed before the kernel. + se::DeviceAddressBase* copy_dst = nullptr; + se::DeviceAddressBase copy_src; + int64_t copy_size = 0; + }; + + // Bumps the per-stream invocation count and builds the launch plan. + absl::StatusOr PrepareLaunch(const ExecuteParams& params); + // Whether the one-shot kernel is enabled. const bool collective_kernel_enabled_; // Collective config being used. Copied over to avoid lifetime issues. diff --git a/third_party/xla/xla/backends/gpu/runtime/command_buffer_conversion_pass.cc b/third_party/xla/xla/backends/gpu/runtime/command_buffer_conversion_pass.cc index ba4312fcbaafb5..fbc9a936b4d288 100644 --- a/third_party/xla/xla/backends/gpu/runtime/command_buffer_conversion_pass.cc +++ b/third_party/xla/xla/backends/gpu/runtime/command_buffer_conversion_pass.cc @@ -79,6 +79,7 @@ std::optional GetCollectiveOpType( case Thunk::kAllGather: return DebugOptions::ALLGATHER; case Thunk::kAllReduce: + case Thunk::kCollectiveKernel: return DebugOptions::ALLREDUCE; case Thunk::kAllToAll: return DebugOptions::ALLTOALL; @@ -213,6 +214,7 @@ std::optional GetCommandBufferCmdType( case Thunk::kAllReduce: case Thunk::kAllToAll: case Thunk::kCollectiveBroadcast: + case Thunk::kCollectiveKernel: case Thunk::kCollectivePermute: case Thunk::kCollectiveReduce: case Thunk::kGroup: diff --git a/third_party/xla/xla/backends/gpu/runtime/topk_test.cc b/third_party/xla/xla/backends/gpu/runtime/topk_test.cc index 0977af2a877258..228585032e1c8a 100644 --- a/third_party/xla/xla/backends/gpu/runtime/topk_test.cc +++ b/third_party/xla/xla/backends/gpu/runtime/topk_test.cc @@ -143,6 +143,9 @@ TEST_P(TopKKernelTest, TopKFloat) { EXPECT_THAT(got, ::testing::ElementsAreArray(slice)) << " k=" << k << ", batch_size=" << batch_size << " i=" << i; } + executor->Deallocate(&input_buffer); + executor->Deallocate(&output_values); + executor->Deallocate(&output_indices); } TEST_P(TopKKernelTest, TopKPackedNegative) { @@ -204,6 +207,9 @@ TEST_P(TopKKernelTest, TopKPackedNegative) { EXPECT_THAT(got, ::testing::ElementsAreArray(slice)) << " k=" << k << ", batch_size=" << batch_size << " i=" << i; } + executor->Deallocate(&input_buffer); + executor->Deallocate(&output_values); + executor->Deallocate(&output_indices); } TEST_P(TopKKernelTest, EnsureSerializable) { diff --git a/third_party/xla/xla/backends/gpu/tests/BUILD b/third_party/xla/xla/backends/gpu/tests/BUILD index 1f14685f3ec03d..73e0ea8e1b3810 100644 --- a/third_party/xla/xla/backends/gpu/tests/BUILD +++ b/third_party/xla/xla/backends/gpu/tests/BUILD @@ -1169,15 +1169,18 @@ cc_library( "//xla/service/gpu:backend_configs_cc", "//xla/stream_executor:device_description", "//xla/stream_executor/cuda:cuda_compute_capability", + "//xla/tsl/platform:env", "//xla/tsl/platform:statusor", "@com_google_absl//absl/log", "@com_google_absl//absl/log:check", "@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_googletest//:gtest_for_library", "@com_google_protobuf//:protobuf", + "@tsl//tsl/platform:path", ] + if_google([ # In OSS this dependency is added automatically for xla_test targets. See # third_party/tensorflow/compiler/xla/xla.default.bzl. @@ -1675,6 +1678,7 @@ xla_test( "//xla/tests:literal_test_util", "//xla/tsl/platform:statusor", "//xla/tsl/platform:test", + "//xla/tsl/testing:temporary_directory", "@com_google_absl//absl/log", "@com_google_absl//absl/log:check", "@com_google_absl//absl/status", @@ -1707,6 +1711,7 @@ xla_test( "//xla/service/gpu:backend_configs_cc", "//xla/tests:literal_test_util", "//xla/tsl/platform:test", + "//xla/tsl/testing:temporary_directory", "@com_google_absl//absl/strings:string_view", "@com_google_googletest//:gtest_main", ], diff --git a/third_party/xla/xla/backends/gpu/tests/all_gather_e2e_test.cc b/third_party/xla/xla/backends/gpu/tests/all_gather_e2e_test.cc index d0a3428223eaaa..0a118b107155db 100644 --- a/third_party/xla/xla/backends/gpu/tests/all_gather_e2e_test.cc +++ b/third_party/xla/xla/backends/gpu/tests/all_gather_e2e_test.cc @@ -30,6 +30,7 @@ limitations under the License. #include "xla/service/gpu/backend_configs.pb.h" #include "xla/tests/literal_test_util.h" #include "xla/tsl/platform/test.h" +#include "xla/tsl/testing/temporary_directory.h" #include "xla/xla.pb.h" namespace xla { @@ -250,5 +251,88 @@ TEST_F(AllGatherTest, TwoDimensional2Gpu) { } } +// Runs 10 consequent all-gathers with FUSION command buffers enabled and +// verifies that all of them use XLA's one-shot all-gather kernel and are placed +// into command buffers. +TEST_F(AllGatherTest, TenConsequentAllGathersWithCudaGraphs) { + constexpr int32_t kNumReplicas = 2; + constexpr int kNumAllGathers = 10; + if (!CheckDeviceCount(kNumReplicas)) { + return; + } + + constexpr absl::string_view kModuleStr = R"( + HloModule test + ENTRY test_computation { + p0 = f32[16] parameter(0) + ag0 = f32[32] all-gather(p0), dimensions={0}, replica_groups={{0,1}} + ag1 = f32[64] all-gather(ag0), dimensions={0}, replica_groups={{0,1}} + ag2 = f32[128] all-gather(ag1), dimensions={0}, replica_groups={{0,1}} + ag3 = f32[256] all-gather(ag2), dimensions={0}, replica_groups={{0,1}} + ag4 = f32[512] all-gather(ag3), dimensions={0}, replica_groups={{0,1}} + ag5 = f32[1024] all-gather(ag4), dimensions={0}, replica_groups={{0,1}} + ag6 = f32[2048] all-gather(ag5), dimensions={0}, replica_groups={{0,1}} + ag7 = f32[4096] all-gather(ag6), dimensions={0}, replica_groups={{0,1}} + ag8 = f32[8192] all-gather(ag7), dimensions={0}, replica_groups={{0,1}} + ROOT ag9 = f32[16384] all-gather(ag8), dimensions={0}, replica_groups={{0,1}} + } + )"; + + ASSERT_OK_AND_ASSIGN( + tsl::testing::TemporaryDirectory dump_dir, + tsl::testing::TemporaryDirectory::CreateForCurrentTestcase()); + ASSERT_OK_AND_ASSIGN(auto module, + ParseAndReturnVerifiedModule(kModuleStr, kNumReplicas)); + + DebugOptions& debug_options = + module->mutable_config().mutable_debug_options(); + // The fixture enables COLLECTIVE_KERNEL_ALL_GATHER, so the 10 all-gathers + // lower to XLA's Triton one-shot all-gather kernel (CollectiveKernelThunk) + // instead of NCCL. Collective kernels are recorded into command buffers + // with the FUSION command buffer type, so COLLECTIVES isn't needed. + debug_options.clear_xla_gpu_enable_command_buffer(); + debug_options.add_xla_gpu_enable_command_buffer(DebugOptions::COLLECTIVES); + debug_options.set_xla_gpu_graph_min_graph_size(1); + debug_options.set_xla_gpu_all_gather_combine_threshold_bytes(0); + // Command buffers are formed at the thunk level, so dump the thunk sequence + // to check which collectives were placed into them. + debug_options.set_xla_dump_to(dump_dir.path()); + + Literal input_r0 = LiteralUtil::CreateR1(std::vector(16, 1.0f)); + Literal input_r1 = LiteralUtil::CreateR1(std::vector(16, 2.0f)); + std::vector> args = {{&input_r0}, {&input_r1}}; + ASSERT_OK_AND_ASSIGN(ExecutionResult result, + ExecuteReplicated(std::move(module), args)); + + // Every all-gather must use KERNEL_STRATEGY_TRITON_ONE_SHOT (XLA one-shot + // kernel, not NCCL). + VerifyOneShotAllGather(result.optimized_module); + + // All one-shot all-gathers must be recorded into command buffers, and none + // may fall back to NCCL. + ASSERT_OK_AND_ASSIGN(CommandBufferThunkCounts one_shot, + CountThunksInDump(dump_dir.path(), "kCollectiveKernel")); + EXPECT_EQ(one_shot.in_command_buffer, kNumAllGathers); + EXPECT_EQ(one_shot.outside_command_buffer, 0); + ASSERT_OK_AND_ASSIGN(CommandBufferThunkCounts nccl, + CountThunksInDump(dump_dir.path(), "kAllGather")); + EXPECT_EQ(nccl.in_command_buffer + nccl.outside_command_buffer, 0); + + // Each all-gather concatenates the two ranks' buffers, so the result is 512 + // repetitions of [16 x 1.0, 16 x 2.0] on both replicas. + ASSERT_EQ(result.results.size(), kNumReplicas); + std::vector expected_data; + expected_data.reserve(16384); + for (int rep = 0; rep < 512; ++rep) { + expected_data.insert(expected_data.end(), 16, 1.0f); + expected_data.insert(expected_data.end(), 16, 2.0f); + } + Literal expected = LiteralUtil::CreateR1(expected_data); + for (int i = 0; i < kNumReplicas; ++i) { + EXPECT_TRUE(LiteralTestUtil::Equal(expected, result.results[i])) + << "Mismatch at replica " << i; + } +} + } // namespace } // namespace xla diff --git a/third_party/xla/xla/backends/gpu/tests/all_reduce_e2e_test.cc b/third_party/xla/xla/backends/gpu/tests/all_reduce_e2e_test.cc index af8da73e11e7ac..fa2089ef918828 100644 --- a/third_party/xla/xla/backends/gpu/tests/all_reduce_e2e_test.cc +++ b/third_party/xla/xla/backends/gpu/tests/all_reduce_e2e_test.cc @@ -40,6 +40,7 @@ limitations under the License. #include "xla/core/collectives/reduction_kind.h" #include "xla/error_spec.h" #include "xla/hlo/ir/hlo_casting_utils.h" +#include "xla/hlo/ir/hlo_computation.h" #include "xla/hlo/ir/hlo_input_output_alias_config.h" #include "xla/hlo/ir/hlo_instruction.h" #include "xla/hlo/ir/hlo_instructions.h" @@ -57,6 +58,7 @@ limitations under the License. #include "xla/tests/literal_test_util.h" #include "xla/tsl/platform/statusor.h" #include "xla/tsl/platform/test.h" +#include "xla/tsl/testing/temporary_directory.h" #include "xla/types.h" #include "xla/util.h" #include "xla/xla.pb.h" @@ -508,6 +510,47 @@ class AllReduceLayoutAwareTest } }; +// Async all-reduce test with XLA's Triton all-reduce kernel enabled. +class AllReduceCollectiveKernelTest : public AllReduceTestNoParams { + public: + AllReduceCollectiveKernelTest() : AllReduceTestNoParams(/*is_async=*/true) {} + + protected: + DebugOptions GetDebugOptionsForTest() const override { + DebugOptions opts = CollectiveOpsWithFlagsBase::GetDebugOptionsForTest(); + opts.clear_xla_gpu_experimental_use_collective_kernels(); + opts.add_xla_gpu_experimental_use_collective_kernels( + DebugOptions::COLLECTIVE_KERNEL_ALL_REDUCE); + return opts; + } + + // Verifies that every all-reduce in `optimized_module` uses XLA's one-shot + // all-reduce kernel. + void VerifyOneShotAllReduce(const HloModule* optimized_module) { + ASSERT_NE(optimized_module, nullptr); + bool found_all_reduce = false; + for (const HloComputation* comp : optimized_module->computations()) { + for (const HloInstruction* instr : comp->instructions()) { + if (instr->opcode() != HloOpcode::kAllReduce && + instr->opcode() != HloOpcode::kAllReduceStart) { + continue; + } + found_all_reduce = true; + ASSERT_OK_AND_ASSIGN(gpu::GpuBackendConfig gpu_config, + instr->backend_config()); + EXPECT_EQ(gpu_config.collective_backend_config().kernel_strategy(), + gpu::CollectiveBackendConfig::KERNEL_STRATEGY_TRITON_ONE_SHOT) + << "Expected AllReduce instruction " << instr->name() + << " to use KERNEL_STRATEGY_TRITON_ONE_SHOT, but got: " + << gpu::CollectiveBackendConfig::CollectiveKernelStrategy_Name( + gpu_config.collective_backend_config().kernel_strategy()); + } + } + EXPECT_TRUE(found_all_reduce) + << "Expected to find an AllReduce instruction in optimized HLO."; + } +}; + INSTANTIATE_TEST_SUITE_P( AllReduceTest, AllReduceTest, ::testing::ValuesIn(AllReduceTestParams::Generate()), @@ -1058,5 +1101,113 @@ TEST_F(AllReduceTestNoParams, TestGlobalDeviceIdMappingWithAsyncAllReduce) { } } +// Runs 10 consequent all-reduces with FUSION command buffers enabled and +// verifies that all of them use XLA's one-shot all-reduce kernel and are placed +// into command buffers. +TEST_F(AllReduceCollectiveKernelTest, TenConsequentAllReducesWithCudaGraphs) { + constexpr int64_t kNumReplicas = 2; + constexpr int kNumAllReduces = 10; + if (!CheckDeviceCount(kNumReplicas)) { + return; + } + + // Adding the per-replica parameter before every all-reduce keeps its input + // non-replicated, so the all-reduce simplifier can't fold the chain. + constexpr absl::string_view kModuleStr = R"( + HloModule test + + add { + x = f32[] parameter(0) + y = f32[] parameter(1) + ROOT add = f32[] add(x, y) + } + + ENTRY test_computation { + p0 = f32[1024] parameter(0) + ar0 = f32[1024] all-reduce(p0), to_apply=add, replica_groups={{0,1}} + in1 = f32[1024] add(ar0, p0) + ar1 = f32[1024] all-reduce(in1), to_apply=add, + replica_groups={{0,1}} + in2 = f32[1024] add(ar1, p0) + ar2 = f32[1024] all-reduce(in2), to_apply=add, + replica_groups={{0,1}} + in3 = f32[1024] add(ar2, p0) + ar3 = f32[1024] all-reduce(in3), to_apply=add, + replica_groups={{0,1}} + in4 = f32[1024] add(ar3, p0) + ar4 = f32[1024] all-reduce(in4), to_apply=add, + replica_groups={{0,1}} + in5 = f32[1024] add(ar4, p0) + ar5 = f32[1024] all-reduce(in5), to_apply=add, + replica_groups={{0,1}} + in6 = f32[1024] add(ar5, p0) + ar6 = f32[1024] all-reduce(in6), to_apply=add, + replica_groups={{0,1}} + in7 = f32[1024] add(ar6, p0) + ar7 = f32[1024] all-reduce(in7), to_apply=add, + replica_groups={{0,1}} + in8 = f32[1024] add(ar7, p0) + ar8 = f32[1024] all-reduce(in8), to_apply=add, + replica_groups={{0,1}} + in9 = f32[1024] add(ar8, p0) + ROOT ar9 = f32[1024] all-reduce(in9), to_apply=add, + replica_groups={{0,1}} + } + )"; + + ASSERT_OK_AND_ASSIGN( + tsl::testing::TemporaryDirectory dump_dir, + tsl::testing::TemporaryDirectory::CreateForCurrentTestcase()); + ASSERT_OK_AND_ASSIGN(auto module, + ParseAndReturnVerifiedModule(kModuleStr, kNumReplicas)); + + DebugOptions& debug_options = + module->mutable_config().mutable_debug_options(); + // The fixture enables COLLECTIVE_KERNEL_ALL_REDUCE, so the 10 all-reduces + // lower to XLA's Triton one-shot all-reduce kernel (CollectiveKernelThunk) + // instead of NCCL. Collective kernels are recorded into command buffers + // with the FUSION command buffer type, so COLLECTIVES isn't needed. + debug_options.clear_xla_gpu_enable_command_buffer(); + debug_options.add_xla_gpu_enable_command_buffer(DebugOptions::COLLECTIVES); + debug_options.set_xla_gpu_graph_min_graph_size(1); + debug_options.set_xla_gpu_all_reduce_combine_threshold_bytes(0); + // Command buffers are formed at the thunk level, so dump the thunk sequence + // to check which collectives were placed into them. + debug_options.set_xla_dump_to(dump_dir.path()); + + Literal input_r0 = + LiteralUtil::CreateR1(std::vector(1024, 1.0f)); + Literal input_r1 = + LiteralUtil::CreateR1(std::vector(1024, 2.0f)); + std::vector> args = {{&input_r0}, {&input_r1}}; + ASSERT_OK_AND_ASSIGN(ExecutionResult result, + ExecuteReplicated(std::move(module), args)); + + // Every all-reduce must use KERNEL_STRATEGY_TRITON_ONE_SHOT (XLA one-shot + // kernel, not NCCL). + VerifyOneShotAllReduce(result.optimized_module); + + // All one-shot all-reduces must be recorded into command buffers, and none + // may fall back to NCCL. + ASSERT_OK_AND_ASSIGN(CommandBufferThunkCounts one_shot, + CountThunksInDump(dump_dir.path(), "kCollectiveKernel")); + EXPECT_EQ(one_shot.in_command_buffer, kNumAllReduces); + EXPECT_EQ(one_shot.outside_command_buffer, 0); + ASSERT_OK_AND_ASSIGN(CommandBufferThunkCounts nccl, + CountThunksInDump(dump_dir.path(), "kAllReduce")); + EXPECT_EQ(nccl.in_command_buffer + nccl.outside_command_buffer, 0); + + // The first all-reduce sums 1 + 2 = 3 and each following one computes + // (a + 1) + (a + 2) = 2a + 3, so both replicas end up with + // 3 * (2^10 - 1) = 3069. + ASSERT_EQ(result.results.size(), kNumReplicas); + Literal expected = + LiteralUtil::CreateR1(std::vector(1024, 3069.0f)); + for (int i = 0; i < kNumReplicas; ++i) { + EXPECT_TRUE(LiteralTestUtil::Equal(expected, result.results[i])) + << "Mismatch at replica " << i; + } +} + } // namespace } // namespace xla diff --git a/third_party/xla/xla/backends/gpu/tests/collective_ops_e2e_test_base.cc b/third_party/xla/xla/backends/gpu/tests/collective_ops_e2e_test_base.cc index 54c5d8edee0a58..ab11d842402e53 100644 --- a/third_party/xla/xla/backends/gpu/tests/collective_ops_e2e_test_base.cc +++ b/third_party/xla/xla/backends/gpu/tests/collective_ops_e2e_test_base.cc @@ -19,14 +19,20 @@ limitations under the License. #include #include #include +#include #include #include #include #include "absl/log/check.h" #include "absl/log/log.h" +#include "absl/status/status.h" #include "absl/status/status_macros.h" #include "absl/status/statusor.h" +#include "absl/strings/ascii.h" +#include "absl/strings/match.h" +#include "absl/strings/str_cat.h" +#include "absl/strings/str_split.h" #include "absl/strings/string_view.h" #include "xla/backends/gpu/tests/hlo_pjrt_gpu_test_base.h" #include "xla/hlo/ir/hlo_instruction.h" @@ -41,9 +47,11 @@ limitations under the License. #include "xla/service/gpu/backend_configs.pb.h" #include "xla/service/hlo_module_config.h" #include "xla/service/hlo_runner_interface.h" +#include "xla/tsl/platform/env.h" #include "xla/tsl/platform/statusor.h" #include "xla/xla.pb.h" #include "xla/xla_data.pb.h" +#include "tsl/platform/path.h" namespace xla { namespace { @@ -190,4 +198,43 @@ CollectiveOpsWithFlagsBase::CreateExecutable(absl::string_view hlo_string, /*run_hlo_passes=*/true); } +absl::StatusOr CountThunksInDump( + absl::string_view dump_dir, absl::string_view thunk_kind_prefix) { + std::vector dump_files; + ABSL_RETURN_IF_ERROR(tsl::Env::Default()->GetMatchingPaths( + tsl::io::JoinPath(dump_dir, "*thunk_sequence_after_thunk_passes*.txt"), + &dump_files)); + if (dump_files.size() != 1) { + return absl::FailedPreconditionError( + absl::StrCat("Expected exactly one thunk sequence dump in ", dump_dir, + ", found ", dump_files.size())); + } + std::string dump; + ABSL_RETURN_IF_ERROR( + tsl::ReadFileToString(tsl::Env::Default(), dump_files[0], &dump)); + + // Each line of the dump describes one thunk as ": ...". + // Thunks nested in a kCommandBuffer thunk are indented under it. + CommandBufferThunkCounts counts; + bool in_command_buffer = false; + for (absl::string_view line : absl::StrSplit(dump, '\n')) { + absl::string_view thunk = absl::StripAsciiWhitespace(line); + if (thunk.empty()) { + continue; + } + const bool is_top_level = thunk.data() == line.data(); + if (size_t pos = thunk.find(": "); pos != absl::string_view::npos) { + thunk.remove_prefix(pos + 2); + } + if (is_top_level) { + in_command_buffer = absl::StartsWith(thunk, "kCommandBuffer"); + } + if (absl::StartsWith(thunk, thunk_kind_prefix)) { + ++(in_command_buffer ? counts.in_command_buffer + : counts.outside_command_buffer); + } + } + return counts; +} + } // namespace xla diff --git a/third_party/xla/xla/backends/gpu/tests/collective_ops_e2e_test_base.h b/third_party/xla/xla/backends/gpu/tests/collective_ops_e2e_test_base.h index dc57ae1ef6faac..7bbd51b3213dbb 100644 --- a/third_party/xla/xla/backends/gpu/tests/collective_ops_e2e_test_base.h +++ b/third_party/xla/xla/backends/gpu/tests/collective_ops_e2e_test_base.h @@ -123,6 +123,19 @@ class CollectiveOpsWithFlagsBase : public CollectiveOpsE2ETestBase { const bool enable_symmetric_buffer_; }; +// Number of thunks of a given kind placed inside and outside of command +// buffers (kCommandBuffer thunks). +struct CommandBufferThunkCounts { + int in_command_buffer = 0; + int outside_command_buffer = 0; +}; + +// Reads the thunk sequence that XLA dumps to `dump_dir` (when compiled with +// --xla_dump_to=`dump_dir`) and counts the thunks whose kind starts with +// `thunk_kind_prefix` (e.g. "kCollectiveKernel" or "kAllGather"). +absl::StatusOr CountThunksInDump( + absl::string_view dump_dir, absl::string_view thunk_kind_prefix); + } // namespace xla #endif // XLA_BACKENDS_GPU_TESTS_COLLECTIVE_OPS_E2E_TEST_BASE_H_ diff --git a/third_party/xla/xla/pjrt/abstract_tracked_device_buffer.cc b/third_party/xla/xla/pjrt/abstract_tracked_device_buffer.cc index 5fec7b2e813ee9..ac58cd13626af5 100644 --- a/third_party/xla/xla/pjrt/abstract_tracked_device_buffer.cc +++ b/third_party/xla/xla/pjrt/abstract_tracked_device_buffer.cc @@ -172,8 +172,8 @@ absl::Status AbstractTrackedDeviceBuffer::WaitUntilBufferReadyOnStream( PjRtMemorySpace* memory_space, std::intptr_t stream) { auto* client = absl::down_cast(memory_space->client()); for (const auto& event : definition_events()) { - ABSL_RETURN_IF_ERROR( - client->raw_client()->WaitOnStream(memory_space, event, stream)); + ABSL_RETURN_IF_ERROR(client->raw_client()->WaitOnStream( + memory_space->devices()[0]->local_device_id(), event, stream)); } return absl::OkStatus(); } diff --git a/third_party/xla/xla/pjrt/common_pjrt_client.cc b/third_party/xla/xla/pjrt/common_pjrt_client.cc index dff4ad82416700..6610bf56efdabe 100644 --- a/third_party/xla/xla/pjrt/common_pjrt_client.cc +++ b/third_party/xla/xla/pjrt/common_pjrt_client.cc @@ -1303,7 +1303,8 @@ CommonPjRtClient::CreateViewOfDeviceBuffer( } ABSL_ASSIGN_OR_RETURN( PjRtDeviceEventRef stream_event, - raw_client()->CreateDeviceEventForStream(memory_space, *stream)); + raw_client()->CreateDeviceEventForStream( + memory_space->devices()[0]->local_device_id(), *stream)); definition_events.emplace_back(std::move(stream_event)); } @@ -2450,6 +2451,9 @@ CommonPjRtLoadedExecutable::LookupDeviceAndAssignment( } CHECK_EQ(device->process_index(), client()->process_index()); result.device = device; + result.local_device_id = device->local_device_id(); + result.global_device_id = device->global_device_id(); + result.process_index = device->process_index(); result.replica = replica; result.partition = partition; return result; @@ -2462,11 +2466,11 @@ absl::Status CommonPjRtLoadedExecutable::ExecutePrepare( size_t host_callback_idx, PjRtDevice* device, int attempt) const { tsl::profiler::TraceMe traceme("CommonPjRtLoadedExecutable::ExecutePrepare"); ABSL_ASSIGN_OR_RETURN( - auto device_and_assign, + DeviceAndAssignment device_and_assign, LookupDeviceAndAssignment(options, replica, partition, device)); + device = device_and_assign.device; // Fill in device to launch_args so it will be present even if ExecutePrepare // fails with OOM. - device = device_and_assign.device; launch_args.device = device; // Execute takes `extra_deps` and waits for those to be diff --git a/third_party/xla/xla/pjrt/common_pjrt_client.h b/third_party/xla/xla/pjrt/common_pjrt_client.h index c25558ff7a8433..cfd960f46b0a02 100644 --- a/third_party/xla/xla/pjrt/common_pjrt_client.h +++ b/third_party/xla/xla/pjrt/common_pjrt_client.h @@ -259,7 +259,9 @@ class CommonPjRtClient : public PjRtClient { std::pair> CreateLinkedEventPromise(PjRtMemorySpace* memory_space, absl::string_view debug_info) { - return raw_client()->CreateLinkedEventPromise(memory_space, debug_info); + return raw_client()->CreateLinkedEventPromise( + memory_space->devices()[0]->local_device_id(), memory_space->kind_id(), + debug_info); } // Track a user-provided future with attached debug_info (if @@ -301,7 +303,9 @@ class CommonPjRtClient : public PjRtClient { virtual absl::StatusOr CreateDeviceEvent( PjRtMemorySpace* memory_space, Future<> dependency) { - return raw_client()->CreateDeviceEvent(memory_space, std::move(dependency)); + return raw_client()->CreateDeviceEvent( + memory_space->devices()[0]->local_device_id(), memory_space->kind_id(), + std::move(dependency)); } absl::StatusOr> CreateErrorBuffer( @@ -314,7 +318,8 @@ class CommonPjRtClient : public PjRtClient { tsl::AsyncValueRef CreateAllocationEventForTransfers( PjRtMemorySpace* memory_space, const std::optional& debug_info) { - if (raw_client()->ShouldCreateAsyncAllocationEvent(memory_space)) { + if (raw_client()->ShouldCreateAsyncAllocationEvent( + memory_space->kind_id())) { tsl::AsyncValueRef result = tsl::MakeConstructedAsyncValueRef(); if (event_tracker()) { @@ -1097,9 +1102,10 @@ class CommonPjRtClientImpl : public CommonPjRtClient { PjRtDeviceEventPromiseRef usage_event_promise, Future serialized_descriptor, PjRtBuffer::RemoteSendCallback on_done) override { - raw_client()->ScheduleRemoteSend(memory_space, raw_buffer, - definition_events, usage_event_promise, - serialized_descriptor, std::move(on_done)); + raw_client()->ScheduleRemoteSend( + memory_space->devices()[0]->local_device_id(), memory_space->kind_id(), + raw_buffer, definition_events, usage_event_promise, + serialized_descriptor, std::move(on_done)); } absl::StatusOr CrossHostReceiveBuffersInto( diff --git a/third_party/xla/xla/pjrt/cpu/cpu_client.cc b/third_party/xla/xla/pjrt/cpu/cpu_client.cc index 7ef56465262556..579e7482bc8a07 100644 --- a/third_party/xla/xla/pjrt/cpu/cpu_client.cc +++ b/third_party/xla/xla/pjrt/cpu/cpu_client.cc @@ -1007,7 +1007,8 @@ PjRtCpuRawClient::CompileInternal( } absl::StatusOr PjRtCpuRawClient::CreateDeviceEvent( - PjRtMemorySpace* memory_space, Future dependency) { + LocalDeviceId local_device_id, int memory_kind_id, + Future dependency) { return ToCpuEvent(std::move(dependency)); } @@ -1034,7 +1035,8 @@ absl::StatusOr PjRtCpuExecutable::GetCompiledMemoryStats() } absl::StatusOr> -PjRtCpuRawClient::CreateLinkedEventPromise(PjRtMemorySpace* memory_space, +PjRtCpuRawClient::CreateLinkedEventPromise(LocalDeviceId local_device_id, + int memory_kind_id, absl::string_view debug_info) { auto definition_event_promise = tsl::MakeIndirectAsyncValue(); auto definition_event = PjRtDeviceEventRef( @@ -1350,7 +1352,7 @@ CreateBufferTable(const BufferAssignment& assignment, } tsl::RCReference PjRtCpuRawClient::MakeLoadState() { - return tsl::MakeRef(); + return tsl::MakeRef(this); } absl::StatusOr> @@ -1360,10 +1362,7 @@ CpuExecutableLoadState::LoadRawExecutable( DeviceAndAssignment device_and_assign, int attempt) { auto result = std::make_unique(run_id); result->executable_ = absl::down_cast(&executable.get()); - auto* client = - absl::down_cast(device_and_assign.device->client()); - result->raw_client_ = - absl::down_cast(client->raw_client()); + result->raw_client_ = raw_client_; int num_addressable_devices = 0; if (device_and_assign.device_assignment != nullptr) { for (int r = 0; r < device_and_assign.device_assignment->replica_count(); @@ -1371,7 +1370,8 @@ CpuExecutableLoadState::LoadRawExecutable( for (int p = 0; p < device_and_assign.device_assignment->computation_count(); ++p) { GlobalDeviceId device_id((*device_and_assign.device_assignment)(r, p)); - if (UnpackCpuProcessIndex(device_id) == client->process_index()) { + if (UnpackCpuProcessIndex(device_id) == + device_and_assign.process_index) { ++num_addressable_devices; } } @@ -1379,8 +1379,8 @@ CpuExecutableLoadState::LoadRawExecutable( } result->num_addressable_devices_ = num_addressable_devices; result->device_assignment_ = std::move(device_and_assign.device_assignment); - result->local_device_id_ = device_and_assign.device->local_device_id(); - result->global_device_id_ = device_and_assign.device->global_device_id(); + result->local_device_id_ = device_and_assign.local_device_id; + result->global_device_id_ = device_and_assign.global_device_id; return result; } diff --git a/third_party/xla/xla/pjrt/cpu/cpu_client.h b/third_party/xla/xla/pjrt/cpu/cpu_client.h index b7539a1a9948fd..25b5489cf513c9 100644 --- a/third_party/xla/xla/pjrt/cpu/cpu_client.h +++ b/third_party/xla/xla/pjrt/cpu/cpu_client.h @@ -161,11 +161,12 @@ class PjRtCpuRawClient : public PjRtRawClient { size_t on_device_bytes_count) override; absl::StatusOr> - CreateLinkedEventPromise(PjRtMemorySpace* memory_space, + CreateLinkedEventPromise(LocalDeviceId local_device_id, int memory_kind_id, absl::string_view debug_info) override; absl::StatusOr CreateDeviceEvent( - PjRtMemorySpace* memory_space, Future<> dependency) override; + LocalDeviceId local_device_id, int memory_kind_id, + Future<> dependency) override; absl::StatusOr ImportForeignMemory( PjRtMemorySpace* memory_space, void* device_ptr, size_t size, @@ -339,7 +340,8 @@ class CpuPjRtRawLoadedExecutable : public PjRtRawLoadedExecutable { class CpuExecutableLoadState : public PjRtExecutableLoadState { public: - explicit CpuExecutableLoadState() = default; + explicit CpuExecutableLoadState(PjRtCpuRawClient* raw_client) + : raw_client_(raw_client) {} ~CpuExecutableLoadState() override = default; @@ -353,6 +355,7 @@ class CpuExecutableLoadState : public PjRtExecutableLoadState { int attempt) override; private: + PjRtCpuRawClient* raw_client_; std::atomic is_deleted_{false}; }; diff --git a/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_client.cc b/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_client.cc index ea6e2554a2cbc5..a7b8bc56730be8 100644 --- a/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_client.cc +++ b/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_client.cc @@ -804,8 +804,8 @@ void StreamExecutorGpuRawClient::ScheduleTransfersOnLocalDevice( } void StreamExecutorGpuRawClient::ScheduleRemoteSend( - PjRtMemorySpace* memory_space, PjRtRawBufferRef raw_buffer, - PjRtDeviceEventRefVector definition_events, + LocalDeviceId local_device_id, int memory_kind_id, + PjRtRawBufferRef raw_buffer, PjRtDeviceEventRefVector definition_events, PjRtDeviceEventPromiseRef usage_event_promise, Future serialized_descriptor, PjRtBuffer::RemoteSendCallback on_done) { diff --git a/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_client.h b/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_client.h index 137c26b16e1c5b..d9fc7804ed980e 100644 --- a/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_client.h +++ b/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_client.h @@ -118,7 +118,7 @@ class StreamExecutorGpuRawClient : public PjRtStreamExecutorRawClient { return kv_store_; } - void ScheduleRemoteSend(PjRtMemorySpace* memory_space, + void ScheduleRemoteSend(LocalDeviceId local_device_id, int memory_kind_id, PjRtRawBufferRef raw_buffer, PjRtDeviceEventRefVector definition_events, PjRtDeviceEventPromiseRef usage_event_promise, diff --git a/third_party/xla/xla/pjrt/raw_pjrt_client.cc b/third_party/xla/xla/pjrt/raw_pjrt_client.cc index 793cc0f011d623..3f2579114d4933 100644 --- a/third_party/xla/xla/pjrt/raw_pjrt_client.cc +++ b/third_party/xla/xla/pjrt/raw_pjrt_client.cc @@ -18,14 +18,14 @@ limitations under the License. namespace xla { void PjRtRawClient::ScheduleRemoteSend( - PjRtMemorySpace* memory_space, PjRtRawBufferRef raw_buffer, - PjRtDeviceEventRefVector definition_events, + LocalDeviceId local_device_id, int memory_kind_id, + PjRtRawBufferRef raw_buffer, PjRtDeviceEventRefVector definition_events, PjRtDeviceEventPromiseRef usage_event_promise, Future serialized_descriptor, PjRtBuffer::RemoteSendCallback on_done) { auto error = absl::UnimplementedError( - absl::StrCat("ScheduleRemoteSend is not implemented for %s", - memory_space->DebugString())); + absl::StrCat("ScheduleRemoteSend is not implemented for device ", + local_device_id.value(), " kind ", memory_kind_id)); on_done(error, /*sends_were_enqueued=*/false); usage_event_promise.SetError(error); } diff --git a/third_party/xla/xla/pjrt/raw_pjrt_client.h b/third_party/xla/xla/pjrt/raw_pjrt_client.h index 18f893868741d2..25fd1028ad9c4e 100644 --- a/third_party/xla/xla/pjrt/raw_pjrt_client.h +++ b/third_party/xla/xla/pjrt/raw_pjrt_client.h @@ -119,6 +119,9 @@ class PjRtExecutableLoadState struct DeviceAndAssignment { PjRtDevice* device; + LocalDeviceId local_device_id; + GlobalDeviceId global_device_id; + int process_index; std::shared_ptr device_assignment; std::optional slice_id; int replica; @@ -155,8 +158,7 @@ class PjRtRawClient { // Returns true if memory allocations for async transfers in `memory_space` // should be deferred behind an allocation event. - virtual bool ShouldCreateAsyncAllocationEvent( - PjRtMemorySpace* memory_space) const { + virtual bool ShouldCreateAsyncAllocationEvent(int memory_kind_id) const { return false; } @@ -196,7 +198,7 @@ class PjRtRawClient { // setting an event into the event promise populates the device-event. virtual absl::StatusOr< std::pair> - CreateLinkedEventPromise(PjRtMemorySpace* memory_space, + CreateLinkedEventPromise(LocalDeviceId local_device_id, int memory_kind_id, absl::string_view debug_info) { return absl::UnimplementedError( "CreateLinkedEventPromise is not supported"); @@ -204,12 +206,13 @@ class PjRtRawClient { // Creates a device event that signals completion of a dependency future. virtual absl::StatusOr CreateDeviceEvent( - PjRtMemorySpace* memory_space, Future<> dependency) = 0; + LocalDeviceId local_device_id, int memory_kind_id, + Future<> dependency) = 0; // Creates a device event that signals completion of work on an external // stream. virtual absl::StatusOr CreateDeviceEventForStream( - PjRtMemorySpace* memory_space, std::intptr_t stream) { + LocalDeviceId local_device_id, std::intptr_t stream) { return absl::UnimplementedError( "CreateDeviceEventForStream is not supported"); } @@ -290,7 +293,7 @@ class PjRtRawClient { return absl::UnimplementedError("TransferToOutfeed is not supported"); } - virtual absl::Status WaitOnStream(PjRtMemorySpace* memory_space, + virtual absl::Status WaitOnStream(LocalDeviceId local_device_id, PjRtDeviceEventRef event, std::intptr_t stream) { return absl::UnimplementedError( @@ -312,7 +315,8 @@ class PjRtRawClient { return absl::UnimplementedError("ClearMemoryStats is not supported."); } - virtual void ScheduleRemoteSend(PjRtMemorySpace* memory_space, + virtual void ScheduleRemoteSend(LocalDeviceId local_device_id, + int memory_kind_id, PjRtRawBufferRef raw_buffer, PjRtDeviceEventRefVector definition_events, PjRtDeviceEventPromiseRef usage_event_promise, diff --git a/third_party/xla/xla/pjrt/se/pjrt_stream_executor_client.cc b/third_party/xla/xla/pjrt/se/pjrt_stream_executor_client.cc index 362c40dfb1a82c..aa7ae99cc0a108 100644 --- a/third_party/xla/xla/pjrt/se/pjrt_stream_executor_client.cc +++ b/third_party/xla/xla/pjrt/se/pjrt_stream_executor_client.cc @@ -590,10 +590,10 @@ PjRtStreamExecutorClient::AllocateLinearizeDest( absl::StatusOr> PjRtStreamExecutorRawClient::CreateLinkedEventPromise( - PjRtMemorySpace* memory_space, absl::string_view debug_info) { - PjRtDevice* device = memory_space->devices()[0]; + LocalDeviceId local_device_id, int memory_kind_id, + absl::string_view debug_info) { ABSL_ASSIGN_OR_RETURN(LocalDeviceState * local_device, - GetLocalDeviceState(device->local_device_id())); + GetLocalDeviceState(local_device_id)); auto result = tsl::MakeRef( this, local_device, async_work_runner()); PjRtDeviceEventRef event = result->event().CopyRef(); @@ -610,12 +610,12 @@ PjRtDeviceEventRef PjRtStreamExecutorRawClient::CreateErrorDeviceEvent( } absl::StatusOr -PjRtStreamExecutorRawClient::CreateDeviceEvent(PjRtMemorySpace* memory_space, +PjRtStreamExecutorRawClient::CreateDeviceEvent(LocalDeviceId local_device_id, + int memory_kind_id, Future<> dependency) { auto definition_event = BufferSequencingEvent::Create(async_work_runner()); - PjRtDevice* device = memory_space->devices()[0]; ABSL_ASSIGN_OR_RETURN(LocalDeviceState * local_device, - GetLocalDeviceState(device->local_device_id())); + GetLocalDeviceState(local_device_id)); dependency.OnReady([definition_event = definition_event.CopyRef(), local_device, this](absl::Status status) mutable { if (!status.ok()) { @@ -682,11 +682,9 @@ PjRtStreamExecutorRawClient::ImportForeignMemory( absl::StatusOr PjRtStreamExecutorRawClient::CreateDeviceEventForStream( - PjRtMemorySpace* memory_space, std::intptr_t stream) { - CHECK_EQ(memory_space->devices().size(), 1); - auto* device = memory_space->devices().front(); + LocalDeviceId local_device_id, std::intptr_t stream) { ABSL_ASSIGN_OR_RETURN(LocalDeviceState * local_device, - GetLocalDeviceState(device->local_device_id())); + GetLocalDeviceState(local_device_id)); auto definition_event = BufferSequencingEvent::Create(this->async_work_runner()); @@ -2113,7 +2111,7 @@ bool PjRtStreamExecutorRawClient::IsHostMemoryPinned(const void* ptr, } absl::Status PjRtStreamExecutorRawClient::WaitOnStream( - PjRtMemorySpace* memory_space, PjRtDeviceEventRef event, + LocalDeviceId local_device_id, PjRtDeviceEventRef event, std::intptr_t stream) { return event.down_cast()->WaitForEventOnExternalStream( stream); @@ -2150,8 +2148,8 @@ PjRtStreamExecutorClient::AllocateForDelinearizationAsync( } void PjRtStreamExecutorRawClient::ScheduleRemoteSend( - PjRtMemorySpace* memory_space, PjRtRawBufferRef raw_buffer, - PjRtDeviceEventRefVector definition_events, + LocalDeviceId local_device_id, int memory_kind_id, + PjRtRawBufferRef raw_buffer, PjRtDeviceEventRefVector definition_events, PjRtDeviceEventPromiseRef usage_event_promise, Future serialized_descriptor, PjRtBuffer::RemoteSendCallback on_done) { diff --git a/third_party/xla/xla/pjrt/se/pjrt_stream_executor_client.h b/third_party/xla/xla/pjrt/se/pjrt_stream_executor_client.h index dcf39112747cdd..ddeabb883dfca4 100644 --- a/third_party/xla/xla/pjrt/se/pjrt_stream_executor_client.h +++ b/third_party/xla/xla/pjrt/se/pjrt_stream_executor_client.h @@ -232,17 +232,17 @@ class PjRtStreamExecutorRawClient : public PjRtRawClient { size_t on_device_bytes_count) override; absl::StatusOr> - CreateLinkedEventPromise(PjRtMemorySpace* memory_space, + CreateLinkedEventPromise(LocalDeviceId local_device_id, int memory_kind_id, absl::string_view debug_info) override; - virtual void ScheduleRemoteSend(PjRtMemorySpace* memory_space, - PjRtRawBufferRef raw_buffer, - PjRtDeviceEventRefVector definition_events, - PjRtDeviceEventPromiseRef usage_event_promise, - Future serialized_descriptor, - PjRtBuffer::RemoteSendCallback on_done); + void ScheduleRemoteSend(LocalDeviceId local_device_id, int memory_kind_id, + PjRtRawBufferRef raw_buffer, + PjRtDeviceEventRefVector definition_events, + PjRtDeviceEventPromiseRef usage_event_promise, + Future serialized_descriptor, + PjRtBuffer::RemoteSendCallback on_done) override; - absl::Status WaitOnStream(PjRtMemorySpace* memory_space, + absl::Status WaitOnStream(LocalDeviceId local_device_id, PjRtDeviceEventRef event, std::intptr_t stream) override; @@ -264,7 +264,8 @@ class PjRtStreamExecutorRawClient : public PjRtRawClient { std::vector transfer_specs); absl::StatusOr CreateDeviceEvent( - PjRtMemorySpace* memory_space, Future<> dependency) override; + LocalDeviceId local_device_id, int memory_kind_id, + Future<> dependency) override; PjRtDeviceEventRef CreateErrorDeviceEvent(absl::Status error); @@ -281,7 +282,7 @@ class PjRtStreamExecutorRawClient : public PjRtRawClient { bool is_mutable) override; absl::StatusOr CreateDeviceEventForStream( - PjRtMemorySpace* memory_space, std::intptr_t stream) override; + LocalDeviceId local_device_id, std::intptr_t stream) override; absl::Status TransferToInfeed(LocalDeviceId local_device_id, const LiteralSlice& literal) override; diff --git a/third_party/xla/xla/pjrt/se/se_raw_buffer.cc b/third_party/xla/xla/pjrt/se/se_raw_buffer.cc index 46f952e0f95e0b..3710029a95e619 100644 --- a/third_party/xla/xla/pjrt/se/se_raw_buffer.cc +++ b/third_party/xla/xla/pjrt/se/se_raw_buffer.cc @@ -286,7 +286,7 @@ PjRtStreamExecutorRawBuffer::MakeAllocationReadyEvent() { ABSL_ASSIGN_OR_RETURN( auto promise_and_event, client_->CreateLinkedEventPromise( - memory_space_, + local_device_->local_device_id(), memory_space_->kind_id(), "PjRtStreamExecutorRawBuffer::MakeAllocationReadyEvent")); auto [promise, event] = std::move(promise_and_event); diff --git a/third_party/xla/xla/python/ifrt/BUILD b/third_party/xla/xla/python/ifrt/BUILD index 6a6fd39440985c..d84003bf10971a 100644 --- a/third_party/xla/xla/python/ifrt/BUILD +++ b/third_party/xla/xla/python/ifrt/BUILD @@ -785,7 +785,8 @@ cc_library( srcs = ["mock.cc"], hdrs = ["mock.h"], visibility = internal_visibility([ - ":mock_users", + ":friends", + ":internal", ]), deps = [ ":attribute_map", diff --git a/third_party/xla/xla/python/ifrt/executable.h b/third_party/xla/xla/python/ifrt/executable.h index 154b5c36336580..e75cf2fbab1df5 100644 --- a/third_party/xla/xla/python/ifrt/executable.h +++ b/third_party/xla/xla/python/ifrt/executable.h @@ -205,7 +205,6 @@ class LoadedExecutable : public RTTIExtends { // prototyping. // TODO(hyeontaek): Factor some of them out as `XlaCompatibleExecutable`. - virtual int num_devices() const = 0; virtual int64_t SizeOfGeneratedCodeInBytes() const = 0; virtual absl::StatusOr GetCompiledMemoryStats() const = 0; @@ -327,13 +326,6 @@ class LoadedExecutable : public RTTIExtends { // particular device list, e.g., portable executables. virtual std::optional devices() const = 0; - // The following APIs are taken from xla::PjRtLoadedExecutable for fast - // prototyping. - // TODO(hyeontaek): Move the following XLA-specific methods to - // pjrt_executable.h and put it in an `XlaCompatibleExecutable`. - - virtual absl::Span addressable_devices() const = 0; - struct DeleteOptions { // Analogous to `ExecuteOptions::execution_stream_id` for any side-effects // of deleting the executable. diff --git a/third_party/xla/xla/python/ifrt/ifrt.bzl b/third_party/xla/xla/python/ifrt/ifrt.bzl index e805d7cf7d7d47..f7217a5df6b3c6 100644 --- a/third_party/xla/xla/python/ifrt/ifrt.bzl +++ b/third_party/xla/xla/python/ifrt/ifrt.bzl @@ -50,8 +50,3 @@ def ifrt_package_groups(name = "ifrt_package_groups"): name = "serdes_week_4_old_version_users", packages = ["//..."], ) - - native.package_group( - name = "mock_users", - packages = ["//..."], - ) diff --git a/third_party/xla/xla/python/ifrt/ir/ifrt_ir_loaded_executable.cc b/third_party/xla/xla/python/ifrt/ir/ifrt_ir_loaded_executable.cc index 14287a70cbe10e..a14a3b2a3d2989 100644 --- a/third_party/xla/xla/python/ifrt/ir/ifrt_ir_loaded_executable.cc +++ b/third_party/xla/xla/python/ifrt/ir/ifrt_ir_loaded_executable.cc @@ -171,11 +171,6 @@ IfrtIrLoadedExecutable::GetHumanReadableProgramText() const { return result; } -int IfrtIrLoadedExecutable::num_devices() const { - DCHECK(this); - return devices_->size(); -} - int64_t IfrtIrLoadedExecutable::SizeOfGeneratedCodeInBytes() const { // TODO(b/261226026): Implement this API or remove it from IFRT. return -1; @@ -338,11 +333,6 @@ IfrtIrLoadedExecutable::GetMpmdCostAnalysis() const { return mpmd_cost_analysis; } -absl::Span IfrtIrLoadedExecutable::addressable_devices() const { - DCHECK(this); - return devices_->devices(); -} - std::optional IfrtIrLoadedExecutable::devices() const { return devices_; } @@ -353,7 +343,13 @@ IfrtIrLoadedExecutable::GetMpmdAddressableDevices() const { mpmd_addressable_devices; mpmd_addressable_devices.reserve(program_->atom_program_executables->size()); for (const auto& [name, executable] : *program_->atom_program_executables) { - mpmd_addressable_devices.insert({name, executable->addressable_devices()}); + if (std::optional devices = executable->devices(); + devices.has_value()) { + mpmd_addressable_devices.insert( + {name, (*devices)->AddressableDeviceList()->devices()}); + } else { + mpmd_addressable_devices.insert({name, {}}); + } } return mpmd_addressable_devices; } diff --git a/third_party/xla/xla/python/ifrt/ir/ifrt_ir_loaded_executable.h b/third_party/xla/xla/python/ifrt/ir/ifrt_ir_loaded_executable.h index 0ea40a3de134e4..af3e874e8a3809 100644 --- a/third_party/xla/xla/python/ifrt/ir/ifrt_ir_loaded_executable.h +++ b/third_party/xla/xla/python/ifrt/ir/ifrt_ir_loaded_executable.h @@ -82,7 +82,6 @@ class IfrtIrLoadedExecutable UserContextRef user_context() const override { return user_context_; } - int num_devices() const override; int64_t SizeOfGeneratedCodeInBytes() const override; absl::StatusOr GetCompiledMemoryStats() const override; @@ -119,8 +118,6 @@ class IfrtIrLoadedExecutable std::optional devices() const override; - absl::Span addressable_devices() const override; - absl::StatusOr>> GetMpmdAddressableDevices() const override; diff --git a/third_party/xla/xla/python/ifrt/ir/transforms/ifrt_verify_bound_external_loaded_executable_pass.cc b/third_party/xla/xla/python/ifrt/ir/transforms/ifrt_verify_bound_external_loaded_executable_pass.cc index 7b2568d77888eb..8bf6747caf6f18 100644 --- a/third_party/xla/xla/python/ifrt/ir/transforms/ifrt_verify_bound_external_loaded_executable_pass.cc +++ b/third_party/xla/xla/python/ifrt/ir/transforms/ifrt_verify_bound_external_loaded_executable_pass.cc @@ -39,6 +39,7 @@ limitations under the License. #include "mlir/Support/TypeID.h" #include "mlir/Support/WalkResult.h" #include "xla/hlo/ir/hlo_sharding.h" +#include "xla/python/ifrt/device_list.h" #include "xla/python/ifrt/ir/atom_program_compiler.h" #include "xla/python/ifrt/ir/ifrt_dialect.h" #include "xla/python/ifrt/ir/ifrt_ops.h" @@ -117,13 +118,19 @@ void IfrtVerifyBoundExternalLoadedExecutablePass::runOnOperation() { const auto exec_it = bound_executable_map_->find(loaded_exec_op.getSymName()); if (exec_it != bound_executable_map_->end()) { - if (loaded_exec_op.getDevices().size() != - exec_it->second->num_devices()) { + std::optional exec_devices = exec_it->second->devices(); + if (!exec_devices.has_value()) { + loaded_exec_op.emitOpError() + << "cannot be bound to a portable executable; IFRT IR does not " + "support portable executables yet"; + return mlir::WalkResult::interrupt(); + } + if (loaded_exec_op.getDevices().size() != (*exec_devices)->size()) { loaded_exec_op.emitOpError() << "expects an executable with " << loaded_exec_op.getDevices().size() << " devices, but was bound to an executable with " - << exec_it->second->num_devices() << " devices"; + << (*exec_devices)->size() << " devices"; return mlir::WalkResult::interrupt(); } diff --git a/third_party/xla/xla/python/ifrt/mock.h b/third_party/xla/xla/python/ifrt/mock.h index 0115b2df541326..ec5eac7f434e49 100644 --- a/third_party/xla/xla/python/ifrt/mock.h +++ b/third_party/xla/xla/python/ifrt/mock.h @@ -403,7 +403,6 @@ class MockLoadedExecutable MOCK_METHOD(absl::StatusOr, GetHumanReadableProgramText, (), (const, final)); MOCK_METHOD(UserContextRef, user_context, (), (const, final)); - MOCK_METHOD(int, num_devices, (), (const, final)); MOCK_METHOD(int64_t, SizeOfGeneratedCodeInBytes, (), (const, final)); MOCK_METHOD(absl::StatusOr, GetCompiledMemoryStats, (), (const, final)); @@ -432,8 +431,6 @@ class MockLoadedExecutable MOCK_METHOD(absl::StatusOr, ExecuteBundle, (absl::Span args, const ExecuteOptions& options), (final)); - MOCK_METHOD(absl::Span, addressable_devices, (), - (const, final)); MOCK_METHOD(std::optional, devices, (), (const, final)); MOCK_METHOD(void, SetDeleteOptions, (const DeleteOptions& options), (final)); @@ -473,7 +470,6 @@ class MockMpmdLoadedExecutable MOCK_METHOD(absl::StatusOr, GetHumanReadableProgramText, (), (const, final)); MOCK_METHOD(UserContextRef, user_context, (), (const, final)); - MOCK_METHOD(int, num_devices, (), (const, final)); MOCK_METHOD(int64_t, SizeOfGeneratedCodeInBytes, (), (const, final)); MOCK_METHOD(absl::StatusOr, GetCompiledMemoryStats, (), (const, final)); @@ -502,8 +498,6 @@ class MockMpmdLoadedExecutable MOCK_METHOD(absl::StatusOr, ExecuteBundle, (absl::Span args, const ExecuteOptions& options), (final)); - MOCK_METHOD(absl::Span, addressable_devices, (), - (const, final)); MOCK_METHOD(std::optional, devices, (), (const, final)); MOCK_METHOD(void, SetDeleteOptions, (const DeleteOptions& options), (final)); diff --git a/third_party/xla/xla/python/ifrt_proxy/client/compiler.cc b/third_party/xla/xla/python/ifrt_proxy/client/compiler.cc index 53ab16b033c8e7..94d5e0825ef4b9 100644 --- a/third_party/xla/xla/python/ifrt_proxy/client/compiler.cc +++ b/third_party/xla/xla/python/ifrt_proxy/client/compiler.cc @@ -231,16 +231,15 @@ Compiler::CreateExecutableFromResponse( if (is_mpmd_executable) { return std::make_unique( client_, rpc_helper_, response->loaded_executable_handle(), - response->name(), response->num_devices(), device_list, - std::move(addressable_devices), std::move(mpmd_addressable_devices), - std::move(fingerprint), std::move(loaded_host_callbacks), + response->name(), device_list, std::move(addressable_devices), + std::move(mpmd_addressable_devices), std::move(fingerprint), + std::move(loaded_host_callbacks), std::move(loaded_host_callback_handles)); } return std::make_unique( client_, rpc_helper_, response->loaded_executable_handle(), - response->name(), response->num_devices(), device_list, - std::move(addressable_devices), std::move(fingerprint), - std::move(loaded_host_callbacks), + response->name(), device_list, std::move(addressable_devices), + std::move(fingerprint), std::move(loaded_host_callbacks), std::move(loaded_host_callback_handles)); } diff --git a/third_party/xla/xla/python/ifrt_proxy/client/compiler_test.cc b/third_party/xla/xla/python/ifrt_proxy/client/compiler_test.cc index bbf98dcd97abe2..ece36588d4e9aa 100644 --- a/third_party/xla/xla/python/ifrt_proxy/client/compiler_test.cc +++ b/third_party/xla/xla/python/ifrt_proxy/client/compiler_test.cc @@ -175,6 +175,7 @@ TEST_F(CompilerTest, Compile) { loaded_executable_handle: 1234 name: "foo-executable" num_devices: 2 + device_ids: [ 0, 1 ] addressable_device_ids: [ 0, 1 ] fingerprint_value: "fingerprint" })pb", @@ -203,9 +204,7 @@ TEST_F(CompilerTest, Compile) { "xla::ifrt::proxy::TestProgram"); EXPECT_EQ(executable->name(), "foo-executable"); - EXPECT_EQ(executable->num_devices(), 2); - EXPECT_THAT(executable->addressable_devices(), - ElementsAre(&devices[0], &devices[1])); + EXPECT_THAT(executable->devices(), Optional(device_list)); EXPECT_THAT(executable->Fingerprint(), absl_testing::IsOkAndHolds(Optional(std::string("fingerprint")))); } diff --git a/third_party/xla/xla/python/ifrt_proxy/client/executable.cc b/third_party/xla/xla/python/ifrt_proxy/client/executable.cc index b16c106054156b..ccb40f327d371f 100644 --- a/third_party/xla/xla/python/ifrt_proxy/client/executable.cc +++ b/third_party/xla/xla/python/ifrt_proxy/client/executable.cc @@ -345,8 +345,7 @@ class LoadedExecutable::OutputSpecCache { LoadedExecutable::LoadedExecutable( xla::ifrt::Client* client, std::shared_ptr rpc_helper, - uint64_t handle, std::string name, int num_devices, - std::optional devices, + uint64_t handle, std::string name, std::optional devices, std::vector addressable_devices, absl::StatusOr> fingerprint, std::vector> @@ -356,9 +355,7 @@ LoadedExecutable::LoadedExecutable( rpc_helper_(std::move(rpc_helper)), handle_(handle), name_(std::move(name)), - num_devices_(num_devices), devices_(devices), - addressable_devices_(std::move(addressable_devices)), fingerprint_(std::move(fingerprint)), user_context_(xla::ifrt::UserContextScope::current()), output_spec_cache_( @@ -551,8 +548,6 @@ absl::StatusOr LoadedExecutable::Serialize() const { "underlying serialization format is not stable"); } -int LoadedExecutable::num_devices() const { return num_devices_; } - int64_t LoadedExecutable::SizeOfGeneratedCodeInBytes() const { tsl::profiler::TraceMe traceme_ifrt_entrypoint( "IfrtProxyEntrypointLoadedExecutableSizeOfGeneratedCodeInBytes"); @@ -882,11 +877,6 @@ std::optional LoadedExecutable::devices() const { return devices_; } -absl::Span LoadedExecutable::addressable_devices() - const { - return addressable_devices_; -} - tsl::Future<> LoadedExecutable::FetchExecuteResult( uint64_t status_handle, std::optional device_time_key) { if (rpc_helper_->protocol_version() < protocol_version::kExecuteResult) { diff --git a/third_party/xla/xla/python/ifrt_proxy/client/executable.h b/third_party/xla/xla/python/ifrt_proxy/client/executable.h index a3d587b8c655e5..fe677bebc58d7b 100644 --- a/third_party/xla/xla/python/ifrt_proxy/client/executable.h +++ b/third_party/xla/xla/python/ifrt_proxy/client/executable.h @@ -59,8 +59,7 @@ class LoadedExecutable final public: LoadedExecutable(xla::ifrt::Client* client, std::shared_ptr rpc_helper, uint64_t handle, - std::string name, int num_devices, - std::optional devices, + std::string name, std::optional devices, std::vector addressable_devices, absl::StatusOr> fingerprint, std::vector> @@ -82,7 +81,6 @@ class LoadedExecutable final return user_context_; } - int num_devices() const override; int64_t SizeOfGeneratedCodeInBytes() const override; absl::StatusOr GetCompiledMemoryStats() const override; @@ -113,7 +111,6 @@ class LoadedExecutable final absl::Span args, const ExecuteOptions& options) override; std::optional devices() const override; - absl::Span addressable_devices() const override; void SetDeleteOptions(const DeleteOptions& options) override; @@ -153,9 +150,7 @@ class LoadedExecutable final const uint64_t handle_; const std::string name_; - const int num_devices_; const std::optional devices_; - const std::vector addressable_devices_; const absl::StatusOr> fingerprint_; const xla::ifrt::UserContextRef user_context_; diff --git a/third_party/xla/xla/python/ifrt_proxy/client/executable_test.cc b/third_party/xla/xla/python/ifrt_proxy/client/executable_test.cc index e2749c8d058b4a..4f6928cc4cef9c 100644 --- a/third_party/xla/xla/python/ifrt_proxy/client/executable_test.cc +++ b/third_party/xla/xla/python/ifrt_proxy/client/executable_test.cc @@ -144,8 +144,7 @@ TEST_F(LoadedExecutableTest, Metadata) { ON_CALL(device2, Id()).WillByDefault(Return(DeviceId(2))); LoadedExecutable executable( &client, rpc_helper_, /*handle=*/1234, /*name=*/"foo", - /*num_devices=*/2, /*devices=*/{}, - /*addressable_devices=*/{}, + /*devices=*/{}, /*addressable_devices=*/{}, /*fingerprint=*/"fingerprint", /*loaded_host_callbacks=*/{}, /*loaded_host_callback_handles=*/{}); @@ -226,7 +225,7 @@ TEST_F(LoadedExecutableTest, Execute) { LoadedExecutable executable( &client, rpc_helper_, /*handle=*/1234, /*name=*/"foo", - /*num_devices=*/2, /*devices=*/{}, /*addressable_devices=*/{}, + /*devices=*/{}, /*addressable_devices=*/{}, /*fingerprint=*/"fingerprint", /*loaded_host_callbacks=*/{}, /*loaded_host_callback_handles=*/{}); @@ -390,7 +389,7 @@ TEST_F(LoadedExecutableTest, DeviceTime) { LoadedExecutable executable( &client, rpc_helper_, /*handle=*/1234, /*name=*/"foo", - /*num_devices=*/1, /*devices=*/{}, /*addressable_devices=*/{}, + /*devices=*/{}, /*addressable_devices=*/{}, /*fingerprint=*/"fingerprint", /*loaded_host_callbacks=*/{}, /*loaded_host_callback_handles=*/{}); diff --git a/third_party/xla/xla/python/ifrt_proxy/client/mpmd_executable.cc b/third_party/xla/xla/python/ifrt_proxy/client/mpmd_executable.cc index 14f855669bc291..f2716076c47365 100644 --- a/third_party/xla/xla/python/ifrt_proxy/client/mpmd_executable.cc +++ b/third_party/xla/xla/python/ifrt_proxy/client/mpmd_executable.cc @@ -56,8 +56,7 @@ namespace proxy { MpmdLoadedExecutable::MpmdLoadedExecutable( xla::ifrt::Client* client, std::shared_ptr rpc_helper, - uint64_t handle, std::string name, int num_devices, - std::optional devices, + uint64_t handle, std::string name, std::optional devices, std::vector addressable_devices, absl::StatusOr< absl::flat_hash_map>> @@ -70,9 +69,9 @@ MpmdLoadedExecutable::MpmdLoadedExecutable( handle_(handle), mpmd_addressable_devices_(std::move(mpmd_addressable_devices)) { loaded_executable_ = std::make_unique( - client, rpc_helper, handle, std::move(name), num_devices, - std::move(devices), std::move(addressable_devices), - std::move(fingerprint), std::move(loaded_host_callbacks), + client, rpc_helper, handle, std::move(name), std::move(devices), + std::move(addressable_devices), std::move(fingerprint), + std::move(loaded_host_callbacks), std::move(loaded_host_callback_handles)); tsl::profiler::TraceMe traceme_ifrt_entrypoint( diff --git a/third_party/xla/xla/python/ifrt_proxy/client/mpmd_executable.h b/third_party/xla/xla/python/ifrt_proxy/client/mpmd_executable.h index 53d83a311667be..3bbcc8738354a0 100644 --- a/third_party/xla/xla/python/ifrt_proxy/client/mpmd_executable.h +++ b/third_party/xla/xla/python/ifrt_proxy/client/mpmd_executable.h @@ -58,8 +58,7 @@ class MpmdLoadedExecutable final public: MpmdLoadedExecutable( xla::ifrt::Client* client, std::shared_ptr rpc_helper, - uint64_t handle, std::string name, int num_devices, - std::optional devices, + uint64_t handle, std::string name, std::optional devices, std::vector addressable_devices, absl::StatusOr< absl::flat_hash_map>> @@ -92,7 +91,6 @@ class MpmdLoadedExecutable final return loaded_executable_->user_context(); } - int num_devices() const override { return loaded_executable_->num_devices(); } int64_t SizeOfGeneratedCodeInBytes() const override { return loaded_executable_->SizeOfGeneratedCodeInBytes(); } @@ -147,9 +145,6 @@ class MpmdLoadedExecutable final std::optional devices() const override { return loaded_executable_->devices(); }; - absl::Span addressable_devices() const override { - return loaded_executable_->addressable_devices(); - } absl::StatusOr< absl::flat_hash_map>> diff --git a/third_party/xla/xla/python/ifrt_proxy/client/mpmd_executable_test.cc b/third_party/xla/xla/python/ifrt_proxy/client/mpmd_executable_test.cc index 5b03765eefa526..6bc11d85e25451 100644 --- a/third_party/xla/xla/python/ifrt_proxy/client/mpmd_executable_test.cc +++ b/third_party/xla/xla/python/ifrt_proxy/client/mpmd_executable_test.cc @@ -161,7 +161,7 @@ TEST_F(MpmdLoadedExecutableTest, GetMpmdAddressableDevicesSuccess) { MpmdLoadedExecutable executable( &client_, rpc_helper_, /*handle=*/1234, /*name=*/"mpmd_foo", - /*num_devices=*/2, /*devices=*/{}, + /*devices=*/{}, /*addressable_devices=*/{}, /*mpmd_addressable_devices=*/devices_map, /*fingerprint=*/"fingerprint", /*loaded_host_callbacks=*/{}, /*loaded_host_callback_handles=*/{}); @@ -177,7 +177,7 @@ TEST_F(MpmdLoadedExecutableTest, GetMpmdAddressableDevicesError) { MpmdLoadedExecutable executable( &client_, rpc_helper_, /*handle=*/1234, /*name=*/"mpmd_foo", - /*num_devices=*/2, /*devices=*/{}, + /*devices=*/{}, /*addressable_devices=*/{}, /*mpmd_addressable_devices=*/absl::InternalError("injected error"), /*fingerprint=*/"fingerprint", @@ -195,7 +195,7 @@ TEST_F(MpmdLoadedExecutableTest, GetMpmdAddressableDevicesVersionCheck) { MpmdLoadedExecutable executable( &client_, rpc_helper_, /*handle=*/1234, /*name=*/"mpmd_foo", - /*num_devices=*/2, /*devices=*/{}, + /*devices=*/{}, /*addressable_devices=*/{}, /*mpmd_addressable_devices=*/absl::InternalError("injected error"), /*fingerprint=*/"fingerprint", @@ -224,7 +224,7 @@ TEST_F(MpmdLoadedExecutableTest, GetMpmdCompiledMemoryStatsSuccess) { MpmdLoadedExecutable executable( &client_, rpc_helper_, /*handle=*/1234, /*name=*/"mpmd_foo", - /*num_devices=*/1, /*devices=*/{}, + /*devices=*/{}, /*addressable_devices=*/{}, /*mpmd_addressable_devices=*/ absl::flat_hash_map>(), /*fingerprint=*/"fingerprint", @@ -247,7 +247,7 @@ TEST_F(MpmdLoadedExecutableTest, GetMpmdCompiledMemoryStatsRpcError) { MpmdLoadedExecutable executable( &client_, rpc_helper_, /*handle=*/1234, /*name=*/"mpmd_foo", - /*num_devices=*/1, /*devices=*/{}, + /*devices=*/{}, /*addressable_devices=*/{}, /*mpmd_addressable_devices=*/ absl::flat_hash_map>(), /*fingerprint=*/"fingerprint", @@ -272,7 +272,7 @@ TEST_F(MpmdLoadedExecutableTest, GetMpmdCompiledMemoryStatsVersionCheck) { MpmdLoadedExecutable executable( &client_, rpc_helper_, /*handle=*/1234, /*name=*/"mpmd_foo", - /*num_devices=*/1, /*devices=*/{}, + /*devices=*/{}, /*addressable_devices=*/{}, /*mpmd_addressable_devices=*/ absl::flat_hash_map>(), /*fingerprint=*/"fingerprint", @@ -288,7 +288,7 @@ TEST_F(MpmdLoadedExecutableTest, GetMpmdHloModules) { MpmdLoadedExecutable executable( &client_, rpc_helper_, /*handle=*/1234, /*name=*/"mpmd_foo", - /*num_devices=*/1, /*devices=*/{}, + /*devices=*/{}, /*addressable_devices=*/{}, /*mpmd_addressable_devices=*/ absl::flat_hash_map>(), /*fingerprint=*/"fingerprint", @@ -327,7 +327,7 @@ TEST_F(MpmdLoadedExecutableTest, GetMpmdCostAnalysisSuccess) { MpmdLoadedExecutable executable( &client_, rpc_helper_, /*handle=*/1234, /*name=*/"mpmd_foo", - /*num_devices=*/1, /*devices=*/{}, + /*devices=*/{}, /*addressable_devices=*/{}, /*mpmd_addressable_devices=*/ absl::flat_hash_map>(), /*fingerprint=*/"fingerprint", @@ -354,7 +354,7 @@ TEST_F(MpmdLoadedExecutableTest, GetMpmdCostAnalysisRpcError) { MpmdLoadedExecutable executable( &client_, rpc_helper_, /*handle=*/1234, /*name=*/"mpmd_foo", - /*num_devices=*/1, /*devices=*/{}, + /*devices=*/{}, /*addressable_devices=*/{}, /*mpmd_addressable_devices=*/ absl::flat_hash_map>(), /*fingerprint=*/"fingerprint", @@ -378,7 +378,7 @@ TEST_F(MpmdLoadedExecutableTest, GetMpmdCostAnalysisVersionCheck) { MpmdLoadedExecutable executable( &client_, rpc_helper_, /*handle=*/1234, /*name=*/"mpmd_foo", - /*num_devices=*/1, /*devices=*/{}, + /*devices=*/{}, /*addressable_devices=*/{}, /*mpmd_addressable_devices=*/ absl::flat_hash_map>(), /*fingerprint=*/"fingerprint", diff --git a/third_party/xla/xla/python/ifrt_proxy/server/ifrt_backend.cc b/third_party/xla/xla/python/ifrt_proxy/server/ifrt_backend.cc index 1d43424a310538..47cab5f4534fbc 100644 --- a/third_party/xla/xla/python/ifrt_proxy/server/ifrt_backend.cc +++ b/third_party/xla/xla/python/ifrt_proxy/server/ifrt_backend.cc @@ -1517,13 +1517,14 @@ tsl::Future IfrtBackend::HandleCompileRequest( // Populate executable metadata. compile_resp->set_name(AsProtoStringData(executable->name())); - compile_resp->set_num_devices(executable->num_devices()); - for (const auto* device : executable->addressable_devices()) { - compile_resp->add_addressable_device_ids(device->Id().value()); - } if (std::optional device_list = executable->devices(); device_list.has_value()) { + compile_resp->set_num_devices((*device_list)->size()); + for (const auto* device : + (*device_list)->AddressableDeviceList()->devices()) { + compile_resp->add_addressable_device_ids(device->Id().value()); + } for (const auto* device : (*device_list)->devices()) { compile_resp->add_device_ids(device->Id().value()); } diff --git a/third_party/xla/xla/python/ifrt_proxy/server/ifrt_backend_test.cc b/third_party/xla/xla/python/ifrt_proxy/server/ifrt_backend_test.cc index c28526914aa497..2669fd1e2bea93 100644 --- a/third_party/xla/xla/python/ifrt_proxy/server/ifrt_backend_test.cc +++ b/third_party/xla/xla/python/ifrt_proxy/server/ifrt_backend_test.cc @@ -1231,6 +1231,7 @@ TEST_P(IfrtBackendHandlerTest, CompileSuccess) { std::vector devices(4); for (int i = 0; i < 4; ++i) { EXPECT_CALL(devices[i], Id()).WillRepeatedly(Return(DeviceId(i))); + EXPECT_CALL(devices[i], IsAddressable()).WillRepeatedly(Return(true)); } std::vector addressable_devices; @@ -1242,11 +1243,8 @@ TEST_P(IfrtBackendHandlerTest, CompileSuccess) { auto executable = std::make_unique(); EXPECT_CALL(*executable, name()).WillOnce(Return("executable_name")); - EXPECT_CALL(*executable, num_devices()).WillOnce(Return(4)); EXPECT_CALL(*executable, devices()) .WillOnce(Return(std::make_optional(device_list))); - EXPECT_CALL(*executable, addressable_devices()) - .WillOnce(Return(absl::MakeSpan(addressable_devices))); EXPECT_CALL(*executable, Fingerprint()).WillOnce(Return("fingerprint")); TF_ASSERT_OK_AND_ASSIGN(CompileResponse response, @@ -1910,11 +1908,8 @@ TEST_P(IfrtBackendHandlerTest, LoadedExecutableMetadataWithMpmd) { MockMpmdLoadedExecutable* executable = e.get(); ON_CALL(*executable, name()).WillByDefault(Return("mpmd_exec")); - ON_CALL(*executable, num_devices()).WillByDefault(Return(1)); auto device_list = BasicDeviceList::Create({}); ON_CALL(*executable, devices()).WillByDefault(Return(device_list)); - ON_CALL(*executable, addressable_devices()) - .WillByDefault(Return(absl::Span({}))); ON_CALL(*executable, Fingerprint()) .WillByDefault(Return("mpmd_fingerprint")); @@ -1973,11 +1968,8 @@ TEST_P(IfrtBackendHandlerTest, LoadedExecutableMpmdCostAnalysis) { MockMpmdLoadedExecutable* executable = e.get(); ON_CALL(*executable, name()).WillByDefault(Return("mpmd_exec")); - ON_CALL(*executable, num_devices()).WillByDefault(Return(1)); auto device_list = BasicDeviceList::Create({}); ON_CALL(*executable, devices()).WillByDefault(Return(device_list)); - ON_CALL(*executable, addressable_devices()) - .WillByDefault(Return(absl::Span())); ON_CALL(*executable, Fingerprint()) .WillByDefault(Return("mpmd_fingerprint")); @@ -2022,11 +2014,8 @@ TEST_P(IfrtBackendHandlerTest, CompileSuccessWithMpmdAddressableDevices) { auto executable = std::make_unique(); ON_CALL(*executable, name()).WillByDefault(Return("mpmd_exec")); - ON_CALL(*executable, num_devices()).WillByDefault(Return(1)); auto empty_device_list = BasicDeviceList::Create({}); ON_CALL(*executable, devices()).WillByDefault(Return(empty_device_list)); - ON_CALL(*executable, addressable_devices()) - .WillByDefault(Return(absl::Span())); ON_CALL(*executable, Fingerprint()).WillByDefault(Return("mpmd_fingerprint")); std::vector mesh1_devices = {mock_devices_[0].get()}; diff --git a/third_party/xla/xla/python/pjrt_ifrt/pjrt_executable.h b/third_party/xla/xla/python/pjrt_ifrt/pjrt_executable.h index a225a622f0e960..0761ae98dbd897 100644 --- a/third_party/xla/xla/python/pjrt_ifrt/pjrt_executable.h +++ b/third_party/xla/xla/python/pjrt_ifrt/pjrt_executable.h @@ -322,11 +322,6 @@ class PjRtLoadedExecutable final }); } - int num_devices() const override { - DCHECK(this); - return pjrt_loaded_executable_->num_replicas() * - pjrt_loaded_executable_->num_partitions(); - } int64_t SizeOfGeneratedCodeInBytes() const override { DCHECK(this); return pjrt_loaded_executable_->SizeOfGeneratedCodeInBytes(); @@ -367,11 +362,6 @@ class PjRtLoadedExecutable final return devices_; } - absl::Span addressable_devices() const override { - DCHECK(this); - return addressable_devices_; - } - absl::StatusOr GetCostAnalysis() const override { ABSL_ASSIGN_OR_RETURN(auto result, pjrt_loaded_executable_->GetCostAnalysis()); return xla::ifrt::FromPjRtAttributeMap(std::move(result)); diff --git a/third_party/xla/xla/python/pjrt_ifrt/xla_executable_impl_test_lib.cc b/third_party/xla/xla/python/pjrt_ifrt/xla_executable_impl_test_lib.cc index 7e0d88f52534ba..ad7061c97fe881 100644 --- a/third_party/xla/xla/python/pjrt_ifrt/xla_executable_impl_test_lib.cc +++ b/third_party/xla/xla/python/pjrt_ifrt/xla_executable_impl_test_lib.cc @@ -268,7 +268,8 @@ module @add_sub attributes { /*replicated=*/false, serialize)); EXPECT_EQ(executable->name(), "add_sub"); - EXPECT_EQ(executable->num_devices(), devices.size()); + ASSERT_TRUE(executable->devices().has_value()); + EXPECT_THAT((*executable->devices())->devices(), ElementsAreArray(devices)); EXPECT_THAT(executable->GetParameterShardings(), Optional(ElementsAre(EquivToProto(R"pb( diff --git a/third_party/xla/xla/python/version.h b/third_party/xla/xla/python/version.h index 9e7b5e5b37d9e3..2ec5ee5303ba41 100644 --- a/third_party/xla/xla/python/version.h +++ b/third_party/xla/xla/python/version.h @@ -18,6 +18,6 @@ limitations under the License. // An increasing version number to protect jax code against breaking changes. // In JAX, reference this via jax._src.lib.ifrt_version. -#define JAX_IFRT_VERSION_NUMBER 72 // unreduced + partial manual shard_map +#define JAX_IFRT_VERSION_NUMBER 74 // Remove LoadedExecutable::num_devices() #endif // XLA_PYTHON_VERSION_H_ diff --git a/third_party/xla/xla/service/BUILD b/third_party/xla/xla/service/BUILD index 8b3ce74bd0d0da..281accbe3b6b0e 100644 --- a/third_party/xla/xla/service/BUILD +++ b/third_party/xla/xla/service/BUILD @@ -2880,6 +2880,8 @@ cc_library( "//xla/hlo/pass:hlo_pass", "@com_google_absl//absl/algorithm:container", "@com_google_absl//absl/container:flat_hash_set", + "@com_google_absl//absl/log", + "@com_google_absl//absl/log:check", "@com_google_absl//absl/status", "@com_google_absl//absl/status:status_macros", "@com_google_absl//absl/status:statusor", @@ -2898,6 +2900,10 @@ xla_cc_test( "//xla/hlo/ir:hlo_instruction_utils", "//xla/hlo/parser:hlo_parser", "//xla/hlo/testlib:hlo_hardware_independent_test_base", + "@com_google_absl//absl/algorithm:container", + "@com_google_absl//absl/log", + "@com_google_absl//absl/status", + "@com_google_absl//absl/strings:string_view", "@com_google_googletest//:gtest_main", ], ) diff --git a/third_party/xla/xla/service/async_collective_custom_call_rewriter.cc b/third_party/xla/xla/service/async_collective_custom_call_rewriter.cc index 10cb775d4e15b6..4a5ea7d10645eb 100644 --- a/third_party/xla/xla/service/async_collective_custom_call_rewriter.cc +++ b/third_party/xla/xla/service/async_collective_custom_call_rewriter.cc @@ -24,6 +24,8 @@ limitations under the License. #include "absl/algorithm/container.h" #include "absl/container/flat_hash_set.h" +#include "absl/log/check.h" +#include "absl/log/log.h" #include "absl/status/status.h" #include "absl/status/status_macros.h" #include "absl/status/statusor.h" @@ -32,9 +34,11 @@ limitations under the License. #include "absl/strings/string_view.h" #include "absl/strings/strip.h" #include "absl/types/span.h" +#include "xla/hlo/ir/hlo_casting_utils.h" #include "xla/hlo/ir/hlo_computation.h" #include "xla/hlo/ir/hlo_instruction.h" #include "xla/hlo/ir/hlo_instruction_utils.h" +#include "xla/hlo/ir/hlo_instructions.h" #include "xla/hlo/ir/hlo_opcode.h" #include "xla/hlo/ir/replica_group.h" #include "xla/service/collective_ops_utils.h" @@ -72,6 +76,18 @@ bool IsCollectiveCustomCall(const HloInstruction* instr, return absl::c_linear_search(kCollectiveBases, target); } +bool IsComputeOnCustomCall(const HloInstruction* instr, + absl::string_view suffix) { + if (instr->opcode() != HloOpcode::kCustomCall) { + return false; + } + absl::string_view target = instr->custom_call_target(); + if (!absl::ConsumeSuffix(&target, suffix)) { + return false; + } + return target == "compute-on"; +} + std::optional MatchingStartTarget(absl::string_view done_target) { absl::string_view base = done_target; if (!absl::ConsumeSuffix(&base, "-done")) { @@ -433,7 +449,131 @@ absl::StatusOr ProcessCollectivePermute( return true; } -absl::StatusOr ProcessPair( +// Validates that `called_comp` is a compute-on computation that this pass knows +// how to rewrite, i.e. it consists only of parameter instructions and exactly +// one all-gather, which must have a single operand. +// Returns that all-gather on success. +absl::StatusOr FindComputeOnAllGather( + HloComputation* called_comp) { + HloAllGatherInstruction* all_gather = nullptr; + for (HloInstruction* inst : called_comp->instructions()) { + if (inst->opcode() == HloOpcode::kAllGather) { + if (all_gather != nullptr) { + return absl::UnimplementedError(absl::StrCat( + "Rewriting compute-on is only supported when the called " + "computation contains a single all-gather, but ", + called_comp->name(), " contains at least two: ", all_gather->name(), + " and ", inst->name())); + } + all_gather = DynCast(inst); + continue; + } + if (inst->opcode() != HloOpcode::kParameter) { + return absl::UnimplementedError(absl::StrCat( + "Rewriting compute-on is only supported when the called computation " + "consists of all-gather and parameter instructions, but ", + called_comp->name(), " contains a ", HloOpcodeString(inst->opcode()), + " instruction: ", inst->name())); + } + } + if (all_gather == nullptr) { + return absl::UnimplementedError(absl::StrCat( + "Rewriting compute-on is only supported when the called computation " + "contains an all-gather, but none was found in ", + called_comp->name())); + } + if (all_gather->operand_count() != 1) { + return absl::UnimplementedError(absl::StrCat( + "Rewriting compute-on is only supported when the all-gather in the " + "called computation has a single operand, but ", + all_gather->name(), " in ", called_comp->name(), " has ", + all_gather->operand_count(), " operands")); + } + return all_gather; +} + +absl::StatusOr SimplifyEmptyComputeOn( + HloComputation* computation, HloComputation* called_comp, + HloInstruction* start_call, HloInstruction* done_call, + absl::Span + forward_path) { + if (absl::c_any_of(called_comp->instructions(), + [](const HloInstruction* inst) { + return inst->opcode() != HloOpcode::kParameter; + })) { + return false; + } + HloInstruction* call_inst = + computation->AddInstruction(HloInstruction::CreateCall( + done_call->shape(), done_call->operands(), called_comp)); + ABSL_RETURN_IF_ERROR(FinishRewrite(computation, start_call, done_call, + /*async_start=*/start_call->mutable_operand(0), + /*async_done=*/call_inst, forward_path, + /*final_result=*/call_inst)); + return true; +} + +} // namespace + +absl::StatusOr AsyncCollectiveCustomCallRewriter::ProcessComputeOn( + HloComputation* computation, HloInstruction* start_call, + HloInstruction* done_call, + absl::Span forward_path, + bool use_legacy_collectives) { + CHECK_EQ(start_call->called_computations().size(), 1); + absl::string_view called_comp_name = + start_call->called_computations().front()->name(); + HloComputation* called_comp = + start_call->GetModule()->GetComputationWithName(called_comp_name); + + ABSL_ASSIGN_OR_RETURN(bool simplified, + SimplifyEmptyComputeOn(computation, called_comp, start_call, + done_call, forward_path)); + if (simplified) { + return true; + } + ABSL_ASSIGN_OR_RETURN(HloAllGatherInstruction * all_gather, + FindComputeOnAllGather(called_comp)); + HloInstruction* async_start; + HloInstruction* async_done; + if (use_legacy_collectives) { + Shape shape = start_call->shape(); + Shape start_shape = + ShapeUtil::MakeTupleShape({start_call->operand(0)->shape(), shape}); + async_start = + computation->AddInstruction(HloInstruction::CreateAllGatherStart( + start_shape, start_call->operands(), + all_gather->all_gather_dimension(), all_gather->device_list(), + /*constrain_layout=*/false, all_gather->channel_id(), + all_gather->use_global_device_ids())); + async_done = computation->AddInstruction(HloInstruction::CreateUnary( + shape, HloOpcode::kAllGatherDone, async_start)); + async_start->set_frontend_attributes(start_call->frontend_attributes()); + } else { + HloInstruction* call_inst = + computation->AddInstruction(HloInstruction::CreateCall( + done_call->shape(), start_call->operands(), called_comp)); + ABSL_ASSIGN_OR_RETURN(async_done, + computation->CreateAsyncInstructions( + call_inst, /*context_shapes=*/{}, + computation->execution_thread(), /*replace=*/false, + /*override_names=*/false)); + ABSL_RETURN_IF_ERROR(computation->RemoveInstruction(call_inst)); + async_start = async_done->mutable_operand(0); + } + + ABSL_RETURN_IF_ERROR(FinishRewrite(computation, start_call, done_call, async_start, + async_done, forward_path, + /*final_result=*/async_done)); + + if (compute_on_helper_ != nullptr) { + ABSL_RETURN_IF_ERROR( + compute_on_helper_->AddBackendSpecializations(start_call, async_start)); + } + return true; +} + +absl::StatusOr AsyncCollectiveCustomCallRewriter::ProcessPair( HloComputation* computation, HloInstruction* start_call, HloInstruction* done_call, absl::Span forward_path, @@ -459,11 +599,13 @@ absl::StatusOr ProcessPair( return ProcessCollectivePermute(computation, start_call, done_call, forward_path, use_legacy_collectives); } + if (absl::StartsWith(target, "compute-on")) { + return ProcessComputeOn(computation, start_call, done_call, forward_path, + use_legacy_collectives); + } return false; } -} // namespace - absl::StatusOr AsyncCollectiveCustomCallRewriter::RunImpl( HloModule* module, const absl::flat_hash_set& execution_threads) { @@ -472,7 +614,9 @@ absl::StatusOr AsyncCollectiveCustomCallRewriter::RunImpl( module->MakeNonfusionComputations(execution_threads)) { std::vector done_calls; for (HloInstruction* instr : computation->MakeInstructionPostOrder()) { - if (IsCollectiveCustomCall(instr, "-done")) { + if (IsCollectiveCustomCall(instr, "-done") || + IsComputeOnCustomCall(instr, "-done")) { + VLOG(1) << "Candidate done: " << instr->name(); done_calls.push_back(instr); } } @@ -481,6 +625,7 @@ absl::StatusOr AsyncCollectiveCustomCallRewriter::RunImpl( auto expected_start = MatchingStartTarget(done_call->custom_call_target()); if (!expected_start.has_value()) { + VLOG(1) << "No matching start for " << done_call->name(); continue; } auto trace = hlo_instruction_utils::async::TraceDataflowPath( @@ -489,6 +634,7 @@ absl::StatusOr AsyncCollectiveCustomCallRewriter::RunImpl( instr->IsCustomCall(*expected_start); }); if (!trace.has_value()) { + VLOG(1) << "No trace found for " << done_call->name(); continue; } auto [start_call, forward_path] = *std::move(trace); diff --git a/third_party/xla/xla/service/async_collective_custom_call_rewriter.h b/third_party/xla/xla/service/async_collective_custom_call_rewriter.h index a251506ee4bf3e..9e6992c314604f 100644 --- a/third_party/xla/xla/service/async_collective_custom_call_rewriter.h +++ b/third_party/xla/xla/service/async_collective_custom_call_rewriter.h @@ -17,34 +17,64 @@ limitations under the License. #define XLA_SERVICE_ASYNC_COLLECTIVE_CUSTOM_CALL_REWRITER_H_ #include "absl/container/flat_hash_set.h" +#include "absl/status/status.h" #include "absl/status/statusor.h" #include "absl/strings/string_view.h" +#include "absl/types/span.h" +#include "xla/hlo/ir/hlo_instruction.h" +#include "xla/hlo/ir/hlo_instruction_utils.h" #include "xla/hlo/ir/hlo_module.h" #include "xla/hlo/pass/hlo_pass_interface.h" namespace xla { +class AsyncComputeOnHelper { + public: + virtual ~AsyncComputeOnHelper() = default; + virtual absl::Status AddBackendSpecializations( + HloInstruction* custom_call_start, HloInstruction* async_start) { + return absl::OkStatus(); + } +}; + // Convert fake async collective custom calls to async collective instructions. // The fake async collective custom calls are temporary workaround to unblock // the development of async collective operations. This pass should be removed -// once the stable HLO async collective are fully supported. +// as soon as possible. class AsyncCollectiveCustomCallRewriter : public HloModulePass { public: explicit AsyncCollectiveCustomCallRewriter( - bool use_legacy_collectives = false) - : use_legacy_collectives_(use_legacy_collectives) {} + bool use_legacy_collectives = false, + AsyncComputeOnHelper* compute_on_helper = nullptr) + : use_legacy_collectives_(use_legacy_collectives), + compute_on_helper_(compute_on_helper) {} absl::string_view name() const override { return "async-collective-custom-call-rewriter"; } - protected: + using HloPassInterface::Run; absl::StatusOr RunImpl( HloModule* module, const absl::flat_hash_set& execution_threads) override; + absl::StatusOr ProcessPair( + HloComputation* computation, HloInstruction* start_call, + HloInstruction* done_call, + absl::Span + forward_path, + bool use_legacy_collectives); + + absl::StatusOr ProcessComputeOn( + HloComputation* computation, HloInstruction* start_call, + HloInstruction* done_call, + absl::Span + forward_path, + bool use_legacy_collectives); + private: bool use_legacy_collectives_; + AsyncComputeOnHelper* compute_on_helper_; }; } // namespace xla diff --git a/third_party/xla/xla/service/async_collective_custom_call_rewriter_test.cc b/third_party/xla/xla/service/async_collective_custom_call_rewriter_test.cc index 52e0e2afb3b096..c2b19168678478 100644 --- a/third_party/xla/xla/service/async_collective_custom_call_rewriter_test.cc +++ b/third_party/xla/xla/service/async_collective_custom_call_rewriter_test.cc @@ -19,9 +19,15 @@ limitations under the License. #include #include +#include "absl/algorithm/container.h" +#include "absl/log/log.h" +#include "absl/status/status.h" +#include "absl/strings/string_view.h" +#include "xla/hlo/ir/hlo_casting_utils.h" #include "xla/hlo/ir/hlo_computation.h" #include "xla/hlo/ir/hlo_instruction.h" #include "xla/hlo/ir/hlo_instruction_utils.h" +#include "xla/hlo/ir/hlo_instructions.h" #include "xla/hlo/ir/hlo_module.h" #include "xla/hlo/ir/hlo_opcode.h" #include "xla/hlo/parser/hlo_parser.h" @@ -215,5 +221,233 @@ TEST_F(AsyncCollectiveCustomCallRewriterTest, } } +TEST_F(AsyncCollectiveCustomCallRewriterTest, ComputeOnStart) { + constexpr absl::string_view hlo_string = R"( +HloModule test +%compute_on.2.2.clone (new_param: f32[256,256]) -> f32[256,2048] { + %new_param = f32[256,256]{1,0} parameter(0) + ROOT %all_gather_reduced.2 = f32[256,2048]{1,0} all-gather(%new_param), channel_id=1, replica_groups=mesh['x'=8] {'x'}, dimensions={1}, use_global_device_ids=true, frontend_attributes={_xla_compute_type="sparseoffload"}, metadata={op_name="shard_map/all_gather_reduced" stack_frame_id=14} +} + +ENTRY %main.0_spmd (param.1: f32[128,2048], param.2: f32[256,256], param: f32[256,256]) -> f32[128,2048] { + %param = f32[256,256]{1,0} parameter(2), sharding={devices=[1,8]<=[8]}, metadata={op_name="w2"} + %compute_on.2 = f32[256,2048]{1,0} custom-call(%param), custom_call_target="compute-on-start", called_computations={%compute_on.2.2.clone}, frontend_attributes={backend_config={"sparse_core_config": {"core_ids": [0], "core_id_mutability": false}}}, metadata={op_name="jit(f)/compute_on" stack_frame_id=13} + %param.1 = f32[128,2048]{1,0} parameter(0), sharding={devices=[8,1]<=[8]}, metadata={op_name="x"} + %param.2 = f32[256,256]{1,0} parameter(1), sharding={devices=[8,1]<=[8]}, metadata={op_name="w1"} + %tuple.2 = (f32[256,2048]{1,0}, f32[128,2048]{1,0}, f32[256,256]{1,0}) tuple(%compute_on.2, %param.1, %param.2), metadata={op_name="jit(f)/optimization_barrier" stack_frame_id=15} + %optimization_barrier.8 = (f32[256,2048]{1,0}, f32[128,2048]{1,0}, f32[256,256]{1,0}) opt-barrier(%tuple.2), metadata={op_name="jit(f)/optimization_barrier" stack_frame_id=15} + %all-gather.2 = f32[2048,256]{1,0} all-gather(%optimization_barrier.8#2), channel_id=4, replica_groups=mesh['axis_0'=8] {'axis_0'}, dimensions={0}, use_global_device_ids=true, frontend_attributes={is_spmd_generated="true"}, metadata={op_name="jit(f)/optimization_barrier" stack_frame_id=15} + %dot.1 = f32[128,256]{1,0} dot(%optimization_barrier.8#1, %all-gather.2), lhs_contracting_dims={1}, rhs_contracting_dims={0}, metadata={op_name="jit(f)/dot_general" stack_frame_id=15} + %tuple.4 = (f32[128,256]{1,0}, f32[256,2048]{1,0}) tuple(%dot.1, %optimization_barrier.8#0), metadata={op_name="jit(f)/optimization_barrier" stack_frame_id=16} + %optimization_barrier.9 = (f32[128,256]{1,0}, f32[256,2048]{1,0}) opt-barrier(%tuple.4), metadata={op_name="jit(f)/optimization_barrier" stack_frame_id=16} + %compute_on_done.2 = f32[256,2048]{1,0} custom-call(%optimization_barrier.9#1), custom_call_target="compute-on-done", metadata={op_name="jit(f)/compute_on_done" stack_frame_id=16} + %tuple.6 = (f32[128,256]{1,0}, f32[256,2048]{1,0}) tuple(%optimization_barrier.9#0, %compute_on_done.2), metadata={op_name="jit(f)/optimization_barrier" stack_frame_id=17} + %optimization_barrier.10 = (f32[128,256]{1,0}, f32[256,2048]{1,0}) opt-barrier(%tuple.6), metadata={op_name="jit(f)/optimization_barrier" stack_frame_id=17} + ROOT %dot.3 = f32[128,2048]{1,0} dot(%optimization_barrier.10#0, %optimization_barrier.10#1), lhs_contracting_dims={1}, rhs_contracting_dims={0}, metadata={op_name="jit(f)/dot_general" stack_frame_id=17} +} +)"; + + // The rewritten done op only reaches the root through %tuple.6 and + // %optimization_barrier.10, so trace back through those intermediaries. + auto find_async_done = [](HloModule* module) { + return hlo_instruction_utils::async::TraceAsyncDataflow( + module->entry_computation()->root_instruction()->mutable_operand(1), + [](const HloInstruction* instr) { return instr->IsAsyncDone(); }); + }; + + // Legacy collectives mode: the compute-on pair becomes an all-gather + // start/done pair configured from the all-gather inside the compute-on + // computation. + { + ASSERT_OK_AND_ASSIGN(std::unique_ptr module, + ParseAndReturnUnverifiedModule(hlo_string)); + AsyncCollectiveCustomCallRewriter rewriter(/*use_legacy_collectives=*/true); + ASSERT_OK_AND_ASSIGN(bool changed, rewriter.Run(module.get())); + EXPECT_TRUE(changed); + EXPECT_TRUE(HloVerifier(HloVerifierOpts{}).Run(module.get()).status().ok()); + + HloInstruction* done = find_async_done(module.get()); + ASSERT_NE(done, nullptr); + EXPECT_EQ(done->opcode(), HloOpcode::kAllGatherDone); + const HloInstruction* producer = + hlo_instruction_utils::async::FindAsyncProducer(done->operand(0)); + ASSERT_NE(producer, nullptr); + ASSERT_EQ(producer->opcode(), HloOpcode::kAllGatherStart); + + const auto* all_gather_start = Cast(producer); + EXPECT_EQ(all_gather_start->all_gather_dimension(), 1); + EXPECT_EQ(all_gather_start->channel_id(), 1); + EXPECT_TRUE(all_gather_start->use_global_device_ids()); + } + + // Generic async mode: the compute-on pair becomes an async start/done pair + // wrapping a call to the compute-on computation. + { + ASSERT_OK_AND_ASSIGN(std::unique_ptr module, + ParseAndReturnUnverifiedModule(hlo_string)); + AsyncCollectiveCustomCallRewriter rewriter( + /*use_legacy_collectives=*/false); + ASSERT_OK_AND_ASSIGN(bool changed, rewriter.Run(module.get())); + EXPECT_TRUE(changed); + EXPECT_TRUE(HloVerifier(HloVerifierOpts{}).Run(module.get()).status().ok()); + + HloInstruction* done = find_async_done(module.get()); + ASSERT_NE(done, nullptr); + EXPECT_EQ(done->opcode(), HloOpcode::kAsyncDone); + const HloInstruction* producer = + hlo_instruction_utils::async::FindAsyncProducer(done->operand(0)); + ASSERT_NE(producer, nullptr); + ASSERT_EQ(producer->opcode(), HloOpcode::kAsyncStart); + + const HloInstruction* wrapped = producer->async_wrapped_instruction(); + EXPECT_EQ(wrapped->opcode(), HloOpcode::kCall); + ASSERT_EQ(wrapped->called_computations().size(), 1); + EXPECT_EQ(wrapped->called_computations()[0]->name(), + "compute_on.2.2.clone"); + } +} + +TEST_F(AsyncCollectiveCustomCallRewriterTest, ComputeOnStartToSimplify) { + constexpr absl::string_view hlo_string = R"( +HloModule test + %compute_on.2.2.clone { + ROOT %new_param = f32[256,256]{1,0} parameter(0) + } + +ENTRY %main.0_spmd { + %param = f32[256,256]{1,0} parameter(2), sharding={devices=[1,8]<=[8]}, metadata={op_name="w2"} + %compute_on.2 = f32[256,256]{1,0} custom-call(%param), custom_call_target="compute-on-start", called_computations={%compute_on.2.2.clone}, frontend_attributes={backend_config={"sparse_core_config": {"core_ids": [0], "core_id_mutability": false}}}, metadata={op_name="jit(f)/compute_on" stack_frame_id=13} + %param.1 = f32[128,2048]{1,0} parameter(0), sharding={devices=[8,1]<=[8]}, metadata={op_name="x"} + %param.2 = f32[256,256]{1,0} parameter(1), sharding={devices=[8,1]<=[8]}, metadata={op_name="w1"} + %tuple.2 = (f32[256,256]{1,0}, f32[128,2048]{1,0}, f32[256,256]{1,0}) tuple(%compute_on.2, %param.1, %param.2), metadata={op_name="jit(f)/optimization_barrier" stack_frame_id=15} + %optimization_barrier.8 = (f32[256,256]{1,0}, f32[128,2048]{1,0}, f32[256,256]{1,0}) opt-barrier(%tuple.2), metadata={op_name="jit(f)/optimization_barrier" stack_frame_id=15} + %all-gather.2 = f32[2048,256]{1,0} all-gather(%optimization_barrier.8#2), channel_id=4, replica_groups=mesh['axis_0'=8] {'axis_0'}, dimensions={0}, use_global_device_ids=true, frontend_attributes={is_spmd_generated="true"}, metadata={op_name="jit(f)/optimization_barrier" stack_frame_id=15} + %dot.1 = f32[128,256]{1,0} dot(%optimization_barrier.8#1, %all-gather.2), lhs_contracting_dims={1}, rhs_contracting_dims={0}, metadata={op_name="jit(f)/dot_general" stack_frame_id=15} + %tuple.4 = (f32[128,256]{1,0}, f32[256,256]{1,0}) tuple(%dot.1, %optimization_barrier.8#0), metadata={op_name="jit(f)/optimization_barrier" stack_frame_id=16} + %optimization_barrier.9 = (f32[128,256]{1,0}, f32[256,256]{1,0}) opt-barrier(%tuple.4), metadata={op_name="jit(f)/optimization_barrier" stack_frame_id=16} + %compute_on_done.2 = f32[256,256]{1,0} custom-call(%optimization_barrier.9#1), custom_call_target="compute-on-done", metadata={op_name="jit(f)/compute_on_done" stack_frame_id=16} + %tuple.6 = (f32[128,256]{1,0}, f32[256,256]{1,0}) tuple(%optimization_barrier.9#0, %compute_on_done.2), metadata={op_name="jit(f)/optimization_barrier" stack_frame_id=17} + %optimization_barrier.10 = (f32[128,256]{1,0}, f32[256,2048]{1,0}) opt-barrier(%tuple.6), metadata={op_name="jit(f)/optimization_barrier" stack_frame_id=17} + ROOT %dot.3 = f32[128,256]{1,0} dot(%optimization_barrier.10#0, %optimization_barrier.10#1), lhs_contracting_dims={1}, rhs_contracting_dims={0}, metadata={op_name="jit(f)/dot_general" stack_frame_id=17} +} +)"; + + ASSERT_OK_AND_ASSIGN(std::unique_ptr module, + ParseAndReturnUnverifiedModule(hlo_string)); + AsyncCollectiveCustomCallRewriter rewriter(/*use_legacy_collectives=*/true); + ASSERT_OK_AND_ASSIGN(bool changed, rewriter.Run(module.get())); + EXPECT_TRUE(changed); + auto instructions = module->entry_computation()->instructions(); + ASSERT_TRUE(absl::c_none_of(instructions, [](const HloInstruction* instr) { + return instr->opcode() == HloOpcode::kCustomCall; + })); +} + +TEST_F(AsyncCollectiveCustomCallRewriterTest, + ComputeOnWithUnsupportedComputationIsUnimplemented) { + // The compute-on computation contains an instruction that is neither an + // all-gather nor a parameter. + constexpr absl::string_view unsupported_instruction_hlo = R"( + HloModule test + + %compute_on_computation (new_param: f32[256,256]) -> f32[256,256] { + %new_param = f32[256,256]{1,0} parameter(0) + ROOT %negate = f32[256,256]{1,0} negate(%new_param) + } + + ENTRY %main (param: f32[256,256]) -> f32[256,256] { + %param = f32[256,256]{1,0} parameter(0) + %compute_on = f32[256,256]{1,0} custom-call(%param), + custom_call_target="compute-on-start", + called_computations={%compute_on_computation} + ROOT %compute_on_done = f32[256,256]{1,0} custom-call(%compute_on), + custom_call_target="compute-on-done" + } + )"; + + // The compute-on computation contains an all-gather with more than one + // operand, which the rewrite does not know how to handle. + constexpr absl::string_view variadic_all_gather_hlo = R"( + HloModule test + + %compute_on_computation (p0: f32[256,256], p1: f32[256,256]) -> (f32[256,2048], f32[256,2048]) { + %p0 = f32[256,256]{1,0} parameter(0) + %p1 = f32[256,256]{1,0} parameter(1) + ROOT %all_gather = (f32[256,2048]{1,0}, f32[256,2048]{1,0}) all-gather(%p0, %p1), channel_id=1, replica_groups={{0,1,2,3,4,5,6,7}}, dimensions={1}, use_global_device_ids=true + } + + ENTRY %main (param: f32[256,256]) -> (f32[256,2048], f32[256,2048]) { + %param = f32[256,256]{1,0} parameter(0) + %compute_on = (f32[256,2048]{1,0}, f32[256,2048]{1,0}) custom-call(%param, %param), + custom_call_target="compute-on-start", + called_computations={%compute_on_computation} + ROOT %compute_on_done = (f32[256,2048]{1,0}, f32[256,2048]{1,0}) custom-call(%compute_on), + custom_call_target="compute-on-done" + } + )"; + + constexpr absl::string_view two_all_gathers_hlo = R"( + HloModule test + + %compute_on_computation (p0: f32[256,256]) -> f32[256,16384] { + %p0 = f32[256,256]{1,0} parameter(0) + %all_gather.0 = f32[256,2048]{1,0} all-gather(%p0), channel_id=1, replica_groups={{0,1,2,3,4,5,6,7}}, dimensions={1}, use_global_device_ids=true + ROOT %all_gather.1 = f32[256,16384]{1,0} all-gather(%all_gather.0), channel_id=2, replica_groups={{0,1,2,3,4,5,6,7}}, dimensions={1}, use_global_device_ids=true + } + + ENTRY %main (param: f32[256,256]) -> f32[256,16384] { + %param = f32[256,256]{1,0} parameter(0) + %compute_on = f32[256,16384]{1,0} custom-call(%param), + custom_call_target="compute-on-start", + called_computations={%compute_on_computation} + ROOT %compute_on_done = f32[256,16384]{1,0} custom-call(%compute_on), + custom_call_target="compute-on-done" + } + )"; + + // Both modes only wrap computations made up of all-gather and parameter + // instructions, so the negate is rejected in either mode. + for (bool use_legacy_collectives : {true, false}) { + SCOPED_TRACE(use_legacy_collectives ? "use_legacy_collectives=true" + : "use_legacy_collectives=false"); + ASSERT_OK_AND_ASSIGN( + std::unique_ptr module, + ParseAndReturnUnverifiedModule(unsupported_instruction_hlo)); + AsyncCollectiveCustomCallRewriter rewriter(use_legacy_collectives); + EXPECT_THAT(rewriter.Run(module.get()), + ::testing::status::StatusIs( + absl::StatusCode::kUnimplemented, + ::testing::HasSubstr("contains a negate instruction"))); + } + + // Both modes require the all-gather to have exactly one operand. + for (bool use_legacy_collectives : {true, false}) { + SCOPED_TRACE(use_legacy_collectives ? "use_legacy_collectives=true" + : "use_legacy_collectives=false"); + ASSERT_OK_AND_ASSIGN( + std::unique_ptr module, + ParseAndReturnUnverifiedModule(variadic_all_gather_hlo)); + AsyncCollectiveCustomCallRewriter rewriter(use_legacy_collectives); + EXPECT_THAT( + rewriter.Run(module.get()), + ::testing::status::StatusIs(absl::StatusCode::kUnimplemented, + ::testing::HasSubstr("has 2 operands"))); + } + + // Both modes require the compute-on computation to contain exactly one + // all-gather. + for (bool use_legacy_collectives : {true, false}) { + SCOPED_TRACE(use_legacy_collectives ? "use_legacy_collectives=true" + : "use_legacy_collectives=false"); + ASSERT_OK_AND_ASSIGN(std::unique_ptr module, + ParseAndReturnUnverifiedModule(two_all_gathers_hlo)); + AsyncCollectiveCustomCallRewriter rewriter(use_legacy_collectives); + EXPECT_THAT(rewriter.Run(module.get()), + ::testing::status::StatusIs( + absl::StatusCode::kUnimplemented, + ::testing::HasSubstr("contains at least two"))); + } +} + } // namespace } // namespace xla diff --git a/third_party/xla/xla/service/gpu/BUILD b/third_party/xla/xla/service/gpu/BUILD index 6b875d38711fb8..73d9ad7dc320c8 100644 --- a/third_party/xla/xla/service/gpu/BUILD +++ b/third_party/xla/xla/service/gpu/BUILD @@ -3899,13 +3899,16 @@ cc_library( hdrs = ["gpu_hlo_ordering.h"], deps = [ "//xla:shape_util", + "//xla/ffi:ffi_registry", "//xla/hlo/analysis:hlo_ordering", "//xla/hlo/analysis:hlo_reachability", "//xla/hlo/ir:hlo", + "//xla/service:hlo_proto_cc", "@com_google_absl//absl/algorithm:container", "@com_google_absl//absl/container:flat_hash_map", "@com_google_absl//absl/log", "@com_google_absl//absl/log:check", + "@com_google_absl//absl/status:statusor", "@com_google_absl//absl/strings", ], ) @@ -3918,9 +3921,13 @@ xla_cc_test( ":gpu_hlo_ordering", "//xla:shape_util", "//xla:xla_data_proto_cc", + "//xla/ffi", + "//xla/ffi:ffi_api", # buildcleaner: keep "//xla/hlo/ir:hlo", "//xla/hlo/testlib:hlo_hardware_independent_test_base", + "//xla/service:hlo_proto_cc", "//xla/tsl/lib/core:status_test_util", + "@com_google_absl//absl/status", "@com_google_absl//absl/status:status_matchers", "@com_google_absl//absl/strings:string_view", "@com_google_googletest//:gtest_main", diff --git a/third_party/xla/xla/service/gpu/gpu_hlo_ordering.cc b/third_party/xla/xla/service/gpu/gpu_hlo_ordering.cc index ccce380544c757..8e028fe556f33b 100644 --- a/third_party/xla/xla/service/gpu/gpu_hlo_ordering.cc +++ b/third_party/xla/xla/service/gpu/gpu_hlo_ordering.cc @@ -24,12 +24,17 @@ limitations under the License. #include "absl/algorithm/container.h" #include "absl/log/check.h" #include "absl/log/log.h" +#include "absl/status/statusor.h" #include "absl/strings/str_cat.h" +#include "xla/ffi/ffi_registry.h" #include "xla/hlo/analysis/hlo_ordering.h" #include "xla/hlo/analysis/hlo_reachability.h" +#include "xla/hlo/ir/hlo_casting_utils.h" #include "xla/hlo/ir/hlo_instruction.h" +#include "xla/hlo/ir/hlo_instructions.h" #include "xla/hlo/ir/hlo_opcode.h" #include "xla/hlo/ir/hlo_schedule.h" +#include "xla/service/hlo.pb.h" #include "xla/shape_util.h" namespace xla { @@ -79,10 +84,21 @@ bool CanRunConcurrently(const HloInstruction* hlo) { } } } - // Custom calls (cuBLAS, cuDNN, etc.) are not considered for concurrent - // execution. + // Legacy custom calls (cuBLAS, cuDNN, etc.) are not considered for + // concurrent execution, whereas typed FFI custom calls registered with + // kCmdBufferCompatible can run concurrently. if (hlo->opcode() == HloOpcode::kCustomCall) { - return false; + const auto* custom_call = Cast(hlo); + if (custom_call->api_version() != + CustomCallApiVersion::API_VERSION_TYPED_FFI) { + return false; + } + absl::StatusOr registration = + ffi::FindHandler(custom_call->custom_call_target(), "gpu"); + if (!registration.ok() || + !ffi::IsCommandBufferCompatible(registration->metadata)) { + return false; + } } // Very large HLO fusions compile into many device instructions. Concurrent // execution can lead to instruction cache thrashing and is not beneficial. diff --git a/third_party/xla/xla/service/gpu/gpu_hlo_ordering_test.cc b/third_party/xla/xla/service/gpu/gpu_hlo_ordering_test.cc index 3d958513946402..f763565b5e6890 100644 --- a/third_party/xla/xla/service/gpu/gpu_hlo_ordering_test.cc +++ b/third_party/xla/xla/service/gpu/gpu_hlo_ordering_test.cc @@ -16,13 +16,18 @@ limitations under the License. #include "xla/service/gpu/gpu_hlo_ordering.h" #include +#include "absl/status/status.h" #include "absl/status/status_matchers.h" #include "absl/strings/string_view.h" +#include "xla/ffi/ffi.h" +#include "xla/hlo/ir/hlo_casting_utils.h" #include "xla/hlo/ir/hlo_computation.h" #include "xla/hlo/ir/hlo_instruction.h" +#include "xla/hlo/ir/hlo_instructions.h" #include "xla/hlo/ir/hlo_opcode.h" #include "xla/hlo/ir/hlo_schedule.h" #include "xla/hlo/testlib/hlo_hardware_independent_test_base.h" +#include "xla/service/hlo.pb.h" #include "xla/shape.h" #include "xla/shape_util.h" #include "xla/tsl/lib/core/status_test_util.h" @@ -32,6 +37,16 @@ namespace xla { namespace gpu { namespace { +absl::Status DummyCmdBufferCompatibleFfiHandler() { return absl::OkStatus(); } + +XLA_FFI_DEFINE_HANDLER(kDummyCmdBufferCompatibleFfiHandler, + DummyCmdBufferCompatibleFfiHandler, ffi::Ffi::Bind(), + {ffi::Traits::kCmdBufferCompatible}); + +XLA_FFI_REGISTER_HANDLER(ffi::GetXlaFfiApi(), + "__test_cmd_buffer_compatible_ffi", "GPU", + kDummyCmdBufferCompatibleFfiHandler); + class ConcurrentRegionsHloOrderingTest : public HloHardwareIndependentTestBase { }; @@ -174,6 +189,40 @@ TEST_F(ConcurrentRegionsHloOrderingTest, GemmSeparatesConcurrentRegions) { EXPECT_EQ(ordering.GetConcurrentRegionId(d), 2); } +TEST_F(ConcurrentRegionsHloOrderingTest, + CmdBufferCompatibleFfiCustomCallStaysInSameRegion) { + auto module = CreateNewVerifiedModule(); + const Shape small_shape = ShapeUtil::MakeShape(xla::F32, {1024, 1024}); + + auto builder = HloComputation::Builder(TestName()); + HloInstruction* param = builder.AddInstruction( + HloInstruction::CreateParameter(0, small_shape, "param")); + HloInstruction* ffi_call = + builder.AddInstruction(HloInstruction::CreateCustomCall( + small_shape, {param}, "__test_cmd_buffer_compatible_ffi")); + Cast(ffi_call)->set_api_version( + CustomCallApiVersion::API_VERSION_TYPED_FFI); + HloInstruction* a = builder.AddInstruction( + HloInstruction::CreateUnary(small_shape, HloOpcode::kNegate, param)); + HloInstruction* legacy_call = builder.AddInstruction( + HloInstruction::CreateCustomCall(small_shape, {param}, "legacy_target")); + HloInstruction* root = builder.AddInstruction( + HloInstruction::CreateTuple({ffi_call, a, legacy_call})); + HloComputation* entry = + module->AddEntryComputation(builder.Build(/*root_instruction=*/root)); + + HloSchedule schedule(module.get()); + schedule.set_sequence(entry, {param, ffi_call, a, legacy_call, root}); + ASSERT_OK(schedule.Verify()); + ConcurrentRegionsHloOrdering ordering(schedule); + + EXPECT_FALSE(ordering.ExecutesBefore(ffi_call, a)); + EXPECT_EQ(ordering.GetConcurrentRegionId(ffi_call), 0); + EXPECT_EQ(ordering.GetConcurrentRegionId(a), 0); + EXPECT_TRUE(ordering.ExecutesBefore(a, legacy_call)); + EXPECT_EQ(ordering.GetConcurrentRegionId(legacy_call), 1); +} + } // namespace } // namespace gpu } // namespace xla diff --git a/third_party/xla/xla/stream_executor/cuda/cudart_kernel_registry_test.cc b/third_party/xla/xla/stream_executor/cuda/cudart_kernel_registry_test.cc index 660aedd4d1b494..f2feff6efa307b 100644 --- a/third_party/xla/xla/stream_executor/cuda/cudart_kernel_registry_test.cc +++ b/third_party/xla/xla/stream_executor/cuda/cudart_kernel_registry_test.cc @@ -39,6 +39,7 @@ limitations under the License. namespace stream_executor::cuda { namespace { +using ::testing::AnyOf; using ::testing::Optional; using Write42Kernel = TypedKernelFactory, int32_t>; @@ -147,9 +148,19 @@ TEST(CudaRuntimeKernelRegistryTest, ExtractsFuncAttributes) { EXPECT_EQ(attrs.compute_capability.major, cc.major); EXPECT_EQ(attrs.compute_capability.minor, cc.minor); EXPECT_EQ(attrs.num_regs, runtime_attrs.numRegs); - EXPECT_EQ(attrs.static_shared_size_bytes, + int reserved_shared_bytes = 0; + ASSERT_EQ(cudaDeviceGetAttribute(&reserved_shared_bytes, + cudaDevAttrReservedSharedMemoryPerBlock, 0), + cudaSuccess); + EXPECT_EQ(runtime_attrs.sharedSizeBytes, kWrite42SharedElements * sizeof(int32_t)); - EXPECT_EQ(attrs.static_shared_size_bytes, runtime_attrs.sharedSizeBytes); + // On Hopper+ (sm_90+), the CUBIN ELF `.nv.shared.` section size + // includes architecture-reserved shared memory per block (e.g., 1024 bytes), + // whereas `cudaFuncGetAttributes` reports only user-declared static shared + // memory in `sharedSizeBytes`. + EXPECT_THAT(attrs.static_shared_size_bytes, + AnyOf(runtime_attrs.sharedSizeBytes, + runtime_attrs.sharedSizeBytes + reserved_shared_bytes)); EXPECT_THAT(attrs.max_threads_per_block, Optional(kWrite42LaunchBoundsMaxThreads)); EXPECT_THAT(attrs.max_threads_per_block,