From 25453d5a90bb1fbb890c89409edf97a12bba7f32 Mon Sep 17 00:00:00 2001 From: Sohaib Iftikhar Date: Mon, 28 Sep 2026 10:44:16 -0700 Subject: [PATCH 01/21] [XLA:GPU] Add asymmetric mode to BlockBarrierOp Extend `triton_xla.block_barrier` with optional `signal_slot` operand and `signal_stride` attributes to support asymmetric block synchronization across GPUs. Before this change Block B on R0 waited for Block B on other GPUs. With this now we can do something like Block B waits for a strided subset of R (world_size) blocks from another GPU. PiperOrigin-RevId: 989715923 --- .../gpu/codegen/triton/collective_emitter.cc | 10 ++-- .../gpu/codegen/triton/ir/tests/invalid.mlir | 17 ++++++ .../gpu/codegen/triton/ir/triton_xla_ops.cc | 13 +++++ .../gpu/codegen/triton/ir/triton_xla_ops.td | 57 ++++++++++++++----- .../tests/triton_xla_block_barrier.mlir | 43 ++++++++++++++ .../triton_xla_lower_block_barrier_pass.cc | 47 ++++++++++++++- 6 files changed, 165 insertions(+), 22 deletions(-) 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/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( From e2957f29941cbe209c42a265044099039e26fbde Mon Sep 17 00:00:00 2001 From: Joshua Lang Date: Mon, 28 Sep 2026 12:06:01 -0700 Subject: [PATCH 02/21] Deallocate GPU buffers in TopKKernelTest to prevent OOM across parameterized runs PiperOrigin-RevId: 989761006 --- third_party/xla/xla/backends/gpu/runtime/topk_test.cc | 6 ++++++ 1 file changed, 6 insertions(+) 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) { From 7b8368c889ba4efd0d25b8473126a2e42ae9c359 Mon Sep 17 00:00:00 2001 From: Terry Heo Date: Mon, 28 Sep 2026 12:09:34 -0700 Subject: [PATCH 03/21] Build the TFLite C library from LiteRT sources in CMake tflite/c/CMakeLists.txt still expected a TensorFlow checkout layout. - Derived TFLITE_SOURCE_DIR from the tflite/c directory and dropped TF_SOURCE_DIR. - Took tflite_types.h from tflite/converter. - Moved the TensorFlow header download pin from v2.21.0-rc0 to v2.21.0. The file hashes are unchanged. PiperOrigin-RevId: 989763450 --- tensorflow/lite/c/CMakeLists.txt | 15 ++++----------- .../cmake/modules/Findtensorflow_headers.cmake | 2 +- 2 files changed, 5 insertions(+), 12 deletions(-) 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/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") # || From 46a616668aa2f285132df5186a9bea6879f1a5a2 Mon Sep 17 00:00:00 2001 From: Mohammed Anany Date: Mon, 28 Sep 2026 12:26:11 -0700 Subject: [PATCH 04/21] [XLA:GPU] Add a test for a tiled reduce duplicated across concatenate regions. This case is different from existing tests that cover duplications in nested regions. This one adds coverage for sibling regions, which shouldn't see each other during emission. PiperOrigin-RevId: 989774195 --- .../tests/reduce/reduce_concat_diamond.hlo | 46 +++++++++++++++++++ 1 file changed, 46 insertions(+) create mode 100644 third_party/xla/xla/backends/gpu/codegen/triton/tests/reduce/reduce_concat_diamond.hlo 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 From dde838be435d132b20941b807893d7825dc160b2 Mon Sep 17 00:00:00 2001 From: Joshua Lang Date: Mon, 28 Sep 2026 12:26:46 -0700 Subject: [PATCH 05/21] Configure DTensor GPU test defaults for Ampere+ GPUs and smaller GPU memory partitions - Set `NVIDIA_TF32_OVERRIDE=0` for DTensor GPU tests so FP32 tests run with IEEE FP32 precision rather than TF32 on Ampere+ GPUs. - Disable `NCCL_CUMEM_ENABLE`, `NCCL_CUMEM_HOST_ENABLE`, and `NCCL_NVLS_ENABLE` in 2-GPU tests for multi-process NCCL compatibility in containerized test environments. - Reduce `_DEFAULT_GPU_MEMORY_LIMIT` from 1024MB to 512MB per virtual GPU in `test_util.py` so 8 virtual devices plus CUDA/cuSolver runtime overhead fit within smaller GPU memory partitions. PiperOrigin-RevId: 989774536 --- tensorflow/dtensor/python/tests/test_util.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) 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): From 53de455e62326e850ee2b06aaeed65ef34ddb665 Mon Sep 17 00:00:00 2001 From: Matthias Guenther Date: Mon, 28 Sep 2026 12:29:30 -0700 Subject: [PATCH 06/21] Update deprecated MLIR API call sites Also, fix a minor formatting issue in `Api.cpp`. PiperOrigin-RevId: 989776315 --- .../xla/third_party/stablehlo/temporary.patch | 259 +++++++++++++++++- 1 file changed, 258 insertions(+), 1 deletion(-) 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 From 463eb9e021bb128460bedc9d8989039ba6a0304d Mon Sep 17 00:00:00 2001 From: "A. Unique TensorFlower" Date: Mon, 28 Sep 2026 12:30:26 -0700 Subject: [PATCH 07/21] Offload EagerContext destructor to a detached thread when it is the last reference. Prevents a thread pool starvation deadlock during `EagerContext` destruction by unconditionally offloading the `Unref()` call via `env_->StartDetachedThread` when called from an executor thread pool. PiperOrigin-RevId: 989776895 --- tensorflow/core/common_runtime/eager/BUILD | 4 + .../core/common_runtime/eager/context.cc | 24 ++++++ .../core/common_runtime/eager/context.h | 7 +- .../core/common_runtime/eager/context_test.cc | 80 ++++++++++++++++++- 4 files changed, 111 insertions(+), 4 deletions(-) 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 From 403dcdb5ee965fe7273bab90fe5be4d4354382c1 Mon Sep 17 00:00:00 2001 From: Joshua Lang Date: Mon, 28 Sep 2026 12:34:34 -0700 Subject: [PATCH 08/21] Replace hardcoded requires-gpu-nvidia tags in tf_cuda_cc_test targets PiperOrigin-RevId: 989779426 --- tensorflow/compiler/jit/BUILD | 6 ++---- tensorflow/compiler/tests/BUILD | 2 -- 2 files changed, 2 insertions(+), 6 deletions(-) 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"], ) From 0b51168be8612d820bfc66a9f8641fc0286686fd Mon Sep 17 00:00:00 2001 From: Ibrahim Umit Akgun Date: Mon, 28 Sep 2026 12:35:17 -0700 Subject: [PATCH 09/21] `ConcurrentRegionsHloOrdering::CanRunConcurrently` previously rejected all `HloOpcode::kCustomCall` instructions, isolating every typed FFI custom kernel into its own singleton concurrent region with barrier dependencies before and after. Additionally, `CommandBufferConversionPass` did not map `Thunk::kCollectiveKernel` to `DebugOptions::COLLECTIVES`, causing one-shot collective kernels to break and split surrounding CUDA Graphs. Update `CanRunConcurrently` in `gpu_hlo_ordering.cc` to allow `API_VERSION_TYPED_FFI` custom calls whose registered GPU FFI handler specifies `xla::ffi::Traits::kCmdBufferCompatible`, while continuing to isolate legacy custom calls (such as cuBLAS and cuDNN). Map `Thunk::kCollectiveKernel` to `DebugOptions::COLLECTIVES` and `DebugOptions::ALLREDUCE` in `command_buffer_conversion_pass.cc` so one-shot collective kernels are captured inside command buffers. PiperOrigin-RevId: 989779912 --- .../runtime/command_buffer_conversion_pass.cc | 2 + third_party/xla/xla/service/gpu/BUILD | 7 +++ .../xla/xla/service/gpu/gpu_hlo_ordering.cc | 22 +++++++-- .../xla/service/gpu/gpu_hlo_ordering_test.cc | 49 +++++++++++++++++++ 4 files changed, 77 insertions(+), 3 deletions(-) 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/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 From a4b4c6fbdece6ab127d1aa70a0fbea8dd301e42f Mon Sep 17 00:00:00 2001 From: Joshua Lang Date: Mon, 28 Sep 2026 12:39:09 -0700 Subject: [PATCH 10/21] Account for Hopper reserved shared memory in CudaRuntimeKernelRegistryTest.ExtractsFuncAttributes PiperOrigin-RevId: 989782323 --- .../cuda/cudart_kernel_registry_test.cc | 15 +++++++++++++-- 1 file changed, 13 insertions(+), 2 deletions(-) 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, From eebea9be23661439d3bf54212418ab775f8ca73b Mon Sep 17 00:00:00 2001 From: "A. Unique TensorFlower" Date: Mon, 28 Sep 2026 12:40:13 -0700 Subject: [PATCH 11/21] Add async support for onednn on aarch64 (the changes are integrated from https://github.com/openxla/xla/pull/48616) PiperOrigin-RevId: 989782947 --- .../mkl_dnn/mkldnn_v1_async_aarch64.BUILD | 196 ++++++++++ .../onednn_async_aarch64_pr_5790.patch | 339 ++++++++++++++++++ .../onednn_async_aarch64_pr_5828.patch | 257 +++++++++++++ .../onednn_async_aarch64_pr_6033.patch | 206 +++++++++++ .../xla/third_party/mkl_dnn/workspace.bzl | 15 + 5 files changed, 1013 insertions(+) create mode 100644 third_party/xla/third_party/mkl_dnn/mkldnn_v1_async_aarch64.BUILD create mode 100644 third_party/xla/third_party/mkl_dnn/onednn_async_aarch64_pr_5790.patch create mode 100644 third_party/xla/third_party/mkl_dnn/onednn_async_aarch64_pr_5828.patch create mode 100644 third_party/xla/third_party/mkl_dnn/onednn_async_aarch64_pr_6033.patch 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", From dabe608b0b30c318a27e2965c489f8ba8b223334 Mon Sep 17 00:00:00 2001 From: Joshua Lang Date: Mon, 28 Sep 2026 12:43:11 -0700 Subject: [PATCH 12/21] Relax float16 tolerance to 1e-2 in StatefulScatterNdTest._VariableRankTest PiperOrigin-RevId: 989784865 --- .../python/kernel_tests/array_ops/scatter_nd_ops_test.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) 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): From 4d007feae90e3caa452ac0a70270d22fea81100c Mon Sep 17 00:00:00 2001 From: "A. Unique TensorFlower" Date: Mon, 28 Sep 2026 12:57:22 -0700 Subject: [PATCH 13/21] Reverts changelist 988511947 PiperOrigin-RevId: 989793164 --- .../xla/third_party/tsl/tsl/platform/BUILD | 1 + .../tsl/platform/unbounded_work_queue_test.cc | 22 +------------------ 2 files changed, 2 insertions(+), 21 deletions(-) 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 From e20101c4ec97978402f0fb77634839b970c4fde9 Mon Sep 17 00:00:00 2001 From: Joshua Lang Date: Mon, 28 Sep 2026 13:12:23 -0700 Subject: [PATCH 14/21] Disable TensorFloat-32 in JacobianTest.test_grad_jacobian_conv PiperOrigin-RevId: 989802032 --- tensorflow/python/eager/backprop_test.py | 3 +++ 1 file changed, 3 insertions(+) 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]) From 4d7aef9a4e4f05e850fd32e93863b0ec858d2777 Mon Sep 17 00:00:00 2001 From: Junwhan Ahn Date: Mon, 28 Sep 2026 13:19:25 -0700 Subject: [PATCH 15/21] Remove LoadedExecutable::addressable_devices() and migrate callers to devices(). `xla::ifrt::LoadedExecutable::addressable_devices()` is superseded by `xla::ifrt::LoadedExecutable::devices()`, which returns `std::optional` and yields `std::nullopt` for portable executables. Update remaining callers across JAX, IFRT IR, and IFRT Proxy to query `LoadedExecutable::devices()` and derive the addressable device list (`(*devices)->AddressableDeviceList()->devices()`) from the returned `DeviceListRef`. When `devices()` returns `std::nullopt`, callers fall back to an empty addressable device list (or return `UnimplementedError` where portable executables are unsupported). Remove `LoadedExecutable::addressable_devices()` from the base interface, mocks, and all subclass implementations. PiperOrigin-RevId: 989806232 --- .../xla/xla/backends/cpu/nanort/ifrt_client.cc | 4 ---- third_party/xla/xla/python/ifrt/executable.h | 7 ------- .../xla/python/ifrt/ir/ifrt_ir_loaded_executable.cc | 13 +++++++------ .../xla/python/ifrt/ir/ifrt_ir_loaded_executable.h | 2 -- third_party/xla/xla/python/ifrt/mock.h | 4 ---- .../xla/python/ifrt_proxy/client/compiler_test.cc | 4 ++-- .../xla/xla/python/ifrt_proxy/client/executable.cc | 6 ------ .../xla/xla/python/ifrt_proxy/client/executable.h | 2 -- .../xla/python/ifrt_proxy/client/mpmd_executable.h | 3 --- .../xla/python/ifrt_proxy/server/ifrt_backend.cc | 7 ++++--- .../python/ifrt_proxy/server/ifrt_backend_test.cc | 9 +-------- .../xla/xla/python/pjrt_ifrt/pjrt_executable.h | 5 ----- third_party/xla/xla/python/version.h | 2 +- 13 files changed, 15 insertions(+), 53 deletions(-) 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..8ad71ec014f0cf 100644 --- a/third_party/xla/xla/backends/cpu/nanort/ifrt_client.cc +++ b/third_party/xla/xla/backends/cpu/nanort/ifrt_client.cc @@ -1087,10 +1087,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/python/ifrt/executable.h b/third_party/xla/xla/python/ifrt/executable.h index 154b5c36336580..efd5418cf3b17c 100644 --- a/third_party/xla/xla/python/ifrt/executable.h +++ b/third_party/xla/xla/python/ifrt/executable.h @@ -327,13 +327,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/ir/ifrt_ir_loaded_executable.cc b/third_party/xla/xla/python/ifrt/ir/ifrt_ir_loaded_executable.cc index 14287a70cbe10e..33f4e0978996d4 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 @@ -338,11 +338,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 +348,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..01ab7a7af80816 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 @@ -119,8 +119,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/mock.h b/third_party/xla/xla/python/ifrt/mock.h index 0115b2df541326..539d6d834a9b68 100644 --- a/third_party/xla/xla/python/ifrt/mock.h +++ b/third_party/xla/xla/python/ifrt/mock.h @@ -432,8 +432,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)); @@ -502,8 +500,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_test.cc b/third_party/xla/xla/python/ifrt_proxy/client/compiler_test.cc index bbf98dcd97abe2..5e7b42b0883e71 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", @@ -204,8 +205,7 @@ TEST_F(CompilerTest, Compile) { 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..c93f8085dfff62 100644 --- a/third_party/xla/xla/python/ifrt_proxy/client/executable.cc +++ b/third_party/xla/xla/python/ifrt_proxy/client/executable.cc @@ -358,7 +358,6 @@ LoadedExecutable::LoadedExecutable( 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_( @@ -882,11 +881,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..9fc414ca4cbde8 100644 --- a/third_party/xla/xla/python/ifrt_proxy/client/executable.h +++ b/third_party/xla/xla/python/ifrt_proxy/client/executable.h @@ -113,7 +113,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; @@ -155,7 +154,6 @@ class LoadedExecutable final 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/mpmd_executable.h b/third_party/xla/xla/python/ifrt_proxy/client/mpmd_executable.h index 53d83a311667be..fd5f52c68f0698 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 @@ -147,9 +147,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/server/ifrt_backend.cc b/third_party/xla/xla/python/ifrt_proxy/server/ifrt_backend.cc index 1d43424a310538..82a472f514fa3d 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 @@ -1518,12 +1518,13 @@ 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()) { + 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..966dcbe79c8fa4 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; @@ -1245,8 +1246,6 @@ TEST_P(IfrtBackendHandlerTest, CompileSuccess) { 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, @@ -1913,8 +1912,6 @@ TEST_P(IfrtBackendHandlerTest, LoadedExecutableMetadataWithMpmd) { 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")); @@ -1976,8 +1973,6 @@ TEST_P(IfrtBackendHandlerTest, LoadedExecutableMpmdCostAnalysis) { 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")); @@ -2025,8 +2020,6 @@ TEST_P(IfrtBackendHandlerTest, CompileSuccessWithMpmdAddressableDevices) { 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..df346225e5f110 100644 --- a/third_party/xla/xla/python/pjrt_ifrt/pjrt_executable.h +++ b/third_party/xla/xla/python/pjrt_ifrt/pjrt_executable.h @@ -367,11 +367,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/version.h b/third_party/xla/xla/python/version.h index 9e7b5e5b37d9e3..6ae610dd6c32d6 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 73 // LoadedExecutable::devices() #endif // XLA_PYTHON_VERSION_H_ From c1643f86cc0b42ed5cdc50cfd66c0a070c308516 Mon Sep 17 00:00:00 2001 From: Junwhan Ahn Date: Mon, 28 Sep 2026 13:31:17 -0700 Subject: [PATCH 16/21] internal visibility changes PiperOrigin-RevId: 989812768 --- third_party/xla/xla/python/ifrt/BUILD | 3 ++- third_party/xla/xla/python/ifrt/ifrt.bzl | 5 ----- 2 files changed, 2 insertions(+), 6 deletions(-) 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/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 = ["//..."], - ) From 7b0145625b1398023a695bac36915b00eba21425 Mon Sep 17 00:00:00 2001 From: Seher Ellis Date: Mon, 28 Sep 2026 13:48:07 -0700 Subject: [PATCH 17/21] [XLA] (2/N User-specified async compute-on) Rewrite compute-on-start/done custom call pairs into async instructions. In legacy mode, the pair becomes an `all-gather-start` / `all-gather-done` pair configured from the all-gather inside the called computation, and in generic mode it becomes an async call wrapping that computation. Only computations built from all-gather and parameter instructions with exactly one single-operand all-gather are supported; anything else returns `UNIMPLEMENTED`. More support will be added after we agree on the general architecture in this CL. Backends can attach backend-specific metadata to the rewritten start op via the new `AsyncComputeOnHelper` extension point. PiperOrigin-RevId: 989821690 --- third_party/xla/xla/service/BUILD | 6 + .../async_collective_custom_call_rewriter.cc | 154 +++++++++++- .../async_collective_custom_call_rewriter.h | 38 ++- ...nc_collective_custom_call_rewriter_test.cc | 234 ++++++++++++++++++ 4 files changed, 424 insertions(+), 8 deletions(-) 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 From fabfe2aba63d7547eb638b80e4c873cc20c49198 Mon Sep 17 00:00:00 2001 From: Levon Ter-Grigoryan Date: Mon, 28 Sep 2026 14:14:23 -0700 Subject: [PATCH 18/21] [XLA:GPU] Record one-shot collective kernels into CUDA graphs with explicit nodes. Currently `CollectiveKernelThunk` can be recorded only with stream capture which takes around 0.25 ms to update at each iteration. With explicit API usage an update time can be reduced to 0.05 ms. This change is also a first step to move invocation counter to the device memory which will make graph update time almost zero. Add e2e tests that run 10 consecutive all-gathers and 10 consecutive all-reduces. PiperOrigin-RevId: 989838065 --- .../xla/xla/backends/gpu/runtime/BUILD | 5 + .../gpu/runtime/collective_kernel_thunk.cc | 134 ++++++++++++++-- .../gpu/runtime/collective_kernel_thunk.h | 36 +++++ third_party/xla/xla/backends/gpu/tests/BUILD | 5 + .../backends/gpu/tests/all_gather_e2e_test.cc | 84 ++++++++++ .../backends/gpu/tests/all_reduce_e2e_test.cc | 151 ++++++++++++++++++ .../gpu/tests/collective_ops_e2e_test_base.cc | 47 ++++++ .../gpu/tests/collective_ops_e2e_test_base.h | 13 ++ 8 files changed, 460 insertions(+), 15 deletions(-) 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/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_ From 0086dd6cb9f72ac4f25d7f1f88d0d2a84008a74b Mon Sep 17 00:00:00 2001 From: Parker Schuh Date: Mon, 28 Sep 2026 14:21:07 -0700 Subject: [PATCH 19/21] Update PjRtRawClient apis to use local_device_id instead of PjRtDevice and PjRtMemorySpace PiperOrigin-RevId: 989842055 --- .../pjrt/abstract_tracked_device_buffer.cc | 4 ++-- .../xla/xla/pjrt/common_pjrt_client.cc | 10 +++++--- third_party/xla/xla/pjrt/common_pjrt_client.h | 18 +++++++++----- third_party/xla/xla/pjrt/cpu/cpu_client.cc | 20 ++++++++-------- third_party/xla/xla/pjrt/cpu/cpu_client.h | 9 ++++--- .../xla/xla/pjrt/gpu/se_gpu_pjrt_client.cc | 4 ++-- .../xla/xla/pjrt/gpu/se_gpu_pjrt_client.h | 2 +- third_party/xla/xla/pjrt/raw_pjrt_client.cc | 8 +++---- third_party/xla/xla/pjrt/raw_pjrt_client.h | 18 ++++++++------ .../pjrt/se/pjrt_stream_executor_client.cc | 24 +++++++++---------- .../xla/pjrt/se/pjrt_stream_executor_client.h | 21 ++++++++-------- third_party/xla/xla/pjrt/se/se_raw_buffer.cc | 2 +- 12 files changed, 78 insertions(+), 62 deletions(-) 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); From 3cdbc2461dd54b9e469bca445e39003e7ab2b4d4 Mon Sep 17 00:00:00 2001 From: Junwhan Ahn Date: Mon, 28 Sep 2026 14:24:47 -0700 Subject: [PATCH 20/21] Remove LoadedExecutable::num_devices() and migrate callers to devices(). `xla::ifrt::LoadedExecutable::num_devices()` is superseded by `xla::ifrt::LoadedExecutable::devices()`. Callers have been updated: - In JAX (`pjit.cc`), when `devices()` returns `std::nullopt` (portable executables), `num_global_devices` evaluates to 1. This is guarded with `#if JAX_IFRT_VERSION_NUMBER >= 74`. - In IFRT IR pass and tests, device counts are derived directly from `(*devices)->size()`. - Bump `JAX_IFRT_VERSION_NUMBER` to 74 in `version.h`. - Remove `LoadedExecutable::num_devices()` from the base interface, mocks, and all subclass implementations. PiperOrigin-RevId: 989844342 --- .../xla/backends/cpu/nanort/ifrt_client.cc | 2 -- third_party/xla/xla/python/ifrt/executable.h | 1 - .../ifrt/ir/ifrt_ir_loaded_executable.cc | 5 ----- .../ifrt/ir/ifrt_ir_loaded_executable.h | 1 - ...y_bound_external_loaded_executable_pass.cc | 13 +++++++++--- third_party/xla/xla/python/ifrt/mock.h | 2 -- .../xla/python/ifrt_proxy/client/compiler.cc | 11 +++++----- .../python/ifrt_proxy/client/compiler_test.cc | 1 - .../python/ifrt_proxy/client/executable.cc | 6 +----- .../xla/python/ifrt_proxy/client/executable.h | 5 +---- .../ifrt_proxy/client/executable_test.cc | 7 +++---- .../ifrt_proxy/client/mpmd_executable.cc | 9 ++++----- .../ifrt_proxy/client/mpmd_executable.h | 4 +--- .../ifrt_proxy/client/mpmd_executable_test.cc | 20 +++++++++---------- .../python/ifrt_proxy/server/ifrt_backend.cc | 2 +- .../ifrt_proxy/server/ifrt_backend_test.cc | 4 ---- .../xla/python/pjrt_ifrt/pjrt_executable.h | 5 ----- .../pjrt_ifrt/xla_executable_impl_test_lib.cc | 3 ++- third_party/xla/xla/python/version.h | 2 +- 19 files changed, 39 insertions(+), 64 deletions(-) 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 8ad71ec014f0cf..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 { diff --git a/third_party/xla/xla/python/ifrt/executable.h b/third_party/xla/xla/python/ifrt/executable.h index efd5418cf3b17c..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; 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 33f4e0978996d4..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; 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 01ab7a7af80816..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; 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 539d6d834a9b68..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)); @@ -471,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)); 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 5e7b42b0883e71..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 @@ -204,7 +204,6 @@ TEST_F(CompilerTest, Compile) { "xla::ifrt::proxy::TestProgram"); EXPECT_EQ(executable->name(), "foo-executable"); - EXPECT_EQ(executable->num_devices(), 2); 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 c93f8085dfff62..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,7 +355,6 @@ LoadedExecutable::LoadedExecutable( rpc_helper_(std::move(rpc_helper)), handle_(handle), name_(std::move(name)), - num_devices_(num_devices), devices_(devices), fingerprint_(std::move(fingerprint)), user_context_(xla::ifrt::UserContextScope::current()), @@ -550,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"); 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 9fc414ca4cbde8..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; @@ -152,7 +150,6 @@ class LoadedExecutable final const uint64_t handle_; const std::string name_; - const int num_devices_; const std::optional 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 fd5f52c68f0698..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(); } 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 82a472f514fa3d..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,10 +1517,10 @@ tsl::Future IfrtBackend::HandleCompileRequest( // Populate executable metadata. compile_resp->set_name(AsProtoStringData(executable->name())); - compile_resp->set_num_devices(executable->num_devices()); 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()); 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 966dcbe79c8fa4..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 @@ -1243,7 +1243,6 @@ 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, Fingerprint()).WillOnce(Return("fingerprint")); @@ -1909,7 +1908,6 @@ 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, Fingerprint()) @@ -1970,7 +1968,6 @@ 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, Fingerprint()) @@ -2017,7 +2014,6 @@ 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, Fingerprint()).WillByDefault(Return("mpmd_fingerprint")); 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 df346225e5f110..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(); 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 6ae610dd6c32d6..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 73 // LoadedExecutable::devices() +#define JAX_IFRT_VERSION_NUMBER 74 // Remove LoadedExecutable::num_devices() #endif // XLA_PYTHON_VERSION_H_ From 96138a6b38ef3364c3c6645bacbf2ab6b421ccd2 Mon Sep 17 00:00:00 2001 From: Fengwu Yao Date: Mon, 28 Sep 2026 14:41:27 -0700 Subject: [PATCH 21/21] Add TfLiteBlockwiseQuantizationV2 for block_shape in the TFLite C API. PiperOrigin-RevId: 989853661 --- tensorflow/lite/core/c/common.cc | 25 ++++ tensorflow/lite/core/c/common.h | 21 ++++ tensorflow/lite/core/c/common_test.cc | 31 +++++ tensorflow/lite/core/interpreter_builder.cc | 65 +++++++++-- tensorflow/lite/core/model_test.cc | 123 ++++++++++++++++++++ 5 files changed, 258 insertions(+), 7 deletions(-) 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");