diff --git a/tensorflow/compiler/mlir/lite/tests/optimize_batch_matmul.mlir b/tensorflow/compiler/mlir/lite/tests/optimize_batch_matmul.mlir index 8c15335815bf4b..8b1829faf66acc 100644 --- a/tensorflow/compiler/mlir/lite/tests/optimize_batch_matmul.mlir +++ b/tensorflow/compiler/mlir/lite/tests/optimize_batch_matmul.mlir @@ -342,4 +342,68 @@ func.func @Batchmatmul2FullyconnectedReshapeTransposeChainConstY(%arg0: tensor<1 // CHECK-NOT: "tfl.batch_matmul" %2 = "tfl.batch_matmul"(%arg0, %1) {adj_x = false, adj_y = false, asymmetric_quantize_inputs = false} : (tensor<16x1024xf32>, tensor<1024x128xf32>) -> tensor<16x128xf32> func.return %2 : tensor<16x128xf32> +} + +// A bf16 constant widened to f32 is still foldable, since flatbuffer export +// materializes the cast. The transpose goes below the cast so it lands +// directly on the constant, where it can be folded. +// CHECK-LABEL: Batchmatmul2FullyconnectedBf16CastConstY +func.func @Batchmatmul2FullyconnectedBf16CastConstY(%arg0: tensor<16x1024xf32>) -> tensor<16x128xf32> { + %weight = arith.constant dense_resource<__elided__> : tensor<1024x128xbf16> + %0 = "tfl.cast"(%weight) : (tensor<1024x128xbf16>) -> tensor<1024x128xf32> + // CHECK: %[[W:.*]] = arith.constant dense_resource<__elided__> : tensor<1024x128xbf16> + // CHECK: %[[T:.*]] = "tfl.transpose"(%[[W]], %{{.*}}) : (tensor<1024x128xbf16>, tensor<2xi32>) -> tensor<128x1024xbf16> + // CHECK: %[[C:.*]] = "tfl.cast"(%[[T]]) : (tensor<128x1024xbf16>) -> tensor<128x1024xf32> + // CHECK: "tfl.fully_connected"(%arg0, %[[C]], %{{.*}}) + // CHECK-NOT: "tfl.batch_matmul" + %1 = "tfl.batch_matmul"(%arg0, %0) {adj_x = false, adj_y = false, asymmetric_quantize_inputs = false} : (tensor<16x1024xf32>, tensor<1024x128xf32>) -> tensor<16x128xf32> + func.return %1 : tensor<16x128xf32> +} + +// With adj_y the rhs is already in filter layout, so the cast is kept as is. +// CHECK-LABEL: Batchmatmul2FullyconnectedBf16CastConstYAdjY +func.func @Batchmatmul2FullyconnectedBf16CastConstYAdjY(%arg0: tensor<16x1024xf32>) -> tensor<16x128xf32> { + %weight = arith.constant dense_resource<__elided__> : tensor<128x1024xbf16> + %0 = "tfl.cast"(%weight) : (tensor<128x1024xbf16>) -> tensor<128x1024xf32> + // CHECK: %[[W:.*]] = arith.constant dense_resource<__elided__> : tensor<128x1024xbf16> + // CHECK: %[[C:.*]] = "tfl.cast"(%[[W]]) : (tensor<128x1024xbf16>) -> tensor<128x1024xf32> + // CHECK-NOT: "tfl.transpose" + // CHECK: "tfl.fully_connected"(%arg0, %[[C]], %{{.*}}) + // CHECK-NOT: "tfl.batch_matmul" + %1 = "tfl.batch_matmul"(%arg0, %0) {adj_x = false, adj_y = true, asymmetric_quantize_inputs = false} : (tensor<16x1024xf32>, tensor<128x1024xf32>) -> tensor<16x128xf32> + func.return %1 : tensor<16x128xf32> +} + +// CHECK-LABEL: NotFuseTransposeFCRhsToBatchMatmulBf16CastConstant +func.func @NotFuseTransposeFCRhsToBatchMatmulBf16CastConstant(%arg0: tensor<16x1024xf32>, %arg1: none) -> tensor<16x128xf32> { + %cst = arith.constant dense<[1, 0]> : tensor<2xi32> + %weight = arith.constant dense_resource<__elided__> : tensor<1024x128xbf16> + %0 = "tfl.cast"(%weight) : (tensor<1024x128xbf16>) -> tensor<1024x128xf32> + %1 = "tfl.transpose"(%0, %cst) : (tensor<1024x128xf32>, tensor<2xi32>) -> tensor<128x1024xf32> + // CHECK: "tfl.fully_connected" + // CHECK-NOT: "tfl.batch_matmul" + %2 = "tfl.fully_connected"(%arg0, %1, %arg1) {asymmetric_quantize_inputs = false, fused_activation_function = "NONE", keep_num_dims = false, weights_format = "DEFAULT"} : (tensor<16x1024xf32>, tensor<128x1024xf32>, none) -> tensor<16x128xf32> + func.return %2 : tensor<16x128xf32> +} + +// A widening cast of a non-constant is not foldable. +// CHECK-LABEL: NotBatchmatmul2FullyconnectedBf16CastNonConstY +func.func @NotBatchmatmul2FullyconnectedBf16CastNonConstY(%arg0: tensor<16x1024xf32>, %arg1: tensor<1024x128xbf16>) -> tensor<16x128xf32> { + %0 = "tfl.cast"(%arg1) : (tensor<1024x128xbf16>) -> tensor<1024x128xf32> + // CHECK: "tfl.batch_matmul" + // CHECK-NOT: "tfl.fully_connected" + %1 = "tfl.batch_matmul"(%arg0, %0) {adj_x = false, adj_y = false, asymmetric_quantize_inputs = false} : (tensor<16x1024xf32>, tensor<1024x128xf32>) -> tensor<16x128xf32> + func.return %1 : tensor<16x128xf32> +} + +// Only float widening is looked through; flatbuffer export does not +// materialize other casts. +// CHECK-LABEL: NotBatchmatmul2FullyconnectedI8CastConstY +func.func @NotBatchmatmul2FullyconnectedI8CastConstY(%arg0: tensor<16x1024xf32>) -> tensor<16x128xf32> { + %weight = arith.constant dense_resource<__elided__> : tensor<1024x128xi8> + %0 = "tfl.cast"(%weight) : (tensor<1024x128xi8>) -> tensor<1024x128xf32> + // CHECK: "tfl.batch_matmul" + // CHECK-NOT: "tfl.fully_connected" + %1 = "tfl.batch_matmul"(%arg0, %0) {adj_x = false, adj_y = false, asymmetric_quantize_inputs = false} : (tensor<16x1024xf32>, tensor<1024x128xf32>) -> tensor<16x128xf32> + func.return %1 : tensor<16x128xf32> } \ No newline at end of file diff --git a/tensorflow/compiler/mlir/lite/transforms/optimize_batch_matmul_pass.cc b/tensorflow/compiler/mlir/lite/transforms/optimize_batch_matmul_pass.cc index c8e4ce5c6077d9..35bdffefc50ff2 100644 --- a/tensorflow/compiler/mlir/lite/transforms/optimize_batch_matmul_pass.cc +++ b/tensorflow/compiler/mlir/lite/transforms/optimize_batch_matmul_pass.cc @@ -44,9 +44,26 @@ namespace mlir { namespace TFL { namespace { +// Returns `op` as a bf16/f16 -> f32 cast, or null if it is anything else. +// +// Weights kept in a narrow float type reach the matmul behind such a cast, +// since the kernels only run in f32. Flatbuffer export materializes this cast +// when its input is a constant, so the cast does not stop the value from being +// a constant filter. +CastOp AsFloatWideningCast(mlir::Operation* op) { + auto cast_op = mlir::dyn_cast_or_null(op); + if (!cast_op) return nullptr; + mlir::Type in_elem = getElementTypeOrSelf(cast_op.getInput().getType()); + mlir::Type out_elem = getElementTypeOrSelf(cast_op.getType()); + if (!(in_elem.isBF16() || in_elem.isF16()) || !out_elem.isF32()) { + return nullptr; + } + return cast_op; +} + // Checks whether `value` is, or can be folded into, a constant. This covers a // plain constant, the dequantize wrappers a quantized constant weight hides -// behind, and chains of ops that only rearrange one of those. +// behind, and chains of ops that only rearrange or widen one of those. bool IsFromFoldableChain(mlir::Value value) { mlir::Operation* defining_op = value.getDefiningOp(); @@ -66,6 +83,8 @@ bool IsFromFoldableChain(mlir::Value value) { defining_op = split_op.getValue().getDefiningOp(); } else if (auto transpose_op = mlir::dyn_cast(defining_op)) { defining_op = transpose_op.getInput().getDefiningOp(); + } else if (auto cast_op = AsFloatWideningCast(defining_op)) { + defining_op = cast_op.getInput().getDefiningOp(); } else { // Stop if the op does not preserve constant nature. break; @@ -193,8 +212,24 @@ struct ConvertBatchMatMulOp2FullyConnectedOp_Rank2ConstantRhs // The rhs need to be transposed if adj_y == false AND this matmul will be // legalized to tfl.fully_connected - Value output_rhs = - !bmm_op.getAdjY() ? create_z_x_transpose_op(input_rhs) : input_rhs; + Value output_rhs = input_rhs; + if (!bmm_op.getAdjY()) { + // Transpose below a widening cast rather than above it. Only a + // transpose directly on a constant gets folded away, and the cast is + // then left directly on the constant too, which flatbuffer export + // materializes. + if (CastOp cast_op = AsFloatWideningCast(input_rhs.getDefiningOp())) { + Value transposed = create_z_x_transpose_op(cast_op.getInput()); + auto transposed_type = + mlir::cast(transposed.getType()); + output_rhs = TFL::CastOp::create( + rewriter, cast_op.getLoc(), + transposed_type.clone(getElementTypeOrSelf(cast_op.getType())), + transposed); + } else { + output_rhs = create_z_x_transpose_op(input_rhs); + } + } Type output_type = bmm_op.getResult().getType(); auto no_input = diff --git a/tensorflow/compiler/mlir/tfr/ir/tfr_ops.td b/tensorflow/compiler/mlir/tfr/ir/tfr_ops.td index 2cad28974ab8ec..be5401d5fabc79 100644 --- a/tensorflow/compiler/mlir/tfr/ir/tfr_ops.td +++ b/tensorflow/compiler/mlir/tfr/ir/tfr_ops.td @@ -41,6 +41,7 @@ def TFR_Dialect : Dialect { }]; let cppNamespace = "::mlir::TFR"; + let useStrictPropertiesInAssemblyFormat = 0; } //===----------------------------------------------------------------------===// diff --git a/tensorflow/compiler/mlir/tfrt/ir/mlrt/mlrt_dialect.td b/tensorflow/compiler/mlir/tfrt/ir/mlrt/mlrt_dialect.td index 13409c3ece1f3d..2650a412198272 100644 --- a/tensorflow/compiler/mlir/tfrt/ir/mlrt/mlrt_dialect.td +++ b/tensorflow/compiler/mlir/tfrt/ir/mlrt/mlrt_dialect.td @@ -26,6 +26,7 @@ def Mlrt_Dialect : Dialect { }]; let cppNamespace = "::mlrt::compiler"; + let useStrictPropertiesInAssemblyFormat = 0; } def MlrtFutureType : DialectType traits = []> : diff --git a/tensorflow/compiler/mlir/tfrt/ir/mlrt/tf_mlrt_tpu_ops.td b/tensorflow/compiler/mlir/tfrt/ir/mlrt/tf_mlrt_tpu_ops.td index a207b83c7e5207..f545786faead7a 100644 --- a/tensorflow/compiler/mlir/tfrt/ir/mlrt/tf_mlrt_tpu_ops.td +++ b/tensorflow/compiler/mlir/tfrt/ir/mlrt/tf_mlrt_tpu_ops.td @@ -28,6 +28,7 @@ def TensorflowMlrtTpu_Dialect : Dialect { }]; let cppNamespace = "::tensorflow::tf_mlrt_tpu"; + let useStrictPropertiesInAssemblyFormat = 0; } class TensorflowMlrtTpu_Op traits = []> : diff --git a/tensorflow/compiler/mlir/tfrt/ir/tfrt_fallback_async.td b/tensorflow/compiler/mlir/tfrt/ir/tfrt_fallback_async.td index 3209f2f253cc2a..7d73e733ffcb97 100644 --- a/tensorflow/compiler/mlir/tfrt/ir/tfrt_fallback_async.td +++ b/tensorflow/compiler/mlir/tfrt/ir/tfrt_fallback_async.td @@ -33,6 +33,7 @@ def FallbackAsync_Dialect : Dialect { }]; let cppNamespace = "::tfrt::fallback_async"; + let useStrictPropertiesInAssemblyFormat = 0; } class FallbackAsync_Op traits = []> : diff --git a/tensorflow/compiler/mlir/tfrt/ir/tfrt_fallback_sync.td b/tensorflow/compiler/mlir/tfrt/ir/tfrt_fallback_sync.td index bba8a0219217c6..4575efae94871f 100644 --- a/tensorflow/compiler/mlir/tfrt/ir/tfrt_fallback_sync.td +++ b/tensorflow/compiler/mlir/tfrt/ir/tfrt_fallback_sync.td @@ -32,6 +32,7 @@ def FallbackSync_Dialect : Dialect { }]; let cppNamespace = "::tfrt::fallback_sync"; + let useStrictPropertiesInAssemblyFormat = 0; } class FallbackSync_Op traits = []> : diff --git a/tensorflow/compiler/mlir/tools/kernel_gen/ir/tf_framework_ops.td b/tensorflow/compiler/mlir/tools/kernel_gen/ir/tf_framework_ops.td index 64f782d02346e8..0f54049616b5ad 100644 --- a/tensorflow/compiler/mlir/tools/kernel_gen/ir/tf_framework_ops.td +++ b/tensorflow/compiler/mlir/tools/kernel_gen/ir/tf_framework_ops.td @@ -40,6 +40,7 @@ def TFFramework_Dialect : Dialect { }]; let useDefaultTypePrinterParser = 1; + let useStrictPropertiesInAssemblyFormat = 1; } def TFFramework_OpKernelContextType : DialectType !tf_framework.jit_callable attributes {tf_entry} { // CHECK: %[[RES:.*]] = tf_framework.jit_compile_from_str %[[CTX]], "placeholder" // CHECK: return %[[RES]] - %0 = tf_framework.jit_compile_from_str "placeholder" { - architectures = ["sm_123", "sm_456"], tileSizes = [1, 2, 3], - unrollFactors = [4], enableFtz = false, - index64Bit = false, cpuCodegen = false } + %0 = tf_framework.jit_compile_from_str "placeholder" < + tileSizes = [1, 2, 3], unrollFactors = [4], enableFtz = false, + index64Bit = false, cpuCodegen = false> { + architectures = ["sm_123", "sm_456"] } func.return %0 : !tf_framework.jit_callable } diff --git a/tensorflow/compiler/mlir/tools/kernel_gen/tests/func_to_jit_invocations.mlir b/tensorflow/compiler/mlir/tools/kernel_gen/tests/func_to_jit_invocations.mlir index d19e8180973b76..b9abece1fee264 100644 --- a/tensorflow/compiler/mlir/tools/kernel_gen/tests/func_to_jit_invocations.mlir +++ b/tensorflow/compiler/mlir/tools/kernel_gen/tests/func_to_jit_invocations.mlir @@ -41,12 +41,12 @@ func.func @unary_tanh(%arg : tensor) -> tensor { // CHECK-SAME: } // CHECK-SAME: } // CHECK-SAME: " -// CHECK-SAME: { -// CHECK-SAME: cpuCodegen = false -// CHECK-SAME: enableFtz = false +// CHECK-SAME: < // CHECK-SAME: tileSizes = [1, 2, 3] // CHECK-SAME: unrollFactors = [3, 2, 1] -// CHECK-SAME: } +// CHECK-SAME: enableFtz = false +// CHECK-SAME: cpuCodegen = false +// CHECK-SAME: > // CHECK: %[[RES:.*]] = tf_framework.jit_execute %[[CALLABLE]](%[[ARG]]) // CHECK: return %[[RES]] @@ -59,11 +59,11 @@ func.func @unary_tanh(%arg : tensor) -> tensor { // CHECK-JFLT: %[[IF:.*]] = scf.if %[[CMPI]] // CHECK-JFLT: %[[JIT:.*]] = tf_framework.jit_compile_from_str // CHECK-JFLT-SAME: "module -// CHECK-JFLT-SAME: cpuCodegen = false -// CHECK-JFLT-SAME: enableFtz = false -// CHECK-JFLT-SAME: index64Bit = true // CHECK-JFLT-SAME: tileSizes = [1, 2, 3] // CHECK-JFLT-SAME: unrollFactors = [3, 2, 1] +// CHECK-JFLT-SAME: enableFtz = false +// CHECK-JFLT-SAME: index64Bit = true +// CHECK-JFLT-SAME: cpuCodegen = false // CHECK-JFLT: %[[JIT_0:.*]] = tf_framework.jit_execute %[[JIT]](%[[ARG0]]) // CHECK-JFLT: scf.yield %[[JIT_0]] // CHECK-JFLT: else @@ -91,12 +91,12 @@ func.func @binary_sub(%arg0 : tensor<*xf32>, %arg1 : tensor<*xf32>) -> tensor<*x // CHECK-SAME: } // CHECK-SAME: } // CHECK-SAME: " -// CHECK-SAME: { -// CHECK-SAME: cpuCodegen = false -// CHECK-SAME: enableFtz = false +// CHECK-SAME: < // CHECK-SAME: tileSizes = [1, 2, 3] // CHECK-SAME: unrollFactors = [3, 2, 1] -// CHECK-SAME: } +// CHECK-SAME: enableFtz = false +// CHECK-SAME: cpuCodegen = false +// CHECK-SAME: > // CHECK: %[[RES:.*]] = tf_framework.jit_execute %[[CALLABLE]](%[[ARG0]], %[[ARG1]]) // CHECK: return %[[RES]] @@ -122,12 +122,12 @@ func.func @binary_sub(%arg0 : tensor<*xf32>, %arg1 : tensor<*xf32>) -> tensor<*x // CHECK-JFLT-SAME: } // CHECK-JFLT-SAME: } // CHECK-JFLT-SAME: " -// CHECK-JFLT-SAME: { -// CHECK-JFLT-SAME: cpuCodegen = false -// CHECK-JFLT-SAME: enableFtz = false +// CHECK-JFLT-SAME: < // CHECK-JFLT-SAME: tileSizes = [1, 2, 3] // CHECK-JFLT-SAME: unrollFactors = [3, 2, 1] -// CHECK-JFLT-SAME: } +// CHECK-JFLT-SAME: enableFtz = false +// CHECK-JFLT-SAME: cpuCodegen = false +// CHECK-JFLT-SAME: > // CHECK-JFLT: %[[RES:.*]] = tf_framework.jit_execute %[[CALLABLE]](%[[ARG0]], %[[ARG1]]) // CHECK-JFLT: scf.yield %[[RES]] : tensor<*xf32> // CHECK-JFLT: } else { @@ -155,13 +155,13 @@ func.func @reciprocal(%arg0: tensor<*xf32>) // CHECK=SAME: return %[[RES_JIT]] : tensor<*xf32> // CHECK-SAME: } // CHECK-SAME: } -// CHECK-SAME: { -// CHECK-SAME: cpuCodegen = false, +// CHECK-SAME: < +// CHECK-SAME: tileSizes = [1, 2, 3], +// CHECK-SAME: unrollFactors = [3, 2, 1], // CHECK-SAME: enableFtz = false, // CHECK-SAME: index64Bit = false, -// CHECK-SAME: tileSizes = [1, 2, 3], -// CHECK-SAME: unrollFactors = [3, 2, 1] -// CHECK-SAME: } +// CHECK-SAME: cpuCodegen = false +// CHECK-SAME: > // CHECK: %[[RES:.*]] = tf_framework.jit_execute %[[CALLABLE]](%[[ARG]]) : tensor<*xf32> -> tensor<*xf32> // CHECK: return %[[RES]] : tensor<*xf32> @@ -174,11 +174,11 @@ func.func @reciprocal(%arg0: tensor<*xf32>) // CHECK-JFLT: %[[IF:.*]] = scf.if %[[CMPI]] // CHECK-JFLT: %[[JIT:.*]] = tf_framework.jit_compile_from_str // CHECK-JFLT-SAME: "module -// CHECK-JFLT-SAME: cpuCodegen = false -// CHECK-JFLT-SAME: enableFtz = false -// CHECK-JFLT-SAME: index64Bit = true // CHECK-JFLT-SAME: tileSizes = [1, 2, 3] // CHECK-JFLT-SAME: unrollFactors = [3, 2, 1] +// CHECK-JFLT-SAME: enableFtz = false +// CHECK-JFLT-SAME: index64Bit = true +// CHECK-JFLT-SAME: cpuCodegen = false // CHECK-JFLT: %[[JIT_0:.*]] = tf_framework.jit_execute %[[JIT]](%[[ARG0]]) // CHECK-JFLT: scf.yield %[[JIT_0]] // CHECK-JFLT: else diff --git a/tensorflow/compiler/mlir/tools/kernel_gen/tests/ops.mlir b/tensorflow/compiler/mlir/tools/kernel_gen/tests/ops.mlir index 8d988e2eae11bc..5829355053fa35 100644 --- a/tensorflow/compiler/mlir/tools/kernel_gen/tests/ops.mlir +++ b/tensorflow/compiler/mlir/tools/kernel_gen/tests/ops.mlir @@ -29,10 +29,10 @@ func.func @alloc(%ctx: !tf_framework.op_kernel_context, // CHECK-LABEL: func @forwarding_alloc func.func @forwarding_alloc(%ctx: !tf_framework.op_kernel_context, %size_0 : index , %size_2 : index) { - %buf = tf_framework.alloc(%ctx, %size_0, %size_2) { + %buf = tf_framework.alloc(%ctx, %size_0, %size_2) < input_indices = [0 : i32, 1 : i32], - output_index = 0 : i32 - } : memref + output_index = 0 + > : memref func.return } @@ -94,20 +94,20 @@ func.func @jit_compile(%ctx : !tf_framework.op_kernel_context) // CHECK-LABEL: func @jit_compile_from_str_wo_ctx func.func @jit_compile_from_str_wo_ctx() -> !tf_framework.jit_callable { - %callable = tf_framework.jit_compile_from_str "placeholder" { - architectures = ["sm_123", "sm_456"], tileSizes = [1, 2, 3], - unrollFactors = [4], enableFtz = false, - index64Bit = false, cpuCodegen = false } + %callable = tf_framework.jit_compile_from_str "placeholder" < + tileSizes = [1, 2, 3], unrollFactors = [4], enableFtz = false, + index64Bit = false, cpuCodegen = false> { + architectures = ["sm_123", "sm_456"] } func.return %callable : !tf_framework.jit_callable } // CHECK-LABEL: func @jit_compile_from_str func.func @jit_compile_from_str(%ctx : !tf_framework.op_kernel_context) -> !tf_framework.jit_callable { - %callable = tf_framework.jit_compile_from_str %ctx , "placeholder" { - architectures = ["sm_123", "sm_456"], tileSizes = [1, 2, 3], - unrollFactors = [4], enableFtz = false, - index64Bit = false, cpuCodegen = false } + %callable = tf_framework.jit_compile_from_str %ctx , "placeholder" < + tileSizes = [1, 2, 3], unrollFactors = [4], enableFtz = false, + index64Bit = false, cpuCodegen = false> { + architectures = ["sm_123", "sm_456"] } func.return %callable : !tf_framework.jit_callable } diff --git a/tensorflow/compiler/mlir/tools/kernel_gen/tests/tf_framework_legalize_to_llvm.mlir b/tensorflow/compiler/mlir/tools/kernel_gen/tests/tf_framework_legalize_to_llvm.mlir index 15e0284d89eb20..e6794dcf88936a 100644 --- a/tensorflow/compiler/mlir/tools/kernel_gen/tests/tf_framework_legalize_to_llvm.mlir +++ b/tensorflow/compiler/mlir/tools/kernel_gen/tests/tf_framework_legalize_to_llvm.mlir @@ -206,9 +206,9 @@ func.func @jit_compile_from_str(%ctx: !tf_framework.op_kernel_context) // CHECK-SAME: %[[NUM_UNROLL_FACTORS]], %[[UNROLL_FACTORS]], // CHECK-SAME: %[[ENABLE_FTZ]], %[[CPU_CODEGEN]] // CHECK: llvm.return %[[RES]] - %0 = tf_framework.jit_compile_from_str %ctx, "placeholder" { + %0 = tf_framework.jit_compile_from_str %ctx, "placeholder" < tileSizes = [1, 2, 3], unrollFactors = [4], - enableFtz = false, index64Bit = false, cpuCodegen = false } + enableFtz = false, index64Bit = false, cpuCodegen = false> func.return %0 : !tf_framework.jit_callable } diff --git a/tensorflow/core/common_runtime/rendezvous_util.cc b/tensorflow/core/common_runtime/rendezvous_util.cc index 5e28d6ac5d70bc..dbfabfb9c1dcf1 100644 --- a/tensorflow/core/common_runtime/rendezvous_util.cc +++ b/tensorflow/core/common_runtime/rendezvous_util.cc @@ -75,14 +75,14 @@ void RecvOutputsFromRendezvousAsync( keys.size(), "; alloc_attrs.size() = ", alloc_attrs.size()))); } - received_tensors->reserve(keys.size()); + received_tensors->resize(keys.size()); std::vector> arguments; + arguments.reserve(keys.size()); for (int i = 0; i < keys.size(); ++i) { Rendezvous::ParsedKey parsed; absl::Status s = Rendezvous::ParseKey(keys[i], &parsed); - received_tensors->push_back(Tensor()); if (!s.ok()) { done(s); return; diff --git a/tensorflow/core/distributed_runtime/graph_mgr.cc b/tensorflow/core/distributed_runtime/graph_mgr.cc index d50a0f1d67f6b6..75375e0f5f47b3 100644 --- a/tensorflow/core/distributed_runtime/graph_mgr.cc +++ b/tensorflow/core/distributed_runtime/graph_mgr.cc @@ -427,10 +427,8 @@ void GraphMgr::RecvOutputsAsync(const int64_t step_id, NamedTensors* out, std::vector keys; std::vector* received_keys = new std::vector; keys.reserve(out->size()); - received_keys->reserve(out->size()); for (const auto& p : *out) { keys.push_back(p.first); - received_keys->push_back(p.second); } RecvOutputsFromRendezvousAsync( rendezvous, nullptr, {}, keys, received_keys, diff --git a/tensorflow/core/ir/dialect.td b/tensorflow/core/ir/dialect.td index f3530af1232a16..106a09638f36ef 100644 --- a/tensorflow/core/ir/dialect.td +++ b/tensorflow/core/ir/dialect.td @@ -170,6 +170,7 @@ def TFGraphDialect : Dialect { let cppNamespace = "::mlir::tfg"; let useDefaultAttributePrinterParser = 1; + let useStrictPropertiesInAssemblyFormat = 0; let hasNonDefaultDestructor = 1; let hasOperationInterfaceFallback = 1; } diff --git a/tensorflow/core/kernels/BUILD b/tensorflow/core/kernels/BUILD index 6b25f346677e3a..1bd46c81866407 100644 --- a/tensorflow/core/kernels/BUILD +++ b/tensorflow/core/kernels/BUILD @@ -335,6 +335,21 @@ tf_kernel_library( alwayslink = 1, ) +tf_cuda_only_cc_test( + name = "conv_2d_gpu_test", + size = "small", + srcs = [ + "conv_2d_gpu.h", + "conv_2d_gpu_test.cu.cc", + ], + features = if_cuda(["-layering_check"]), + deps = [ + ":conv_2d", + "//tensorflow/core:test", + "//tensorflow/core:test_main", + ], +) + cc_library( name = "conv_2d_hdrs", hdrs = ["conv_2d.h"], diff --git a/tensorflow/core/kernels/conv_2d_gpu.h b/tensorflow/core/kernels/conv_2d_gpu.h index 3c9ccb5e6330f8..89e53268c814a1 100644 --- a/tensorflow/core/kernels/conv_2d_gpu.h +++ b/tensorflow/core/kernels/conv_2d_gpu.h @@ -150,10 +150,10 @@ struct Index : Array { }; // A helper function that converts a tensor index into a flat array index. -template -EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE int TensorIndexToFlat( +template +EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE FlatIndex TensorIndexToFlat( const Index& index, const Dimension& dims) { - int flat_index = index[0]; + FlatIndex flat_index = index[0]; for (int i = 1; i < IndexCount; i++) { flat_index = flat_index * dims[i] + index[i]; } @@ -320,8 +320,8 @@ __global__ void SwapDimension1And2InTensor3UsingTiles( input_tile_index[2] * TileSizeJ, }; - int input_origin_flat_index = - TensorIndexToFlat(input_tile_origin, input_dims); + int64_t input_origin_flat_index = + TensorIndexToFlat(input_tile_origin, input_dims); bool full_tile = true; int tile_width = TileSizeJ; @@ -350,8 +350,10 @@ __global__ void SwapDimension1And2InTensor3UsingTiles( // dimension of the input array. int ti = x / TileSizeJ; int tj = x % TileSizeJ; - int input_index = input_origin_flat_index + ti * input_dims[2] + tj; - int input_increment = ReadRowPerPass * input_dims[2]; + int64_t input_index = + input_origin_flat_index + static_cast(ti) * input_dims[2] + tj; + int64_t input_increment = + static_cast(ReadRowPerPass) * input_dims[2]; if (full_tile) { #pragma unroll @@ -385,8 +387,8 @@ __global__ void SwapDimension1And2InTensor3UsingTiles( output_tile_index[2] * TileSizeI, }; - int output_origin_flat_index = - TensorIndexToFlat(output_tile_origin, output_dims); + int64_t output_origin_flat_index = + TensorIndexToFlat(output_tile_origin, output_dims); constexpr int out_effective_thread_num = NumThreads / TileSizeI * TileSizeI; @@ -396,8 +398,10 @@ __global__ void SwapDimension1And2InTensor3UsingTiles( // dimension of the output array. int ti = x / TileSizeI; int tj = x % TileSizeI; - int output_index = output_origin_flat_index + ti * output_dims[2] + tj; - int output_increment = WriteRowPerPass * output_dims[2]; + int64_t output_index = output_origin_flat_index + + static_cast(ti) * output_dims[2] + tj; + int64_t output_increment = + static_cast(WriteRowPerPass) * output_dims[2]; if (full_tile) { #pragma unroll diff --git a/tensorflow/core/kernels/conv_2d_gpu_test.cu.cc b/tensorflow/core/kernels/conv_2d_gpu_test.cu.cc new file mode 100644 index 00000000000000..112f812a0a51b6 --- /dev/null +++ b/tensorflow/core/kernels/conv_2d_gpu_test.cu.cc @@ -0,0 +1,41 @@ +/* Copyright 2026 The TensorFlow 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. +==============================================================================*/ + +#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM + +#include "tensorflow/core/kernels/conv_2d_gpu.h" + +#include + +#include "tensorflow/core/platform/test.h" + +namespace tensorflow { +namespace functor { +namespace { + +TEST(TensorIndexToFlatTest, HandlesFlatIndexPastInt32Limit) { + constexpr int kDimension = 46341; + const Dimension<3> dims(1, kDimension, kDimension); + const Index<3> index(0, kDimension - 1, kDimension - 1); + + EXPECT_EQ(TensorIndexToFlat(index, dims), + static_cast(kDimension) * kDimension - 1); +} + +} // namespace +} // namespace functor +} // namespace tensorflow + +#endif // GOOGLE_CUDA || TENSORFLOW_USE_ROCM diff --git a/tensorflow/core/kernels/conv_grad_input_ops.h b/tensorflow/core/kernels/conv_grad_input_ops.h index d28077272041b0..f78dba00be3e76 100644 --- a/tensorflow/core/kernels/conv_grad_input_ops.h +++ b/tensorflow/core/kernels/conv_grad_input_ops.h @@ -356,7 +356,10 @@ class Conv2DBackpropInputOp : public OpKernel { OP_REQUIRES( context, out_backprop.dims() == 4, absl::InvalidArgumentError(absl::StrCat( - "input_sizes must be 4-dimensional, got: ", out_backprop.dims()))); + "out_backprop must be 4-dimensional, got: ", out_backprop.dims()))); + OP_REQUIRES(context, filter.dims() == 4, + absl::InvalidArgumentError(absl::StrCat( + "filter must be 4-dimensional, got: ", filter.dims()))); TensorShape input_shape; OP_REQUIRES_OK(context, @@ -468,7 +471,10 @@ class Conv2DCustomBackpropInputOp : public OpKernel { OP_REQUIRES( context, out_backprop.dims() == 4, absl::InvalidArgumentError(absl::StrCat( - "input_sizes must be 4-dimensional, got: ", out_backprop.dims()))); + "out_backprop must be 4-dimensional, got: ", out_backprop.dims()))); + OP_REQUIRES(context, filter.dims() == 4, + absl::InvalidArgumentError(absl::StrCat( + "filter must be 4-dimensional, got: ", filter.dims()))); TensorShape input_shape; OP_REQUIRES_OK(context, diff --git a/tensorflow/core/kernels/conv_grad_shape_utils.cc b/tensorflow/core/kernels/conv_grad_shape_utils.cc index 45a4a6c83291c0..da1b4b569d45f6 100644 --- a/tensorflow/core/kernels/conv_grad_shape_utils.cc +++ b/tensorflow/core/kernels/conv_grad_shape_utils.cc @@ -181,6 +181,17 @@ absl::Status Conv2DBackpropComputeInputShape( const Tensor& input_sizes, const TensorShape& filter_shape, const TensorShape& out_backprop_shape, const TensorFormat& data_format, TensorShape* input_shape) { + if (filter_shape.dims() != 4) { + return absl::InvalidArgumentError( + absl::StrCat("Conv2DBackpropInput: filter must be 4-dimensional, not ", + filter_shape.dims())); + } + if (out_backprop_shape.dims() != 4) { + return absl::InvalidArgumentError(absl::StrCat( + "Conv2DBackpropInput: out_backprop must be 4-dimensional, not ", + out_backprop_shape.dims())); + } + if (!TensorShapeUtils::IsVector(input_sizes.shape())) { return absl::InvalidArgumentError(absl::StrCat( "Conv2DBackpropInput: input_sizes input must be 1-dim, not ", diff --git a/tensorflow/core/kernels/cwise_ops.h b/tensorflow/core/kernels/cwise_ops.h index 547077cb9c3263..5e7be3ef285e46 100644 --- a/tensorflow/core/kernels/cwise_ops.h +++ b/tensorflow/core/kernels/cwise_ops.h @@ -1033,6 +1033,35 @@ struct log1p : base> {}; template struct sign : base> {}; +// Specialization of sign for complex types to avoid underflow in |z|^2. +// Eigen's scalar_sign_op computes |z|^2 = re^2 + im^2 as a float32 +// intermediate. When |z| < sqrt(float32_tiny) ~= 1.08e-19, this underflows +// to 0, causing sign(z) to incorrectly return 0. We fix this by checking +// whether the input is nonzero directly, and using std::abs() (which uses +// hypot internally) to avoid intermediate underflow/overflow. +template +struct safe_complex_sign_op { + EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE std::complex operator()( + const std::complex& z) const { + const RealType re = z.real(); + const RealType im = z.imag(); + if (re == RealType(0) && im == RealType(0)) { + return std::complex(0, 0); + } + // Use hypot directly to avoid intermediate underflow/overflow. + // std::abs(complex) may compute sqrt(re*re + im*im) which underflows + // when FTZ is enabled and |re|^2 or |im|^2 is subnormal. + const RealType mag = std::hypot(re, im); + return std::complex(re / mag, im / mag); + } +}; + +template <> +struct sign : base> {}; + +template <> +struct sign : base> {}; + template struct sinh : base> {}; diff --git a/tensorflow/dtensor/python/tests/test_util.py b/tensorflow/dtensor/python/tests/test_util.py index d8ebcb6b199e4e..d6e875a5b5fe15 100644 --- a/tensorflow/dtensor/python/tests/test_util.py +++ b/tensorflow/dtensor/python/tests/test_util.py @@ -54,9 +54,7 @@ DEFAULT_TOL = 1e-5 -# 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 +_DEFAULT_GPU_MEMORY_LIMIT = 1024 # 1G def get_use_xla_spmd(device_type): diff --git a/tensorflow/lite/core/c/BUILD b/tensorflow/lite/core/c/BUILD index 4b1de6f583bd4b..13e1a20563db8f 100644 --- a/tensorflow/lite/core/c/BUILD +++ b/tensorflow/lite/core/c/BUILD @@ -31,11 +31,7 @@ load( package( # copybara:uncomment default_applicable_licenses = ["//tensorflow:LICENSE"], - default_visibility = [ - # "//litert:__subpackages__", # copybara:uncomment - # copybara:uncomment "//third_party/odml/litert:__subpackages__", - "//tensorflow/lite:__subpackages__", - ], + default_visibility = ["//tensorflow/lite:__subpackages__"], licenses = ["notice"], ) @@ -325,7 +321,6 @@ tflite_cc_library_with_c_headers_test( visibility = [ # "//litert:__subpackages__", # copybara:uncomment # copybara:uncomment "//third_party/odml/litert/litert:__subpackages__", - "//tensorflow/compiler/mlir/lite/experimental/lrt:__subpackages__", "//tensorflow/lite:__subpackages__", ] + c_api_visibility_allowlist(), deps = [ diff --git a/tensorflow/lite/delegates/xnnpack/BUILD b/tensorflow/lite/delegates/xnnpack/BUILD index dd912cf6b12ae9..93dcc92d204914 100644 --- a/tensorflow/lite/delegates/xnnpack/BUILD +++ b/tensorflow/lite/delegates/xnnpack/BUILD @@ -1323,6 +1323,27 @@ cc_library( ], ) +cc_library( + name = "moe_block_scale", + hdrs = ["moe_block_scale.h"], + compatible_with = get_compatible_with_portable(), + copts = tflite_copts(), +) + +cc_test( + name = "moe_block_scale_test", + srcs = ["moe_block_scale_test.cc"], + linkopts = select({ + "//tensorflow:emscripten": EMSCRIPTEN_LINKOPTS, + "//conditions:default": [], + }), + deps = [ + ":moe_block_scale", + ":test_main", + "@com_google_googletest//:gtest", + ], +) + cc_library( name = "moe_delegate_kernel", srcs = ["moe_delegate_kernel.cc"], @@ -1330,6 +1351,7 @@ cc_library( compatible_with = get_compatible_with_portable(), copts = tflite_copts() + ["-DFLATBUFFERS_LOCALE_INDEPENDENT=0"], deps = [ + ":moe_block_scale", "//tensorflow/lite:builtin_ops", "//tensorflow/lite/c:common", "//tensorflow/lite/core/c:common", @@ -1341,6 +1363,25 @@ cc_library( ], ) +cc_test( + name = "moe_delegate_kernel_test", + srcs = ["moe_delegate_kernel_test.cc"], + linkopts = select({ + "//tensorflow:emscripten": EMSCRIPTEN_LINKOPTS, + "//conditions:default": [], + }), + deps = [ + ":moe_delegate_kernel", + ":test_main", + "//tensorflow/lite:builtin_ops", + "//tensorflow/lite:framework", + "//tensorflow/lite/core:framework", + "//tensorflow/lite/core/c:common", + "@com_google_googletest//:gtest", + "@flatbuffers//:runtime_cc", + ], +) + cc_test( name = "average_pool_2d_test", srcs = ["average_pool_2d_test.cc"], diff --git a/tensorflow/lite/delegates/xnnpack/moe_block_scale.h b/tensorflow/lite/delegates/xnnpack/moe_block_scale.h new file mode 100644 index 00000000000000..f62ecfd4589703 --- /dev/null +++ b/tensorflow/lite/delegates/xnnpack/moe_block_scale.h @@ -0,0 +1,120 @@ +/* Copyright 2026 The TensorFlow 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. +==============================================================================*/ +#ifndef TENSORFLOW_LITE_DELEGATES_XNNPACK_MOE_BLOCK_SCALE_H_ +#define TENSORFLOW_LITE_DELEGATES_XNNPACK_MOE_BLOCK_SCALE_H_ + +#include +#include +#include + +// Scale layout and dequantization helpers shared by the custom "moe" XNNPACK +// delegate kernel. They live in their own header so they can be unit tested +// without standing up a delegate. +// +// MoE expert weights are laid out as `[output_channels, num_experts, 1, +// input_channels]`, i.e. one row per (output channel, expert) pair. The +// matching scale tensor is `[output_channels, num_experts, 1, blocks_per_row]`. +// `blocks_per_row == 1` is per-output-channel quantization; larger values split +// the input axis into that many equally sized blocks, each with its own scale. +namespace tflite { +namespace xnnpack { + +// How a single row's scales map onto its input channels. +struct BlockScaleLayout { + // Number of scales per row. 1 means per-output-channel quantization. + size_t groups_per_row; + // Number of input channels covered by one scale. + size_t group_size; +}; + +// Derives the layout from the total number of scale elements. `scale_elements` +// is expected to be an exact multiple of `num_rows`; callers validate that +// before reaching here, and a ragged value degrades to treating the row as +// per-channel rather than reading out of bounds. +inline BlockScaleLayout ResolveBlockScaleLayout(size_t scale_elements, + size_t num_rows, + size_t input_channels) { + const size_t groups_per_row = + (num_rows > 0) ? (scale_elements / num_rows) : 0; + const size_t group_size = + (groups_per_row > 0 && groups_per_row <= input_channels) + ? (input_channels / groups_per_row) + : 1; + return {groups_per_row, group_size}; +} + +// Index of the scale covering input channel `in`. Clamped to the last block so +// an input axis that does not divide evenly cannot read past the row. +inline size_t BlockScaleIndex(const BlockScaleLayout& layout, size_t in) { + if (layout.groups_per_row <= 1 || layout.group_size == 0) { + return 0; + } + return std::min(in / layout.group_size, layout.groups_per_row - 1); +} + +// Dequantizes the rows belonging to `expert` into `dst`, which is a dense +// `[output_channels, input_channels]` buffer. +inline void CopyAndDequantizeExpertWeightRowsInt8( + const int8_t* weight_i8, const float* scale, size_t scale_elements, + size_t num_experts, size_t expert, size_t output_channels, + size_t input_channels, float* dst) { + const size_t num_rows = output_channels * num_experts; + const BlockScaleLayout layout = + ResolveBlockScaleLayout(scale_elements, num_rows, input_channels); + + for (size_t out = 0; out < output_channels; ++out) { + const size_t row_idx = out * num_experts + expert; + const int8_t* src_row = weight_i8 + row_idx * input_channels; + float* dst_row = dst + out * input_channels; + const float* row_scales = scale + row_idx * layout.groups_per_row; + for (size_t in = 0; in < input_channels; ++in) { + const float scale_val = row_scales[BlockScaleIndex(layout, in)]; + dst_row[in] = static_cast(src_row[in]) * scale_val; + } + } +} + +// As above, for int4 weights packed two nibbles per byte, low nibble first. +inline void CopyAndDequantizeExpertWeightRowsInt4( + const int8_t* weight_i4_packed, const float* scale, size_t scale_elements, + size_t num_experts, size_t expert, size_t output_channels, + size_t input_channels, float* dst) { + const size_t num_rows = output_channels * num_experts; + const BlockScaleLayout layout = + ResolveBlockScaleLayout(scale_elements, num_rows, input_channels); + + for (size_t out = 0; out < output_channels; ++out) { + const size_t row_idx = out * num_experts + expert; + const int8_t* src_row_packed = + weight_i4_packed + (row_idx * input_channels) / 2; + float* dst_row = dst + out * input_channels; + const float* row_scales = scale + row_idx * layout.groups_per_row; + + for (size_t in = 0; in < input_channels; ++in) { + const size_t byte_idx = in / 2; + const int8_t byte_val = src_row_packed[byte_idx]; + const int8_t nibble = (in % 2 == 0) + ? static_cast(byte_val << 4) >> 4 + : static_cast(byte_val >> 4); + const float scale_val = row_scales[BlockScaleIndex(layout, in)]; + dst_row[in] = static_cast(nibble) * scale_val; + } + } +} + +} // namespace xnnpack +} // namespace tflite + +#endif // TENSORFLOW_LITE_DELEGATES_XNNPACK_MOE_BLOCK_SCALE_H_ diff --git a/tensorflow/lite/delegates/xnnpack/moe_block_scale_test.cc b/tensorflow/lite/delegates/xnnpack/moe_block_scale_test.cc new file mode 100644 index 00000000000000..45ca9dc0b5b9a0 --- /dev/null +++ b/tensorflow/lite/delegates/xnnpack/moe_block_scale_test.cc @@ -0,0 +1,244 @@ +/* Copyright 2026 The TensorFlow 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. +==============================================================================*/ +#include "tensorflow/lite/delegates/xnnpack/moe_block_scale.h" + +#include +#include +#include + +#include +#include + +namespace tflite { +namespace xnnpack { +namespace { + +using ::testing::ElementsAre; + +TEST(ResolveBlockScaleLayoutTest, OneScalePerRowIsPerChannel) { + const BlockScaleLayout layout = ResolveBlockScaleLayout( + /*scale_elements=*/8, /*num_rows=*/8, /*input_channels=*/16); + EXPECT_EQ(layout.groups_per_row, 1); + EXPECT_EQ(layout.group_size, 16); +} + +TEST(ResolveBlockScaleLayoutTest, SplitsInputAxisIntoEqualBlocks) { + const BlockScaleLayout layout = ResolveBlockScaleLayout( + /*scale_elements=*/32, /*num_rows=*/8, /*input_channels=*/64); + EXPECT_EQ(layout.groups_per_row, 4); + EXPECT_EQ(layout.group_size, 16); +} + +// A scale per input channel is the finest blocking the layout can express. +TEST(ResolveBlockScaleLayoutTest, HandlesOneScalePerInputChannel) { + const BlockScaleLayout layout = ResolveBlockScaleLayout( + /*scale_elements=*/64, /*num_rows=*/8, /*input_channels=*/8); + EXPECT_EQ(layout.groups_per_row, 8); + EXPECT_EQ(layout.group_size, 1); +} + +// More scales than input channels cannot describe a real blocking. The parser +// rejects such tensors, but the layout still has to stay in bounds rather than +// compute a zero group size and divide by it. +TEST(ResolveBlockScaleLayoutTest, ClampsGroupSizeWhenBlocksExceedChannels) { + const BlockScaleLayout layout = ResolveBlockScaleLayout( + /*scale_elements=*/64, /*num_rows=*/8, /*input_channels=*/4); + EXPECT_EQ(layout.groups_per_row, 8); + EXPECT_EQ(layout.group_size, 1); +} + +TEST(ResolveBlockScaleLayoutTest, ToleratesZeroRows) { + const BlockScaleLayout layout = ResolveBlockScaleLayout( + /*scale_elements=*/0, /*num_rows=*/0, /*input_channels=*/16); + EXPECT_EQ(layout.groups_per_row, 0); + EXPECT_EQ(layout.group_size, 1); +} + +TEST(BlockScaleIndexTest, PerChannelLayoutAlwaysUsesTheFirstScale) { + const BlockScaleLayout layout = {/*groups_per_row=*/1, /*group_size=*/16}; + EXPECT_EQ(BlockScaleIndex(layout, 0), 0); + EXPECT_EQ(BlockScaleIndex(layout, 15), 0); +} + +TEST(BlockScaleIndexTest, AdvancesAtBlockBoundaries) { + const BlockScaleLayout layout = {/*groups_per_row=*/4, /*group_size=*/4}; + std::vector indices; + indices.reserve(16); + for (size_t in = 0; in < 16; ++in) { + indices.push_back(BlockScaleIndex(layout, in)); + } + EXPECT_THAT(indices, + ElementsAre(0, 0, 0, 0, 1, 1, 1, 1, 2, 2, 2, 2, 3, 3, 3, 3)); +} + +// Reading past the last block would walk into the next row's scales, so the +// index saturates instead. +TEST(BlockScaleIndexTest, ClampsToTheLastBlock) { + const BlockScaleLayout layout = {/*groups_per_row=*/4, /*group_size=*/4}; + EXPECT_EQ(BlockScaleIndex(layout, 16), 3); + EXPECT_EQ(BlockScaleIndex(layout, 1000), 3); +} + +TEST(BlockScaleIndexTest, ToleratesZeroGroupSize) { + const BlockScaleLayout layout = {/*groups_per_row=*/4, /*group_size=*/0}; + EXPECT_EQ(BlockScaleIndex(layout, 7), 0); +} + +// Two experts, two output channels, four input channels. Rows are interleaved +// as [out, expert], which is what makes the expert stride `num_experts`. +class DequantizeInt8Test : public ::testing::Test { + protected: + static constexpr size_t kNumExperts = 2; + static constexpr size_t kOutputChannels = 2; + static constexpr size_t kInputChannels = 4; + + // Row r holds the values {1, 2, 3, 4} offset by 10 * r, so a mis-indexed row + // is obvious in the output. + std::vector Weights() const { + std::vector weights; + for (size_t row = 0; row < kOutputChannels * kNumExperts; ++row) { + for (size_t in = 0; in < kInputChannels; ++in) { + weights.push_back(static_cast(10 * row + in + 1)); + } + } + return weights; + } +}; + +TEST_F(DequantizeInt8Test, PerChannelScalesApplyToTheWholeRow) { + const std::vector weights = Weights(); + // One scale per row; rows are [out0/e0, out0/e1, out1/e0, out1/e1]. + const std::vector scales = {1.0f, 100.0f, 2.0f, 200.0f}; + std::vector dst(kOutputChannels * kInputChannels, -1.0f); + + CopyAndDequantizeExpertWeightRowsInt8( + weights.data(), scales.data(), scales.size(), kNumExperts, /*expert=*/0, + kOutputChannels, kInputChannels, dst.data()); + + // Expert 0 owns rows 0 and 2, scaled by 1.0 and 2.0. + EXPECT_THAT(dst, ElementsAre(1.0f, 2.0f, 3.0f, 4.0f, // row 0 * 1.0 + 42.0f, 44.0f, 46.0f, 48.0f)); // row 2 * 2.0 +} + +TEST_F(DequantizeInt8Test, BlockwiseScalesChangeHalfwayThroughTheRow) { + const std::vector weights = Weights(); + // Two scales per row over four input channels, so a block is two channels + // wide. Rows are [out0/e0, out0/e1, out1/e0, out1/e1]. + const std::vector scales = {1.0f, 10.0f, // row 0 + -1.0f, -10.0f, // row 1 (other expert) + 2.0f, 20.0f, // row 2 + -2.0f, -20.0f}; // row 3 (other expert) + std::vector dst(kOutputChannels * kInputChannels, -1.0f); + + CopyAndDequantizeExpertWeightRowsInt8( + weights.data(), scales.data(), scales.size(), kNumExperts, /*expert=*/0, + kOutputChannels, kInputChannels, dst.data()); + + EXPECT_THAT(dst, ElementsAre( + // row 0 = {1,2,3,4}, scales {1, 1, 10, 10} + 1.0f, 2.0f, 30.0f, 40.0f, + // row 2 = {21,22,23,24}, scales {2, 2, 20, 20} + 42.0f, 44.0f, 460.0f, 480.0f)); +} + +// The second expert must read its own interleaved rows and its own scales. +TEST_F(DequantizeInt8Test, BlockwiseScalesAreSelectedPerExpert) { + const std::vector weights = Weights(); + const std::vector scales = {1.0f, 10.0f, -1.0f, -10.0f, + 2.0f, 20.0f, -2.0f, -20.0f}; + std::vector dst(kOutputChannels * kInputChannels, 0.0f); + + CopyAndDequantizeExpertWeightRowsInt8( + weights.data(), scales.data(), scales.size(), kNumExperts, /*expert=*/1, + kOutputChannels, kInputChannels, dst.data()); + + EXPECT_THAT(dst, ElementsAre( + // row 1 = {11,12,13,14}, scales {-1, -1, -10, -10} + -11.0f, -12.0f, -130.0f, -140.0f, + // row 3 = {31,32,33,34}, scales {-2, -2, -20, -20} + -62.0f, -64.0f, -660.0f, -680.0f)); +} + +// One scale per input channel is the limiting case of blocking. +TEST_F(DequantizeInt8Test, SupportsOneScalePerInputChannel) { + const std::vector weights = Weights(); + std::vector scales(kOutputChannels * kNumExperts * kInputChannels, + 0.0f); + // Only expert 0's rows (0 and 2) matter here. + for (size_t i = 0; i < kInputChannels; ++i) { + scales[0 * kInputChannels + i] = static_cast(i + 1); + scales[2 * kInputChannels + i] = static_cast(i + 1); + } + std::vector dst(kOutputChannels * kInputChannels, 0.0f); + + CopyAndDequantizeExpertWeightRowsInt8( + weights.data(), scales.data(), scales.size(), kNumExperts, /*expert=*/0, + kOutputChannels, kInputChannels, dst.data()); + + EXPECT_THAT(dst, ElementsAre( + // row 0 = {1,2,3,4} times {1,2,3,4} + 1.0f, 4.0f, 9.0f, 16.0f, + // row 2 = {21,22,23,24} times {1,2,3,4} + 21.0f, 44.0f, 69.0f, 96.0f)); +} + +class DequantizeInt4Test : public ::testing::Test { + protected: + static constexpr size_t kNumExperts = 2; + static constexpr size_t kOutputChannels = 1; + static constexpr size_t kInputChannels = 4; + + // Two nibbles per byte, low nibble first. Row 0 holds {1, 2, 3, 4} and row 1 + // holds {-1, -2, -3, -4}. + std::vector PackedWeights() const { + return { + static_cast(0x21), + static_cast(0x43), // row 0 + static_cast(0xEF), + static_cast(0xCD), // row 1 + }; + } +}; + +TEST_F(DequantizeInt4Test, UnpacksNibblesWithPerChannelScale) { + const std::vector weights = PackedWeights(); + const std::vector scales = {2.0f, 3.0f}; + std::vector dst(kInputChannels, 0.0f); + + CopyAndDequantizeExpertWeightRowsInt4( + weights.data(), scales.data(), scales.size(), kNumExperts, /*expert=*/0, + kOutputChannels, kInputChannels, dst.data()); + + EXPECT_THAT(dst, ElementsAre(2.0f, 4.0f, 6.0f, 8.0f)); +} + +TEST_F(DequantizeInt4Test, AppliesBlockwiseScalesToNibbles) { + const std::vector weights = PackedWeights(); + // Two scales per row over four input channels. + const std::vector scales = {1.0f, 10.0f, // row 0 + 5.0f, 50.0f}; // row 1 + std::vector dst(kInputChannels, 0.0f); + + CopyAndDequantizeExpertWeightRowsInt4( + weights.data(), scales.data(), scales.size(), kNumExperts, /*expert=*/1, + kOutputChannels, kInputChannels, dst.data()); + + // Row 1 = {-1, -2, -3, -4} with scales {5, 5, 50, 50}. + EXPECT_THAT(dst, ElementsAre(-5.0f, -10.0f, -150.0f, -200.0f)); +} + +} // namespace +} // namespace xnnpack +} // namespace tflite diff --git a/tensorflow/lite/delegates/xnnpack/moe_delegate_kernel.cc b/tensorflow/lite/delegates/xnnpack/moe_delegate_kernel.cc index 27421f5dca40fb..bc06fd34069150 100644 --- a/tensorflow/lite/delegates/xnnpack/moe_delegate_kernel.cc +++ b/tensorflow/lite/delegates/xnnpack/moe_delegate_kernel.cc @@ -34,6 +34,7 @@ limitations under the License. #include "tensorflow/lite/c/builtin_op_data.h" #include "tensorflow/lite/core/c/builtin_op_data.h" #include "tensorflow/lite/core/c/common.h" +#include "tensorflow/lite/delegates/xnnpack/moe_block_scale.h" #include "tensorflow/lite/kernels/internal/tensor_ctypes.h" #include "tensorflow/lite/kernels/kernel_util.h" @@ -185,15 +186,11 @@ class MoeExpertsDelegateKernel::Impl { return false; } const int count = NumElements(t); - if (attr.weight_type == MoeExpertsAttributes::WeightType::kInt8) { - if (count != static_cast(num_rows)) { - TF_LITE_KERNEL_LOG( - context, - "%s node #%d %s element count %d does not match rows %zu", - kMoeCustomOp, node_index, name, count, num_rows); - return false; - } - } else if (attr.weight_type == MoeExpertsAttributes::WeightType::kInt4) { + // One scale per row is per-output-channel quantization; any exact + // multiple of the row count is blockwise, with that multiple being the + // number of blocks along the input axis. + if (attr.weight_type == MoeExpertsAttributes::WeightType::kInt8 || + attr.weight_type == MoeExpertsAttributes::WeightType::kInt4) { if (count <= 0 || (static_cast(count) % num_rows != 0)) { TF_LITE_KERNEL_LOG( context, @@ -646,55 +643,6 @@ class MoeExpertsDelegateKernel::Impl { } } - static void CopyAndDequantizeExpertWeightRowsInt8( - const int8_t* weight_i8, const float* scale, size_t num_experts, - size_t expert, size_t output_channels, size_t input_channels, - float* dst) { - for (size_t out = 0; out < output_channels; ++out) { - const size_t row_idx = out * num_experts + expert; - const float row_scale = scale[row_idx]; - const int8_t* src_row = weight_i8 + row_idx * input_channels; - float* dst_row = dst + out * input_channels; - for (size_t in = 0; in < input_channels; ++in) { - dst_row[in] = static_cast(src_row[in]) * row_scale; - } - } - } - - static void CopyAndDequantizeExpertWeightRowsInt4( - const int8_t* weight_i4_packed, const float* scale, size_t scale_elements, - size_t num_experts, size_t expert, size_t output_channels, - size_t input_channels, float* dst) { - const size_t num_rows = output_channels * num_experts; - const size_t groups_per_row = scale_elements / num_rows; - const size_t group_size = - (groups_per_row > 0 && groups_per_row <= input_channels) - ? (input_channels / groups_per_row) - : 1; - - for (size_t out = 0; out < output_channels; ++out) { - const size_t row_idx = out * num_experts + expert; - const int8_t* src_row_packed = - weight_i4_packed + (row_idx * input_channels) / 2; - float* dst_row = dst + out * input_channels; - const float* row_scales = scale + row_idx * groups_per_row; - - for (size_t in = 0; in < input_channels; ++in) { - const size_t byte_idx = in / 2; - const int8_t byte_val = src_row_packed[byte_idx]; - const int8_t nibble = (in % 2 == 0) - ? static_cast(byte_val << 4) >> 4 - : static_cast(byte_val >> 4); - const size_t group_idx = - (groups_per_row > 1 && group_size > 0) - ? std::min(in / group_size, groups_per_row - 1) - : 0; - const float scale_val = row_scales[group_idx]; - dst_row[in] = static_cast(nibble) * scale_val; - } - } - } - void CopyGateUpExpertWeight(const void* gate_weight, const float* gate_scale, size_t gate_scale_elements, const void* ff1_weight, const float* ff1_scale, @@ -715,11 +663,12 @@ class MoeExpertsDelegateKernel::Impl { dst + hidden_dim * model_dim); } else if (attr_.weight_type == MoeExpertsAttributes::WeightType::kInt8) { CopyAndDequantizeExpertWeightRowsInt8( - static_cast(gate_weight), gate_scale, num_experts, - expert, hidden_dim, model_dim, dst); + static_cast(gate_weight), gate_scale, + gate_scale_elements, num_experts, expert, hidden_dim, model_dim, dst); CopyAndDequantizeExpertWeightRowsInt8( - static_cast(ff1_weight), ff1_scale, num_experts, - expert, hidden_dim, model_dim, dst + hidden_dim * model_dim); + static_cast(ff1_weight), ff1_scale, ff1_scale_elements, + num_experts, expert, hidden_dim, model_dim, + dst + hidden_dim * model_dim); } else { CopyExpertWeightRows(static_cast(gate_weight), num_experts, expert, hidden_dim, model_dim, dst); @@ -741,8 +690,9 @@ class MoeExpertsDelegateKernel::Impl { kernel_buffer_.data()); } else if (attr_.weight_type == MoeExpertsAttributes::WeightType::kInt8) { CopyAndDequantizeExpertWeightRowsInt8( - static_cast(weight), scale, num_experts, expert, - output_channels, input_channels, kernel_buffer_.data()); + static_cast(weight), scale, scale_elements, + num_experts, expert, output_channels, input_channels, + kernel_buffer_.data()); } else { CopyExpertWeightRows(static_cast(weight), num_experts, expert, output_channels, input_channels, diff --git a/tensorflow/lite/delegates/xnnpack/moe_delegate_kernel_test.cc b/tensorflow/lite/delegates/xnnpack/moe_delegate_kernel_test.cc new file mode 100644 index 00000000000000..a2f4430833f0d9 --- /dev/null +++ b/tensorflow/lite/delegates/xnnpack/moe_delegate_kernel_test.cc @@ -0,0 +1,217 @@ +/* Copyright 2026 The TensorFlow 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. +==============================================================================*/ +#include "tensorflow/lite/delegates/xnnpack/moe_delegate_kernel.h" + +#include +#include +#include +#include +#include + +#include +#include "flatbuffers/flexbuffers.h" // from @flatbuffers +#include "tensorflow/lite/builtin_ops.h" +#include "tensorflow/lite/core/c/common.h" +#include "tensorflow/lite/interpreter.h" + +namespace tflite { +namespace xnnpack { +namespace { + +constexpr int kNumExperts = 2; +constexpr int kNumActiveExperts = 1; +constexpr int kModelDim = 8; +constexpr int kHiddenDim = 4; +constexpr int kNumTokens = 2; + +// gate/ff1 are [hidden_dim, num_experts, 1, model_dim]; linear is the +// transpose. One scale per row is per-output-channel quantization. +constexpr int kGateFf1Rows = kHiddenDim * kNumExperts; +constexpr int kLinearRows = kModelDim * kNumExperts; + +std::vector MoeCustomData(const std::string& weight_type, + const std::string& activation) { + flexbuffers::Builder fbb; + fbb.Map([&]() { + fbb.Int("num_experts", kNumExperts); + fbb.Int("num_active_experts", kNumActiveExperts); + fbb.Int("model_dim", kModelDim); + fbb.Int("hidden_dim", kHiddenDim); + fbb.String("weight_type", weight_type); + fbb.String("activation", activation); + fbb.Bool("renormalized_top_weights", true); + }); + fbb.Finish(); + return fbb.GetBuffer(); +} + +// Builds a single-node interpreter holding one custom "moe" op so that +// MoeExpertsDelegateKernel::IsSupported can be called against a real +// TfLiteContext. The op is never invoked, so the tensor contents are +// irrelevant; only types, shapes and constness matter. +class MoeNodeBuilder { + public: + MoeNodeBuilder() : interpreter_(std::make_unique()) { + registration_ = {nullptr, nullptr, nullptr, nullptr}; + registration_.builtin_code = kTfLiteBuiltinCustom; + registration_.custom_name = "moe"; + registration_.version = 1; + } + + void AddInput(TfLiteType type, const std::vector& shape) { + interpreter_->AddTensors(1); + interpreter_->SetTensorParametersReadWrite( + next_tensor_, type, "", shape, {kTfLiteNoQuantization, nullptr}); + inputs_.push_back(next_tensor_++); + } + + void AddConstInput(TfLiteType type, const std::vector& shape, + size_t num_bytes) { + buffers_.push_back(std::vector(num_bytes, 0)); + interpreter_->AddTensors(1); + interpreter_->SetTensorParametersReadOnly( + next_tensor_, type, "", shape, {kTfLiteNoQuantization, nullptr}, + buffers_.back().data(), buffers_.back().size()); + inputs_.push_back(next_tensor_++); + } + + void AddOutput(TfLiteType type, const std::vector& shape) { + interpreter_->AddTensors(1); + interpreter_->SetTensorParametersReadWrite( + next_tensor_, type, "", shape, {kTfLiteNoQuantization, nullptr}); + outputs_.push_back(next_tensor_++); + } + + void SetCustomData(const std::vector& data) { + custom_data_.assign(data.begin(), data.end()); + } + + void Finish() { + interpreter_->SetInputs(inputs_); + interpreter_->SetOutputs(outputs_); + interpreter_->AddNodeWithParameters(inputs_, outputs_, custom_data_.data(), + custom_data_.size(), nullptr, + ®istration_); + } + + TfLiteStatus IsSupported() { + const auto* pair = interpreter_->node_and_registration(0); + return MoeExpertsDelegateKernel::IsSupported( + interpreter_->primary_subgraph().context(), &pair->first, &pair->second, + /*node_index=*/0); + } + + private: + std::unique_ptr interpreter_; + TfLiteRegistration registration_; + std::vector inputs_; + std::vector outputs_; + std::vector custom_data_; + // Backing storage for the constant tensors, which are referenced rather than + // copied by the interpreter. + std::vector> buffers_; + int next_tensor_ = 0; +}; + +// Assembles the ten inputs of an int8 moe op. `*_scales` is the total element +// count of each scale tensor: `rows` is per-output-channel and `rows * n` is +// blockwise with n blocks along the input axis. +void BuildInt8Moe(MoeNodeBuilder& builder, int gate_scales, int ff1_scales, + int linear_scales, + const std::string& activation = "gelu_tanh") { + builder.AddInput(kTfLiteFloat32, {1, 1, kNumTokens, kModelDim}); + builder.AddInput(kTfLiteFloat32, {1, 1, kNumTokens, kNumActiveExperts}); + builder.AddInput(kTfLiteInt32, {1, 1, kNumTokens, kNumActiveExperts}); + + builder.AddConstInput(kTfLiteInt8, {kHiddenDim, kNumExperts, 1, kModelDim}, + kGateFf1Rows * kModelDim * sizeof(int8_t)); + builder.AddConstInput(kTfLiteFloat32, {gate_scales}, + gate_scales * sizeof(float)); + builder.AddConstInput(kTfLiteInt8, {kHiddenDim, kNumExperts, 1, kModelDim}, + kGateFf1Rows * kModelDim * sizeof(int8_t)); + builder.AddConstInput(kTfLiteFloat32, {ff1_scales}, + ff1_scales * sizeof(float)); + builder.AddConstInput(kTfLiteInt8, {kModelDim, kNumExperts, 1, kHiddenDim}, + kLinearRows * kHiddenDim * sizeof(int8_t)); + builder.AddConstInput(kTfLiteFloat32, {linear_scales}, + linear_scales * sizeof(float)); + builder.AddConstInput(kTfLiteFloat32, {kNumExperts}, + kNumExperts * sizeof(float)); + + builder.AddOutput(kTfLiteFloat32, {1, 1, kNumTokens, kModelDim}); + builder.SetCustomData(MoeCustomData("int8", activation)); + builder.Finish(); +} + +TEST(MoeDelegateKernelIsSupportedTest, AcceptsPerChannelScales) { + MoeNodeBuilder builder; + BuildInt8Moe(builder, kGateFf1Rows, kGateFf1Rows, kLinearRows); + EXPECT_EQ(builder.IsSupported(), kTfLiteOk); +} + +// Four scales per row, i.e. the input axis split into four blocks. This is the +// case the delegate could not previously represent. +TEST(MoeDelegateKernelIsSupportedTest, AcceptsBlockwiseScales) { + MoeNodeBuilder builder; + BuildInt8Moe(builder, kGateFf1Rows * 4, kGateFf1Rows * 4, kLinearRows * 4); + EXPECT_EQ(builder.IsSupported(), kTfLiteOk); +} + +// Different projections may be blocked differently, since gate/ff1 and linear +// have different input axes. +TEST(MoeDelegateKernelIsSupportedTest, AcceptsMixedBlockCounts) { + MoeNodeBuilder builder; + BuildInt8Moe(builder, kGateFf1Rows * 2, kGateFf1Rows, kLinearRows * 4); + EXPECT_EQ(builder.IsSupported(), kTfLiteOk); +} + +// A scale count that is not a whole number of scales per row cannot describe +// any blocking. +TEST(MoeDelegateKernelIsSupportedTest, RejectsScaleCountNotMultipleOfRows) { + MoeNodeBuilder builder; + BuildInt8Moe(builder, kGateFf1Rows + 1, kGateFf1Rows, kLinearRows); + EXPECT_EQ(builder.IsSupported(), kTfLiteError); +} + +TEST(MoeDelegateKernelIsSupportedTest, RejectsEmptyScaleTensor) { + MoeNodeBuilder builder; + BuildInt8Moe(builder, kGateFf1Rows, kGateFf1Rows, 0); + EXPECT_EQ(builder.IsSupported(), kTfLiteError); +} + +// The exporter now labels the activation "gelu_tanh"; the legacy "gelu" +// spelling has to keep working for already-converted models. +TEST(MoeDelegateKernelIsSupportedTest, AcceptsBothGeluSpellings) { + MoeNodeBuilder tanh_builder; + BuildInt8Moe(tanh_builder, kGateFf1Rows, kGateFf1Rows, kLinearRows, + /*activation=*/"gelu_tanh"); + EXPECT_EQ(tanh_builder.IsSupported(), kTfLiteOk); + + MoeNodeBuilder legacy_builder; + BuildInt8Moe(legacy_builder, kGateFf1Rows, kGateFf1Rows, kLinearRows, + /*activation=*/"gelu"); + EXPECT_EQ(legacy_builder.IsSupported(), kTfLiteOk); +} + +TEST(MoeDelegateKernelIsSupportedTest, RejectsUnknownActivation) { + MoeNodeBuilder builder; + BuildInt8Moe(builder, kGateFf1Rows, kGateFf1Rows, kLinearRows, + /*activation=*/"silu"); + EXPECT_EQ(builder.IsSupported(), kTfLiteError); +} + +} // namespace +} // namespace xnnpack +} // namespace tflite diff --git a/tensorflow/lite/kernels/internal/reference/select.h b/tensorflow/lite/kernels/internal/reference/select.h index 5fe86f174a56ab..408d584c98b20e 100644 --- a/tensorflow/lite/kernels/internal/reference/select.h +++ b/tensorflow/lite/kernels/internal/reference/select.h @@ -109,18 +109,22 @@ void RunSelectOp(const D* cond, const T* x, const T* y, T* output, } else { TFLITE_DCHECK_EQ(cond_stride_0, 1); if (x_stride_0 == 0 && y_stride_0 == 0) { + const T x0 = *x; + const T y0 = *y; for (size_t i = 0; i < output_shape_0; ++i) { - output[i] = cond[i] ? *x : *y; + output[i] = cond[i] ? x0 : y0; } } else if (x_stride_0 == 0) { TFLITE_DCHECK_EQ(y_stride_0, 1); + const T x0 = *x; for (size_t i = 0; i < output_shape_0; ++i) { - output[i] = cond[i] ? *x : y[i]; + output[i] = cond[i] ? x0 : y[i]; } } else if (y_stride_0 == 0) { TFLITE_DCHECK_EQ(x_stride_0, 1); + const T y0 = *y; for (size_t i = 0; i < output_shape_0; ++i) { - output[i] = cond[i] ? x[i] : *y; + output[i] = cond[i] ? x[i] : y0; } } else { TFLITE_DCHECK_EQ(x_stride_0, 1); diff --git a/tensorflow/lite/python/analyzer_wrapper/BUILD b/tensorflow/lite/python/analyzer_wrapper/BUILD index ccf7b77a348340..cd9c91a5f219c2 100644 --- a/tensorflow/lite/python/analyzer_wrapper/BUILD +++ b/tensorflow/lite/python/analyzer_wrapper/BUILD @@ -49,7 +49,6 @@ cc_library( visibility = ["//visibility:public"], deps = [ "//tensorflow/compiler/mlir/lite/schema:schema_utils", - "//tensorflow/core:framework", "//tensorflow/core/public:release_version", "//tensorflow/lite:version", "//tensorflow/lite/core:model_builder", diff --git a/tensorflow/lite/toco/graph_transformations/reorder_reshape_transpose.cc b/tensorflow/lite/toco/graph_transformations/reorder_reshape_transpose.cc index 8d417c8fde2d6e..abb92666625751 100644 --- a/tensorflow/lite/toco/graph_transformations/reorder_reshape_transpose.cc +++ b/tensorflow/lite/toco/graph_transformations/reorder_reshape_transpose.cc @@ -12,7 +12,6 @@ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#include #include #include #include diff --git a/tensorflow/lite/tools/strip_buffers/BUILD b/tensorflow/lite/tools/strip_buffers/BUILD index 723bcb216657d9..1c2ddc1af6b349 100644 --- a/tensorflow/lite/tools/strip_buffers/BUILD +++ b/tensorflow/lite/tools/strip_buffers/BUILD @@ -34,7 +34,6 @@ cc_library( hdrs = ["stripping_lib.h"], copts = tflite_copts(), deps = [ - "//tensorflow/core:tflite_portable_logging", "//tensorflow/lite:framework", "//tensorflow/lite:string", "//tensorflow/lite/core:framework", @@ -54,7 +53,6 @@ cc_binary( linkopts = tflite_linkopts(), deps = [ ":stripping_lib", - "//tensorflow/core:tflite_portable_logging", "//tensorflow/lite:framework", "//tensorflow/lite/core:framework", "//tensorflow/lite/core/c:common", @@ -71,7 +69,6 @@ cc_binary( linkopts = tflite_linkopts(), deps = [ ":stripping_lib", - "//tensorflow/core:tflite_portable_logging", "//tensorflow/lite:framework", "//tensorflow/lite/core:framework", "//tensorflow/lite/core/c:common", diff --git a/tensorflow/lite/tools/strip_buffers/reconstitute_buffers_into_fb.cc b/tensorflow/lite/tools/strip_buffers/reconstitute_buffers_into_fb.cc index e02cdecef4896a..acd9fcfd257931 100644 --- a/tensorflow/lite/tools/strip_buffers/reconstitute_buffers_into_fb.cc +++ b/tensorflow/lite/tools/strip_buffers/reconstitute_buffers_into_fb.cc @@ -20,7 +20,6 @@ limitations under the License. #include "absl/log/log.h" #include "flatbuffers/flatbuffer_builder.h" // from @flatbuffers -#include "tensorflow/core/platform/logging.h" #include "tensorflow/lite/core/c/common.h" #include "tensorflow/lite/model_builder.h" #include "tensorflow/lite/tools/command_line_flags.h" diff --git a/tensorflow/lite/tools/strip_buffers/strip_buffers_from_fb.cc b/tensorflow/lite/tools/strip_buffers/strip_buffers_from_fb.cc index 6915075582927a..7f04a2031e46b9 100644 --- a/tensorflow/lite/tools/strip_buffers/strip_buffers_from_fb.cc +++ b/tensorflow/lite/tools/strip_buffers/strip_buffers_from_fb.cc @@ -20,7 +20,6 @@ limitations under the License. #include "absl/log/log.h" #include "flatbuffers/flatbuffer_builder.h" // from @flatbuffers -#include "tensorflow/core/platform/logging.h" #include "tensorflow/lite/core/c/common.h" #include "tensorflow/lite/model_builder.h" #include "tensorflow/lite/tools/command_line_flags.h" diff --git a/tensorflow/lite/tools/strip_buffers/stripping_lib.cc b/tensorflow/lite/tools/strip_buffers/stripping_lib.cc index d29a757e491a21..121705e6343e56 100644 --- a/tensorflow/lite/tools/strip_buffers/stripping_lib.cc +++ b/tensorflow/lite/tools/strip_buffers/stripping_lib.cc @@ -27,7 +27,6 @@ limitations under the License. #include "absl/strings/string_view.h" #include "flatbuffers/buffer.h" // from @flatbuffers #include "flatbuffers/flatbuffer_builder.h" // from @flatbuffers -#include "tensorflow/core/platform/logging.h" #include "tensorflow/lite/core/c/common.h" #include "tensorflow/lite/model_builder.h" #include "tensorflow/lite/schema/schema_generated.h" diff --git a/tensorflow/python/data/experimental/kernel_tests/service/cross_trainer_cache_test.py b/tensorflow/python/data/experimental/kernel_tests/service/cross_trainer_cache_test.py index 243862993d6b51..bd1b3f3b8b0574 100644 --- a/tensorflow/python/data/experimental/kernel_tests/service/cross_trainer_cache_test.py +++ b/tensorflow/python/data/experimental/kernel_tests/service/cross_trainer_cache_test.py @@ -108,8 +108,8 @@ def testConcurrentReaders(self): num_cpus = multiprocessing.cpu_count() cluster = self._create_cluster( num_workers=1, cross_trainer_cache_size_bytes=(num_cpus + 8) * 423) - num_readers = 20 - num_elements = 50 + num_readers = 8 + num_elements = 20 dataset = dataset_ops.Dataset.range(10000000).repeat() datasets = [] diff --git a/tensorflow/python/kernel_tests/control_flow/control_flow_ops_py_test.py b/tensorflow/python/kernel_tests/control_flow/control_flow_ops_py_test.py index 463b75934f8085..ad85997e0df7e2 100644 --- a/tensorflow/python/kernel_tests/control_flow/control_flow_ops_py_test.py +++ b/tensorflow/python/kernel_tests/control_flow/control_flow_ops_py_test.py @@ -4970,6 +4970,78 @@ def testGuardedAssertDoesNotCopyWhenTrue(self): # No copy was performed for the guarded assert self.assertEqual([], guarded_memcpy_nodestat_names) + @test_util.run_in_graph_and_eager_modes + def testAssertRaisesWithRaggedTensorData(self): + rt = ragged_factory_ops.constant([[1], [2]]) + cond = math_ops.reduce_all(math_ops.equal(rt.row_lengths(), 2)) + with self.assertRaisesRegex(errors_impl.InvalidArgumentError, + r"\[\[1\], \[2\]\]"): + self.evaluate(control_flow_assert.Assert(cond, [rt])) + + @test_util.run_in_graph_and_eager_modes + def testAssertPassesWithRaggedTensorData(self): + rt = ragged_factory_ops.constant([[1], [2]]) + cond = math_ops.reduce_all(math_ops.equal(rt.row_lengths(), 1)) + self.evaluate(control_flow_assert.Assert(cond, [rt])) + + @test_util.run_in_graph_and_eager_modes + def testAssertRaisesWithDenseAndRaggedTensorData(self): + dense = constant_op.constant([7.5, 8.5]) + rt = ragged_factory_ops.constant([[1, 2, 3, 4], [5]]) + with self.assertRaisesRegex(errors_impl.InvalidArgumentError, + r"(?s)7\.5.*\[\[1, 2, 3, 4\], \[5\]\]"): + self.evaluate( + control_flow_assert.Assert(False, [dense, rt], summarize=-1)) + + @test_util.run_in_graph_and_eager_modes + def testAssertSummarizesRaggedTensorData(self): + rt = ragged_factory_ops.constant([[i] for i in range(8)]) + with self.assertRaisesRegex( + errors_impl.InvalidArgumentError, + r"\[\[0\], \[1\], \[2\], \.\.\., \[5\], \[6\], \[7\]\]"): + self.evaluate(control_flow_assert.Assert(False, [rt])) + with self.assertRaisesRegex(errors_impl.InvalidArgumentError, + r"\[\[0\], \[1\], \[2\], \[3\], \[4\]"): + self.evaluate(control_flow_assert.Assert(False, [rt], summarize=-2)) + with self.assertRaisesRegex(errors_impl.InvalidArgumentError, + r"\[\[0\], \[1\], \.\.\., \[6\], \[7\]\]"): + self.evaluate( + control_flow_assert.Assert(False, [rt], summarize=np.int64(2))) + with self.assertRaises(errors_impl.InvalidArgumentError): + self.evaluate(control_flow_assert.Assert(False, [rt], summarize=0)) + + @test_util.run_in_graph_and_eager_modes + def testAssertWithUnknownRankRaggedTensorData(self): + + @eager_def_function.function(input_signature=[ + ragged_tensor.RaggedTensorSpec(ragged_rank=1, dtype=dtypes.int32), + tensor_lib.TensorSpec([], dtypes.bool), + ]) + def f(rt, cond): + with ops.control_dependencies([control_flow_assert.Assert(cond, [rt])]): + return array_ops.identity(rt.flat_values) + + rt = ragged_factory_ops.constant([[1, 2], [3]]) + self.assertAllEqual([1, 2, 3], self.evaluate(f(rt, True))) + with self.assertRaisesRegex( + errors_impl.InvalidArgumentError, + r"flat_values=\[1 2 3\], nested_row_splits=\[\[0 2 3\]\]"): + self.evaluate(f(rt, False)) + + @test_util.run_in_graph_and_eager_modes + def testAssertWithRaggedTensorDataCompilesWithXla(self): + if not test_util.is_xla_enabled(): + self.skipTest("Requires XLA.") + + @eager_def_function.function(jit_compile=True) + def f(rt): + cond = math_ops.reduce_all(math_ops.greater(rt.row_lengths(), 0)) + with ops.control_dependencies([control_flow_assert.Assert(cond, [rt])]): + return math_ops.reduce_sum(rt.flat_values) + + self.assertEqual(6, self.evaluate(f(ragged_factory_ops.constant([[1, 2], + [3]])))) + class WhileOpBenchmark(test.Benchmark): """Evaluate the performance of while_loop op.""" diff --git a/tensorflow/python/kernel_tests/math_ops/bincount_op_test.py b/tensorflow/python/kernel_tests/math_ops/bincount_op_test.py index 293d0c67f4b4e4..36760fd2227b80 100644 --- a/tensorflow/python/kernel_tests/math_ops/bincount_op_test.py +++ b/tensorflow/python/kernel_tests/math_ops/bincount_op_test.py @@ -756,6 +756,30 @@ def test_sparse_bincount_input_validation(self): size=size, weights=[])) + def test_sparse_bincount_rank0_fails(self): + # Regression test for the rank-0 guard in SparseBincountOp::Compute + # ("dense_shape must have at least 1 dimension"). Without it the kernel + # reads dense_shape(0) on an empty dense_shape. Shape inference may reject + # the input first, so both messages are accepted. + with self.assertRaisesRegex( + (ValueError, errors.InvalidArgumentError), + "Input must be less than rank 2|" + "dense_shape must have at least 1 dimension", + ): + self.evaluate( + gen_math_ops.sparse_bincount( + indices=constant_op.constant( + [], shape=[1, 0], dtype=dtypes.int64 + ), + values=[0], + dense_shape=constant_op.constant( + [], shape=[0], dtype=dtypes.int64 + ), + size=10, + weights=[], + ) + ) + class RaggedBincountOpTest(test_util.TensorFlowTestCase, parameterized.TestCase): diff --git a/tensorflow/python/kernel_tests/math_ops/cwise_ops_unary_test.py b/tensorflow/python/kernel_tests/math_ops/cwise_ops_unary_test.py index b45144d6205e25..5e4214eda409e7 100644 --- a/tensorflow/python/kernel_tests/math_ops/cwise_ops_unary_test.py +++ b/tensorflow/python/kernel_tests/math_ops/cwise_ops_unary_test.py @@ -695,6 +695,42 @@ def g(x): *gradient_checker_v2.compute_gradient(g, [ops.convert_to_tensor(2.0)])) self.assertLess(err, 1e-3) + @test_util.run_in_graph_and_eager_modes + def testComplexSignSmallMagnitude(self): + """Regression test for #116945: sign(z) must not return 0 for small |z|. + + Eigen's scalar_sign_op computes |z|^2 = re^2 + im^2 as a float32 + intermediate, which underflows to 0 when |z| < sqrt(float32_tiny) ~= + 1.08e-19, causing sign(z) to incorrectly return 0. + """ + # Test values with magnitudes below the underflow threshold. + for dtype in [np.complex64, np.complex128]: + small_values = np.array( + [ + 1e-20 + 0j, + 1e-30 + 0j, + 0 + 1e-20j, + -1e-20 + 0j, + 1e-20 + 1e-20j, + ], + dtype=dtype, + ) + + result = self.evaluate(math_ops.sign(constant_op.constant(small_values))) + + # sign(z) for nonzero z must be z / |z| (unit magnitude), never 0. + for i, z in enumerate(small_values): + expected = z / np.abs(z) + self.assertNotEqual( + result[i], 0, msg=f"sign({z}) returned 0 for dtype={dtype.__name__}" + ) + self.assertAllClose( + result[i], + expected, + rtol=1e-5, + msg=f"sign({z}) incorrect for dtype={dtype.__name__}", + ) + @test_util.run_in_graph_and_eager_modes def testDigamma(self): try: diff --git a/tensorflow/python/kernel_tests/nn_ops/conv2d_transpose_test.py b/tensorflow/python/kernel_tests/nn_ops/conv2d_transpose_test.py index 32bc59a0800a44..b795c888b4ca24 100644 --- a/tensorflow/python/kernel_tests/nn_ops/conv2d_transpose_test.py +++ b/tensorflow/python/kernel_tests/nn_ops/conv2d_transpose_test.py @@ -369,5 +369,30 @@ def testConv2DTransposeLargeOutputShape(self): filters=[[[[1]]]]) self.evaluate(op) + def testConv2DTransposeInvalidFilterRank(self): + with self.session(): + with self.assertRaises((errors.InvalidArgumentError, ValueError)): + op = nn_ops.conv2d_transpose( + input=constant_op.constant(1.0, shape=[1, 3, 3, 8]), + filters=8, + output_shape=(2, 2), + strides=(2, 2), + padding="SAME", + ) + self.evaluate(op) + + def testConv2DTransposeInvalidOutputShapeRank(self): + with self.session(): + with self.assertRaises((errors.InvalidArgumentError, ValueError)): + op = nn_ops.conv2d_transpose( + input=constant_op.constant(1.0, shape=[1, 4, 4, 3]), + filters=constant_op.constant(1.0, shape=[3, 3, 3, 3]), + output_shape=(2, 4, 2), + strides=[1, 2, 2, 1], + padding="SAME", + ) + self.evaluate(op) + + if __name__ == "__main__": test.main() diff --git a/tensorflow/python/ops/BUILD b/tensorflow/python/ops/BUILD index 2bcda6a1843e58..675b0563574d5d 100644 --- a/tensorflow/python/ops/BUILD +++ b/tensorflow/python/ops/BUILD @@ -1123,12 +1123,15 @@ py_library( ":array_ops", ":cond", ":control_flow_ops_gen", + ":control_flow_util", ":logging_ops_gen", ":math_ops_gen", + ":string_ops", "//tensorflow/python/eager:context", "//tensorflow/python/framework:dtypes", "//tensorflow/python/framework:errors", "//tensorflow/python/framework:ops", + "//tensorflow/python/types:internal", "//tensorflow/python/util:dispatch", "//tensorflow/python/util:tf_export", "//tensorflow/python/util:tf_should_use", diff --git a/tensorflow/python/ops/control_flow_assert.py b/tensorflow/python/ops/control_flow_assert.py index 18928fa6f35754..5dd7ee7bff7902 100644 --- a/tensorflow/python/ops/control_flow_assert.py +++ b/tensorflow/python/ops/control_flow_assert.py @@ -20,9 +20,12 @@ from tensorflow.python.framework import ops from tensorflow.python.ops import array_ops from tensorflow.python.ops import cond +from tensorflow.python.ops import control_flow_util from tensorflow.python.ops import gen_control_flow_ops from tensorflow.python.ops import gen_logging_ops from tensorflow.python.ops import gen_math_ops +from tensorflow.python.ops import string_ops +from tensorflow.python.types import internal as internal_types from tensorflow.python.util import dispatch from tensorflow.python.util import tf_should_use from tensorflow.python.util.tf_export import tf_export @@ -57,6 +60,38 @@ def _summarize_eager(tensor, summarize=None): return ", ".join(lst) +def _format_ragged(x, summarize): + """Returns `x`, or a scalar string `Tensor` to pass in its place if ragged. + + Assert only handles dense tensors, so a `RaggedTensor` of known rank is + formatted the way `tf.print` formats it. + + Args: + x: An entry of the `data` argument to `Assert`. + summarize: The `summarize` argument to `Assert`. + """ + if not isinstance(x, internal_types.RaggedTensor): + return x + # Default to Assert's own 3. Ragged formatting accepts only a Python int + # that is -1 (everything) or positive. + summarize = 3 if summarize is None else int(summarize) + if summarize == 0 or control_flow_util.GraphOrParentsInXlaContext( + ops.get_default_graph()): + # Assert shows no entries for summarize=0, and XLA cannot compile the + # formatting ops. + return ops.convert_to_tensor("") + if summarize < 0: + summarize = -1 + if x.shape.rank is None: + # Ragged formatting needs a known rank; fall back to the components. + splits = x.nested_row_splits + template = "tf.RaggedTensor(flat_values={}, nested_row_splits=[%s])" % ( + ", ".join(["{}"] * len(splits))) + return string_ops.string_format( + template, [x.flat_values, *splits], summarize=summarize) + return string_ops.string_format("{}", [x], summarize=summarize) + + # Assert and Print are special symbols in python, so we must # use an upper-case version of them. @tf_export("debugging.Assert", "Assert") @@ -70,8 +105,10 @@ def Assert(condition, data, summarize=None, name=None): Args: condition: The condition to evaluate. - data: The tensors to print out when condition is false. - summarize: Print this many entries of each tensor. + data: The tensors to print out when condition is false. A `RaggedTensor` + of known rank is printed as `tf.print` prints it. + summarize: Print this many entries of each tensor. For a `RaggedTensor`, + the first and last this many entries of each dimension. name: A name for this operation (optional). Returns: @@ -97,6 +134,7 @@ def Assert(condition, data, summarize=None, name=None): """ if context.executing_eagerly(): if not condition: + data = [_format_ragged(x, summarize) for x in data] xs = ops.convert_n_to_tensor(data) data_str = [_summarize_eager(x, summarize) for x in xs] raise errors.InvalidArgumentError( @@ -107,8 +145,13 @@ def Assert(condition, data, summarize=None, name=None): return with ops.name_scope(name, "Assert", [condition, data]) as name: - xs = ops.convert_n_to_tensor(data) - if all(x.dtype in {dtypes.string, dtypes.int32} for x in xs): + # Ragged data cannot go through convert_n_to_tensor, and its formatting + # ops should only run when the assert fails, so it takes the guarded path. + data = list(data) + has_ragged = any(isinstance(x, internal_types.RaggedTensor) for x in data) + if not has_ragged and all( + x.dtype in {dtypes.string, dtypes.int32} + for x in ops.convert_n_to_tensor(data)): # As a simple heuristic, we assume that string and int32 are # on host to avoid the need to use cond. If it is not case, # we will pay the price copying the tensor to host memory. @@ -118,7 +161,8 @@ def Assert(condition, data, summarize=None, name=None): def true_assert(): return gen_logging_ops._assert( # pylint: disable=protected-access - condition, data, summarize, name="Assert") + condition, [_format_ragged(x, summarize) for x in data], + summarize, name="Assert") guarded_assert = cond.cond( condition, diff --git a/tensorflow/python/ops/nn_ops.py b/tensorflow/python/ops/nn_ops.py index 58c30af22bec4b..70d28b1a186cf2 100644 --- a/tensorflow/python/ops/nn_ops.py +++ b/tensorflow/python/ops/nn_ops.py @@ -2986,6 +2986,13 @@ def conv2d_transpose_v2( data_format = "NHWC" channel_index = 1 if data_format.startswith("NC") else 3 + filters = ops.convert_to_tensor(filters, name="filters") + if filters.shape.rank is not None and filters.shape.rank != 4: + raise ValueError( + "`filters` must be 4-dimensional. " + f"Received: filters with shape {filters.shape}" + ) + strides = _get_sequence(strides, 2, channel_index, "strides") dilations = _get_sequence(dilations, 2, channel_index, "dilations") padding, explicit_paddings = convert_padding(padding) diff --git a/tensorflow/workspace2.bzl b/tensorflow/workspace2.bzl index eacb39dbf8a972..bcfe4415614b9f 100644 --- a/tensorflow/workspace2.bzl +++ b/tensorflow/workspace2.bzl @@ -67,6 +67,7 @@ load("@xla//third_party/tensorrt:tensorrt_configure.bzl", "tensorrt_configure") load("@xla//third_party/tensorrt:workspace.bzl", tensorrt = "repo") load("@xla//third_party/triton:workspace.bzl", triton = "repo") load("@xla//third_party/xnnpack:workspace.bzl", xnnpack = "repo") +load("@xla//third_party/xz:workspace.bzl", xz = "repo") load("@xla//tools/def_file_filter:def_file_filter_configure.bzl", "def_file_filter_configure") load("@xla//tools/toolchains:cpus/aarch64/aarch64_compiler_configure.bzl", "aarch64_compiler_configure") load("@xla//tools/toolchains:cpus/arm/arm_compiler_configure.bzl", "arm_compiler_configure") @@ -255,6 +256,7 @@ def _tf_repositories(): ) xnnpack() + xz() # XNNPack dependency. tf_http_archive( diff --git a/third_party/xla/.github/workflows/ci_multi_device.yml b/third_party/xla/.github/workflows/ci_multi_device.yml index 8a8ab8fa70b477..4ea65669e06f76 100644 --- a/third_party/xla/.github/workflows/ci_multi_device.yml +++ b/third_party/xla/.github/workflows/ci_multi_device.yml @@ -69,12 +69,14 @@ jobs: halt-dispatch-input: ${{ inputs.halt-for-connection }} - name: "Run build.py" working-directory: ${{ matrix.job_info.repo }} + env: + PYTHONPATH: ${{ github.workspace }}/openxla/xla run: | # shellcheck disable=SC2193 if [[ "${{ matrix.job_info.pool }}" == *"windows"* ]]; then python "$GITHUB_WORKSPACE"\\openxla\\xla\\build_tools\\ci\\build.py --build="${{ matrix.job_info.name }}_github_actions" else - "$GITHUB_WORKSPACE"/openxla/xla/build_tools/ci/build.py --build="${{ matrix.job_info.name }}_github_actions" + python3 "$GITHUB_WORKSPACE"/openxla/xla/build_tools/ci/build.py --build="${{ matrix.job_info.name }}_github_actions" fi diff --git a/third_party/xla/AGENTS.md b/third_party/xla/AGENTS.md new file mode 100644 index 00000000000000..d3dcfded10316e --- /dev/null +++ b/third_party/xla/AGENTS.md @@ -0,0 +1,11 @@ +# AI Assistant Guidelines for OpenXLA + +Welcome to the OpenXLA (`openxla/xla`) repository. + +* **Core XLA compiler guidelines**: When generating, reviewing, or modifying + code under [`xla/`](xla/), follow the detailed C++, MLIR, testing, and + Bazel guidelines in [`xla/AGENTS.md`](xla/AGENTS.md). +* [`build_tools/`](build_tools/) contains CI and developer tooling, + [`docs/`](docs/) the user-facing documentation, and + [`third_party/`](third_party/) vendored dependencies and Bazel repository + rules. diff --git a/third_party/xla/docs/custom_call.md b/third_party/xla/docs/custom_call.md index 358fdaa26044e5..2356c39dc657cb 100644 --- a/third_party/xla/docs/custom_call.md +++ b/third_party/xla/docs/custom_call.md @@ -309,7 +309,7 @@ The external FFI API in `xla/ffi/api/ffi.h` returns `Error` from `Verify` and ```c++ ABSL_ASSIGN_OR_RETURN(BufferR2 input, - Match("input", buffer, m::Buffer())); + Match("input", buffer, m::Buffer())); ABSL_RETURN_IF_ERROR(Verify( "input", input, m::Buffer().WithDims(expected_rows, m::Dim()))); diff --git a/third_party/xla/opensource_only.files b/third_party/xla/opensource_only.files index 097ae77aa29598..9c227079557dd5 100644 --- a/third_party/xla/opensource_only.files +++ b/third_party/xla/opensource_only.files @@ -160,4 +160,5 @@ xla/third_party/tensorrt/tensorrt/tensorrt_config.py.tpl: xla/third_party/tensorrt/tensorrt_configure.bzl: xla/third_party/tensorrt/workspace.bzl: xla/third_party/xla_googletest_wrapper/include/gmock/gmock.h: +xla/third_party/xz/config/config_generic.h: xla/third_party/zlib.BUILD: diff --git a/third_party/xla/third_party/cuda_tile/patches/llvm_cl_989877215.patch b/third_party/xla/third_party/cuda_tile/patches/llvm_cl_989877215.patch new file mode 100644 index 00000000000000..821700fa0c8609 --- /dev/null +++ b/third_party/xla/third_party/cuda_tile/patches/llvm_cl_989877215.patch @@ -0,0 +1,11 @@ +diff --git a/include/cuda_tile/Dialect/CudaTile/IR/Dialect.td b/include/cuda_tile/Dialect/CudaTile/IR/Dialect.td +--- a/include/cuda_tile/Dialect/CudaTile/IR/Dialect.td ++++ b/include/cuda_tile/Dialect/CudaTile/IR/Dialect.td +@@ -42,6 +42,7 @@ + let useDefaultTypePrinterParser = 1; + let useDefaultAttributePrinterParser = 1; + ++ let useStrictPropertiesInAssemblyFormat = 0; + let extraClassDeclaration = [{ + template + void addExternalOperations() { diff --git a/third_party/xla/third_party/cuda_tile/patches/series b/third_party/xla/third_party/cuda_tile/patches/series index 9134fcfde4c7bd..46a712b02b5ef8 100644 --- a/third_party/xla/third_party/cuda_tile/patches/series +++ b/third_party/xla/third_party/cuda_tile/patches/series @@ -1,3 +1,4 @@ constructor.patch unused_variable.patch pow_op_alias.patch +llvm_cl_989877215.patch diff --git a/third_party/xla/third_party/extensions/llvm.bzl b/third_party/xla/third_party/extensions/llvm.bzl index 8d3262d8cb498f..b3ca0e0c658fdc 100644 --- a/third_party/xla/third_party/extensions/llvm.bzl +++ b/third_party/xla/third_party/extensions/llvm.bzl @@ -19,6 +19,7 @@ load("@bazel_tools//tools/build_defs/repo:http.bzl", "http_archive") load("@llvm-raw//utils/bazel:configure.bzl", "llvm_configure") load("@llvm-raw//utils/bazel:linux_uapi.bzl", "linux_uapi_setup") load("//third_party:repo.bzl", "tf_mirror_urls") +load("//third_party/xz:workspace.bzl", xz = "repo") _PYYAML_CONTENT = """\ load("@rules_python//python:defs.bzl", "py_library") @@ -64,6 +65,7 @@ _llvm_zstd_compat = repository_rule( def _llvm_extension_impl(mctx): # @unused _llvm_zlib_compat(name = "llvm_zlib") _llvm_zstd_compat(name = "llvm_zstd") + xz() linux_uapi_setup(name = "linux_uapi") http_archive( name = "pyyaml", diff --git a/third_party/xla/third_party/llvm/toolchains.patch b/third_party/xla/third_party/llvm/toolchains.patch index a8c15b85d9235a..1170c7610ab573 100644 --- a/third_party/xla/third_party/llvm/toolchains.patch +++ b/third_party/xla/third_party/llvm/toolchains.patch @@ -13,13 +13,13 @@ # limitations under the License. # ============================================================================== diff --git a/utils/bazel/llvm-project-overlay/llvm/BUILD.bazel b/utils/bazel/llvm-project-overlay/llvm/BUILD.bazel -index b683f86ca386..70819c622402 100644 --- a/utils/bazel/llvm-project-overlay/llvm/BUILD.bazel +++ b/utils/bazel/llvm-project-overlay/llvm/BUILD.bazel -@@ -59,6 +59,36 @@ exports_files([ +@@ -57,6 +57,36 @@ + # This one is needed for building and vendoring out lldb from off tree. "utils/lldbDataFormatters.py", ]) - ++ +config_setting( + name = "macos_arm64", + constraint_values = [ @@ -49,20 +49,9 @@ index b683f86ca386..70819c622402 100644 + "@platforms//cpu:x86_64", + ], +) -+ + config_setting( name = "darwin_arm64", - constraint_values = [ -@@ -533,7 +563,7 @@ cc_library( - # We unconditionally depend on the custom LLVM lzma wrapper. This only - # provides the configuration define unless lzma is enabled, in which - # case it also links the system liblzma. -- "//third-party:lzma", -+ # "//third-party:lzma", - "//libc:shared_math_headers_for_apfloat", - ], - ) - diff --git a/utils/bazel/llvm-project-overlay/llvm/config.bzl b/utils/bazel/llvm-project-overlay/llvm/config.bzl --- a/utils/bazel/llvm-project-overlay/llvm/config.bzl +++ b/utils/bazel/llvm-project-overlay/llvm/config.bzl diff --git a/third_party/xla/third_party/triton/common/series.bzl b/third_party/xla/third_party/triton/common/series.bzl index 3c729811758214..e5a4592a9bec1b 100644 --- a/third_party/xla/third_party/triton/common/series.bzl +++ b/third_party/xla/third_party/triton/common/series.bzl @@ -64,5 +64,6 @@ common_patch_list = [ "//third_party/triton:common/llvm_cl979526570.patch", "//third_party/triton:common/llvm_cl983115707.patch", "//third_party/triton:common/llvm_cl986804242.patch", + "//third_party/triton:common/useStrictPropertiesInAssemblyFormat.patch", # Add new patches just above this line ] diff --git a/third_party/xla/third_party/triton/common/useStrictPropertiesInAssemblyFormat.patch b/third_party/xla/third_party/triton/common/useStrictPropertiesInAssemblyFormat.patch new file mode 100644 index 00000000000000..83639d1da45ea9 --- /dev/null +++ b/third_party/xla/third_party/triton/common/useStrictPropertiesInAssemblyFormat.patch @@ -0,0 +1,114 @@ +Prevents issues due to https://github.com/llvm/llvm-project/commit/c495a9b796c46456e523a5eeb3690415ba99f5db. +This is a temporary fix to unblock the LLVM integration, the opt-out will eventually be removed. +See b/567092262 for more details. + +--- a/include/triton/Dialect/Gluon/IR/GluonDialect.td ++++ b/include/triton/Dialect/Gluon/IR/GluonDialect.td +@@ -16,6 +16,7 @@ + "mlir::gpu::GPUDialect", + ]; + let useDefaultAttributePrinterParser = 1; ++ let useStrictPropertiesInAssemblyFormat = 0; + } + + #endif + +--- a/include/triton/Dialect/Triton/IR/TritonDialect.td ++++ b/include/triton/Dialect/Triton/IR/TritonDialect.td +@@ -51,6 +51,7 @@ + + let hasConstantMaterializer = 1; + let useDefaultTypePrinterParser = 1; ++ let useStrictPropertiesInAssemblyFormat = 0; + } + + include "triton/Dialect/Triton/IR/TritonTypes.td" + +--- a/include/triton/Dialect/TritonGPU/IR/TritonGPUDialect.td ++++ b/include/triton/Dialect/TritonGPU/IR/TritonGPUDialect.td +@@ -35,6 +35,7 @@ + + let useDefaultTypePrinterParser = 1; + let useDefaultAttributePrinterParser = 1; ++ let useStrictPropertiesInAssemblyFormat = 0; + } + + #endif + +--- a/include/triton/Dialect/TritonInstrument/IR/TritonInstrumentDialect.td ++++ b/include/triton/Dialect/TritonInstrument/IR/TritonInstrumentDialect.td +@@ -6,6 +6,7 @@ + def TritonInstrument_Dialect : Dialect { + let name = "tti"; + let cppNamespace = "::mlir::triton::instrument"; ++ let useStrictPropertiesInAssemblyFormat = 0; + } + + #endif // TRITONINSTRUMENT_DIALECT + +--- a/include/triton/Dialect/TritonNvidiaGPU/IR/TritonNvidiaGPUDialect.td ++++ b/include/triton/Dialect/TritonNvidiaGPU/IR/TritonNvidiaGPUDialect.td +@@ -43,6 +43,7 @@ + + let useDefaultAttributePrinterParser = 1; + let useDefaultTypePrinterParser = 1; ++ let useStrictPropertiesInAssemblyFormat = 0; + } + + #endif + +--- a/third_party/amd/include/Dialect/TritonAMDGPU/IR/TritonAMDGPUDialect.td ++++ b/third_party/amd/include/Dialect/TritonAMDGPU/IR/TritonAMDGPUDialect.td +@@ -37,6 +37,7 @@ + let dependentDialects = ["triton::TritonDialect"]; + + let useDefaultAttributePrinterParser = 1; ++ let useStrictPropertiesInAssemblyFormat = 0; + } + + #endif + +--- a/third_party/nvidia/include/Dialect/NVGPU/IR/NVGPUDialect.td ++++ b/third_party/nvidia/include/Dialect/NVGPU/IR/NVGPUDialect.td +@@ -35,6 +35,7 @@ + let dependentDialects = [ + "mlir::LLVM::LLVMDialect" + ]; ++ let useStrictPropertiesInAssemblyFormat = 0; + } + + #endif + +--- a/third_party/nvidia/include/Dialect/NVWS/IR/NVWSDialect.td ++++ b/third_party/nvidia/include/Dialect/NVWS/IR/NVWSDialect.td +@@ -39,6 +39,7 @@ + + let useDefaultTypePrinterParser = 1; + let useDefaultAttributePrinterParser = 1; ++ let useStrictPropertiesInAssemblyFormat = 0; + } + + #endif + +--- a/third_party/proton/Dialect/include/Dialect/Proton/IR/ProtonDialect.td ++++ b/third_party/proton/Dialect/include/Dialect/Proton/IR/ProtonDialect.td +@@ -13,6 +13,7 @@ + }]; + + let dependentDialects = []; ++ let useStrictPropertiesInAssemblyFormat = 0; + } + + #endif + +--- a/third_party/proton/Dialect/include/Dialect/ProtonGPU/IR/ProtonGPUDialect.td ++++ b/third_party/proton/Dialect/include/Dialect/ProtonGPU/IR/ProtonGPUDialect.td +@@ -22,6 +22,7 @@ + + let useDefaultTypePrinterParser = 1; + let useDefaultAttributePrinterParser = 1; ++ let useStrictPropertiesInAssemblyFormat = 0; + } + + #endif // PROTONGPU_DIALECT + diff --git a/third_party/xla/third_party/xz/BUILD.bazel b/third_party/xla/third_party/xz/BUILD.bazel new file mode 100644 index 00000000000000..e69de29bb2d1d6 diff --git a/third_party/xla/third_party/xz/config/BUILD.bazel b/third_party/xla/third_party/xz/config/BUILD.bazel new file mode 100644 index 00000000000000..d0efcb1c3bd83c --- /dev/null +++ b/third_party/xla/third_party/xz/config/BUILD.bazel @@ -0,0 +1 @@ +exports_files(["config_generic.h"]) diff --git a/third_party/xla/third_party/xz/config/config_generic.h b/third_party/xla/third_party/xz/config/config_generic.h new file mode 100644 index 00000000000000..78807e812e8722 --- /dev/null +++ b/third_party/xla/third_party/xz/config/config_generic.h @@ -0,0 +1,557 @@ +/* Copyright 2020 The TensorStore Authors + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +/* config.h. Generated from config.h.in by configure. */ +/* config.h.in. Generated from configure.ac by autoheader. */ + +/* Define if building universal (internal helper macro) */ +/* #undef AC_APPLE_UNIVERSAL_BUILD */ + +/* How many MiB of RAM to assume if the real amount cannot be determined. */ +#define ASSUME_RAM 128 + +/* Define to 1 if translation of program messages to the user's native + language is requested. */ +/* #undef ENABLE_NLS */ + +#ifndef _WIN32 +/* Define to 1 if bswap_16 is available. */ +#define HAVE_BSWAP_16 1 + +/* Define to 1 if bswap_32 is available. */ +#define HAVE_BSWAP_32 1 + +/* Define to 1 if bswap_64 is available. */ +#define HAVE_BSWAP_64 1 +#endif + +/* Define to 1 if you have the header file. */ +#if defined(__APPLE__) || defined(_WIN32) +// byteswap.h does not exist in Apple SDKs. +// This #define must be entirely eliminated (not 0), as usages are checking +// for it being defined. +#undef HAVE_BYTESWAP_H +#else +#define HAVE_BYTESWAP_H 1 +#endif + +/* Define to 1 if Capsicum is available. */ +#undef HAVE_CAPSICUM + +/* Define to 1 if the system has the type `CC_SHA256_CTX'. */ +#if defined(__APPLE__) +#define HAVE_CC_SHA256_CTX 1 +#endif + +/* Define to 1 if you have the `CC_SHA256_Init' function. */ +#if defined(__APPLE__) +#define HAVE_CC_SHA256_INIT 1 +#endif + +/* Define to 1 if you have the MacOS X function CFLocaleCopyCurrent in the + CoreFoundation framework. */ +/* #undef HAVE_CFLOCALECOPYCURRENT */ + +/* Define to 1 if you have the MacOS X function CFPreferencesCopyAppValue in + the CoreFoundation framework. */ +/* #undef HAVE_CFPREFERENCESCOPYAPPVALUE */ + +/* Define to 1 if crc32 integrity check is enabled. */ +#define HAVE_CHECK_CRC32 1 + +/* Define to 1 if crc64 integrity check is enabled. */ +#define HAVE_CHECK_CRC64 1 + +/* Define to 1 if sha256 integrity check is enabled. */ +#define HAVE_CHECK_SHA256 1 + +/* Define to 1 if you have the `clock_gettime' function. */ +#if !defined(__APPLE__) +#define HAVE_CLOCK_GETTIME 1 +#endif + +/* Define to 1 if you have the header file. */ +#if defined(__APPLE__) +#define HAVE_COMMONCRYPTO_COMMONDIGEST_H 1 +#endif + +/* Define if the GNU dcgettext() function is already present or preinstalled. + */ +/* #undef HAVE_DCGETTEXT */ + +/* Define to 1 if you have the declaration of `CLOCK_MONOTONIC', and to 0 if + you don't. */ +#if !defined(__APPLE__) +#define HAVE_DECL_CLOCK_MONOTONIC 1 +#endif + +/* Define to 1 if you have the declaration of `program_invocation_name', and + to 0 if you don't. */ +#if defined(__ANDROID__) || defined(__APPLE__) || defined(_WIN32) +#define HAVE_DECL_PROGRAM_INVOCATION_NAME 0 +#else +#define HAVE_DECL_PROGRAM_INVOCATION_NAME 1 +#endif + +/* Define to 1 if any of HAVE_DECODER_foo have been defined. */ +#define HAVE_DECODERS 1 + +/* Define to 1 if arm decoder is enabled. */ +#define HAVE_DECODER_ARM 1 + +/* Define to 1 if armthumb decoder is enabled. */ +#define HAVE_DECODER_ARMTHUMB 1 + +/* Define to 1 if delta decoder is enabled. */ +#define HAVE_DECODER_DELTA 1 + +/* Define to 1 if ia64 decoder is enabled. */ +#define HAVE_DECODER_IA64 1 + +/* Define to 1 if lzma1 decoder is enabled. */ +#define HAVE_DECODER_LZMA1 1 + +/* Define to 1 if lzma2 decoder is enabled. */ +#define HAVE_DECODER_LZMA2 1 + +/* Define to 1 if powerpc decoder is enabled. */ +#define HAVE_DECODER_POWERPC 1 + +/* Define to 1 if sparc decoder is enabled. */ +#define HAVE_DECODER_SPARC 1 + +/* Define to 1 if x86 decoder is enabled. */ +#define HAVE_DECODER_X86 1 + +/* Define to 1 if you have the header file. */ +#define HAVE_DLFCN_H 1 + +/* Define to 1 if any of HAVE_ENCODER_foo have been defined. */ +#define HAVE_ENCODERS 1 + +/* Define to 1 if arm encoder is enabled. */ +#define HAVE_ENCODER_ARM 1 + +/* Define to 1 if armthumb encoder is enabled. */ +#define HAVE_ENCODER_ARMTHUMB 1 + +/* Define to 1 if delta encoder is enabled. */ +#define HAVE_ENCODER_DELTA 1 + +/* Define to 1 if ia64 encoder is enabled. */ +#define HAVE_ENCODER_IA64 1 + +/* Define to 1 if lzma1 encoder is enabled. */ +#define HAVE_ENCODER_LZMA1 1 + +/* Define to 1 if lzma2 encoder is enabled. */ +#define HAVE_ENCODER_LZMA2 1 + +/* Define to 1 if powerpc encoder is enabled. */ +#define HAVE_ENCODER_POWERPC 1 + +/* Define to 1 if sparc encoder is enabled. */ +#define HAVE_ENCODER_SPARC 1 + +/* Define to 1 if x86 encoder is enabled. */ +#define HAVE_ENCODER_X86 1 + +/* Define to 1 if you have the header file. */ +#define HAVE_FCNTL_H 1 + +/* Define to 1 if you have the `futimens' function. */ +#define HAVE_FUTIMENS 1 + +/* Define to 1 if you have the `futimes' function. */ +/* #undef HAVE_FUTIMES */ + +/* Define to 1 if you have the `futimesat' function. */ +/* #undef HAVE_FUTIMESAT */ + +/* Define to 1 if you have the header file. */ +#define HAVE_GETOPT_H 1 + +/* Define to 1 if you have the `getopt_long' function. */ +#define HAVE_GETOPT_LONG 1 + +/* Define if the GNU gettext() function is already present or preinstalled. */ +/* #undef HAVE_GETTEXT */ + +/* Define if you have the iconv() function and it works. */ +/* #undef HAVE_ICONV */ + +/* Define to 1 if you have the header file. */ +/* #undef HAVE_IMMINTRIN_H */ + +/* Define to 1 if you have the header file. */ +#define HAVE_INTTYPES_H 1 + +/* Define to 1 if you have the header file. */ +#define HAVE_LIMITS_H 1 + +/* Define to 1 if mbrtowc and mbstate_t are properly declared. */ +#define HAVE_MBRTOWC 1 + +/* Define to 1 if you have the header file. */ +#define HAVE_MEMORY_H 1 + +/* Define to 1 to enable bt2 match finder. */ +#define HAVE_MF_BT2 1 + +/* Define to 1 to enable bt3 match finder. */ +#define HAVE_MF_BT3 1 + +/* Define to 1 to enable bt4 match finder. */ +#define HAVE_MF_BT4 1 + +/* Define to 1 to enable hc3 match finder. */ +#define HAVE_MF_HC3 1 + +/* Define to 1 to enable hc4 match finder. */ +#define HAVE_MF_HC4 1 + +/* Define to 1 if getopt.h declares extern int optreset. */ +/* #undef HAVE_OPTRESET */ + +/* Define to 1 if you have the `posix_fadvise' function. */ +#if !defined(__APPLE__) +#define HAVE_POSIX_FADVISE 1 +#endif + +/* Define to 1 if you have the `pthread_condattr_setclock' function. */ +#if !defined(__APPLE__) && !defined(__ANDROID__) +#define HAVE_PTHREAD_CONDATTR_SETCLOCK 1 +#endif + +/* Have PTHREAD_PRIO_INHERIT. */ +#define HAVE_PTHREAD_PRIO_INHERIT 1 + +/* Define to 1 if you have the `SHA256Init' function. */ +/* #undef HAVE_SHA256INIT */ + +/* Define to 1 if the system has the type `SHA256_CTX'. */ +/* #undef HAVE_SHA256_CTX */ + +/* Define to 1 if you have the header file. */ +/* #undef HAVE_SHA256_H */ + +/* Define to 1 if you have the `SHA256_Init' function. */ +/* #undef HAVE_SHA256_INIT */ + +/* Define to 1 if the system has the type `SHA2_CTX'. */ +/* #undef HAVE_SHA2_CTX */ + +/* Define to 1 if you have the header file. */ +/* #undef HAVE_SHA2_H */ + +/* Define to 1 if optimizing for size. */ +#undef HAVE_SMALL + +/* Define to 1 if stdbool.h conforms to C99. */ +#define HAVE_STDBOOL_H 1 + +/* Define to 1 if you have the header file. */ +#define HAVE_STDINT_H 1 + +/* Define to 1 if you have the header file. */ +#define HAVE_STDLIB_H 1 + +/* Define to 1 if you have the header file. */ +/* #define HAVE_STRINGS_H 1 */ + +/* Define to 1 if you have the header file. */ +#define HAVE_STRING_H 1 + +/* Define to 1 if `st_atimensec' is a member of `struct stat'. */ +/* #undef HAVE_STRUCT_STAT_ST_ATIMENSEC */ + +/* Define to 1 if `st_atimespec.tv_nsec' is a member of `struct stat'. */ +/* #undef HAVE_STRUCT_STAT_ST_ATIMESPEC_TV_NSEC */ + +/* Define to 1 if `st_atim.st__tim.tv_nsec' is a member of `struct stat'. */ +/* #undef HAVE_STRUCT_STAT_ST_ATIM_ST__TIM_TV_NSEC */ + +/* Define to 1 if `st_atim.tv_nsec' is a member of `struct stat'. */ +#define HAVE_STRUCT_STAT_ST_ATIM_TV_NSEC 1 + +/* Define to 1 if `st_uatime' is a member of `struct stat'. */ +/* #undef HAVE_STRUCT_STAT_ST_UATIME */ + +/* Define to 1 if you have the header file. */ +/* #undef HAVE_SYS_BYTEORDER_H */ + +/* Define to 1 if you have the header file. */ +#undef HAVE_SYS_CAPSICUM_H + +/* Define to 1 if you have the header file. */ +/* #undef HAVE_SYS_ENDIAN_H */ + +#ifndef _WIN32 +/* Define to 1 if you have the header file. */ +#define HAVE_SYS_PARAM_H 1 + +/* Define to 1 if you have the header file. */ +#define HAVE_SYS_STAT_H 1 + +/* Define to 1 if you have the header file. */ +#define HAVE_SYS_TIME_H 1 + +/* Define to 1 if you have the header file. */ +#define HAVE_SYS_TYPES_H 1 +#endif + +/* Define to 1 if the system has the type `uintptr_t'. */ +#define HAVE_UINTPTR_T 1 + +/* Define to 1 if you have the header file. */ +#define HAVE_UNISTD_H 1 + +/* Define to 1 if you have the `utime' function. */ +/* #undef HAVE_UTIME */ + +/* Define to 1 if you have the `utimes' function. */ +/* #undef HAVE_UTIMES */ + +/* Define to 1 or 0, depending whether the compiler supports simple visibility + declarations. */ +#define HAVE_VISIBILITY 0 + +#ifndef _WIN32 +/* Define to 1 if you have the `wcwidth' function. */ +#define HAVE_WCWIDTH 1 +#endif + +/* Define to 1 if the system has the type `_Bool'. */ +#define HAVE__BOOL 1 + +/* Define to 1 if you have the `_futime' function. */ +#undef HAVE__FUTIME + +/* Define to 1 if _mm_movemask_epi8 is available. */ +/* #undef HAVE__MM_MOVEMASK_EPI8 */ + +/* Define to the sub-directory where libtool stores uninstalled libraries. */ +#define LT_OBJDIR ".libs/" + +#ifdef _WIN32 + +/* Define to 1 when using Windows Vista compatible threads. This uses features + that are not available on Windows XP. */ +#define MYTHREAD_VISTA 1 + +#else + +/* Define to 1 when using POSIX threads (pthreads). */ +#define MYTHREAD_POSIX 1 + +#endif + +/* Define to 1 when using Windows 95 (and thus XP) compatible threads. This + avoids use of features that were added in Windows Vista. */ +/* #undef MYTHREAD_WIN95 */ + +/* Define to 1 to disable debugging code. */ +#define NDEBUG 1 + +/* Name of package */ +#define PACKAGE "xz" + +/* Define to the address where bug reports for this package should be sent. */ +#define PACKAGE_BUGREPORT "lasse.collin@tukaani.org" + +/* Define to the full name of this package. */ +#define PACKAGE_NAME "XZ Utils" + +/* Define to the full name and version of this package. */ +#define PACKAGE_STRING "XZ Utils 5.2.4" + +/* Define to the one symbol short name of this package. */ +#define PACKAGE_TARNAME "xz" + +/* Define to the home page for this package. */ +#define PACKAGE_URL "http://tukaani.org/xz/" + +/* Define to the version of this package. */ +#define PACKAGE_VERSION "5.2.4" + +/* Define to necessary symbol if this constant uses a non-standard name on + your system. */ +/* #undef PTHREAD_CREATE_JOINABLE */ + +/* The size of `size_t', as computed by sizeof. */ +#define SIZEOF_SIZE_T 8 + +/* Define to 1 if you have the ANSI C header files. */ +#define STDC_HEADERS 1 + +/* Define to 1 if the number of available CPU cores can be detected with + cpuset(2). */ +/* #undef TUKLIB_CPUCORES_CPUSET */ + +/* Define to 1 if the number of available CPU cores can be detected with + pstat_getdynamic(). */ +/* #undef TUKLIB_CPUCORES_PSTAT_GETDYNAMIC */ + +/* Define to 1 if the number of available CPU cores can be detected with + sysconf(_SC_NPROCESSORS_ONLN) or sysconf(_SC_NPROC_ONLN). */ +#define TUKLIB_CPUCORES_SYSCONF 1 + +/* Define to 1 if the number of available CPU cores can be detected with + sysctl(). */ +/* #undef TUKLIB_CPUCORES_SYSCTL */ + +/* Define to 1 if the system supports fast unaligned access to 16-bit and + 32-bit integers. */ +#define TUKLIB_FAST_UNALIGNED_ACCESS 1 + +/* Define to 1 if the amount of physical memory can be detected with + _system_configuration.physmem. */ +/* #undef TUKLIB_PHYSMEM_AIX */ + +/* Define to 1 if the amount of physical memory can be detected with + getinvent_r(). */ +/* #undef TUKLIB_PHYSMEM_GETINVENT_R */ + +/* Define to 1 if the amount of physical memory can be detected with + getsysinfo(). */ +/* #undef TUKLIB_PHYSMEM_GETSYSINFO */ + +/* Define to 1 if the amount of physical memory can be detected with + pstat_getstatic(). */ +/* #undef TUKLIB_PHYSMEM_PSTAT_GETSTATIC */ + +/* Define to 1 if the amount of physical memory can be detected with + sysconf(_SC_PAGESIZE) and sysconf(_SC_PHYS_PAGES). */ +#if defined(__APPLE__) +// Apple SDKs (either iOS or OSX) do not define _SC_PHYS_PAGES. +// This #define must be entirely eliminated (not 0), as usages are checking +// for it being defined. +#undef TUKLIB_PHYSMEM_SYSCONF +#else +#define TUKLIB_PHYSMEM_SYSCONF 1 +#endif + +/* Define to 1 if the amount of physical memory can be detected with sysctl(). + */ +/* #undef TUKLIB_PHYSMEM_SYSCTL */ + +/* Define to 1 if the amount of physical memory can be detected with Linux + sysinfo(). */ +/* #undef TUKLIB_PHYSMEM_SYSINFO */ + +/* Enable extensions on AIX 3, Interix. */ +#ifndef _ALL_SOURCE +# define _ALL_SOURCE 1 +#endif +/* Enable GNU extensions on systems that have them. */ +#ifndef _GNU_SOURCE +# define _GNU_SOURCE 1 +#endif +/* Enable threading extensions on Solaris. */ +#ifndef _POSIX_PTHREAD_SEMANTICS +# define _POSIX_PTHREAD_SEMANTICS 1 +#endif +/* Enable extensions on HP NonStop. */ +#ifndef _TANDEM_SOURCE +# define _TANDEM_SOURCE 1 +#endif +/* Enable general extensions on Solaris. */ +#ifndef __EXTENSIONS__ +# define __EXTENSIONS__ 1 +#endif + + +/* Version number of package */ +#define VERSION "5.2.4" + +/* Define WORDS_BIGENDIAN to 1 if your processor stores words with the most + significant byte first (like Motorola and SPARC, unlike Intel). */ +#if defined AC_APPLE_UNIVERSAL_BUILD +# if defined __BIG_ENDIAN__ +# define WORDS_BIGENDIAN 1 +# endif +#else +# ifndef WORDS_BIGENDIAN +/* # undef WORDS_BIGENDIAN */ +# endif +#endif + +/* Enable large inode numbers on Mac OS X 10.5. */ +#ifndef _DARWIN_USE_64_BIT_INODE +# define _DARWIN_USE_64_BIT_INODE 1 +#endif + +/* Number of bits in a file offset, on hosts where this is settable. */ +/* #undef _FILE_OFFSET_BITS */ + +/* Define for large files, on AIX-style hosts. */ +/* #undef _LARGE_FILES */ + +/* Define to 1 if on MINIX. */ +/* #undef _MINIX */ + +/* Define to 2 if the system does not provide POSIX.1 features except with + this defined. */ +/* #undef _POSIX_1_SOURCE */ + +/* Define to 1 if you need to in order for `stat' and other things to work. */ +/* #undef _POSIX_SOURCE */ + +/* Define for Solaris 2.5.1 so the uint32_t typedef from , + , or is not used. If the typedef were allowed, the + #define below would cause a syntax error. */ +/* #undef _UINT32_T */ + +/* Define for Solaris 2.5.1 so the uint64_t typedef from , + , or is not used. If the typedef were allowed, the + #define below would cause a syntax error. */ +/* #undef _UINT64_T */ + +/* Define for Solaris 2.5.1 so the uint8_t typedef from , + , or is not used. If the typedef were allowed, the + #define below would cause a syntax error. */ +/* #undef _UINT8_T */ + +/* Define to rpl_ if the getopt replacement functions and variables should be + used. */ +/* #undef __GETOPT_PREFIX */ + +/* Define to the type of a signed integer type of width exactly 32 bits if + such a type exists and the standard includes do not define it. */ +/* #undef int32_t */ + +/* Define to the type of a signed integer type of width exactly 64 bits if + such a type exists and the standard includes do not define it. */ +/* #undef int64_t */ + +/* Define to the type of an unsigned integer type of width exactly 16 bits if + such a type exists and the standard includes do not define it. */ +/* #undef uint16_t */ + +/* Define to the type of an unsigned integer type of width exactly 32 bits if + such a type exists and the standard includes do not define it. */ +/* #undef uint32_t */ + +/* Define to the type of an unsigned integer type of width exactly 64 bits if + such a type exists and the standard includes do not define it. */ +/* #undef uint64_t */ + +/* Define to the type of an unsigned integer type of width exactly 8 bits if + such a type exists and the standard includes do not define it. */ +/* #undef uint8_t */ + +/* Define to the type of an unsigned integer type wide enough to hold a + pointer, if such a type exists, and if the system does not define it. */ +/* #undef uintptr_t */ diff --git a/third_party/xla/third_party/xz/patches/BUILD.bazel b/third_party/xla/third_party/xz/patches/BUILD.bazel new file mode 100644 index 00000000000000..1d4b626f2c2ffb --- /dev/null +++ b/third_party/xla/third_party/xz/patches/BUILD.bazel @@ -0,0 +1 @@ +exports_files(["remove_have_config.diff"]) diff --git a/third_party/xla/third_party/xz/patches/remove_have_config.diff b/third_party/xla/third_party/xz/patches/remove_have_config.diff new file mode 100644 index 00000000000000..8d6612736c8df6 --- /dev/null +++ b/third_party/xla/third_party/xz/patches/remove_have_config.diff @@ -0,0 +1,14 @@ +diff -Nur a/src/common/sysdefs.h b/src/common/sysdefs.h +--- a/src/common/sysdefs.h 2025-04-03 11:34:58.000000000 +0000 ++++ b/src/common/sysdefs.h 2025-06-20 23:44:13.985352903 +0000 +@@ -19,9 +19,7 @@ + // Includes // + ////////////// + +-#ifdef HAVE_CONFIG_H +-# include +-#endif ++#include "lzma_config/config.h" + + // Choose if MinGW-w64's stdio replacement functions should be used. + // The default has varied slightly in the past so it's clearest to always diff --git a/third_party/xla/third_party/xz/workspace.bzl b/third_party/xla/third_party/xz/workspace.bzl new file mode 100644 index 00000000000000..22e73727d36198 --- /dev/null +++ b/third_party/xla/third_party/xz/workspace.bzl @@ -0,0 +1,37 @@ +# Copyright 2020 The TensorStore Authors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# buildifier: disable=module-docstring + +load("@bazel_tools//tools/build_defs/repo:utils.bzl", "maybe") +load( + "//third_party:repo.bzl", + "tf_http_archive", + "tf_mirror_urls", +) + +def repo(): + maybe( + tf_http_archive, + # The repo name is the same as the one used in llvm_project to avoid the conflict with `xz` + # repository in rules_ml_toolchain. + name = "llvm_xz", + strip_prefix = "xz-5.8.3", + urls = tf_mirror_urls("https://github.com/tukaani-project/xz/releases/download/v5.8.3/xz-5.8.3.tar.gz"), + sha256 = "3d3a1b973af218114f4f889bbaa2f4c037deaae0c8e815eec381c3d546b974a0", + patch_file = [ + "//third_party/xz/patches:remove_have_config.diff", + ], + build_file = "//third_party/xz:xz.BUILD.bazel", + ) diff --git a/third_party/xla/third_party/xz/xz.BUILD.bazel b/third_party/xla/third_party/xz/xz.BUILD.bazel new file mode 100644 index 00000000000000..4aacd302a267c7 --- /dev/null +++ b/third_party/xla/third_party/xz/xz.BUILD.bazel @@ -0,0 +1,172 @@ +load("@bazel_skylib//rules:copy_file.bzl", "copy_file") +load("@rules_cc//cc:cc_library.bzl", "cc_library") + +package(default_visibility = ["//visibility:private"]) + +licenses(["unencumbered"]) + +exports_files(["COPYING"]) + +cc_library( + name = "lzma", + srcs = glob( + include = [ + "src/**/*.h", + "src/liblzma/**/*.c", + ], + exclude = [ + "src/**/*_tablegen.c", + "src/liblzma/check/crc_clmul_consts_gen.c", + "src/liblzma/check/sha256.c", + "src/liblzma/check/crc*_small.c", + "src/liblzma/common/stream_decoder_mt.c", + "src/liblzma/common/stream_encoder_mt.c", + ], + ) + select({ + "@platforms//os:osx": [], + "@platforms//os:ios": [], + "//conditions:default": [ + "src/liblzma/check/sha256.c", + ], + }) + select({ + "@platforms//os:none": [], + "//conditions:default": [ + "src/common/tuklib_cpucores.c", + "src/common/tuklib_physmem.c", + "src/liblzma/common/stream_decoder_mt.c", + "src/liblzma/common/stream_encoder_mt.c", + ], + }), + hdrs = glob(["src/liblzma/api/**/*.h"]), + copts = select({ + "@platforms//os:windows": [], + "//conditions:default": ["-std=c99"], + }), + defines = select({ + "@platforms//os:windows": ["LZMA_API_STATIC"], + "//conditions:default": [], + }), + implementation_deps = [ + ":config", + ":src_common", + ":src_liblzma", + ":src_liblzma_check", + ":src_liblzma_common", + ":src_liblzma_delta", + ":src_liblzma_lz", + ":src_liblzma_lzma", + ":src_liblzma_rangecoder", + ":src_liblzma_simple", + ], + linkopts = select({ + "@platforms//os:android": [], + "@platforms//os:windows": [], + "@platforms//os:none": [], + "//conditions:default": ["-lpthread"], + }), + linkstatic = select({ + "@platforms//os:windows": True, + "//conditions:default": False, + }), + strip_include_prefix = "src/liblzma/api", + visibility = ["//visibility:public"], +) + +# Generate config.h from config_generic.h +copy_file( + name = "config_copy", + src = "@xla//third_party/xz/config:config_generic.h", + out = "lzma_config/config.h", +) + +cc_library( + name = "config", + hdrs = ["lzma_config/config.h"], +) + +cc_library( + name = "src_common", + srcs = select({ + "@platforms//os:none": [], + "//conditions:default": [ + "src/common/tuklib_exit.c", + "src/common/tuklib_progname.c", + ], + }), + hdrs = glob(["src/common/*.h"]), + defines = select({ + "@platforms//os:windows": ["LZMA_API_STATIC"], + "//conditions:default": [], + }), + local_defines = select({ + "@platforms//os:windows": ["TUKLIB_GETTEXT=0"], + "//conditions:default": [], + }), + strip_include_prefix = "src/common", + deps = [ + ":config", + ":src_liblzma", + ":src_liblzma_api", + ":src_liblzma_common", + ], +) + +cc_library( + name = "src_liblzma", + hdrs = glob([ + "src/liblzma/common/*.h", + "src/liblzma/lzma/*.h", + ]), + strip_include_prefix = "src/liblzma", +) + +cc_library( + name = "src_liblzma_api", + hdrs = glob( + include = ["src/liblzma/api/**/*.h"], + exclude = ["src/liblzma/api/lzma.h"], + ), + strip_include_prefix = "src/liblzma/api", +) + +cc_library( + name = "src_liblzma_check", + hdrs = glob(["src/liblzma/check/*.h"]), + strip_include_prefix = "src/liblzma/check", +) + +cc_library( + name = "src_liblzma_common", + hdrs = glob(["src/liblzma/common/*.h"]), + strip_include_prefix = "src/liblzma/common", +) + +cc_library( + name = "src_liblzma_delta", + hdrs = glob(["src/liblzma/delta/*.h"]), + strip_include_prefix = "src/liblzma/delta", +) + +cc_library( + name = "src_liblzma_lz", + hdrs = glob(["src/liblzma/lz/*.h"]), + strip_include_prefix = "src/liblzma/lz", +) + +cc_library( + name = "src_liblzma_lzma", + hdrs = glob(["src/liblzma/lzma/*.h"]), + strip_include_prefix = "src/liblzma/lzma", +) + +cc_library( + name = "src_liblzma_rangecoder", + hdrs = glob(["src/liblzma/rangecoder/*.h"]), + strip_include_prefix = "src/liblzma/rangecoder", +) + +cc_library( + name = "src_liblzma_simple", + hdrs = glob(["src/liblzma/simple/*.h"]), + strip_include_prefix = "src/liblzma/simple", +) diff --git a/third_party/xla/workspace2.bzl b/third_party/xla/workspace2.bzl index 8a99921e4d611f..00a196dcefbfa2 100644 --- a/third_party/xla/workspace2.bzl +++ b/third_party/xla/workspace2.bzl @@ -90,6 +90,7 @@ load("//third_party/transformer_engine:workspace.bzl", transformer_engine = "rep load("//third_party/triton:workspace.bzl", triton = "repo") load("//third_party/uv:workspace.bzl", uv = "repo") load("//third_party/xnnpack:workspace.bzl", xnnpack = "repo") +load("//third_party/xz:workspace.bzl", xz = "repo") load("//tools/def_file_filter:def_file_filter_configure.bzl", "def_file_filter_configure") load("//tools/toolchains:cpus/aarch64/aarch64_compiler_configure.bzl", "aarch64_compiler_configure") load("//tools/toolchains:cpus/arm/arm_compiler_configure.bzl", "arm_compiler_configure") @@ -157,6 +158,7 @@ def _initialize_third_party(): triton() uv() xnnpack() + xz() cuda_tile() cutlass() cutlass_cutedsl_runtime() diff --git a/third_party/xla/xla/AGENTS.md b/third_party/xla/xla/AGENTS.md index 4b59d200fc7bdf..58e5d6c4077e71 100644 --- a/third_party/xla/xla/AGENTS.md +++ b/third_party/xla/xla/AGENTS.md @@ -1,8 +1,8 @@ # AI Assistant Guidelines for OpenXLA Development This document provides guidelines for AI code assistants when generating, -suggesting, or modifying code within the -`third_party/tensorflow/compiler/xla` (OpenXLA) codebase. +suggesting, or modifying code within the OpenXLA codebase (everything under +`xla/` in this repository). ## General Context @@ -10,12 +10,27 @@ suggesting, or modifying code within the * **Code Quality:** Adhere to [Google C++ Style Guide](https://google.github.io/styleguide/cppguide.html) and OpenXLA-specific conventions. * **Portability:** This code is open-sourced and runs on a number of host platforms (e.g. Linux, Windows, etc.) +## Repository Layout and Build + +* Source lives under `xla/`. Include paths are relative to the repository + root, e.g. `#include "xla/hlo/ir/hlo_module.h"`. +* Bazel labels use the same prefix, e.g. `//xla/hlo/ir:hlo_module`. +* Build and test with Bazel, e.g. `bazel test //xla/service:hlo_verifier_test`. +* Common external dependencies and their labels: + `@com_google_absl//absl/...`, `@llvm-project//mlir/...`, + `@llvm-project//llvm/...`, `@com_google_googletest//:gtest_main`, + `@com_google_protobuf//:protobuf`. TSL (the platform layer) lives in-tree + under `//xla/tsl/...`. +* Test macros: `xla_cc_test` from `//xla:xla.default.bzl` (hardware- + independent unit tests) and `xla_test` from `//xla/tests:build_defs.bzl` + (tests that run on configured hardware backends). + ## Coding Guidelines for AI Assistance 1. **Error Handling (`absl::Status`, `absl::StatusOr`)**: * **Always** use `absl::Status` or `absl::StatusOr` for functions that can encounter recoverable errors. * **Macros**: - * Use header `tsl/platform/status_macros.h`. + * Use header `absl/status/status_macros.h`. * Use `ABSL_RETURN_IF_ERROR` for error propagation. * Use `ABSL_ASSIGN_OR_RETURN` for `StatusOr` assignments. * **Safely access `StatusOr` values**: Check `.ok()` before accessing. @@ -35,13 +50,14 @@ suggesting, or modifying code within the return absl::OkStatus(); } ``` + 3. **Decision making**: * **Avoid `bool`** for returning a decision to do or not to do something. * **Prefer `Decision`**: - * Located in `third_party/tensorflow/compiler/xla/service/decision.h` + * Located in `xla/service/decision.h` * Example: ```cpp - #include "third_party/tensorflow/compiler/xla/service/decision.h" + #include "xla/service/decision.h" using AutotunerDecision = Decision; @@ -50,7 +66,7 @@ suggesting, or modifying code within the return AutotunerDecision::Forbid("Cublas autotuning was explicitly disabled"); } - voud AutotuneCublas(const AutotunerDecision& decision) { + void AutotuneCublas(const AutotunerDecision& decision) { if (decision.IsForbidden()) { return; } @@ -63,24 +79,28 @@ suggesting, or modifying code within the 5. **Testing**: * Write unit tests using `EXPECT_EQ`, `EXPECT_TRUE`, etc. - * Use macros in `tsl/platform/status_matchers.h` instead of their TF_* - counterparts. For example: - * Use `ASSERT_OK_AND_ASSIGN`, `ASSERT_OK`, and `EXPECT_OK`. + * Use the status testing macros `ASSERT_OK_AND_ASSIGN`, `ASSERT_OK`, and + `EXPECT_OK` from `` and the status matchers + (`absl_testing::IsOk`, `absl_testing::IsOkAndHolds`, + `absl_testing::StatusIs`) from `absl/status/status_matchers.h` instead + of their legacy `TF_*` or `tsl::testing::*` counterparts: * DO NOT USE `TF_ASSERT_OK_AND_ASSIGN`, `TF_ASSERT_OK`, and `TF_EXPECT_OK`. - * When you refactor code that uses the TF_* macros., replace them, but - do not touch unrelated code. - * Put tests into an anonymous namespace + * When you refactor code that uses the `TF_*` macros, replace them, + but do not touch unrelated code. + * Put tests into an anonymous namespace. * Use `HloPjRtInterpreterReferenceMixin` or `HloHardwareIndependentTestBase` for compiler pass tests locally where possible. + * When running GPU tests, pass `--config=cuda` to the test command. * Ensure tests are deterministic and do not flake. 6. **BUILD targets**: - * When defining BUILD targets prefer these XLA specific rules: + * When defining BUILD targets prefer these XLA specific rules: * Instead of `proto_library` use `tf_proto_library`. There is no need to define language specific targets with `tf_proto_library`. - * Instead of `cc_test` use `xla_cc_test`. + * Instead of `cc_test` use `xla_cc_test` (or `xla_test` when the test + requires a real hardware backend). 7. **Explicit Typing**: * **Avoid `auto`** in public headers or complex logic chains. @@ -92,8 +112,39 @@ suggesting, or modifying code within the * *Example*: Do not rely on layout information before Layout Assignment. 9. **Namespaces**: - * Prefer xla::gpu over nested namespaces. + * Prefer `xla::gpu` over nested namespaces. 10. **MLIR Operation Creation**: * **Always** use the static `OpTy::create(rewriter, ...)` method when creating MLIR operations. * **Avoid** using `rewriter.create(...)`. This syntax is deprecated. + +11. **License Headers**: + * Any new source or build files created in this codebase must include the + OpenXLA Apache 2.0 copyright header at the top of the file. + * Replace `` in the template below with the current 4-digit calendar + year (e.g., `2026`—do not leave `` literally or copy an outdated + year): + ```cpp + /* Copyright The OpenXLA Authors. + + Licensed under the Apache License, Version 2.0 (the "License"); + you may not use this file except in compliance with the License. + You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + + Unless required by applicable law or agreed to in writing, software + distributed under the License is distributed on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + See the License for the specific language governing permissions and + limitations under the License. + ==============================================================================*/ + ``` + (Use `#` comment prefixes for Python and `BUILD` files.) + +## Editing This File + +This file is shared between the open-source repository and Google-internal +development. Keep it in open-source terms only: repository-relative paths +(`xla/...`), OSS macro names (`ABSL_RETURN_IF_ERROR`, `ABSL_ASSIGN_OR_RETURN`), +Bazel labels (`//xla/...`, `@repo//...`), and `bazel` commands. diff --git a/third_party/xla/xla/BUILD b/third_party/xla/xla/BUILD index b720b59902e79c..3da1bb89bd9bc2 100644 --- a/third_party/xla/xla/BUILD +++ b/third_party/xla/xla/BUILD @@ -1353,7 +1353,6 @@ cc_library( "//xla/backends/autotuner:backends_proto_cc", "//xla/hlo/pass:hlo_pass_filter", "//xla/service:collective_utils", - "//xla/stream_executor/cuda:nvjitlink_support", "//xla/stream_executor/cuda:ptx_compiler_support", "//xla/tsl/platform:logging", "//xla/tsl/util:command_line_flags", @@ -1426,11 +1425,23 @@ xla_cc_test( cc_library( name = "union_find", hdrs = ["union_find.h"], + deps = ["@com_google_absl//absl/log:check"], +) + +xla_cc_test( + name = "union_find_test", + srcs = ["union_find_test.cc"], + deps = [ + ":union_find", + "//xla/tsl/platform:test", + "//xla/tsl/platform:test_benchmark", + "//xla/tsl/platform:test_main", + ], ) cc_library( - name = "inlined_bit_set", - hdrs = ["inlined_bit_set.h"], + name = "bit_set", + hdrs = ["bit_set.h"], deps = [ "@com_google_absl//absl/container:inlined_vector", "@com_google_absl//absl/log:check", @@ -1438,11 +1449,12 @@ cc_library( ) xla_cc_test( - name = "inlined_bit_set_test", - srcs = ["inlined_bit_set_test.cc"], + name = "bit_set_test", + srcs = ["bit_set_test.cc"], deps = [ - ":inlined_bit_set", + ":bit_set", "//xla/tsl/platform:test", + "//xla/tsl/platform:test_benchmark", "//xla/tsl/platform:test_main", "@com_google_absl//absl/container:flat_hash_set", ], diff --git a/third_party/xla/xla/backends/cpu/codegen/tiled/tests/elementwise/strided_overhanging_tiles.hlo b/third_party/xla/xla/backends/cpu/codegen/tiled/tests/elementwise/strided_overhanging_tiles.hlo new file mode 100644 index 00000000000000..2cdfa33a4c0ab5 --- /dev/null +++ b/third_party/xla/xla/backends/cpu/codegen/tiled/tests/elementwise/strided_overhanging_tiles.hlo @@ -0,0 +1,41 @@ +// RUN: test_correctness %s + +// RUN: fusion_to_xtile %s | fusion_compiler_opt \ +// RUN: --xtile-cpu-new-xtile-to-vector --xla-simplify-arith \ +// RUN: | FileCheck %s --check-prefix=NEW-VECTOR +// RUN: fusion_to_xtile %s | fusion_compiler_opt \ +// RUN: --xtile-cpu-new-xtile-to-vector --xtile-cpu-new-vector-to-llvm \ +// RUN: | FileCheck %s --check-prefix=NEW-VECTOR-LLVM +// RUN: test_correctness %s --xla_cpu_use_new_xtile_lowering=true + +fusion { + p0 = f32[224] parameter(0) + ROOT slice = f32[112] slice(p0), slice={[0:224:2]} +} + +ENTRY main { + p0 = f32[224] parameter(0) + ROOT fusion = f32[112] fusion(p0), kind=kLoop, calls=fusion, + backend_config={ + "fusion_backend_config":{ + "block_level_fusion_config":{ + "output_tiles":[{"sizes":["16"]}], + "num_ctas":"3" + } + } + } +} + +// NEW-VECTOR: #[[$UB_MAP:.*]] = #xla.indexing_map<"(pid, tid) -> (min(max(pid * -48 - tid * 16 + 112, 0), 16)), domain: pid in [0, 2], tid in [0, 2]"> +// NEW-VECTOR-LABEL: func.func @wrapped_fusion_impl( +// NEW-VECTOR-SAME: %[[PID:[^:]+]]: index {xla.range = [0 : index, 2 : index]} +// NEW-VECTOR: scf.for %[[TID:.*]] = +// NEW-VECTOR: %[[ALLOCA:.*]] = memref.alloca() : memref<16xf32> +// NEW-VECTOR: %[[UB:.*]] = xla.apply_indexing #[[$UB_MAP]](%[[PID]], %[[TID]]) +// NEW-VECTOR: scf.for %[[IV:.*]] = %{{.*}} to %[[UB]] +// NEW-VECTOR: memref.load +// NEW-VECTOR: memref.store %{{.*}}, %[[ALLOCA]][%[[IV]]] + +// NEW-VECTOR-LLVM-NOT: xla.apply_indexing +// NEW-VECTOR-LLVM-LABEL: llvm.func internal @wrapped_fusion_impl +// NEW-VECTOR-LLVM: llvm.intr.smax diff --git a/third_party/xla/xla/backends/cpu/tests/ynn_fusion_test.cc b/third_party/xla/xla/backends/cpu/tests/ynn_fusion_test.cc index 6ed7c4a6260beb..8bf18a9fb03912 100644 --- a/third_party/xla/xla/backends/cpu/tests/ynn_fusion_test.cc +++ b/third_party/xla/xla/backends/cpu/tests/ynn_fusion_test.cc @@ -113,6 +113,43 @@ TEST_P(YnnFusionTest, Pad) { RunTest(kModuleStr); } +TEST_P(YnnFusionTest, Copy) { + constexpr absl::string_view kModuleStr = R"( + HloModule copy + + ynn_fusion { + %input = $dtype[8, 10] parameter(0) + ROOT %copy = $dtype[8, 10] copy(%input) + } + + ENTRY entry { + %p0 = $dtype[8, 10] parameter(0) + ROOT %fusion = $dtype[8, 10] fusion(%p0), kind=kCustom, calls=ynn_fusion, + backend_config={"fusion_config": {kind: "__ynn_fusion"}} + })"; + + RunTest(kModuleStr); +} + +TEST_P(YnnFusionTest, CopyDegenerateLayout) { + constexpr absl::string_view kModuleStr = R"( + HloModule copy_degenerate_layout + + ynn_fusion { + %input = $dtype[8, 10, 1]{1,0,2} parameter(0) + ROOT %copy = $dtype[8, 10, 1]{2,1,0} copy(%input) + } + + ENTRY entry { + %p0 = $dtype[8, 10, 1]{1,0,2} parameter(0) + ROOT %fusion = $dtype[8, 10, 1]{2,1,0} fusion(%p0), kind=kCustom, + calls=ynn_fusion, + backend_config={"fusion_config": {kind: "__ynn_fusion"}} + })"; + + RunTest(kModuleStr); +} + TEST_P(YnnFusionTest, Transpose) { constexpr absl::string_view kModuleStr = R"( HloModule transpose diff --git a/third_party/xla/xla/backends/cpu/transforms/ynn_matcher.h b/third_party/xla/xla/backends/cpu/transforms/ynn_matcher.h index e5d18a75601cc1..f43b25f2e60df9 100644 --- a/third_party/xla/xla/backends/cpu/transforms/ynn_matcher.h +++ b/third_party/xla/xla/backends/cpu/transforms/ynn_matcher.h @@ -50,7 +50,7 @@ class YnnMatcher : public LibraryMatcher { HloOpcode::kConvolution, HloOpcode::kReshape, HloOpcode::kBitcast, HloOpcode::kBroadcast, HloOpcode::kTranspose, HloOpcode::kPad, - HloOpcode::kIota}; + HloOpcode::kCopy, HloOpcode::kIota}; for (const auto& [op, _] : GetYnnUnaryOpMap()) { supported_ops.insert(op); } @@ -85,6 +85,9 @@ class YnnMatcher : public LibraryMatcher { if (instr->opcode() == HloOpcode::kPad) { return IsPadOpSupportedByYnn(instr); } + if (instr->opcode() == HloOpcode::kCopy) { + return IsCopyOpSupportedByYnn(instr); + } if (!IsInstructionPreferredByYnn(instr)) { // TODO: It might make sense sometimes that even though an instruction is // not preferred by YNNPACK, that we should still fuse it, if it lies diff --git a/third_party/xla/xla/backends/cpu/transforms/ynn_matcher_test.cc b/third_party/xla/xla/backends/cpu/transforms/ynn_matcher_test.cc index 79afc3055657ea..6d3607737c29de 100644 --- a/third_party/xla/xla/backends/cpu/transforms/ynn_matcher_test.cc +++ b/third_party/xla/xla/backends/cpu/transforms/ynn_matcher_test.cc @@ -201,6 +201,33 @@ TEST_F(YnnReduceTest, ConvertReduce) { )"); } +TEST_F(YnnReduceTest, CopyDegenerateLayout) { + const char* hlo_text = R"( + HloModule copy_degenerate_layout + + add { + lhs = f32[] parameter(0) + rhs = f32[] parameter(1) + ROOT add = f32[] add(lhs, rhs) + } + + ENTRY main { + input = f32[128,64,1]{1,0,2} parameter(0) + copied = f32[128,64,1]{2,1,0} copy(input) + init = f32[] constant(0) + ROOT result = f32[128] reduce(copied, init), dimensions={1,2}, to_apply=add + } + )"; + + MatchOptimizedHlo(hlo_text, R"( + CHECK: copy + CHECK: reduce + CHECK: ENTRY + CHECK: kind=kCustom + CHECK: "kind":"__ynn_fusion" + )"); +} + struct DotTestConfig { absl::string_view lhs_dtype; absl::string_view rhs_dtype; diff --git a/third_party/xla/xla/backends/cpu/ynn_emitter.cc b/third_party/xla/xla/backends/cpu/ynn_emitter.cc index 395e8ff24177d6..25ba2f528528cc 100644 --- a/third_party/xla/xla/backends/cpu/ynn_emitter.cc +++ b/third_party/xla/xla/backends/cpu/ynn_emitter.cc @@ -217,6 +217,21 @@ absl::StatusOr DefineReshapeOp(ynn_subgraph_t subgraph, return DefineBitcastOp(subgraph, tensor_ids, instr); } +absl::StatusOr DefineCopyOp(ynn_subgraph_t subgraph, + TensorIdMap& tensor_ids, + const HloInstruction* instr) { + VLOG(3) << absl::StreamFormat("Define tensor value for copy op: %s", + instr->ToString()); + CHECK_EQ(instr->opcode(), HloOpcode::kCopy); + const HloInstruction* input = instr->operand(0); + CHECK_EQ(input->shape().element_type(), instr->shape().element_type()); + ABSL_ASSIGN_OR_RETURN(auto in, FindTensorValue(tensor_ids, input)); + ABSL_ASSIGN_OR_RETURN(auto out, DefineTensorValue(subgraph, instr)); + + YNN_RETURN_IF_ERROR(ynn_define_copy(subgraph, in, &out, /*flags=*/0)); + return out; +} + absl::StatusOr DefineTransposeOp(ynn_subgraph_t subgraph, TensorIdMap& tensor_ids, const HloInstruction* instr) { @@ -845,7 +860,7 @@ absl::StatusOr EmitYnnSubgraph( continue; } - if (instr->IsElementwise()) { + if (instr->IsElementwise() && instr->opcode() != HloOpcode::kCopy) { if (!IsElementwiseOpSupportedByYnn(instr)) { return InvalidArgument( "Unsupported elementwise instruction in YNN fusion: %s", @@ -894,6 +909,16 @@ absl::StatusOr EmitYnnSubgraph( DefineReshapeOp(subgraph.get(), tensor_ids, instr)); } break; + case HloOpcode::kCopy: { + if (!IsCopyOpSupportedByYnn(instr)) { + return InvalidArgument( + "Unsupported copy instruction in YNN fusion: %s", + instr->ToString()); + } + ABSL_ASSIGN_OR_RETURN(tensor_ids[instr], + DefineCopyOp(subgraph.get(), tensor_ids, instr)); + } break; + case HloOpcode::kTranspose: { if (!IsTransposeOpSupportedByYnn(instr)) { return InvalidArgument( diff --git a/third_party/xla/xla/backends/cpu/ynn_support.cc b/third_party/xla/xla/backends/cpu/ynn_support.cc index 38bbfe050ab1ff..2ae056f728b000 100644 --- a/third_party/xla/xla/backends/cpu/ynn_support.cc +++ b/third_party/xla/xla/backends/cpu/ynn_support.cc @@ -17,7 +17,9 @@ limitations under the License. #include #include +#include #include +#include #include "ynnpack/include/ynnpack.h" #include "absl/algorithm/container.h" @@ -130,7 +132,16 @@ bool IsLayoutSupportedByYnn(const Shape& shape) { // TODO(b/460602165): We should eliminate this limitation. return false; } - return !shape.has_layout() || LayoutUtil::HasDescendingLayout(shape.layout()); + if (!shape.has_layout()) { + return true; + } + std::vector minor_to_major; + for (int64_t dim : shape.layout().minor_to_major()) { + if (shape.dimensions(dim) != 1) { + minor_to_major.push_back(dim); + } + } + return absl::c_is_sorted(minor_to_major, std::greater()); } namespace { @@ -165,6 +176,28 @@ bool IsBitcastOpSupportedByYnn(const HloInstruction* hlo) { return hlo->shape().element_type() == input->shape().element_type(); } +bool IsCopyOpSupportedByYnn(const HloInstruction* hlo) { + CHECK_EQ(hlo->opcode(), HloOpcode::kCopy); + if (!CheckOperandCount(hlo, 1)) { + return false; + } + if (!YnnType(hlo->shape().element_type()).ok()) { + return false; + } + const HloInstruction* input = hlo->operand(0); + if (hlo->shape().element_type() != input->shape().element_type()) { + return false; + } + if (hlo->shape().dimensions() != input->shape().dimensions()) { + return false; + } + if (!IsLayoutSupportedByYnn(hlo->shape()) || + !IsLayoutSupportedByYnn(input->shape())) { + return false; + } + return true; +} + bool IsReshapeOpSupportedByYnn(const HloInstruction* hlo) { CHECK_EQ(hlo->opcode(), HloOpcode::kReshape); if (!CheckOperandCount(hlo, 1)) { diff --git a/third_party/xla/xla/backends/cpu/ynn_support.h b/third_party/xla/xla/backends/cpu/ynn_support.h index de5ad5f2696eda..2eb3c3417d72ee 100644 --- a/third_party/xla/xla/backends/cpu/ynn_support.h +++ b/third_party/xla/xla/backends/cpu/ynn_support.h @@ -56,6 +56,9 @@ bool IsLayoutSupportedByYnn(const Shape& shape); // Returns true if the bitcast op is supported by YNNPACK. bool IsBitcastOpSupportedByYnn(const HloInstruction* hlo); +// Returns true if the copy op is supported by YNNPACK. +bool IsCopyOpSupportedByYnn(const HloInstruction* hlo); + // Returns true if the reshape op is supported by YNNPACK. bool IsReshapeOpSupportedByYnn(const HloInstruction* hlo); diff --git a/third_party/xla/xla/backends/gpu/autotuner/block_level_emitter.cc b/third_party/xla/xla/backends/gpu/autotuner/block_level_emitter.cc index 8329f81c5bceb4..88ec1af6050f4b 100644 --- a/third_party/xla/xla/backends/gpu/autotuner/block_level_emitter.cc +++ b/third_party/xla/xla/backends/gpu/autotuner/block_level_emitter.cc @@ -123,9 +123,11 @@ BlockLevelEmitterBackend::GetSupportedConfigs(const HloInstruction& instr) { ->config() .debug_options() .xla_gpu_fusion_autotune_top_k_configs(); - ABSL_ASSIGN_OR_RETURN(TopKTiledRunTimeDataOrError tiled_runtime_data, - indexing_performance_model_.TryFindTopKBestTilingsForFusion( - *fusion_adaptor, num_configs)); + ABSL_ASSIGN_OR_RETURN( + TopKTiledRunTimeDataOrError tiled_runtime_data, + indexing_performance_model_ + .TryFindTopKBestTilingsForFusionAsync(*fusion_adaptor, num_configs) + .Await()); if (std::holds_alternative(tiled_runtime_data)) { return std::vector>(); @@ -152,9 +154,10 @@ BlockLevelEmitterBackend::GetCostModelConfig(const HloInstruction& instr) { auto fusion_adaptor = HloFusionAdaptor::ForInstruction(Cast(&instr)); - ABSL_ASSIGN_OR_RETURN( - TiledRunTimeDataOrError tiled_runtime_data_or_error, - indexing_performance_model_.TryFindBestTilingForFusion(*fusion_adaptor)); + ABSL_ASSIGN_OR_RETURN(TiledRunTimeDataOrError tiled_runtime_data_or_error, + indexing_performance_model_ + .TryFindBestTilingForFusionAsync(*fusion_adaptor) + .Await()); if (const auto* fusion_decision = std::get_if(&tiled_runtime_data_or_error)) { diff --git a/third_party/xla/xla/backends/gpu/autotuner/fission_backend.cc b/third_party/xla/xla/backends/gpu/autotuner/fission_backend.cc index 0813a96f1e2904..d814f140920cd7 100644 --- a/third_party/xla/xla/backends/gpu/autotuner/fission_backend.cc +++ b/third_party/xla/xla/backends/gpu/autotuner/fission_backend.cc @@ -306,7 +306,8 @@ absl::StatusOr FissionBackend::EstimateFissionPrologueEpilogue( auto fusion_adaptor = HloFusionAdaptor::ForInstruction(instr); ABSL_ASSIGN_OR_RETURN( TiledRunTimeDataOrError tiled_data_or_decision, - indexing_cost_model.TryFindBestTilingForFusion(*fusion_adaptor)); + indexing_cost_model.TryFindBestTilingForFusionAsync(*fusion_adaptor) + .Await()); if (const auto* tiled_data = std::get_if(&tiled_data_or_decision)) { diff --git a/third_party/xla/xla/backends/gpu/autotuner/miopen.cc b/third_party/xla/xla/backends/gpu/autotuner/miopen.cc index a2e925c27c0d30..0b9e59b4f58a86 100644 --- a/third_party/xla/xla/backends/gpu/autotuner/miopen.cc +++ b/third_party/xla/xla/backends/gpu/autotuner/miopen.cc @@ -276,22 +276,15 @@ bool MIOpenBackend::IsSupported(const HloInstruction& instr) { absl::StatusOr> MIOpenBackend::GetDefaultConfig( const HloInstruction& instr) { - if (IsSupported(instr)) { - auto config = std::make_unique(); - config->mutable_algorithm()->set_algo_id(0); - return config; - } - return absl::InvalidArgumentError( - "MIOpen backend doesn't support getting a default config for this " - "instruction."); + return absl::UnimplementedError( + "MIOpen backend doesn't support getting a default config"); } absl::StatusOr>> -GetConvolutionCustomCallConfigs(const HloCustomCallInstruction* instr, - const HloModule* module, - se::StreamExecutor* stream_executor, - se::DeviceAddressAllocator* allocator, - se::Stream* stream) { +MIOpenBackend::GetConvolutionCustomCallConfigs( + const HloCustomCallInstruction* instr, const HloModule* module, + se::StreamExecutor* stream_executor, se::DeviceAddressAllocator* allocator, + se::Stream* stream) { if (stream_executor == nullptr) { return absl::InvalidArgumentError("Null stream executor is not supported."); } @@ -357,7 +350,7 @@ GetConvolutionCustomCallConfigs(const HloCustomCallInstruction* instr, gpu_conv_config.filter_descriptor, gpu_conv_params.filter_buf, gpu_conv_config.output_descriptor, gpu_conv_params.output_buf, gpu_conv_config.conv_desc, - /* use_fallback = */ false, &scratch_allocator, engine_options, + /* use_fallback = */ do_not_autotune_, &scratch_allocator, engine_options, &conv_runners)); std::vector> configs; @@ -373,10 +366,10 @@ GetConvolutionCustomCallConfigs(const HloCustomCallInstruction* instr, } absl::StatusOr>> -GetFusedConvolutionCustomCallConfigs(const HloCustomCallInstruction* instr, - const HloModule* module, - se::StreamExecutor* stream_executor, - se::DeviceAddressAllocator* allocator) { +MIOpenBackend::GetFusedConvolutionCustomCallConfigs( + const HloCustomCallInstruction* instr, const HloModule* module, + se::StreamExecutor* stream_executor, + se::DeviceAddressAllocator* allocator) { if (stream_executor == nullptr) { return absl::InvalidArgumentError("Null stream executor is not supported."); } @@ -453,13 +446,6 @@ MIOpenBackend::GetSupportedConfigs(const HloInstruction& instr) { stream_executor(), allocator_); } - if (do_not_autotune_) { - ABSL_ASSIGN_OR_RETURN(auto default_config, GetDefaultConfig(instr)); - std::vector> configs; - configs.push_back(std::move(default_config)); - return std::move(configs); - } - return GetConvolutionCustomCallConfigs( custom_call_instr, custom_call_instr->GetModule(), stream_executor(), allocator_, /* stream */ nullptr); diff --git a/third_party/xla/xla/backends/gpu/autotuner/miopen.h b/third_party/xla/xla/backends/gpu/autotuner/miopen.h index aa42f86b25c55d..67d18c5c3ff132 100644 --- a/third_party/xla/xla/backends/gpu/autotuner/miopen.h +++ b/third_party/xla/xla/backends/gpu/autotuner/miopen.h @@ -26,6 +26,8 @@ limitations under the License. #include "xla/backends/autotuner/codegen_backend.h" #include "xla/backends/gpu/autotuner/gpu_codegen_backend.h" #include "xla/hlo/ir/hlo_instruction.h" +#include "xla/hlo/ir/hlo_instructions.h" +#include "xla/hlo/ir/hlo_module.h" #include "xla/service/compiler.h" #include "xla/stream_executor/device_address_allocator.h" #include "xla/stream_executor/stream_executor.h" @@ -62,6 +64,19 @@ class MIOpenBackend : public GpuCodegenBackend { private: bool IsSupported(const HloInstruction& instr) override; + absl::StatusOr>> + GetConvolutionCustomCallConfigs( + const HloCustomCallInstruction* instr, const HloModule* module, + stream_executor::StreamExecutor* stream_executor, + stream_executor::DeviceAddressAllocator* allocator, + stream_executor::Stream* stream); + + absl::StatusOr>> + GetFusedConvolutionCustomCallConfigs( + const HloCustomCallInstruction* instr, const HloModule* module, + stream_executor::StreamExecutor* stream_executor, + stream_executor::DeviceAddressAllocator* allocator); + bool do_not_autotune_; stream_executor::DeviceAddressAllocator* allocator_; }; diff --git a/third_party/xla/xla/backends/gpu/autotuner/miopen_test.cc b/third_party/xla/xla/backends/gpu/autotuner/miopen_test.cc index 500935b952009c..a7e1bc0ef63232 100644 --- a/third_party/xla/xla/backends/gpu/autotuner/miopen_test.cc +++ b/third_party/xla/xla/backends/gpu/autotuner/miopen_test.cc @@ -142,10 +142,7 @@ TEST_F(MIOpenBackendTest, GetDefaultConfigFromMIOpenCustomCall) { absl::StatusOr> config = backend_.GetDefaultConfig( (*hlo_module->entry_computation()->root_instruction()->operand(0))); - TF_ASSERT_OK(config); - ASSERT_TRUE((*config)->has_algorithm()); - MIOpenBackendConfig algorithm_config = (*config)->algorithm(); - EXPECT_EQ(algorithm_config.algo_id(), 0); + EXPECT_THAT(config, absl_testing::StatusIs(absl::StatusCode::kUnimplemented)); } TEST_F(MIOpenBackendTest, ApplyConfigToMIOpenCustomCall) { 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 a40bcbd0aed1b3..b6cdb9217cf98a 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 @@ -47,7 +47,7 @@ limitations under the License. #include "mlir/Dialect/Tensor/IR/Tensor.h" #include "mlir/IR/Attributes.h" #include "mlir/IR/Builders.h" -#include "mlir/IR/BuiltinOps.h" +#include "mlir/IR/BuiltinAttributes.h" #include "mlir/IR/BuiltinTypeInterfaces.h" #include "mlir/IR/BuiltinTypes.h" #include "mlir/IR/Location.h" @@ -285,11 +285,14 @@ absl::Status ValidateBlockLevelFusionConfig( return absl::OkStatus(); } -absl::StatusOr> GetAllReduceUnmanagedKernelArguments( +// Unmanaged kernel arguments of collectives that use the AllReduce argument +// layout: rank, signal value, signal buffers and one remote buffer per +// parameter. +absl::StatusOr> GetRemoteBufferUnmanagedKernelArguments( const HloComputation* computation, - const HloAllReduceInstruction* all_reduce) { + const HloCollectiveInstruction* collective) { const int32_t num_devices = - all_reduce->device_list()->num_devices_per_group(); + collective->device_list()->num_devices_per_group(); std::vector unmanaged_arguments; unmanaged_arguments.reserve(computation->num_parameters() + kNumCollectiveMetadataArgs); @@ -422,6 +425,80 @@ mlir::Value EmitDeviceInvocationCount(mlir::ImplicitLocOpBuilder& b, b, counter, arith::ConstantOp::create(b, b.getI32IntegerAttr(1))); } +// Returns the element offset of the active half of a double buffer that holds +// `num_elements` elements of `element_type` per half. The lowest bit of +// `signal_value` selects the half. +mlir::Value EmitDoubleBufferOffset(mlir::ImplicitLocOpBuilder& b, + mlir::Value signal_value, + int64_t num_elements, + PrimitiveType element_type) { + const mlir::Type i64_type = b.getI64Type(); + mlir::Value signal_value_i64 = + arith::ExtSIOp::create(b, i64_type, signal_value); + mlir::Value one = + arith::ConstantOp::create(b, i64_type, b.getI64IntegerAttr(1)); + mlir::Value buffer_index = arith::AndIOp::create(b, signal_value_i64, one); + // The allocated double buffer size in bytes is aligned to + // kXlaAllocatedBufferAlignBytes. To get the offset to the second buffer, we + // divide the buffer size by 2 and then divide by the element size to get + // the number of elements in the buffer. Its important to do this to make + // sure that start of the buffers are always aligned to 16 bytes. + const int64_t element_size = ShapeUtil::ByteSizeOfPrimitiveType(element_type); + const int64_t buffer_size = xla::RoundUpTo( + num_elements * element_size, kXlaAllocatedBufferAlignBytes); + const int64_t elements_per_buffer = buffer_size / element_size; + mlir::Value elements_per_buffer_op = arith::ConstantOp::create( + b, i64_type, b.getI64IntegerAttr(elements_per_buffer)); + return arith::MulIOp::create(b, buffer_index, elements_per_buffer_op); +} + +// Emits instructions to get the pointer to the remote buffer of the given +// rank. +// We have a !tt.ptr pointing to the base of the remote +// buffers. We add the rank index to the base pointer and load to get to the +// base pointer of the remote buffer of the given rank. Then we add the buffer +// offset to get the pointer to the correct buffer inside (double buffering). +// Note that the returned pointer is of the storage type not the logical type. +mlir::Value EmitRemoteBufferPtr(mlir::ImplicitLocOpBuilder& b, + mlir::Value remote_buffers, mlir::Value rank, + mlir::Value buffer_offset, + mlir::Type elem_storage_type) { + const auto ptr_to_i64_type = + ttir::PointerType::get(b.getI64Type(), kGlobalAddressSpace); + const auto ptr_to_elem_type = + ttir::PointerType::get(elem_storage_type, kGlobalAddressSpace); + mlir::Value remote_buf_ptr_addr = + ttir::AddPtrOp::create(b, ptr_to_i64_type, remote_buffers, rank); + mlir::Value remote_buf_i64 = + ttir::LoadOp::create(b, // + remote_buf_ptr_addr, // + ttir::CacheModifier::NONE, // + ttir::EvictionPolicy::NORMAL, // + /*isVolatile=*/false); // + mlir::Value remote_buf_ptr_base = ttir::IntToPtrOp::create( + b, ptr_to_elem_type, remote_buf_i64, + llvm::ArrayRef{xtile::GetDivisibilityAttr(b)}); + return ttir::AddPtrOp::create(b, ptr_to_elem_type, remote_buf_ptr_base, + buffer_offset); +} + +// Emits a barrier between this block and the matching blocks of the other +// ranks. All threads in the block should have completed their writes before we +// proceed with a block barrier. +// Otherwise, remote ranks might start reading the data before it is ready. +void EmitBlockBarrier(mlir::ImplicitLocOpBuilder& b, mlir::Value signal_buffers, + mlir::Value rank, mlir::Value signal_value, + int64_t world_size, mlir::Value signal_slot = nullptr, + int64_t signal_stride = 0) { + mlir::triton::gpu::BarrierOp::create(b, mlir::triton::gpu::AddrSpace::Local); + mtx::BlockBarrierOp::create( + b, signal_buffers, rank, signal_value, signal_slot, + b.getI32IntegerAttr(static_cast(world_size)), + signal_stride > 0 + ? b.getI32IntegerAttr(static_cast(signal_stride)) + : mlir::IntegerAttr()); +} + class AllReduceEmitter { public: static mlir::LogicalResult Emit(AllReduceEmitterContext ctx, @@ -519,8 +596,6 @@ class AllReduceEmitter { remote_input_buffers_i64_ = ttir::BitcastOp::create( builder_, ptr_to_i64_type_, remote_input_buffers_); } - const mlir::Type i64_type = builder_.getI64Type(); - // Check if last bit of signal_value is 0 or 1. mlir::Value signal_value = signal_value_; // For two-shot signal value is always incremented by 2. // So we need to right shift it by 1 to get the buffer index. @@ -530,55 +605,18 @@ class AllReduceEmitter { arith::ConstantOp::create(builder_, signal_value.getType(), builder_.getI32IntegerAttr(1))); } - mlir::Value buffer_index = arith::AndIOp::create( - builder_, i64_type, - arith::ExtSIOp::create(builder_, i64_type, signal_value), // i32->i64 - arith::ConstantOp::create(builder_, i64_type, - builder_.getI64IntegerAttr(1))); - // The allocated double buffer size in bytes is aligned to - // kXlaAllocatedBufferAlignBytes. To get the offset to the second buffer, we - // divide the buffer size by 2 and then divide by the element size to get - // the number of elements in the buffer. Its important to do this to make - // sure that start of the buffers are always aligned to 16 bytes. - const int64_t buffer_size = xla::RoundUpTo( - ctx_.num_elements * - ShapeUtil::ByteSizeOfPrimitiveType(ctx_.element_type), - kXlaAllocatedBufferAlignBytes); - const int64_t elements_per_buffer = - buffer_size / ShapeUtil::ByteSizeOfPrimitiveType(ctx_.element_type); - buffer_offset_ = arith::MulIOp::create( - builder_, i64_type, buffer_index, - arith::ConstantOp::create( - builder_, i64_type, - builder_.getI64IntegerAttr(elements_per_buffer))); + buffer_offset_ = EmitDoubleBufferOffset( + builder_, signal_value, ctx_.num_elements, ctx_.element_type); initialized_ = true; return absl::OkStatus(); } - // Emits instructions to get the pointer to the remote buffer of the given - // rank. - // We have a !tt.ptr pointing to the base of the remote - // buffers. We add the rank index to the base pointer and load to get to the - // base pointer of the remote buffer of the given rank. Then we add the buffer - // offset to get the pointer to the correct buffer inside (double buffering). - // Note that the returned pointer is of the storage type not the logical type. + // Returns the pointer to the active remote buffer of the given rank. See + // EmitRemoteBufferPtr. mlir::Value GetRemoteBufferPtr(mlir::Value rank_idx) { CHECK(initialized_); - mlir::Value remote_buf_ptr_addr = ttir::AddPtrOp::create( - builder_, ptr_to_i64_type_, remote_input_buffers_i64_, rank_idx); - mlir::Value remote_buf_i64 = - ttir::LoadOp::create(builder_, // - remote_buf_ptr_addr, // - ttir::CacheModifier::NONE, // - ttir::EvictionPolicy::NORMAL, // - /*isVolatile=*/false); // - mlir::Value remote_buf_ptr_base = - ttir::IntToPtrOp::create(builder_, ptr_to_elem_type_, remote_buf_i64, - llvm::ArrayRef{ - xtile::GetDivisibilityAttr(builder_)}); - mlir::Value remote_buf_ptr = ttir::AddPtrOp::create( - builder_, ptr_to_elem_type_, remote_buf_ptr_base, buffer_offset_); - return remote_buf_ptr; + return EmitRemoteBufferPtr(builder_, remote_input_buffers_i64_, rank_idx, + buffer_offset_, elem_storage_type_); } // Loads a tile from the remote buffer of the given rank. @@ -675,15 +713,8 @@ class AllReduceEmitter { mlir::LogicalResult EmitSync(mlir::Value signal_value) { CHECK(initialized_); - // All threads in the block should have completed their writes before we - // proceed with a block barrier. - // Otherwise, remote ranks might start reading the data before it is ready. - mlir::triton::gpu::BarrierOp::create(builder_, - mlir::triton::gpu::AddrSpace::Local); - mtx::BlockBarrierOp::create(builder_, signal_buffers_, device_rank_, - signal_value, /*signal_slot=*/nullptr, - builder_.getI32IntegerAttr(ctx_.world_size), - /*signal_stride=*/nullptr); + EmitBlockBarrier(builder_, signal_buffers_, device_rank_, signal_value, + ctx_.world_size); return mlir::success(); } @@ -1155,8 +1186,8 @@ absl::StatusOr> GetCollectiveUnmanagedKernelArguments( const HloInstruction* root = computation->root_instruction(); switch (root->opcode()) { case HloOpcode::kAllReduce: - return GetAllReduceUnmanagedKernelArguments( - computation, Cast(root)); + return GetRemoteBufferUnmanagedKernelArguments( + computation, Cast(root)); case HloOpcode::kAllGather: { // AllGather only needs the 3 metadata args (rank, signal_value, // signal_buffers). No per-parameter scratch buffer args because the @@ -1248,81 +1279,109 @@ absl::StatusOr CreateCollectiveKernelSpec( } } -CollectiveCodegenConfig CreateCollectiveCodegenConfig( - const HloInstruction* instr) { - const HloInstruction* collective = instr; - if (instr->opcode() == HloOpcode::kFusion) { - collective = instr->fused_instructions_computation()->root_instruction(); - } - CollectiveCodegenConfig config; - auto gpu_config = collective->backend_config(); - if (!gpu_config.ok()) { - return config; - } - const auto strategy = - gpu_config->collective_backend_config().kernel_strategy(); - // One-shot AllGather: the runtime copies input to symmetric scratch and the - // kernel needs a barrier before reading peers' data. - if (collective->opcode() == HloOpcode::kAllGather && - strategy == CollectiveBackendConfig::KERNEL_STRATEGY_TRITON_ONE_SHOT) { - config.copy_input_to_scratch = true; - config.emit_entry_barrier = true; - } - return config; -} +namespace { -absl::Status EmitCollectiveEntryBarrier(mlir::ModuleOp module, - int32_t world_size) { - // Find the xtile::EntryFuncOp in the module. - xtile::EntryFuncOp entry_func = nullptr; - for (auto fn : module.getOps()) { - entry_func = fn; - break; +struct AllGatherRewriteContext { + mlir::stablehlo::AllGatherOp op; + xtile::EntryFuncOp entry_func; + xtile::ExtractTileOp input_extract; + xtile::SelectBufferOp pull_select; + mlir::Value rank_arg; + mlir::Value signal_buffers_arg; + int64_t world_size = 0; +}; + +absl::StatusOr BuildAllGatherContext( + mlir::stablehlo::AllGatherOp op) { + if (op.getOperands().size() != 1) { + return absl::InvalidArgumentError( + "AllGather op must have exactly one operand."); } + auto entry_func = op->getParentOfType(); if (!entry_func) { - return absl::InternalError( - "No xtile::EntryFuncOp found in module for barrier insertion."); + return absl::InvalidArgumentError( + "AllGather op must be in an XTile entry function."); } - - // The opaque args are appended after the regular input/output args. - // Their count is stored in the "num_opaque_args" attribute. auto num_opaque_attr = entry_func->getAttrOfType("num_opaque_args"); if (!num_opaque_attr || num_opaque_attr.getInt() < kNumCollectiveMetadataArgs) { - return absl::InternalError( - absl::StrCat("Expected at least ", kNumCollectiveMetadataArgs, - " opaque args for collective entry barrier, got ", - num_opaque_attr ? num_opaque_attr.getInt() : 0)); + return absl::InvalidArgumentError( + "EntryFuncOp does not have collective metadata arguments."); } - // The opaque args start at (total_args - num_opaque_args). + mlir::Value input_tile = op.getOperand(0); + auto input_extract = llvm::dyn_cast_if_present( + input_tile.getDefiningOp()); + if (!input_extract && input_tile.getDefiningOp() && + input_tile.getDefiningOp()->getNumOperands() > 0) { + // Workaround(i1_to_i8_workaround): booleans are stored as i8 and cast to + // i1 after ExtractTileOp. + input_extract = llvm::dyn_cast_if_present( + input_tile.getDefiningOp()->getOperand(0).getDefiningOp()); + } + if (!input_extract) { + return absl::InvalidArgumentError( + "AllGather operand must be defined by xtile::ExtractTileOp."); + } + + auto pull_select = llvm::dyn_cast_if_present( + input_extract.getSource().getDefiningOp()); + if (!pull_select) { + return absl::InvalidArgumentError( + "AllGather ExtractTileOp source must be defined by " + "xtile::SelectBufferOp."); + } + + ABSL_ASSIGN_OR_RETURN(auto replica_groups, + xla::ConvertReplicaGroups(op.getReplicaGroups(), op)); + if (replica_groups->replica_groups().empty()) { + return absl::InvalidArgumentError( + "AllGather replica groups must not be empty."); + } + + AllGatherRewriteContext ctx; + ctx.op = op; + ctx.entry_func = entry_func; + ctx.input_extract = input_extract; + ctx.pull_select = pull_select; + ctx.world_size = replica_groups->num_devices_per_group(); + int32_t total_args = entry_func.getNumArguments(); - int32_t opaque_start = + int32_t metadata_args_start = total_args - num_opaque_attr.getInt() - kNumTileIndexArgs; // Layout: opaque[0]=rank, opaque[1]=invocation count (unused, the signal // value comes from a counter in device memory), opaque[2]=signal_buffers. - mlir::Value rank_arg = entry_func.getArgument(opaque_start); - mlir::Value signal_buffers_arg = entry_func.getArgument(opaque_start + 2); + ctx.rank_arg = entry_func.getArgument(metadata_args_start); + ctx.signal_buffers_arg = entry_func.getArgument(metadata_args_start + 2); + return ctx; +} + +} // namespace + +mlir::LogicalResult RewriteAllGather(mlir::stablehlo::AllGatherOp op, + mlir::PatternRewriter& rewriter) { + absl::StatusOr maybe_ctx = BuildAllGatherContext(op); + if (!maybe_ctx.ok()) { + return rewriter.notifyMatchFailure(op, maybe_ctx.status().message()); + } + AllGatherRewriteContext& ctx = *maybe_ctx; - // Insert at the beginning of the entry block, right after program_id - // extraction (which is the first op). We want the barrier before any - // scf.for tile loop. - mlir::Block& entry_block = entry_func.front(); - auto loc = entry_func.getLoc(); - mlir::ImplicitLocOpBuilder builder(loc, &entry_block, entry_block.begin()); + mlir::ImplicitLocOpBuilder builder(ctx.pull_select.getLoc(), rewriter); + builder.setInsertionPoint(ctx.pull_select); mlir::Value block_id = ttir::GetProgramIdOp::create(builder, 0); mlir::Value signal_value = EmitDeviceInvocationCount( - builder, signal_buffers_arg, rank_arg, block_id, world_size); + builder, ctx.signal_buffers_arg, ctx.rank_arg, block_id, ctx.world_size); // Inter-block barrier via signal flags. This blocks until all // remote ranks have also signaled. - mtx::BlockBarrierOp::create(builder, signal_buffers_arg, rank_arg, + mtx::BlockBarrierOp::create(builder, ctx.signal_buffers_arg, ctx.rank_arg, signal_value, /*signal_slot=*/nullptr, - builder.getI32IntegerAttr(world_size), + builder.getI32IntegerAttr(ctx.world_size), /*signal_stride=*/nullptr); - return absl::OkStatus(); + rewriter.replaceOp(op, op.getOperand(0)); + return mlir::success(); } } // namespace xla::gpu diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/collective_emitter.h b/third_party/xla/xla/backends/gpu/codegen/triton/collective_emitter.h index 094e529a0ded1b..5088623ca49d13 100644 --- a/third_party/xla/xla/backends/gpu/codegen/triton/collective_emitter.h +++ b/third_party/xla/xla/backends/gpu/codegen/triton/collective_emitter.h @@ -24,7 +24,6 @@ limitations under the License. #include "absl/status/statusor.h" #include "llvm/ADT/SmallVector.h" #include "mlir/IR/Builders.h" -#include "mlir/IR/BuiltinOps.h" #include "mlir/IR/PatternMatch.h" #include "mlir/IR/Types.h" #include "mlir/Support/LLVM.h" @@ -101,23 +100,13 @@ absl::StatusOr> GetCollectiveUnmanagedKernelArguments( mlir::LogicalResult RewriteAllReduce(mlir::stablehlo::AllReduceOp op, mlir::PatternRewriter& rewriter); -// Creates a lightweight codegen config from the collective HLO instruction. -// The returned config drives codegen decisions (e.g. whether the runtime must -// copy the input to scratch before kernel launch). -CollectiveCodegenConfig CreateCollectiveCodegenConfig( - const HloInstruction* instr); +// Rewrites stablehlo all-gather op to a triton implementation. +mlir::LogicalResult RewriteAllGather(mlir::stablehlo::AllGatherOp op, + mlir::PatternRewriter& rewriter); // Creates a CollectiveKernelSpec for a given collective or fusion instruction. absl::StatusOr CreateCollectiveKernelSpec( const HloInstruction* instr, const LaunchDimensions& launch_dimensions); -// Emits a collective entry barrier at the start of the entry function in -// |module|. The barrier ensures all ranks have completed their D2D copies -// before any rank starts reading from the symmetric scratch buffers. -// Uses the opaque metadata args (rank, signal_value, signal_buffers) already -// present in the EntryFuncOp. -absl::Status EmitCollectiveEntryBarrier(mlir::ModuleOp module, - int32_t world_size); - } // namespace xla::gpu #endif // XLA_BACKENDS_GPU_CODEGEN_TRITON_COLLECTIVE_EMITTER_H_ 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 bf680cada9aaeb..603a011b0af778 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 @@ -79,16 +79,36 @@ func.func @insert_wrong_shape(%arg0: tensor<8xbf16>, %arg1: !tt.ptr) { // ----- 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) -> () + // expected-error @+1 {{'triton_xla.block_barrier' op symmetric barrier_mode requires signal_slot to be omitted and signal_stride to be 0}} + triton_xla.block_barrier %ptr, %rank, %sig : (!tt.ptr, i32, i32) -> () + return +} + +func.func @block_barrier_symmetric_with_stride(%ptr: !tt.ptr, %rank: i32, %sig: i32) { + // expected-error @+1 {{'triton_xla.block_barrier' op symmetric barrier_mode requires signal_slot to be omitted and signal_stride to be 0}} + triton_xla.block_barrier %ptr, %rank, %sig, : (!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) -> () + // expected-error @+1 {{'triton_xla.block_barrier' op symmetric barrier_mode requires signal_slot to be omitted and signal_stride to be 0}} + triton_xla.block_barrier %ptr, %rank, %sig, %slot : (!tt.ptr, i32, i32, i32) -> () return } +func.func @block_barrier_asymmetric_without_slot(%ptr: !tt.ptr, %rank: i32, %sig: i32) { + // expected-error @+1 {{'triton_xla.block_barrier' op asymmetric barrier_mode requires signal_slot to be specified}} + triton_xla.block_barrier %ptr, %rank, %sig : (!tt.ptr, i32, i32) -> () + return +} + +// ----- + +func.func @block_barrier_asymmetric_without_stride(%ptr: !tt.ptr, %rank: i32, %sig: i32, %slot: i32) { + // expected-error @+1 {{'triton_xla.block_barrier' op asymmetric barrier_mode requires a positive signal_stride}} + triton_xla.block_barrier %ptr, %rank, %sig, %slot : (!tt.ptr, i32, i32, i32) -> () + return +} diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/ir/triton_xla_attrs.td b/third_party/xla/xla/backends/gpu/codegen/triton/ir/triton_xla_attrs.td index ab18493965e4d7..de54d77026f7da 100644 --- a/third_party/xla/xla/backends/gpu/codegen/triton/ir/triton_xla_attrs.td +++ b/third_party/xla/xla/backends/gpu/codegen/triton/ir/triton_xla_attrs.td @@ -86,5 +86,19 @@ def TTXLA_ComparatorAttr : EnumAttr let assemblyFormat = "$value"; } +def TTXLA_BarrierModeEnum : I32EnumAttr<"BarrierMode", + "Synchronization mode for triton_xla.block_barrier.", + [ + I32EnumAttrCase<"kSymmetric", 0, "symmetric">, + I32EnumAttrCase<"kConsumerSymmetric", 1, "consumer_symmetric">, + I32EnumAttrCase<"kProducerSymmetric", 2, "producer_symmetric">, + ]> { + let cppNamespace = "::mlir::triton::xla"; + let genSpecializedAttr = 0; +} +def TTXLA_BarrierModeAttr + : EnumAttr { + let assemblyFormat = "`<` $value `>`"; +} #endif // XLA_BACKENDS_GPU_CODEGEN_TRITON_IR_TRITON_XLA_ATTRS_TD_ diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/ir/triton_xla_dialect.td b/third_party/xla/xla/backends/gpu/codegen/triton/ir/triton_xla_dialect.td index e978703b7ccdf5..d5c5e683ab6961 100644 --- a/third_party/xla/xla/backends/gpu/codegen/triton/ir/triton_xla_dialect.td +++ b/third_party/xla/xla/backends/gpu/codegen/triton/ir/triton_xla_dialect.td @@ -28,6 +28,7 @@ def XlaTritonDialect : Dialect { let cppNamespace = "::mlir::triton::xla"; let useDefaultAttributePrinterParser = 1; + let useStrictPropertiesInAssemblyFormat = 1; let dependentDialects = [ "mlir::triton::TritonDialect", "mlir::triton::gpu::TritonGPUDialect", 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 98ae71a0bb3c3f..608eecd7de12e2 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 @@ -324,14 +324,21 @@ LogicalResult PtrToMemrefOp::verify() { } LogicalResult BlockBarrierOp::verify() { - if (getSignalStride() < 0) { - return emitOpError("signal_stride must be non-negative"); + if (getBarrierMode() == BarrierMode::kSymmetric) { + if (getSignalSlot() || getSignalStride() != 0) { + return emitOpError( + "symmetric barrier_mode requires signal_slot to be omitted and " + "signal_stride to be 0"); + } + return success(); } - if (getSignalStride() > 0 && !getSignalSlot()) { - return emitOpError("signal_stride requires signal_slot to be specified"); + if (getSignalStride() <= 0) { + return emitOpError( + "asymmetric barrier_mode requires a positive signal_stride"); } - if (getSignalSlot() && getSignalStride() == 0) { - return emitOpError("signal_slot requires a positive signal_stride"); + if (!getSignalSlot()) { + return emitOpError( + "asymmetric barrier_mode requires signal_slot to be specified"); } return success(); } 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 27bbca78c5aad7..b8ebaee318e174 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 @@ -204,7 +204,7 @@ def TTXLA_SqueezeDimsOp : TTXLA_Op<"squeeze_dims", [ let summary = "Remove a size-1 dimension from a tensor."; let arguments = (ins TT_Tensor:$src, I32Attr:$axis); let results = (outs TT_Tensor:$result); - let assemblyFormat = "$src attr-dict `:` type($src) `->` type($result)"; + let assemblyFormat = "$src prop-dict attr-dict `:` type($src) `->` type($result)"; } def TTXLA_GetTidOp : TTXLA_Op<"get_tid", [Pure]> { @@ -337,9 +337,9 @@ def TTXLA_BlockBarrierOp : TTXLA_Op<"block_barrier", []> { of `i32` flags with shape `[num_blocks][world_size]`, flattened as `block_idx * world_size + slot_idx`. - This barrier supports two synchronization modes: + This barrier supports three synchronization modes via `$barrier_mode`: - ### 1. Symmetric Mode (`$signal_stride == 0`, `$signal_slot` omitted) + ### 1. `symmetric` (default; `$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 @@ -349,7 +349,7 @@ def TTXLA_BlockBarrierOp : TTXLA_Op<"block_barrier", []> { 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). + ### 2. `consumer_symmetric` (`$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: @@ -366,6 +366,19 @@ def TTXLA_BlockBarrierOp : TTXLA_Op<"block_barrier", []> { slots at `block_idx * world_size + [0..world_size)` reach `>= $signal_value`. + ### 3. `producer_symmetric` (`$signal_stride > 0`, `$signal_slot` provided) + Used when block `block_idx` on every GPU produces data at its own + `block_idx`, and consumer block `block_idx` on GPU `r` waits for data + produced by `target_block` across all peer GPUs: + - Let `base_block = block_idx - $signal_slot * $signal_stride` and + `target_block = base_block + r * $signal_stride`. + - **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 + `target_block * 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 than the number of SMs on a GPU. That is to say that if BlockN is launched @@ -404,14 +417,17 @@ def TTXLA_BlockBarrierOp : TTXLA_Op<"block_barrier", []> { // same as this value. I32Attr:$world_size, // Optional block-ID stride between consumer blocks for asymmetric barriers. - DefaultValuedOptionalAttr:$signal_stride + DefaultValuedOptionalAttr:$signal_stride, + // Synchronization mode for the barrier. + DefaultValuedAttr:$barrier_mode ); // This is a barrier op and produces no results. let results = (outs); let assemblyFormat = [{ - $signal_buffers `,` $device_rank `,` $signal_value `,` ($signal_slot^ `,`)? attr-dict `:` + $signal_buffers `,` $device_rank `,` $signal_value (`,` $signal_slot^)? prop-dict attr-dict `:` functional-type(operands, results) }]; @@ -449,7 +465,7 @@ def TTXLA_GetPeerPtrOp : TTXLA_Op<"get_peer_ptr", [Pure]> { let results = (outs Arg]>:$result); let assemblyFormat = [{ - $address `,` $peer_id `,` $metadata `,` attr-dict `:` + $address `,` $peer_id `,` $metadata prop-dict attr-dict `:` functional-type(operands, results) }]; } diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/BUILD b/third_party/xla/xla/backends/gpu/codegen/triton/tests/BUILD index 22cfd08b2cd3b6..1532db909a751e 100644 --- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/BUILD +++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/BUILD @@ -68,7 +68,6 @@ xla_cc_binary( deviceless_tests = [ "collectives/all_reduce.hlo", "collectives/all_gather.hlo", - "collectives/all_gather_one_shot.hlo", "concat/concat_active_scalar_constant.hlo", "concat/concat_offset_not_a_multiple_of_tile_size.hlo", "constraints/max_tensor_elements_violation.hlo", diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/collectives/all_gather.hlo b/third_party/xla/xla/backends/gpu/codegen/triton/tests/collectives/all_gather.hlo index 39b3bc64e1b787..b065647800ae85 100644 --- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/collectives/all_gather.hlo +++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/collectives/all_gather.hlo @@ -33,14 +33,27 @@ ENTRY entry { // CHECK: xtile.entry_func @triton_fn( // CHECK-SAME: %[[ARG0:[a-zA-Z0-9_]+]]: memref<2xi64>, // CHECK-SAME: %[[ARG1:[a-zA-Z0-9_]+]]: memref<256x128xf32>, -// CHECK-SAME: %[[ARG2:[a-zA-Z0-9_]+]]: i32, -// CHECK-SAME: %[[ARG3:[a-zA-Z0-9_]+]]: i32, -// CHECK-SAME: %[[ARG4:[a-zA-Z0-9_]+]]: !tt.ptr, +// CHECK-SAME: %[[RANK:[a-zA-Z0-9_]+]]: i32, +// CHECK-SAME: %[[INVOCATION_COUNT:[a-zA-Z0-9_]+]]: i32, +// CHECK-SAME: %[[SIGNAL_BUFFERS:[a-zA-Z0-9_]+]]: !tt.ptr, // CHECK-SAME: %[[PID:[a-zA-Z0-9_]+]]: index {xla.range = [0 : index, 127 : index]}) -// CHECK: %[[IDX0:.*]] = xla.apply_indexing #[[MAP0]](%[[PID]]) -// CHECK: %[[IDX1:.*]] = xla.apply_indexing #[[MAP1]](%[[PID]]) -// CHECK: %[[IDX2:.*]] = xla.apply_indexing #[[MAP2]](%[[PID]]) +// CHECK-DAG: %[[IDX0:.*]] = xla.apply_indexing #[[MAP0]](%[[PID]]) +// CHECK-DAG: %[[IDX1:.*]] = xla.apply_indexing #[[MAP1]](%[[PID]]) +// CHECK-DAG: %[[IDX2:.*]] = xla.apply_indexing #[[MAP2]](%[[PID]]) +// CHECK-DAG: %[[C1:.*]] = arith.constant 1 : i32 +// CHECK-DAG: %[[WORLD_SIZE:.*]] = arith.constant 2 : i32 +// CHECK: %[[BLOCK_ID:.*]] = tt.get_program_id x +// CHECK: %[[SIGNAL_BUFFER_ADDR:.*]] = tt.addptr %[[SIGNAL_BUFFERS]], %[[RANK]] +// CHECK: %[[SIGNAL_BUFFER_INT:.*]] = tt.load %[[SIGNAL_BUFFER_ADDR]] +// CHECK: %[[SIGNAL_BUFFER:.*]] = tt.int_to_ptr %[[SIGNAL_BUFFER_INT]] +// CHECK: %[[BLOCK_OFFSET:.*]] = arith.muli %[[BLOCK_ID]], %[[WORLD_SIZE]] +// CHECK: %[[COUNTER_IDX:.*]] = arith.addi %[[BLOCK_OFFSET]], %[[RANK]] +// CHECK: %[[COUNTER_PTR:.*]] = tt.addptr %[[SIGNAL_BUFFER]], %[[COUNTER_IDX]] +// CHECK: %[[COUNTER:.*]] = tt.load %[[COUNTER_PTR]] {isVolatile = true} +// CHECK: %[[SIGNAL_VALUE:.*]] = arith.addi %[[COUNTER]], %[[C1]] +// CHECK: triton_xla.block_barrier %[[SIGNAL_BUFFERS]], %[[RANK]], %[[SIGNAL_VALUE]] +// CHECK-SAME: // CHECK: %[[BUF:.*]] = xtile.select_buffer %[[ARG0]][%[[IDX2]]] : memref<2xi64> -> memref<128x128xf32> // CHECK: %[[EXT:.*]] = xtile.extract %[[BUF]][%[[IDX0]], %[[IDX1]]] [16, 16] [1, 1] : memref<128x128xf32> -> tensor<16x16xf32> // CHECK: %[[IDX3:.*]] = xla.apply_indexing #[[MAP3]](%[[PID]]) diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/collectives/all_gather_one_shot.hlo b/third_party/xla/xla/backends/gpu/codegen/triton/tests/collectives/all_gather_one_shot.hlo deleted file mode 100644 index f04d9c1b766fe6..00000000000000 --- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/collectives/all_gather_one_shot.hlo +++ /dev/null @@ -1,58 +0,0 @@ -// RUN: fusion_to_triton %s | FileCheck %s -// RUN: fusion_to_triton %s --xla_gpu_experimental_enable_tiling_propagation | FileCheck %s -f { - %param0 = f32[128,128]{1,0} parameter(0) - ROOT result = f32[256,128]{1,0} all-gather(%param0), - replica_groups={{0,1}}, dimensions={0}, - backend_config={"collective_backend_config":{ - "kernel_strategy":"KERNEL_STRATEGY_TRITON_ONE_SHOT"}} -} - -ENTRY entry { - p0 = f32[128,128]{1,0} parameter(0) - ROOT %fusion = f32[256,128]{1,0} fusion(p0), - kind=kCustom, calls=f, - backend_config={ - "fusion_backend_config": { - "kind": "__triton_collective", - "block_level_fusion_config": { - "num_warps": "4", - "output_tiles": [{sizes: [16,16]}], - "num_ctas": 1, - "num_stages": 1, - "is_tma_allowed": false, - "is_warp_specialization_allowed": false - } - } - } -} - -// One-shot all-gather waits on an entry barrier before reading the peers' -// scratch buffers. The signal value of the barrier is derived from this block's -// own slot in the local signal buffer (`SignalBuffers[rank] -// [block_id * world_size + rank]`). - -// CHECK: xtile.entry_func @triton_fn( -// CHECK-SAME: %[[ARG0:[a-zA-Z0-9_]+]]: memref<2xi64>, -// CHECK-SAME: %[[ARG1:[a-zA-Z0-9_]+]]: memref<256x128xf32>, -// CHECK-SAME: %[[RANK:[a-zA-Z0-9_]+]]: i32, -// CHECK-SAME: %[[INVOCATION_COUNT:[a-zA-Z0-9_]+]]: i32, -// CHECK-SAME: %[[SIGNAL_BUFFERS:[a-zA-Z0-9_]+]]: !tt.ptr, -// CHECK-SAME: %[[PID:[a-zA-Z0-9_]+]]: index{{( \{xla.range = \[0 : index, 127 : index\]\})?}}) -// CHECK-DAG: %[[C1:.*]] = arith.constant 1 : i32 -// CHECK-DAG: %[[WORLD_SIZE:.*]] = arith.constant 2 : i32 -// CHECK: %[[BLOCK_ID:.*]] = tt.get_program_id x -// CHECK: %[[SIGNAL_BUFFER_ADDR:.*]] = tt.addptr %[[SIGNAL_BUFFERS]], %[[RANK]] -// CHECK: %[[SIGNAL_BUFFER_INT:.*]] = tt.load %[[SIGNAL_BUFFER_ADDR]] -// CHECK: %[[SIGNAL_BUFFER:.*]] = tt.int_to_ptr %[[SIGNAL_BUFFER_INT]] -// CHECK: %[[BLOCK_OFFSET:.*]] = arith.muli %[[BLOCK_ID]], %[[WORLD_SIZE]] -// CHECK: %[[COUNTER_IDX:.*]] = arith.addi %[[BLOCK_OFFSET]], %[[RANK]] -// CHECK: %[[COUNTER_PTR:.*]] = tt.addptr %[[SIGNAL_BUFFER]], %[[COUNTER_IDX]] -// CHECK: %[[COUNTER:.*]] = tt.load %[[COUNTER_PTR]] {isVolatile = true} -// CHECK: %[[SIGNAL_VALUE:.*]] = arith.addi %[[COUNTER]], %[[C1]] -// CHECK: triton_xla.block_barrier %[[SIGNAL_BUFFERS]], %[[RANK]], %[[SIGNAL_VALUE]], -// CHECK-SAME: {world_size = 2 : i32} -// CHECK: xtile.select_buffer %[[ARG0]] -// CHECK: xtile.extract -// CHECK: xtile.insert -// CHECK: xtile.return diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/collectives/all_reduce.hlo b/third_party/xla/xla/backends/gpu/codegen/triton/tests/collectives/all_reduce.hlo index 8ae71f7f0d3ba6..731791d3cbe9b0 100644 --- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/collectives/all_reduce.hlo +++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/collectives/all_reduce.hlo @@ -43,8 +43,8 @@ fusion{ // CHECK-SAME: : tensor<4096x!tt.ptr> // CHECK: ttg.barrier local -// CHECK: triton_xla.block_barrier %[[ARG4]], %[[ARG2]], %[[SIGNAL_VALUE]], -// CHECK-SAME: {world_size = 2 : i32} : (!tt.ptr, i32, i32) +// CHECK: triton_xla.block_barrier %[[ARG4]], %[[ARG2]], %[[SIGNAL_VALUE]] +// CHECK-SAME: : (!tt.ptr, i32, i32) // CHECK: xtile.insert // CHECK-SAME: : tensor<4096xf32> -> memref<65536xf32> diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/fusion_emitter_shared_dialect_test.cc b/third_party/xla/xla/backends/gpu/codegen/triton/tests/fusion_emitter_shared_dialect_test.cc index ce1a4f839e0350..4b12d7be11a5fe 100644 --- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/fusion_emitter_shared_dialect_test.cc +++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/fusion_emitter_shared_dialect_test.cc @@ -292,7 +292,7 @@ ENTRY e { // CHECK: %[[EXTRACT1:.*]] = xtile.extract %arg1[] [] [] : memref -> tensor // CHECK: %[[INIT:.*]] = stablehlo.reshape %[[EXTRACT1]] : (tensor) -> tensor<1xf32> // CHECK: %[[OUTPUT:.*]], %{{.*}} = xtile.scan(%[[EXTRACT0]]) inits(%[[INIT]]) -// CHECK-SAME: dimension = 0 {scan_dim_size = 1024 : i64} +// CHECK-SAME: dimension = 0 // CHECK-SAME: : (tensor<1024xf32>), (tensor<1xf32>) -> (tensor<1024xf32>), (tensor<1xf32>) { // CHECK: ^bb0(%[[INPUT:.*]]: tensor, %[[CARRY:.*]]: tensor): // CHECK: %[[ADD:.*]] = stablehlo.add %[[INPUT]], %[[CARRY]] : tensor @@ -411,7 +411,7 @@ ENTRY e { EXPECT_OK(CreateXTileIrAndFileCheck( *module->GetComputationWithName("triton_dot"), block_level_parameters, R"( - CHECK: %[[DOT:.*]] = xtile.dot_scaled %[[LHS:.*]] scale %[[LHS_SCALE:.*]], %[[RHS:.*]] scale %[[RHS_SCALE:.*]] {dot_dimension_numbers = #stablehlo.dot, fastMath = true, lhs_elem_type = f8E5M2, rhs_elem_type = f8E5M2} : tensor<128x128xf8E5M2>, tensor<128x4xi8> * tensor<128x256xf8E5M2>, tensor<256x4xi8> -> tensor<128x256xf32> + CHECK: %[[DOT:.*]] = xtile.dot_scaled %[[LHS:.*]] scale %[[LHS_SCALE:.*]], %[[RHS:.*]] scale %[[RHS_SCALE:.*]] > : tensor<128x128xf8E5M2>, tensor<128x4xi8> * tensor<128x256xf8E5M2>, tensor<256x4xi8> -> tensor<128x256xf32> CHECK: %[[RES:.*]] = arith.addf %{{.*}}, %[[DOT]] : tensor<128x256xf32> )")); } @@ -584,8 +584,10 @@ TEST_F(XTileDialectTest, HloAllGatherDotLowering) { CHECK: %[[SELECT2:.*]] = xtile.select_buffer %[[SELECT1]][%{{.*}}] CHECK-SAME: : memref<2xi64> -> memref<128x128xf32> CHECK: %[[LHS_TILE:.*]] = xtile.extract %[[SELECT2]] + CHECK: %[[AG1:.*]] = "stablehlo.all_gather"(%[[LHS_TILE]]) + CHECK: %[[AG2:.*]] = "stablehlo.all_gather"(%[[AG1]]) CHECK: %[[RHS_TILE:.*]] = xtile.extract %arg1 - CHECK: stablehlo.dot_general %[[LHS_TILE]], %[[RHS_TILE]] + CHECK: stablehlo.dot_general %[[AG2]], %[[RHS_TILE]] )")); } @@ -678,7 +680,7 @@ ENTRY e { CHECK: scf.for %[[IV:.*]] = %{{.*}} to %{{.*}} step %{{.*}} iter_args(%[[CARRY:.*]] = %{{.*}}) CHECK: %[[REVERSE_IDX:.*]] = arith.subi %{{.*}}, %[[IV]] : index CHECK: %[[INPUT_TILE:.*]] = xtile.extract %arg0 - CHECK: %[[SCAN_OUT:.*]], %[[NEW_CARRY:.*]] = xtile.scan(%[[INPUT_TILE]]) inits(%[[CARRY]]) dimension = 0 {is_reverse = true + CHECK: %[[SCAN_OUT:.*]], %[[NEW_CARRY:.*]] = xtile.scan(%[[INPUT_TILE]]) inits(%[[CARRY]]) dimension = 0 CHECK: xtile.insert %[[SCAN_OUT]] into %arg2 CHECK: scf.yield %[[NEW_CARRY]] )")); diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/transforms/stablehlo_lower_to_triton.cc b/third_party/xla/xla/backends/gpu/codegen/triton/transforms/stablehlo_lower_to_triton.cc index c47c144a25d843..09747bc73d0261 100644 --- a/third_party/xla/xla/backends/gpu/codegen/triton/transforms/stablehlo_lower_to_triton.cc +++ b/third_party/xla/xla/backends/gpu/codegen/triton/transforms/stablehlo_lower_to_triton.cc @@ -1030,6 +1030,18 @@ class LowerDotGeneral : public mlir::OpRewritePattern { bool warp_specialization_allowed_; }; +class LowerAllGather : public mlir::OpRewritePattern { + public: + using OpRewritePattern::OpRewritePattern; + + private: + mlir::LogicalResult matchAndRewrite( + stablehlo::AllGatherOp op, + mlir::PatternRewriter& rewriter) const override { + return ::xla::gpu::RewriteAllGather(op, rewriter); + } +}; + class LowerAllReduce : public mlir::OpRewritePattern { public: using OpRewritePattern::OpRewritePattern; @@ -1054,7 +1066,8 @@ class StableHLOLowerToTritonPass { mlir::RewritePatternSet patterns(mlir_context); patterns.add(mlir_context); + LowerReduce, LowerReshape, LowerAllGather, LowerAllReduce>( + mlir_context); patterns.add(mlir_context); patterns.add(mlir_context, warp_specialization_allowed_); if (mlir::failed(mlir::applyPatternsGreedily(getOperation(), 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 7fcbfbf52b4d08..d587b6c3788d1f 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 @@ -45,17 +45,17 @@ tt.func @block_barrier_kernel( // CHECK-NEXT: } // CHECK-NEXT: ttg.barrier local // CHECK-NEXT: tt.return - triton_xla.block_barrier %ptr, %rank, %signal_value, { world_size = 8 : i32 } : + triton_xla.block_barrier %ptr, %rank, %signal_value : (!tt.ptr, i32, i32) -> () tt.return } -// CHECK-LABEL: block_barrier_asymmetric_kernel +// CHECK-LABEL: block_barrier_consumer_symmetric_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( +tt.func @block_barrier_consumer_symmetric_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 @@ -88,7 +88,52 @@ tt.func @block_barrier_asymmetric_kernel( // 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 } : + triton_xla.block_barrier %ptr, %rank, %signal_value, %signal_slot + : (!tt.ptr, i32, i32, i32) -> () tt.return -} \ No newline at end of file +} + +// CHECK-LABEL: block_barrier_producer_symmetric_kernel +// CHECK-SAME: %[[PTR:[^:]+]]: !tt.ptr, +// CHECK-SAME: %[[SIGNAL_VALUE:[^:]+]]: i32, +// CHECK-SAME: %[[RANK:[^:]+]]: i32, +// CHECK-SAME: %[[SIGNAL_SLOT:[^:]+]]: i32 +tt.func @block_barrier_producer_symmetric_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: %[[WAIT_BLOCK:.+]] = arith.addi %[[BASE_T]], %[[RANK_OFFSET]] + // CHECK-NEXT: %[[WAIT_BLOCK_OFFSET:.+]] = arith.muli %[[WAIT_BLOCK]], %[[WORLD_SIZE]] + // CHECK-NEXT: %[[WRITE_OFFSET:.+]] = arith.addi %[[MULI]], %[[RANK]] + // 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]], %[[WAIT_BLOCK_OFFSET]] + // 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 + : + (!tt.ptr, i32, i32, i32) -> () + tt.return +} diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/transforms/tests/triton_xla_fold_transpose.mlir b/third_party/xla/xla/backends/gpu/codegen/triton/transforms/tests/triton_xla_fold_transpose.mlir index ac41b44c900344..e9009950d7949d 100644 --- a/third_party/xla/xla/backends/gpu/codegen/triton/transforms/tests/triton_xla_fold_transpose.mlir +++ b/third_party/xla/xla/backends/gpu/codegen/triton/transforms/tests/triton_xla_fold_transpose.mlir @@ -93,3 +93,59 @@ func.func @push_transpose_up_through_mask(%arg0: tensor<4x8xf32>, %arg1: f32) -> %1 = tt.trans %0 {order = array} : tensor<4x8xf32> -> tensor<8x4xf32> return %1 : tensor<8x4xf32> } + +// CHECK-LABEL: func @push_transpose_up_through_diamond_dag +func.func @push_transpose_up_through_diamond_dag( + %input: memref<4x8xf32, #xtile.layout<[1, 0]>>, + %offset0: index, %offset1: index) -> tensor<8x4xf32> { + // CHECK: %[[TRANSPOSE:.*]] = memref.transpose %arg0 (d0, d1) -> (d1, d0) + // CHECK-NEXT: %[[EXTRACT:.*]] = xtile.extract %[[TRANSPOSE]][%arg2, %arg1] [8, 4] [1, 1] + // CHECK-NEXT: %[[SQ:.*]] = arith.mulf %[[EXTRACT]], %[[EXTRACT]] : tensor<8x4xf32> + // CHECK-NEXT: %[[ADD:.*]] = arith.addf %[[EXTRACT]], %[[SQ]] : tensor<8x4xf32> + // CHECK-NEXT: %[[MUL:.*]] = arith.mulf %[[SQ]], %[[ADD]] : tensor<8x4xf32> + // CHECK-NOT: tt.trans + // CHECK-NEXT: return %[[MUL]] : tensor<8x4xf32> + %tile = xtile.extract %input[%offset0, %offset1][4, 8][1, 1] + : memref<4x8xf32, #xtile.layout<[1, 0]>> -> tensor<4x8xf32> + %sq = arith.mulf %tile, %tile : tensor<4x8xf32> + %add = arith.addf %tile, %sq : tensor<4x8xf32> + %mul = arith.mulf %sq, %add : tensor<4x8xf32> + %res = tt.trans %mul {order = array} : tensor<4x8xf32> -> tensor<8x4xf32> + return %res : tensor<8x4xf32> +} + +// CHECK-LABEL: func @push_transpose_up_through_if_multiple_uses +func.func @push_transpose_up_through_if_multiple_uses( + %arg0: tensor<4x8xf32>, %arg1: tensor<4x8xf32>, %cond: i1) -> tensor<8x4xf32> { + // CHECK-DAG: %[[TRANS0:.*]] = tt.trans %arg0 {order = array} : tensor<4x8xf32> -> tensor<8x4xf32> + // CHECK-DAG: %[[TRANS1:.*]] = tt.trans %arg1 {order = array} : tensor<4x8xf32> -> tensor<8x4xf32> + // CHECK: %[[IF:.*]] = scf.if %arg2 -> (tensor<8x4xf32>) { + // CHECK: scf.yield %[[TRANS0]] : tensor<8x4xf32> + // CHECK: } else { + // CHECK: scf.yield %[[TRANS1]] : tensor<8x4xf32> + // CHECK: } + // CHECK-NEXT: %[[SQ:.*]] = arith.mulf %[[IF]], %[[IF]] : tensor<8x4xf32> + // CHECK-NEXT: %[[ADD:.*]] = arith.addf %[[IF]], %[[SQ]] : tensor<8x4xf32> + // CHECK-NOT: tt.trans + // CHECK-NEXT: return %[[ADD]] : tensor<8x4xf32> + %0 = scf.if %cond -> tensor<4x8xf32> { + scf.yield %arg0 : tensor<4x8xf32> + } else { + scf.yield %arg1 : tensor<4x8xf32> + } + %sq = arith.mulf %0, %0 : tensor<4x8xf32> + %add = arith.addf %0, %sq : tensor<4x8xf32> + %res = tt.trans %add {order = array} : tensor<4x8xf32> -> tensor<8x4xf32> + return %res : tensor<8x4xf32> +} + +// CHECK-LABEL: func @do_not_push_transpose_when_producer_has_non_transpose_user +func.func @do_not_push_transpose_when_producer_has_non_transpose_user( + %arg0: tensor<4x8xf32>) -> (tensor<4x8xf32>, tensor<8x4xf32>) { + // CHECK-NEXT: %[[NEG:.*]] = arith.negf %arg0 : tensor<4x8xf32> + // CHECK-NEXT: %[[TRANS:.*]] = tt.trans %[[NEG]] {order = array} : tensor<4x8xf32> -> tensor<8x4xf32> + // CHECK-NEXT: return %[[NEG]], %[[TRANS]] : tensor<4x8xf32>, tensor<8x4xf32> + %0 = arith.negf %arg0 : tensor<4x8xf32> + %1 = tt.trans %0 {order = array} : tensor<4x8xf32> -> tensor<8x4xf32> + return %0, %1 : tensor<4x8xf32>, tensor<8x4xf32> +} diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/transforms/tests/triton_xla_remote_access.mlir b/third_party/xla/xla/backends/gpu/codegen/triton/transforms/tests/triton_xla_remote_access.mlir index 180108fe7f2b0a..313a6fd6265595 100644 --- a/third_party/xla/xla/backends/gpu/codegen/triton/transforms/tests/triton_xla_remote_access.mlir +++ b/third_party/xla/xla/backends/gpu/codegen/triton/transforms/tests/triton_xla_remote_access.mlir @@ -62,8 +62,8 @@ tt.func @get_peer_ptr( // Load metadata->buffer_root_ptrs[argument_offset + peer_id] + offset. // CHECK-NEXT: %11 = arith.addi %10, %6 : i64 // CHECK-NEXT: %12 = tt.int_to_ptr %11 : i64 -> !tt.ptr - %arg_0_peer_ptr = triton_xla.get_peer_ptr %arg0, %peer_id, %metadata, - { argument_index = 0 : i32, world_size = 2 : i32 } : + %arg_0_peer_ptr = triton_xla.get_peer_ptr %arg0, %peer_id, %metadata + : (!tt.ptr, i64, !tt.ptr) -> !tt.ptr // Load metadata->rank @@ -90,8 +90,8 @@ tt.func @get_peer_ptr( // CHECK-NEXT: %24 = arith.addi %23, %19 : i64 // CHECK-NEXT: %25 = tt.int_to_ptr %24 : i64 -> !tt.ptr - %arg_1_peer_ptr = triton_xla.get_peer_ptr %arg1, %peer_id, %metadata, - { argument_index = 1 : i32, world_size = 2 : i32 } : + %arg_1_peer_ptr = triton_xla.get_peer_ptr %arg1, %peer_id, %metadata + : (!tt.ptr, i64, !tt.ptr) -> !tt.ptr // Avoid optimizing away the get_peer_ptr calls, by returning xor of the two diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/transforms/tests/triton_xla_squeeze_dims.mlir b/third_party/xla/xla/backends/gpu/codegen/triton/transforms/tests/triton_xla_squeeze_dims.mlir index 63d0cbc0faa245..4e87ba0dc1048d 100644 --- a/third_party/xla/xla/backends/gpu/codegen/triton/transforms/tests/triton_xla_squeeze_dims.mlir +++ b/third_party/xla/xla/backends/gpu/codegen/triton/transforms/tests/triton_xla_squeeze_dims.mlir @@ -22,7 +22,7 @@ func.func @push_squeeze_dims_up_through_elementwise(%arg0: tensor<4x1x8xf32>) -> tensor<4x8xf32> { // CHECK: arith.negf {{.*}} : tensor<4x8xf32> %0 = arith.negf %arg0 : tensor<4x1x8xf32> - %1 = triton_xla.squeeze_dims %0 {axis = 1 : i32} : tensor<4x1x8xf32> -> tensor<4x8xf32> + %1 = triton_xla.squeeze_dims %0 : tensor<4x1x8xf32> -> tensor<4x8xf32> return %1 : tensor<4x8xf32> } @@ -32,7 +32,7 @@ func.func @push_squeeze_dims_up_through_elementwise(%arg0: tensor<4x1x8xf32>) -> func.func @push_squeeze_dims_up_through_multiple_results(%arg0: tensor<4x1x8xf32>) -> tensor<4x8xf32> { // CHECK: tt.elementwise_inline_asm {{.*}} : tensor<4x8xf32> -> tensor<4x8xf32>, tensor<4x8xf32> %0:2 = tt.elementwise_inline_asm "" {constraints = "", packed_element = 1 : i32, pure = true} %arg0 : tensor<4x1x8xf32> -> tensor<4x1x8xf32>, tensor<4x1x8xf32> - %1 = triton_xla.squeeze_dims %0#0 {axis = 1 : i32} : tensor<4x1x8xf32> -> tensor<4x8xf32> + %1 = triton_xla.squeeze_dims %0#0 : tensor<4x1x8xf32> -> tensor<4x8xf32> return %1 : tensor<4x8xf32> } @@ -42,7 +42,7 @@ func.func @push_squeeze_dims_up_through_multiple_results(%arg0: tensor<4x1x8xf32 func.func @push_squeeze_dims_up_through_broadcast(%arg0: tensor<1x4x1x8xf32>) -> tensor<4x16x8xf32> { // CHECK: tt.broadcast {{.*}} : tensor<4x1x8xf32> -> tensor<4x16x8xf32> %0 = tt.broadcast %arg0 : tensor<1x4x1x8xf32> -> tensor<1x4x16x8xf32> - %1 = triton_xla.squeeze_dims %0 {axis = 0 : i32} : tensor<1x4x16x8xf32> -> tensor<4x16x8xf32> + %1 = triton_xla.squeeze_dims %0 : tensor<1x4x16x8xf32> -> tensor<4x16x8xf32> return %1 : tensor<4x16x8xf32> } @@ -52,7 +52,7 @@ func.func @push_squeeze_dims_up_through_broadcast(%arg0: tensor<1x4x1x8xf32>) -> func.func @push_squeeze_dims_up_through_trans(%arg0: tensor<4x1x8xf32>) -> tensor<8x4xf32> { // CHECK: tt.trans {{.*}} {order = array} : tensor<4x8xf32> -> tensor<8x4xf32> %0 = tt.trans %arg0 {order = array} : tensor<4x1x8xf32> -> tensor<8x4x1xf32> - %1 = triton_xla.squeeze_dims %0 {axis = 2 : i32} : tensor<8x4x1xf32> -> tensor<8x4xf32> + %1 = triton_xla.squeeze_dims %0 : tensor<8x4x1xf32> -> tensor<8x4xf32> return %1 : tensor<8x4xf32> } @@ -62,7 +62,7 @@ func.func @push_squeeze_dims_up_through_trans(%arg0: tensor<4x1x8xf32>) -> tenso func.func @push_squeeze_dims_up_through_join(%arg0: tensor<1x4xf32>, %arg1: tensor<1x4xf32>) -> tensor<4x2xf32> { // CHECK-DAG: tt.join {{.*}} : tensor<4xf32> -> tensor<4x2xf32> %0 = tt.join %arg0, %arg1 : tensor<1x4xf32> -> tensor<1x4x2xf32> - %1 = triton_xla.squeeze_dims %0 {axis = 0 : i32} : tensor<1x4x2xf32> -> tensor<4x2xf32> + %1 = triton_xla.squeeze_dims %0 : tensor<1x4x2xf32> -> tensor<4x2xf32> return %1 : tensor<4x2xf32> } @@ -77,7 +77,7 @@ func.func @push_squeeze_dims_up_through_reduce(%arg0: tensor<8x4x1xf32>) -> tens tt.reduce.return %1 : f32 // CHECK: }) : (tensor<8x4xf32>) -> tensor<8xf32> }) : (tensor<8x4x1xf32>) -> tensor<8x1xf32> - %2 = triton_xla.squeeze_dims %0 {axis = 1 : i32} : tensor<8x1xf32> -> tensor<8xf32> + %2 = triton_xla.squeeze_dims %0 : tensor<8x1xf32> -> tensor<8xf32> return %2 : tensor<8xf32> } @@ -88,7 +88,7 @@ func.func @fold_squeeze_of_expand_cancelling(%arg0: tensor<4x8xf32>) -> tensor<4 // CHECK-NOT: tt.expand_dims // CHECK-NOT: triton_xla.squeeze_dims %0 = tt.expand_dims %arg0 {axis = 1 : i32} : tensor<4x8xf32> -> tensor<4x1x8xf32> - %1 = triton_xla.squeeze_dims %0 {axis = 1 : i32} : tensor<4x1x8xf32> -> tensor<4x8xf32> + %1 = triton_xla.squeeze_dims %0 : tensor<4x1x8xf32> -> tensor<4x8xf32> return %1 : tensor<4x8xf32> } @@ -96,10 +96,10 @@ func.func @fold_squeeze_of_expand_cancelling(%arg0: tensor<4x8xf32>) -> tensor<4 // CHECK-LABEL: func @fold_squeeze_of_expand_swapping func.func @fold_squeeze_of_expand_swapping(%arg0: tensor<4x1x8xf32>) -> tensor<1x4x8xf32> { - // CHECK: triton_xla.squeeze_dims {{.*}} {axis = 1 : i32} : tensor<4x1x8xf32> -> tensor<4x8xf32> + // CHECK: triton_xla.squeeze_dims {{.*}} : tensor<4x1x8xf32> -> tensor<4x8xf32> // CHECK: tt.expand_dims {{.*}} {axis = 0 : i32} : tensor<4x8xf32> -> tensor<1x4x8xf32> %0 = tt.expand_dims %arg0 {axis = 0 : i32} : tensor<4x1x8xf32> -> tensor<1x4x1x8xf32> - %1 = triton_xla.squeeze_dims %0 {axis = 2 : i32} : tensor<1x4x1x8xf32> -> tensor<1x4x8xf32> + %1 = triton_xla.squeeze_dims %0 : tensor<1x4x1x8xf32> -> tensor<1x4x8xf32> return %1 : tensor<1x4x8xf32> } @@ -107,8 +107,8 @@ func.func @fold_squeeze_of_expand_swapping(%arg0: tensor<4x1x8xf32>) -> tensor<1 // CHECK-LABEL: func @squeeze_reshape func.func @squeeze_reshape(%arg0: tensor<4x1x1xf32>) -> tensor<4xf32> { - // CHECK: triton_xla.squeeze_dims {{.*}} {axis = 1 : i32} : tensor<4x1x1xf32> -> tensor<4x1xf32> - // CHECK: triton_xla.squeeze_dims {{.*}} {axis = 1 : i32} : tensor<4x1xf32> -> tensor<4xf32> + // CHECK: triton_xla.squeeze_dims {{.*}} : tensor<4x1x1xf32> -> tensor<4x1xf32> + // CHECK: triton_xla.squeeze_dims {{.*}} : tensor<4x1xf32> -> tensor<4xf32> %0 = tt.reshape %arg0 : tensor<4x1x1xf32> -> tensor<4xf32> return %0 : tensor<4xf32> } @@ -152,7 +152,7 @@ module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 1 : i32} { // CHECK-LABEL: func @reshape_with_encoding // CHECK-SAME: tensor<1x32xf32, #[[ARG_ENC:.+]]>) -> tensor<32xf32, #[[RES_ENC:.+]]> func.func @reshape_with_encoding(%arg0: tensor<1x32xf32, #arg_enc>) -> tensor<32xf32, #res_enc> { - // CHECK: triton_xla.squeeze_dims {{.*}} {axis = 0 : i32} : + // CHECK: triton_xla.squeeze_dims {{.*}} : // CHECK-SAME: tensor<1x32xf32, #[[ARG_ENC]]> -> tensor<32xf32, #[[RES_ENC]]> %0 = tt.reshape %arg0 : tensor<1x32xf32, #arg_enc> -> tensor<32xf32, #res_enc> return %0 : tensor<32xf32, #res_enc> @@ -171,7 +171,7 @@ func.func @fold_squeeze_dims_of_extract( %tile = xtile.extract %input[%offset, %offset, %offset][4, 1, 8][1, 1, 1] : memref<4x16x8xf32, #xtile.layout<[2, 1, 0]>> -> tensor<4x1x8xf32> // CHECK-NOT: triton_xla.squeeze_dims - %squeezed = triton_xla.squeeze_dims %tile {axis = 1 : i32} : tensor<4x1x8xf32> -> tensor<4x8xf32> + %squeezed = triton_xla.squeeze_dims %tile : tensor<4x1x8xf32> -> tensor<4x8xf32> // CHECK: return %[[EXTRACT]] return %squeezed : tensor<4x8xf32> } @@ -187,7 +187,7 @@ func.func @squeeze_insert( %arg1: tensor<4x1x8xf32>, %offset: index) { // CHECK: %[[REDUCED:.*]] = triton_xla.squeeze_dims %[[TILE]] - // CHECK-SAME: {axis = 1 : i32} : tensor<4x1x8xf32> -> tensor<4x8xf32> + // CHECK-SAME: : tensor<4x1x8xf32> -> tensor<4x8xf32> // CHECK: xtile.insert %[[REDUCED]] into %[[BUFFER]] // CHECK-SAME: tensor<4x8xf32> -> memref<4x16x8xf32, #xtile.layout<[2, 1, 0]>> xtile.insert %arg1 into %arg0[%offset, %offset, %offset][4, 1, 8][1, 1, 1] @@ -213,10 +213,10 @@ func.func @squeeze_insert_unit_tensor( // CHECK-LABEL: func @reorder_squeeze_dims func.func @reorder_squeeze_dims(%arg0: tensor<4x1x8x1xf32>) -> tensor<4x8xf32> { - // CHECK: triton_xla.squeeze_dims {{.*}} {axis = 1 : i32} : tensor<4x1x8x1xf32> -> tensor<4x8x1xf32> - // CHECK: triton_xla.squeeze_dims {{.*}} {axis = 2 : i32} : tensor<4x8x1xf32> -> tensor<4x8xf32> - %0 = triton_xla.squeeze_dims %arg0 {axis = 3 : i32} : tensor<4x1x8x1xf32> -> tensor<4x1x8xf32> - %1 = triton_xla.squeeze_dims %0 {axis = 1 : i32} : tensor<4x1x8xf32> -> tensor<4x8xf32> + // CHECK: triton_xla.squeeze_dims {{.*}} : tensor<4x1x8x1xf32> -> tensor<4x8x1xf32> + // CHECK: triton_xla.squeeze_dims {{.*}} : tensor<4x8x1xf32> -> tensor<4x8xf32> + %0 = triton_xla.squeeze_dims %arg0 : tensor<4x1x8x1xf32> -> tensor<4x1x8xf32> + %1 = triton_xla.squeeze_dims %0 : tensor<4x1x8xf32> -> tensor<4x8xf32> return %1 : tensor<4x8xf32> } @@ -227,7 +227,7 @@ func.func @diamond(%arg0: tensor<4x1x8xf32>) -> tensor<4x8xf32> { // CHECK-NOT: arith.negf {{.*}} : tensor<4x1x8xf32> %0 = arith.negf %arg0 : tensor<4x1x8xf32> %1 = arith.addf %0, %0 : tensor<4x1x8xf32> - %2 = triton_xla.squeeze_dims %1 {axis = 1 : i32} : tensor<4x1x8xf32> -> tensor<4x8xf32> + %2 = triton_xla.squeeze_dims %1 : tensor<4x1x8xf32> -> tensor<4x8xf32> return %2 : tensor<4x8xf32> } @@ -246,7 +246,7 @@ func.func @insert_expand_dims(%arg0: tensor<4x1xf32>) -> tensor<4xf32> { }) : (tensor<4x1xf32>) -> tensor<4xf32> %3 = tt.expand_dims %1 {axis = 1 : i32} : tensor<4xf32> -> tensor<4x1xf32> %4 = arith.addf %0, %3 : tensor<4x1xf32> - %5 = triton_xla.squeeze_dims %4 {axis = 1 : i32} : tensor<4x1xf32> -> tensor<4xf32> + %5 = triton_xla.squeeze_dims %4 : tensor<4x1xf32> -> tensor<4xf32> return %5 : tensor<4xf32> } @@ -267,7 +267,7 @@ func.func @push_squeeze_dims_up_through_if( scf.yield %arg0, %arg1 : tensor<16xf32>, tensor<4x1xf32> } // CHECK-NOT: triton_xla.squeeze_dims - %squeeze = triton_xla.squeeze_dims %if#1 {axis = 1 : i32} + %squeeze = triton_xla.squeeze_dims %if#1 : tensor<4x1xf32> -> tensor<4xf32> // CHECK: return return %if#0, %squeeze : tensor<16xf32>, tensor<4xf32> @@ -278,7 +278,7 @@ func.func @push_squeeze_dims_up_through_if( // FINALIZE-LABEL: func @squeeze_dims_to_reshape func.func @squeeze_dims_to_reshape(%arg0: tensor<4x1x8xf32>) -> tensor<4x8xf32> { // FINALIZE: tt.reshape {{.*}} : tensor<4x1x8xf32> -> tensor<4x8xf32> - %0 = triton_xla.squeeze_dims %arg0 {axis = 1 : i32} : tensor<4x1x8xf32> -> tensor<4x8xf32> + %0 = triton_xla.squeeze_dims %arg0 : tensor<4x1x8xf32> -> tensor<4x8xf32> return %0 : tensor<4x8xf32> } @@ -289,6 +289,6 @@ func.func @push_squeeze_dims_up_through_mask( %arg0: tensor<4x1x8xf32>, %arg1: f32) -> tensor<4x8xf32> { // CHECK: xtile.mask %{{.*}} bounds [4, 6], %arg1 : tensor<4x8xf32> %0 = xtile.mask %arg0 bounds [4, 1, 6], %arg1 : tensor<4x1x8xf32> - %1 = triton_xla.squeeze_dims %0 {axis = 1 : i32} : tensor<4x1x8xf32> -> tensor<4x8xf32> + %1 = triton_xla.squeeze_dims %0 : tensor<4x1x8xf32> -> tensor<4x8xf32> return %1 : tensor<4x8xf32> } diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/transforms/tests/xtile_to_triton_lowering.mlir b/third_party/xla/xla/backends/gpu/codegen/triton/transforms/tests/xtile_to_triton_lowering.mlir index 1657db56a13535..3bf693607595ff 100644 --- a/third_party/xla/xla/backends/gpu/codegen/triton/transforms/tests/xtile_to_triton_lowering.mlir +++ b/third_party/xla/xla/backends/gpu/codegen/triton/transforms/tests/xtile_to_triton_lowering.mlir @@ -26,8 +26,8 @@ func.func @lower_dot_scaled_add_to_triton( // CHECK: %[[RES:.*]] = tt.dot_scaled %[[LHS]] scale %[[LHS_SCALE]], %[[RHS]] scale %[[RHS_SCALE]], %[[ACC]] lhs = e5m2 rhs = e5m2 {fastMath = true} : tensor<128x128xf8E5M2>, tensor<128x4xi8> * tensor<128x256xf8E5M2>, tensor<256x4xi8> -> tensor<128x256xf32> // CHECK-NOT: arith.addf %0 = xtile.dot_scaled %lhs scale %lhs_scale, %rhs scale %rhs_scale - {dot_dimension_numbers = #stablehlo.dot, - fastMath = true, lhs_elem_type = f8E5M2, rhs_elem_type = f8E5M2} : tensor<128x128xf8E5M2>, + , + fastMath = true, lhs_elem_type = f8E5M2, rhs_elem_type = f8E5M2> : tensor<128x128xf8E5M2>, tensor<128x4xi8> * tensor<128x256xf8E5M2>, tensor<256x4xi8> -> tensor<128x256xf32> %1 = arith.addf %acc, %0 : tensor<128x256xf32> // CHECK: return %[[RES]] : tensor<128x256xf32> @@ -44,8 +44,8 @@ func.func @lower_packed_dot_scaled_add_to_triton( // CHECK: %[[RES:.*]] = tt.dot_scaled %[[LHS]] scale %[[LHS_SCALE]], %[[RHS]] scale %[[RHS_SCALE]], %[[ACC]] lhs = e2m1 rhs = e2m1 {fastMath = true} : tensor<128x32xi8>, tensor<128x2xi8> * tensor<32x128xi8>, tensor<128x2xi8> -> tensor<128x128xf32> // CHECK-NOT: arith.addf %0 = xtile.dot_scaled %lhs scale %lhs_scale, %rhs scale %rhs_scale - {dot_dimension_numbers = #stablehlo.dot, - fastMath = true, lhs_elem_type = f4E2M1FN, rhs_elem_type = f4E2M1FN} : tensor<128x32xi8>, + , + fastMath = true, lhs_elem_type = f4E2M1FN, rhs_elem_type = f4E2M1FN> : tensor<128x32xi8>, tensor<128x2xi8> * tensor<32x128xi8>, tensor<128x2xi8> -> tensor<128x128xf32> %1 = arith.addf %acc, %0 : tensor<128x128xf32> // CHECK: return %[[RES]] : tensor<128x128xf32> @@ -67,8 +67,8 @@ func.func @lower_dot_scaled_in_loop_non_canonical( // CHECK: %[[DOT:.*]] = tt.dot_scaled %{{.*}} scale %{{.*}}, %{{.*}} scale %{{.*}}, %[[ACC_2D]] // CHECK: scf.yield %[[DOT]] : tensor<128x256xf32> %0 = xtile.dot_scaled %lhs scale %lhs_scale, %rhs scale %rhs_scale - {dot_dimension_numbers = #stablehlo.dot, - fastMath = true, lhs_elem_type = f8E5M2, rhs_elem_type = f8E5M2} : tensor<1x128x128xf8E5M2>, tensor<1x128x4xi8> * tensor<1x128x256xf8E5M2>, tensor<1x256x4xi8> -> tensor<1x128x256xf32> + , + fastMath = true, lhs_elem_type = f8E5M2, rhs_elem_type = f8E5M2> : tensor<1x128x128xf8E5M2>, tensor<1x128x4xi8> * tensor<1x128x256xf8E5M2>, tensor<1x256x4xi8> -> tensor<1x128x256xf32> %1 = arith.addf %accum, %0 : tensor<1x128x256xf32> scf.yield %1 : tensor<1x128x256xf32> } @@ -92,7 +92,7 @@ func.func @scan_lowering(%input0: tensor<16x16x16xf32>, %init0: tensor<16x16x1xf // CHECK-DAG: %[[RES1:.*]] = arith.addf %[[BCAST_INIT0]], %[[SCAN]] : tensor<16x16x16xf32> // CHECK: return %[[RES1]] : tensor<16x16x16xf32> - %0, %1 = xtile.scan(%input0) inits(%init0) dimension = 2 {scan_dim_size = 16 : i64} : (tensor<16x16x16xf32>), (tensor<16x16x1xf32>) -> (tensor<16x16x16xf32>), (tensor<16x16x1xf32>) { + %0, %1 = xtile.scan(%input0) inits(%init0) dimension = 2 : (tensor<16x16x16xf32>), (tensor<16x16x1xf32>) -> (tensor<16x16x16xf32>), (tensor<16x16x1xf32>) { ^bb0(%arg0: f32, %arg1: f32): %add = arith.addf %arg0, %arg1 : f32 xtile.yield %add : f32 @@ -131,7 +131,7 @@ func.func @scan_lowering_carry(%input0: tensor<16x16x16xf32>, %init0: tensor<16x // CHECK-NOT: tt.reshape // CHECK: return %[[RES0]], %[[RES1]] - %0, %1 = xtile.scan(%input0) inits(%init0) dimension = 2 {scan_dim_size = 16 : i64} : (tensor<16x16x16xf32>), (tensor<16x16x1xf32>) -> (tensor<16x16x16xf32>), (tensor<16x16x1xf32>) { + %0, %1 = xtile.scan(%input0) inits(%init0) dimension = 2 : (tensor<16x16x16xf32>), (tensor<16x16x1xf32>) -> (tensor<16x16x16xf32>), (tensor<16x16x1xf32>) { ^bb0(%arg0: f32, %arg1: f32): %add = arith.addf %arg0, %arg1 : f32 xtile.yield %add : f32 @@ -167,7 +167,7 @@ func.func @scan_lowering_1d_carry(%input0: tensor<16xf32>, %init0: tensor<1xf32> // CHECK-NOT: tt.gather // CHECK-NOT: tt.reshape // CHECK: return %[[RES0]], %[[RES1]] - %0, %1 = xtile.scan(%input0) inits(%init0) dimension = 0 {scan_dim_size = 16 : i64} : (tensor<16xf32>), (tensor<1xf32>) -> (tensor<16xf32>), (tensor<1xf32>) { + %0, %1 = xtile.scan(%input0) inits(%init0) dimension = 0 : (tensor<16xf32>), (tensor<1xf32>) -> (tensor<16xf32>), (tensor<1xf32>) { ^bb0(%arg0: f32, %arg1: f32): %add = arith.addf %arg0, %arg1 : f32 xtile.yield %add : f32 @@ -198,7 +198,7 @@ func.func @reverse_scan_lowering_carry(%input0: tensor<16xf32>, %init0: tensor<1 // CHECK-NOT: tt.gather // CHECK-NOT: tt.reshape // CHECK: return %[[RES0]], %[[RES1]] - %0, %1 = xtile.scan(%input0) inits(%init0) dimension = 0 {scan_dim_size = 16 : i64, is_reverse = true} : (tensor<16xf32>), (tensor<1xf32>) -> (tensor<16xf32>), (tensor<1xf32>) { + %0, %1 = xtile.scan(%input0) inits(%init0) dimension = 0 : (tensor<16xf32>), (tensor<1xf32>) -> (tensor<16xf32>), (tensor<1xf32>) { ^bb0(%arg0: f32, %arg1: f32): %add = arith.addf %arg0, %arg1 : f32 xtile.yield %add : f32 @@ -215,8 +215,8 @@ func.func @lower_dot_scaled_without_add_falls_back_to_xtile( -> tensor<128x256xf32> { // CHECK: %[[RES:.*]] = xtile.dot_scaled %[[LHS]] scale %[[LHS_SCALE]], %[[RHS]] scale %[[RHS_SCALE]] {{.*}}fastMath = true{{.*}} : tensor<128x128xf8E5M2>, tensor<128x4xi8> * tensor<128x256xf8E5M2>, tensor<256x4xi8> -> tensor<128x256xf32> %0 = xtile.dot_scaled %lhs scale %lhs_scale, %rhs scale %rhs_scale - {dot_dimension_numbers = #stablehlo.dot, - fastMath = true, lhs_elem_type = f8E5M2, rhs_elem_type = f8E5M2} : tensor<128x128xf8E5M2>, + , + fastMath = true, lhs_elem_type = f8E5M2, rhs_elem_type = f8E5M2> : tensor<128x128xf8E5M2>, tensor<128x4xi8> * tensor<128x256xf8E5M2>, tensor<256x4xi8> -> tensor<128x256xf32> // CHECK: return %[[RES]] : tensor<128x256xf32> return %0 : tensor<128x256xf32> diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/transforms/triton_xla_fold_transpose_pass.cc b/third_party/xla/xla/backends/gpu/codegen/triton/transforms/triton_xla_fold_transpose_pass.cc index 80d77ef1fcf3cc..790918f99f2f31 100644 --- a/third_party/xla/xla/backends/gpu/codegen/triton/transforms/triton_xla_fold_transpose_pass.cc +++ b/third_party/xla/xla/backends/gpu/codegen/triton/transforms/triton_xla_fold_transpose_pass.cc @@ -20,6 +20,7 @@ limitations under the License. #include "absl/algorithm/container.h" #include "llvm/ADT/ArrayRef.h" +#include "llvm/ADT/DenseMap.h" #include "llvm/ADT/STLExtras.h" #include "llvm/ADT/SmallVector.h" #include "llvm/Support/ErrorHandling.h" @@ -57,6 +58,24 @@ namespace { return guard; } +bool CanFoldTransposeThroughProducer(TransOp op) { + return absl::c_all_of(op.getSrc().getUsers(), [&](Operation* user) { + auto other_trans = dyn_cast(user); + return other_trans && other_trans.getOrder() == op.getOrder(); + }); +} + +void ReplaceAllTransposeUsersWith(PatternRewriter& rewriter, Value old_src, + Value new_src) { + for (Operation* user : llvm::make_early_inc_range(old_src.getUsers())) { + rewriter.replaceOp(user, new_src); + } + if (Operation* producer = old_src.getDefiningOp(); + producer && producer->use_empty()) { + rewriter.eraseOp(producer); + } +} + // Push the transpose up through the extract tile, this will then be folded into // MemrefToPtr at the lowering stage. LogicalResult PushTransposeThroughExtractTile(TransOp op, @@ -65,6 +84,10 @@ LogicalResult PushTransposeThroughExtractTile(TransOp op, if (!extract) { return rewriter.notifyMatchFailure(op, "Transpose source is not extract."); } + if (!CanFoldTransposeThroughProducer(op)) { + return rewriter.notifyMatchFailure( + op, "Transpose source has users that are not matching transposes."); + } SmallVector reduced_dims = to_vector(extract.getReducedDimensions()); @@ -107,19 +130,18 @@ LogicalResult PushTransposeThroughExtractTile(TransOp op, return result; }; + OpBuilder::InsertionGuard guard = SetInsertionPoint(rewriter, extract); auto permutation_map = mlir::AffineMapAttr::get( mlir::AffineMap::getPermutationMap(permutation, rewriter.getContext())); // TODO(willfroom): Return a permutation layout (b/455478641). auto pushed_transpose = mlir::memref::TransposeOp::create( rewriter, extract.getLoc(), extract.getSource(), permutation_map); - rewriter.replaceOpWithNewOp<::xla::xtile::ExtractTileOp>( - op, op.getType(), pushed_transpose, permute(extract.getOffsets()), - permute(extract.getFullTileShape()), permute(extract.getStrides())); - - if (extract->use_empty()) { - rewriter.eraseOp(extract); - } + Value new_extract = ::xla::xtile::ExtractTileOp::create( + rewriter, extract.getLoc(), op.getType(), pushed_transpose, + permute(extract.getOffsets()), permute(extract.getFullTileShape()), + permute(extract.getStrides())); + ReplaceAllTransposeUsersWith(rewriter, op.getSrc(), new_extract); return success(); } @@ -131,9 +153,16 @@ LogicalResult PushTransposeUpThroughBroadcast(TransOp op, return rewriter.notifyMatchFailure( // op, "Transpose source is not a broadcast."); } - Value new_trans = TransOp::create(rewriter, op.getLoc(), broadcast.getSrc(), - op.getOrderAttr()); - rewriter.replaceOpWithNewOp(op, op.getType(), new_trans); + if (!CanFoldTransposeThroughProducer(op)) { + return rewriter.notifyMatchFailure( + op, "Transpose source has users that are not matching transposes."); + } + OpBuilder::InsertionGuard guard = SetInsertionPoint(rewriter, broadcast); + Value new_trans = TransOp::create(rewriter, broadcast.getLoc(), + broadcast.getSrc(), op.getOrderAttr()); + Value new_broadcast = BroadcastOp::create(rewriter, broadcast.getLoc(), + op.getType(), new_trans); + ReplaceAllTransposeUsersWith(rewriter, op.getSrc(), new_broadcast); return success(); } @@ -144,6 +173,10 @@ LogicalResult PushTransposeUpThroughExpandDims(TransOp op, return rewriter.notifyMatchFailure( op, "Transpose source is not an expand_dims."); } + if (!CanFoldTransposeThroughProducer(op)) { + return rewriter.notifyMatchFailure( + op, "Transpose source has users that are not matching transposes."); + } unsigned new_axis = [&] { for (auto [i, dim] : llvm::enumerate(op.getOrder())) { @@ -160,10 +193,12 @@ LogicalResult PushTransposeUpThroughExpandDims(TransOp op, dim -= dim > expand_dims.getAxis(); } - Value new_trans = - TransOp::create(rewriter, op.getLoc(), expand_dims.getSrc(), new_order); - rewriter.replaceOpWithNewOp(op, op.getType(), new_trans, - new_axis); + OpBuilder::InsertionGuard guard = SetInsertionPoint(rewriter, expand_dims); + Value new_trans = TransOp::create(rewriter, expand_dims.getLoc(), + expand_dims.getSrc(), new_order); + Value new_expand_dims = ExpandDimsOp::create( + rewriter, expand_dims.getLoc(), op.getType(), new_trans, new_axis); + ReplaceAllTransposeUsersWith(rewriter, op.getSrc(), new_expand_dims); return success(); } @@ -175,13 +210,23 @@ LogicalResult PushTransposeUpThroughElementwise(TransOp op, return rewriter.notifyMatchFailure( op, "source is not a single-result elementwise op"); } + if (!CanFoldTransposeThroughProducer(op)) { + return rewriter.notifyMatchFailure( + op, "Transpose source has users that are not matching transposes."); + } + OpBuilder::InsertionGuard guard = SetInsertionPoint(rewriter, elementwise); + llvm::SmallDenseMap transposed_operands; SmallVector new_operands; new_operands.reserve(elementwise->getNumOperands()); for (Value operand : elementwise->getOperands()) { - if (auto tensor_type = dyn_cast(operand.getType())) { - operand = TransOp::create(rewriter, elementwise->getLoc(), operand, - op.getOrderAttr()); + if (isa(operand.getType())) { + auto [it, inserted] = transposed_operands.try_emplace(operand, nullptr); + if (inserted) { + it->second = TransOp::create(rewriter, elementwise->getLoc(), operand, + op.getOrderAttr()); + } + operand = it->second; } new_operands.push_back(operand); } @@ -189,7 +234,7 @@ LogicalResult PushTransposeUpThroughElementwise(TransOp op, Operation* new_op = rewriter.clone(*elementwise); new_op->setOperands(new_operands); new_op->getResult(0).setType(op.getType()); - rewriter.replaceOp(op, new_op->getResults()); + ReplaceAllTransposeUsersWith(rewriter, op.getSrc(), new_op->getResult(0)); return success(); } @@ -213,19 +258,24 @@ LogicalResult PushTransposeUpThroughElementwise(TransOp op, LogicalResult PushTransposeUpIntoIf(TransOp op, PatternRewriter& rewriter) { Value src = op.getSrc(); auto if_op = src.getDefiningOp(); - if (!if_op || !src.hasOneUse()) { + if (!if_op) { return rewriter.notifyMatchFailure(op, "Expected scf.if producer."); } + if (!CanFoldTransposeThroughProducer(op)) { + return rewriter.notifyMatchFailure( + op, "Transpose source has users that are not matching transposes."); + } // Compute the new types for the if op. unsigned result_number = cast(src).getResultNumber(); auto new_types = llvm::to_vector(if_op.getResultTypes()); new_types[result_number] = op.getType(); - auto new_if_op = - scf::IfOp::create(rewriter, op.getLoc(), new_types, if_op.getCondition(), - /*addThenBlock=*/false, - /*addElseBlock=*/false); + OpBuilder::InsertionGuard guard = SetInsertionPoint(rewriter, if_op); + auto new_if_op = scf::IfOp::create(rewriter, if_op.getLoc(), new_types, + if_op.getCondition(), + /*addThenBlock=*/false, + /*addElseBlock=*/false); // Update then and else regions. for (auto [old_region, new_region] : @@ -235,13 +285,16 @@ LogicalResult PushTransposeUpIntoIf(TransOp op, PatternRewriter& rewriter) { continue; } auto yield_op = new_region->front().getTerminator(); - OpBuilder::InsertionGuard guard = SetInsertionPoint(rewriter, yield_op); + OpBuilder::InsertionGuard yield_guard = + SetInsertionPoint(rewriter, yield_op); auto trans_op = - TransOp::create(rewriter, op.getLoc(), op.getType(), + TransOp::create(rewriter, if_op.getLoc(), op.getType(), yield_op->getOperand(result_number), op.getOrderAttr()); yield_op->setOperand(result_number, trans_op); } - rewriter.replaceOp(op, new_if_op.getResult(result_number)); + for (Operation* user : llvm::make_early_inc_range(src.getUsers())) { + rewriter.replaceOp(user, new_if_op.getResult(result_number)); + } rewriter.replaceOp(if_op, new_if_op); return success(); } @@ -255,7 +308,7 @@ LogicalResult HoistTransposeUpFromIf(TransOp op, PatternRewriter& rewriter) { return rewriter.notifyMatchFailure(op, "Operand defined inside scf.if."); } - op->moveBefore(if_op); + rewriter.modifyOpInPlace(op, [&] { op->moveBefore(if_op); }); return success(); } @@ -274,6 +327,10 @@ LogicalResult PushTransposeUpThroughReshape(TransOp op, return rewriter.notifyMatchFailure(op, "Transpose source is not a reshape."); } + if (!CanFoldTransposeThroughProducer(op)) { + return rewriter.notifyMatchFailure( + op, "Transpose source has users that are not matching transposes."); + } auto operand_shape = reshape.getSrc().getType().getShape(); auto reshape_shape = reshape.getType().getShape(); @@ -311,9 +368,12 @@ LogicalResult PushTransposeUpThroughReshape(TransOp op, } } + OpBuilder::InsertionGuard guard = SetInsertionPoint(rewriter, reshape); auto new_trans = TransOp::create(rewriter, reshape.getLoc(), reshape.getSrc(), new_order); - rewriter.replaceOpWithNewOp(op, op.getType(), new_trans); + Value new_reshape = + ReshapeOp::create(rewriter, reshape.getLoc(), op.getType(), new_trans); + ReplaceAllTransposeUsersWith(rewriter, op.getSrc(), new_reshape); return success(); } @@ -323,17 +383,24 @@ LogicalResult PushTransposeUpThroughMask(TransOp op, if (!mask_op) { return rewriter.notifyMatchFailure(op, "source is not a mask op"); } + if (!CanFoldTransposeThroughProducer(op)) { + return rewriter.notifyMatchFailure( + op, "Transpose source has users that are not matching transposes."); + } llvm::SmallVector new_bounds(op.getOrder().size()); for (auto [idx, dim] : llvm::enumerate(op.getOrder())) { new_bounds[idx] = mask_op.getBounds()[dim]; } - auto new_transpose = TransOp::create(rewriter, op.getLoc(), + OpBuilder::InsertionGuard guard = SetInsertionPoint(rewriter, mask_op); + auto new_transpose = TransOp::create(rewriter, mask_op.getLoc(), mask_op.getSource(), op.getOrderAttr()); - rewriter.replaceOpWithNewOp<::xla::xtile::MaskOp>( - op, op.getType(), new_transpose, new_bounds, mask_op.getValue()); + Value new_mask = ::xla::xtile::MaskOp::create(rewriter, mask_op.getLoc(), + op.getType(), new_transpose, + new_bounds, mask_op.getValue()); + ReplaceAllTransposeUsersWith(rewriter, op.getSrc(), new_mask); return success(); } 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 41f48e48010301..f53e2aa9a7f030 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 @@ -52,6 +52,7 @@ LogicalResult LowerBlockBarrierOp(BlockBarrierOp block_barrier, block_barrier.getSignalValue(); const int32_t world_size = block_barrier.getWorldSize(); const int32_t signal_stride = block_barrier.getSignalStride(); + const BarrierMode barrier_mode = block_barrier.getBarrierMode(); // Triton magic constant. constexpr auto kGlobalAddressSpace = mlir::triton::PtrAddrSpace::Global; @@ -128,7 +129,8 @@ LogicalResult LowerBlockBarrierOp(BlockBarrierOp block_barrier, auto block_offset = mlir::arith::MulIOp::create( builder, i32_type, block_id, world_size_op); mlir::Value write_offset; - if (signal_stride > 0) { + mlir::Value wait_block_offset = block_offset; + if (barrier_mode != BarrierMode::kSymmetric) { // Suppose we have 4 blocks (0, 1, 2, 3), 2 ranks (world_size = 2), // and signal_stride = 2. // @@ -144,12 +146,14 @@ LogicalResult LowerBlockBarrierOp(BlockBarrierOp block_barrier, // - 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`. + // 2. Find `target_block` (`base_block + rank * signal_stride`): + // - In `kConsumerSymmetric`, `target_block` is the block on peer + // GPUs that reads this rank's data, so we write to `signal_slot` + // inside `target_block`'s slots + // (`target_block * world_size + signal_slot`). + // - In `kProducerSymmetric`, `target_block` is the producer block + // whose slots (`target_block * world_size + [0..world_size)`) + // this consumer block waits on. mlir::Value signal_slot = block_barrier.getSignalSlot(); auto stride_op = mlir::arith::ConstantOp::create( builder, builder.getI32IntegerAttr(signal_stride)); @@ -163,8 +167,14 @@ LogicalResult LowerBlockBarrierOp(BlockBarrierOp block_barrier, 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); + if (barrier_mode == BarrierMode::kConsumerSymmetric) { + 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); + wait_block_offset = target_block_offset; + } } else { write_offset = mlir::arith::AddIOp::create(builder, i32_type, block_offset, rank); @@ -201,10 +211,10 @@ LogicalResult LowerBlockBarrierOp(BlockBarrierOp block_barrier, // -> !tt.ptr auto read_address = mlir::triton::IntToPtrOp::create( builder, ptr_to_i32_type, read_address_i64); - // Pointer to SignalBuffers[rank][block_id] + // Pointer to SignalBuffers[rank][wait_block_id] // -> !tt.ptr auto read_address_at_block_offset = mlir::triton::AddPtrOp::create( - builder, ptr_to_i32_type, read_address, block_offset); + builder, ptr_to_i32_type, read_address, wait_block_offset); // -> tensor> auto read_address_at_block_offset_tensor = mlir::triton::SplatOp::create(builder, tensor_of_ptr_to_i32_type, diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/xtile_compiler.cc b/third_party/xla/xla/backends/gpu/codegen/triton/xtile_compiler.cc index d52453328a497e..7cf600930a5697 100644 --- a/third_party/xla/xla/backends/gpu/codegen/triton/xtile_compiler.cc +++ b/third_party/xla/xla/backends/gpu/codegen/triton/xtile_compiler.cc @@ -447,14 +447,6 @@ absl::StatusOr CreateTritonModule( absl::MakeSpan(opaque_args_types), mlir_context, use_experimental_tiling, enable_same_shape_multi_output_fusion)); - if (fusion_kind == kTritonCollectiveFusionKind && - CreateCollectiveCodegenConfig(&fusion).emit_entry_barrier) { - const HloInstruction* root = hlo_computation->root_instruction(); - int32_t world_size = root->replica_groups()[0].replica_ids_size(); - ABSL_RETURN_IF_ERROR( - EmitCollectiveEntryBarrier(triton_module.get(), world_size)); - } - if (DumpingEnabledForHloModule(*hlo_computation->parent()) && DumpingEnabledForEmitter("triton-fusion", debug_options)) { auto suffix = absl::StrCat(fusion.name(), ".before_validation.ttir.txt"); diff --git a/third_party/xla/xla/backends/gpu/runtime/BUILD b/third_party/xla/xla/backends/gpu/runtime/BUILD index 8fcdd36e3b9583..f5c3acd8945ea8 100644 --- a/third_party/xla/xla/backends/gpu/runtime/BUILD +++ b/third_party/xla/xla/backends/gpu/runtime/BUILD @@ -5433,6 +5433,64 @@ xla_cc_test( ], ) +cc_library( + name = "device_slot", + hdrs = ["device_slot.h"], + visibility = ["//visibility:private"], + deps = [ + "@com_google_absl//absl/base:core_headers", + "@com_google_absl//absl/functional:function_ref", + ], +) + +cc_library( + name = "vector_storage", + srcs = ["vector_storage.cc"], + hdrs = ["vector_storage.h"], + visibility = ["//visibility:private"], + deps = [":device_slot"], +) + +xla_cc_test( + name = "vector_storage_test", + srcs = ["vector_storage_test.cc"], + deps = [ + ":device_slot", + ":vector_storage", + "@com_google_absl//absl/base:core_headers", + "@com_google_absl//absl/container:flat_hash_set", + "@com_google_googletest//:gtest_main", + ], +) + +cc_library( + name = "cow_storage", + srcs = ["cow_storage.cc"], + hdrs = ["cow_storage.h"], + visibility = ["//visibility:private"], + deps = [ + ":device_slot", + "@com_google_absl//absl/base:core_headers", + "@com_google_absl//absl/container:inlined_vector", + "@com_google_absl//absl/synchronization", + ], +) + +xla_cc_test( + name = "cow_storage_test", + srcs = ["cow_storage_test.cc"], + deps = [ + ":cow_storage", + ":device_slot", + "//xla/tsl/platform:env", + "@com_google_absl//absl/base:core_headers", + "@com_google_absl//absl/container:flat_hash_set", + "@com_google_absl//absl/synchronization", + "@com_google_absl//absl/time", + "@com_google_googletest//:gtest_main", + ], +) + cc_library( name = "custom_kernel_thunk", srcs = ["custom_kernel_thunk.cc"], diff --git a/third_party/xla/xla/backends/gpu/runtime/all_gather.cc b/third_party/xla/xla/backends/gpu/runtime/all_gather.cc index 7ac59489518d7a..f0e781ec3bb726 100644 --- a/third_party/xla/xla/backends/gpu/runtime/all_gather.cc +++ b/third_party/xla/xla/backends/gpu/runtime/all_gather.cc @@ -232,7 +232,6 @@ absl::StatusOr CreateAllGatherKernelSpec( CollectiveKernelSpec kernel_spec = { /* .codegen_config= */ { /* .copy_input_to_scratch= */ true, - /* .emit_entry_barrier= */ true, /* .input_buffer_specs= */ {{/*requires_multimem=*/false, SymmetricMemoryType::kNone}}, /* .output_buffer_specs= */ diff --git a/third_party/xla/xla/backends/gpu/runtime/all_reduce.cc b/third_party/xla/xla/backends/gpu/runtime/all_reduce.cc index 8887f57cea1780..627337b25436bc 100644 --- a/third_party/xla/xla/backends/gpu/runtime/all_reduce.cc +++ b/third_party/xla/xla/backends/gpu/runtime/all_reduce.cc @@ -466,7 +466,6 @@ absl::StatusOr CreateAllReduceKernelSpec( CollectiveKernelSpec kernel_spec = { /* .codegen_config= */ { /* .copy_input_to_scratch= */ false, - /* .emit_entry_barrier= */ false, /* .input_buffer_specs= */ {{/*requires_multimem=*/false, SymmetricMemoryType::kNone}}, /* .output_buffer_specs= */ diff --git a/third_party/xla/xla/backends/gpu/runtime/collective_kernel_thunk_test.cc b/third_party/xla/xla/backends/gpu/runtime/collective_kernel_thunk_test.cc index 79192758a8ddb9..03b534f1da4d8f 100644 --- a/third_party/xla/xla/backends/gpu/runtime/collective_kernel_thunk_test.cc +++ b/third_party/xla/xla/backends/gpu/runtime/collective_kernel_thunk_test.cc @@ -203,7 +203,6 @@ CollectiveKernelSpec CreateCollectiveKernelSpec( return { /*codegen_config=*/{ /*copy_input_to_scratch=*/false, - /*emit_entry_barrier=*/false, /*input_buffer_specs=*/ {{/*requires_multimem=*/false, SymmetricMemoryType::kNone}}, /*output_buffer_specs=*/ diff --git a/third_party/xla/xla/backends/gpu/runtime/collective_params.h b/third_party/xla/xla/backends/gpu/runtime/collective_params.h index 2c2e97224e1f13..99fc121d4319b8 100644 --- a/third_party/xla/xla/backends/gpu/runtime/collective_params.h +++ b/third_party/xla/xla/backends/gpu/runtime/collective_params.h @@ -165,9 +165,6 @@ struct CollectiveCodegenConfig { // buffer before kernel launch. The kernel receives the scratch buffer as // its input argument. bool copy_input_to_scratch = false; - // If true, a cross-rank barrier is emitted before the tile loop. The barrier - // waits until all ranks have populated the symmetric scratch buffers. - bool emit_entry_barrier = false; // Specs for input operand buffers. std::vector input_buffer_specs; // Specs for output result buffers. 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 fbc9a936b4d288..35618e2b64f268 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 @@ -204,6 +204,8 @@ std::optional GetCommandBufferCmdType( case Thunk::kPartitionId: case Thunk::kReplicaId: return DebugOptions::FUSION; + case Thunk::kCollectiveKernel: + return DebugOptions::COLLECTIVES_KERNEL; case Thunk::kWhile: return DebugOptions::WHILE; case Thunk::kConditional: @@ -214,7 +216,6 @@ std::optional GetCommandBufferCmdType( case Thunk::kAllReduce: case Thunk::kAllToAll: case Thunk::kCollectiveBroadcast: - case Thunk::kCollectiveKernel: case Thunk::kCollectivePermute: case Thunk::kCollectiveReduce: case Thunk::kGroup: diff --git a/third_party/xla/xla/backends/gpu/runtime/cow_storage.cc b/third_party/xla/xla/backends/gpu/runtime/cow_storage.cc new file mode 100644 index 00000000000000..9b86fc78b28fcb --- /dev/null +++ b/third_party/xla/xla/backends/gpu/runtime/cow_storage.cc @@ -0,0 +1,78 @@ +/* Copyright 2026 The OpenXLA Authors. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#include "xla/backends/gpu/runtime/cow_storage.h" + +#include +#include +#include +#include + +#include "absl/container/inlined_vector.h" +#include "absl/synchronization/mutex.h" +#include "xla/backends/gpu/runtime/device_slot.h" + +namespace xla::gpu { + +DeviceSlot* CowStorage::Find(int device_ordinal) const { + if (device_ordinal < 0) { + return nullptr; + } + const Snapshot* snapshot = snapshot_.load(std::memory_order_acquire); + if (snapshot == nullptr) { + return nullptr; + } + return FindIn(*snapshot, device_ordinal); +} + +DeviceSlot* CowStorage::GetOrCreate(int device_ordinal, + DeviceSlotFactoryRef factory) { + if (device_ordinal < 0) { + return nullptr; + } + if (DeviceSlot* slot = Find(device_ordinal)) { + return slot; + } + absl::MutexLock lock(mutex_); + // Re-check under `mutex_` in case another thread concurrently created the + // slot for `device_ordinal` after the lock-free `Find` above. Snapshots are + // only published under `mutex_`, so relaxed is enough. + const Snapshot* current = snapshot_.load(std::memory_order_relaxed); + if (current != nullptr) { + if (DeviceSlot* slot = FindIn(*current, device_ordinal)) { + return slot; + } + } + slots_.push_back(factory()); + std::unique_ptr next = current == nullptr + ? std::make_unique() + : std::make_unique(*current); + next->push_back(Entry{device_ordinal, slots_.back().get()}); + snapshots_.push_back(std::move(next)); + const Snapshot* published = snapshots_.back().get(); + snapshot_.store(published, std::memory_order_release); + return published->back().slot; +} + +DeviceSlot* CowStorage::FindIn(const Snapshot& snapshot, int device_ordinal) { + for (const Entry& entry : snapshot) { + if (entry.device_ordinal == device_ordinal) { + return entry.slot; + } + } + return nullptr; +} + +} // namespace xla::gpu diff --git a/third_party/xla/xla/backends/gpu/runtime/cow_storage.h b/third_party/xla/xla/backends/gpu/runtime/cow_storage.h new file mode 100644 index 00000000000000..d7eae435c08554 --- /dev/null +++ b/third_party/xla/xla/backends/gpu/runtime/cow_storage.h @@ -0,0 +1,80 @@ +/* Copyright 2026 The OpenXLA Authors. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#ifndef XLA_BACKENDS_GPU_RUNTIME_COW_STORAGE_H_ +#define XLA_BACKENDS_GPU_RUNTIME_COW_STORAGE_H_ + +#include +#include +#include +#include + +#include "absl/base/thread_annotations.h" +#include "absl/container/inlined_vector.h" +#include "absl/synchronization/mutex.h" +#include "xla/backends/gpu/runtime/device_slot.h" + +namespace xla::gpu { + +// Copy-on-write storage for holding per-device state slots built on first +// touch. The storage itself is thread-safe for concurrent lookups (`Find` is +// lock-free) and insertions (`GetOrCreate` serializes via a mutex and publishes +// a new immutable snapshot). However, the returned state slots are not +// internally synchronized; it is only safe to mutate them across threads if +// each thread is dedicated to a distinct device ordinal/index. +class CowStorage { + public: + CowStorage() = default; + + CowStorage(const CowStorage&) = delete; + CowStorage& operator=(const CowStorage&) = delete; + + // Lock-free. Returns nullptr if `device_ordinal` is negative or has no slot + // yet. + DeviceSlot* Find(int device_ordinal) const; + + // Returns the slot for `device_ordinal`, building it on the first call. + // Lock-free once the slot exists. Returns nullptr if `device_ordinal` is + // negative. + DeviceSlot* GetOrCreate(int device_ordinal, DeviceSlotFactoryRef factory) + ABSL_LOCKS_EXCLUDED(mutex_); + + private: + struct Entry { + int device_ordinal; + DeviceSlot* slot; + }; + + // Sized for a typical 8-GPU host so snapshots stay inline in the common + // case. + static constexpr size_t kInlinedSnapshotSize = 8; + + // Immutable once published. + using Snapshot = absl::InlinedVector; + + static DeviceSlot* FindIn(const Snapshot& snapshot, int device_ordinal); + + std::atomic snapshot_{nullptr}; + absl::Mutex mutex_; + // Heap slots keep their address across snapshot copies and stay cache-line + // aligned. + std::vector> slots_ ABSL_GUARDED_BY(mutex_); + // Freed only in the destructor, so readers never see a freed snapshot. + std::vector> snapshots_ ABSL_GUARDED_BY(mutex_); +}; + +} // namespace xla::gpu + +#endif // XLA_BACKENDS_GPU_RUNTIME_COW_STORAGE_H_ diff --git a/third_party/xla/xla/backends/gpu/runtime/cow_storage_test.cc b/third_party/xla/xla/backends/gpu/runtime/cow_storage_test.cc new file mode 100644 index 00000000000000..f5e24d835914f0 --- /dev/null +++ b/third_party/xla/xla/backends/gpu/runtime/cow_storage_test.cc @@ -0,0 +1,304 @@ +/* Copyright 2026 The OpenXLA Authors. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#include "xla/backends/gpu/runtime/cow_storage.h" + +#include +#include +#include +#include +#include + +#include +#include +#include "absl/base/optimization.h" +#include "absl/container/flat_hash_set.h" +#include "absl/synchronization/notification.h" +#include "absl/time/time.h" +#include "xla/backends/gpu/runtime/device_slot.h" +#include "xla/tsl/platform/env.h" +#include "xla/tsl/platform/threadpool.h" + +namespace xla::gpu { +namespace { + +using ::testing::Contains; +using ::testing::Each; +using ::testing::Not; + +// Counts constructions and destructions. Neither copyable nor movable, so a +// storage that copies or moves its states fails to compile. +struct Tracked { + static constexpr int kInitialValue = 42; + + Tracked() { constructed.fetch_add(1); } + ~Tracked() { destroyed.fetch_add(1); } + Tracked(const Tracked&) = delete; + Tracked& operator=(const Tracked&) = delete; + + static void Reset() { + constructed.store(0); + destroyed.store(0); + } + + static inline std::atomic constructed{0}; + static inline std::atomic destroyed{0}; + + int value = kInitialValue; +}; +static_assert(!std::is_copy_constructible_v); +static_assert(!std::is_copy_assignable_v); +static_assert(!std::is_move_constructible_v); +static_assert(!std::is_move_assignable_v); + +// Notifies `entered`, then waits for `gate`, for whichever is set. Lets a test +// keep a CowStorage insert, and so its mutex, held. +struct GatedState { + GatedState() { + if (absl::Notification* notification = entered.load()) { + notification->Notify(); + } + if (absl::Notification* notification = gate.load()) { + notification->WaitForNotification(); + } + } + + static void Reset() { + entered.store(nullptr); + gate.store(nullptr); + } + + static inline std::atomic entered{nullptr}; + static inline std::atomic gate{nullptr}; +}; + +bool IsCacheLineAligned(const void* ptr) { + return reinterpret_cast(ptr) % ABSL_CACHELINE_SIZE == 0; +} + +template +size_t CountDistinct(const std::vector& ptrs) { + return absl::flat_hash_set(ptrs.begin(), ptrs.end()).size(); +} + +class CowStorageTest : public ::testing::Test { + protected: + void SetUp() override { + Tracked::Reset(); + GatedState::Reset(); + } + void TearDown() override { GatedState::Reset(); } +}; + +TEST_F(CowStorageTest, FindReturnsNullForOrdinalsNotCreated) { + CowStorage storage; + EXPECT_EQ(storage.Find(-1), nullptr); + EXPECT_EQ(storage.GetOrCreate(-1, &DeviceSlot::Create), nullptr); + EXPECT_EQ(storage.Find(0), nullptr); + EXPECT_NE(storage.GetOrCreate(5, &DeviceSlot::Create), nullptr); + EXPECT_EQ(storage.Find(-1), nullptr); + EXPECT_EQ(storage.GetOrCreate(-1, &DeviceSlot::Create), nullptr); + EXPECT_EQ(storage.Find(4), nullptr); + EXPECT_EQ(storage.Find(6), nullptr); + EXPECT_EQ(Tracked::constructed.load(), 1); +} + +TEST_F(CowStorageTest, SparseOrdinals) { + CowStorage storage; + DeviceSlot* s0 = storage.GetOrCreate(0, &DeviceSlot::Create); + DeviceSlot* s7 = storage.GetOrCreate(7, &DeviceSlot::Create); + DeviceSlot* s31 = storage.GetOrCreate(31, &DeviceSlot::Create); + + EXPECT_NE(s0, nullptr); + EXPECT_NE(s7, nullptr); + EXPECT_NE(s31, nullptr); + EXPECT_EQ(storage.Find(0), s0); + EXPECT_EQ(storage.Find(7), s7); + EXPECT_EQ(storage.Find(31), s31); + for (int gap : {1, 6, 8, 30, 32}) { + EXPECT_EQ(storage.Find(gap), nullptr) << "ordinal " << gap; + } + EXPECT_EQ(Tracked::constructed.load(), 3); +} + +TEST_F(CowStorageTest, GetOrCreateBuildsOnce) { + CowStorage storage; + DeviceSlot* slot = storage.GetOrCreate(3, &DeviceSlot::Create); + EXPECT_NE(slot, nullptr); + EXPECT_EQ(storage.GetOrCreate(3, &DeviceSlot::Create), slot); + EXPECT_EQ(storage.Find(3), slot); + EXPECT_EQ(Tracked::constructed.load(), 1); +} + +TEST_F(CowStorageTest, StatesKeepAddressAcrossInserts) { + // More ordinals than a snapshot stores inline. + constexpr int kNumOrdinals = 16; + CowStorage storage; + std::vector states; + for (int ordinal = 0; ordinal < kNumOrdinals; ++ordinal) { + states.push_back(DeviceSlot::Unwrap( + storage.GetOrCreate(ordinal, &DeviceSlot::Create))); + } + EXPECT_THAT(states, Not(Contains(nullptr))); + EXPECT_EQ(CountDistinct(states), kNumOrdinals); + for (int ordinal = 0; ordinal < kNumOrdinals; ++ordinal) { + DeviceSlot* slot = storage.Find(ordinal); + EXPECT_EQ(DeviceSlot::Unwrap(slot), states[ordinal]) + << "ordinal " << ordinal; + EXPECT_TRUE(IsCacheLineAligned(slot)) << "ordinal " << ordinal; + } + EXPECT_EQ(Tracked::constructed.load(), kNumOrdinals); +} + +TEST_F(CowStorageTest, DestroysEveryState) { + { + CowStorage storage; + for (int ordinal = 0; ordinal < 3; ++ordinal) { + storage.GetOrCreate(ordinal, &DeviceSlot::Create); + } + EXPECT_EQ(Tracked::destroyed.load(), 0); + } + EXPECT_EQ(Tracked::constructed.load(), 3); + EXPECT_EQ(Tracked::destroyed.load(), 3); +} + +TEST_F(CowStorageTest, ConcurrentGetOrCreateBuildsOneState) { + constexpr int kNumThreads = 8; + CowStorage storage; + std::vector states(kNumThreads, nullptr); + absl::Notification start; + { + tsl::thread::ThreadPool pool(tsl::Env::Default(), "cow_storage_test", + kNumThreads); + for (int i = 0; i < kNumThreads; ++i) { + pool.Schedule([&, i] { + start.WaitForNotification(); + states[i] = DeviceSlot::Unwrap( + storage.GetOrCreate(0, &DeviceSlot::Create)); + }); + } + start.Notify(); + } // Joins the pool. + EXPECT_EQ(Tracked::constructed.load(), 1); + EXPECT_NE(states[0], nullptr); + EXPECT_THAT(states, Each(states[0])); +} + +TEST_F(CowStorageTest, ConcurrentReadersSeeStableFullyBuiltStates) { + constexpr int kNumOrdinals = 16; + constexpr int kNumReaders = 4; + CowStorage storage; + std::vector created(kNumOrdinals, nullptr); + std::vector> seen( + kNumReaders, std::vector(kNumOrdinals, nullptr)); + absl::Notification start; + std::atomic writers_done{false}; + std::atomic errors{0}; + { + tsl::thread::ThreadPool readers(tsl::Env::Default(), "readers", + kNumReaders); + for (int reader = 0; reader < kNumReaders; ++reader) { + readers.Schedule([&, reader] { + start.WaitForNotification(); + std::vector& last_seen = seen[reader]; + bool done = false; + while (!done) { + // Loaded before the scan, so the last scan starts after all writers + // have finished. + done = writers_done.load(); + for (int ordinal = 0; ordinal < kNumOrdinals; ++ordinal) { + Tracked* state = DeviceSlot::Unwrap(storage.Find(ordinal)); + if (state == nullptr) { + continue; + } + // A published state is fully built and never replaced. + if (state->value != Tracked::kInitialValue || + (last_seen[ordinal] != nullptr && + last_seen[ordinal] != state)) { + errors.fetch_add(1); + } + last_seen[ordinal] = state; + } + } + }); + } + { + tsl::thread::ThreadPool writers(tsl::Env::Default(), "writers", + kNumOrdinals); + for (int ordinal = 0; ordinal < kNumOrdinals; ++ordinal) { + writers.Schedule([&, ordinal] { + start.WaitForNotification(); + created[ordinal] = DeviceSlot::Unwrap( + storage.GetOrCreate(ordinal, &DeviceSlot::Create)); + }); + } + start.Notify(); + } // Joins the writers. + writers_done.store(true); + } // Joins the readers. + EXPECT_EQ(errors.load(), 0); + EXPECT_EQ(Tracked::constructed.load(), kNumOrdinals); + EXPECT_THAT(created, Not(Contains(nullptr))); + for (int reader = 0; reader < kNumReaders; ++reader) { + EXPECT_EQ(seen[reader], created) << "reader " << reader; + } +} + +// An insert holds the mutex while it builds the state. `Find` must not wait +// for it. +TEST_F(CowStorageTest, FindDoesNotWaitForInsert) { + CowStorage storage; + DeviceSlot* existing = + storage.GetOrCreate(0, &DeviceSlot::Create); + ASSERT_NE(existing, nullptr); + + absl::Notification entered; + absl::Notification release; + absl::Notification read_done; + GatedState::entered.store(&entered); + GatedState::gate.store(&release); + + bool insert_in_progress = false; + bool read_finished = false; + DeviceSlot* found = nullptr; + DeviceSlot* uncommitted = nullptr; + { + tsl::thread::ThreadPool pool(tsl::Env::Default(), "cow_storage_test", 2); + pool.Schedule( + [&] { storage.GetOrCreate(1, &DeviceSlot::Create); }); + insert_in_progress = + entered.WaitForNotificationWithTimeout(absl::Seconds(10)); + if (insert_in_progress) { + pool.Schedule([&] { + found = storage.Find(0); + uncommitted = storage.Find(1); + read_done.Notify(); + }); + read_finished = + read_done.WaitForNotificationWithTimeout(absl::Seconds(10)); + } + // Always unblock the insert, so the pool can join. + release.Notify(); + } // Joins the pool. + GatedState::Reset(); + + EXPECT_TRUE(insert_in_progress); + EXPECT_TRUE(read_finished); + EXPECT_EQ(found, existing); + EXPECT_EQ(uncommitted, nullptr); +} + +} // namespace +} // namespace xla::gpu diff --git a/third_party/xla/xla/backends/gpu/runtime/device_slot.h b/third_party/xla/xla/backends/gpu/runtime/device_slot.h new file mode 100644 index 00000000000000..02640185ab6f97 --- /dev/null +++ b/third_party/xla/xla/backends/gpu/runtime/device_slot.h @@ -0,0 +1,59 @@ +/* Copyright 2026 The OpenXLA Authors. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#ifndef XLA_BACKENDS_GPU_RUNTIME_DEVICE_SLOT_H_ +#define XLA_BACKENDS_GPU_RUNTIME_DEVICE_SLOT_H_ + +#include + +#include "absl/base/optimization.h" +#include "absl/functional/function_ref.h" + +namespace xla::gpu { + +// Type-erased, cache-line-aligned base for a per-device state slot. +// Aligned so that states of different devices never share a cache line. +struct alignas(ABSL_CACHELINE_SIZE) DeviceSlot { + virtual ~DeviceSlot() = default; + + template + static std::unique_ptr Create(); + + template + static T* Unwrap(DeviceSlot* slot); +}; + +template +struct TypedDeviceSlot final : DeviceSlot { + T value; +}; + +template +std::unique_ptr DeviceSlot::Create() { + return std::make_unique>(); +} + +template +T* DeviceSlot::Unwrap(DeviceSlot* slot) { + return slot != nullptr ? &(static_cast*>(slot)->value) + : nullptr; +} + +// No argument function that creates a DeviceSlot. +using DeviceSlotFactoryRef = absl::FunctionRef()>; + +} // namespace xla::gpu + +#endif // XLA_BACKENDS_GPU_RUNTIME_DEVICE_SLOT_H_ diff --git a/third_party/xla/xla/backends/gpu/runtime/vector_storage.cc b/third_party/xla/xla/backends/gpu/runtime/vector_storage.cc new file mode 100644 index 00000000000000..d0f3b2112db59b --- /dev/null +++ b/third_party/xla/xla/backends/gpu/runtime/vector_storage.cc @@ -0,0 +1,29 @@ +/* Copyright 2026 The OpenXLA Authors. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#include "xla/backends/gpu/runtime/vector_storage.h" + +#include "xla/backends/gpu/runtime/device_slot.h" + +namespace xla::gpu { + +DeviceSlot* VectorStorage::Find(int device_ordinal) const { + if (device_ordinal < 0 || device_ordinal >= size()) { + return nullptr; + } + return slots_[device_ordinal].get(); +} + +} // namespace xla::gpu diff --git a/third_party/xla/xla/backends/gpu/runtime/vector_storage.h b/third_party/xla/xla/backends/gpu/runtime/vector_storage.h new file mode 100644 index 00000000000000..3be2705cab164a --- /dev/null +++ b/third_party/xla/xla/backends/gpu/runtime/vector_storage.h @@ -0,0 +1,49 @@ +/* Copyright 2026 The OpenXLA Authors. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#ifndef XLA_BACKENDS_GPU_RUNTIME_VECTOR_STORAGE_H_ +#define XLA_BACKENDS_GPU_RUNTIME_VECTOR_STORAGE_H_ + +#include +#include +#include + +#include "xla/backends/gpu/runtime/device_slot.h" + +namespace xla::gpu { + +// Pre-sized storage for holding per-device state slots in [0, size), which can +// be accessed across threads safely as long as each thread only accesses a +// single ordinal/index. `Find` is a bounds check plus an array index. +class VectorStorage { + public: + explicit VectorStorage(std::vector> slots) + : slots_(std::move(slots)) {} + + VectorStorage(const VectorStorage&) = delete; + VectorStorage& operator=(const VectorStorage&) = delete; + + int size() const { return static_cast(slots_.size()); } + + // Returns nullptr if `device_ordinal` is outside [0, size). + DeviceSlot* Find(int device_ordinal) const; + + private: + std::vector> slots_; +}; + +} // namespace xla::gpu + +#endif // XLA_BACKENDS_GPU_RUNTIME_VECTOR_STORAGE_H_ diff --git a/third_party/xla/xla/backends/gpu/runtime/vector_storage_test.cc b/third_party/xla/xla/backends/gpu/runtime/vector_storage_test.cc new file mode 100644 index 00000000000000..2ee5c802fae2ef --- /dev/null +++ b/third_party/xla/xla/backends/gpu/runtime/vector_storage_test.cc @@ -0,0 +1,138 @@ +/* Copyright 2026 The OpenXLA Authors. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#include "xla/backends/gpu/runtime/vector_storage.h" + +#include +#include +#include +#include +#include +#include + +#include +#include +#include "absl/base/optimization.h" +#include "absl/container/flat_hash_set.h" +#include "xla/backends/gpu/runtime/device_slot.h" + +namespace xla::gpu { +namespace { + +using ::testing::Contains; +using ::testing::Not; + +// Counts constructions and destructions. Neither copyable nor movable, so a +// storage that copies or moves its states fails to compile. +struct Tracked { + static constexpr int kInitialValue = 42; + + Tracked() { constructed.fetch_add(1); } + ~Tracked() { destroyed.fetch_add(1); } + Tracked(const Tracked&) = delete; + Tracked& operator=(const Tracked&) = delete; + + static void Reset() { + constructed.store(0); + destroyed.store(0); + } + + static inline std::atomic constructed{0}; + static inline std::atomic destroyed{0}; + + int value = kInitialValue; +}; +static_assert(!std::is_copy_constructible_v); +static_assert(!std::is_copy_assignable_v); +static_assert(!std::is_move_constructible_v); +static_assert(!std::is_move_assignable_v); + +template +std::vector> MakeSlots(int size) { + std::vector> slots; + slots.reserve(size); + for (int i = 0; i < size; ++i) { + slots.push_back(DeviceSlot::Create()); + } + return slots; +} + +bool IsCacheLineAligned(const void* ptr) { + return reinterpret_cast(ptr) % ABSL_CACHELINE_SIZE == 0; +} + +template +size_t CountDistinct(const std::vector& ptrs) { + return absl::flat_hash_set(ptrs.begin(), ptrs.end()).size(); +} + +class VectorStorageTest : public ::testing::Test { + protected: + void SetUp() override { Tracked::Reset(); } +}; + +TEST_F(VectorStorageTest, StatesExistAfterConstruction) { + { + VectorStorage storage(MakeSlots(4)); + EXPECT_EQ(storage.size(), 4); + EXPECT_EQ(Tracked::constructed.load(), 4); + for (int ordinal = 0; ordinal < 4; ++ordinal) { + Tracked* state = DeviceSlot::Unwrap(storage.Find(ordinal)); + ASSERT_NE(state, nullptr) << "ordinal " << ordinal; + EXPECT_EQ(state->value, Tracked::kInitialValue) << "ordinal " << ordinal; + } + } + EXPECT_EQ(Tracked::destroyed.load(), 4); +} + +TEST_F(VectorStorageTest, FindOutsideRangeReturnsNull) { + VectorStorage storage(MakeSlots(4)); + EXPECT_EQ(storage.Find(-1), nullptr); + EXPECT_EQ(storage.Find(4), nullptr); +} + +TEST_F(VectorStorageTest, EmptyStorageFindsNothing) { + VectorStorage empty({}); + EXPECT_EQ(empty.size(), 0); + EXPECT_EQ(empty.Find(0), nullptr); + EXPECT_EQ(Tracked::constructed.load(), 0); +} + +TEST_F(VectorStorageTest, StatesAreDistinctAndStable) { + VectorStorage storage(MakeSlots(4)); + std::vector states; + for (int ordinal = 0; ordinal < 4; ++ordinal) { + states.push_back(DeviceSlot::Unwrap(storage.Find(ordinal))); + } + EXPECT_THAT(states, Not(Contains(nullptr))); + EXPECT_EQ(CountDistinct(states), 4); + for (int ordinal = 0; ordinal < 4; ++ordinal) { + EXPECT_EQ(DeviceSlot::Unwrap(storage.Find(ordinal)), + states[ordinal]) + << "ordinal " << ordinal; + } +} + +TEST_F(VectorStorageTest, StatesAreCacheLineAligned) { + VectorStorage storage(MakeSlots(4)); + for (int ordinal = 0; ordinal < 4; ++ordinal) { + DeviceSlot* slot = storage.Find(ordinal); + EXPECT_NE(slot, nullptr) << "ordinal " << ordinal; + EXPECT_TRUE(IsCacheLineAligned(slot)) << "ordinal " << ordinal; + } +} + +} // namespace +} // namespace xla::gpu diff --git a/third_party/xla/xla/backends/gpu/tests/BUILD b/third_party/xla/xla/backends/gpu/tests/BUILD index 73e0ea8e1b3810..c7fd9b12c2af57 100644 --- a/third_party/xla/xla/backends/gpu/tests/BUILD +++ b/third_party/xla/xla/backends/gpu/tests/BUILD @@ -296,6 +296,7 @@ xla_test( deps = [ ":hlo_pjrt_gpu_test_base", "//xla:error_spec", + "//xla:xla_proto_cc", "//xla/hlo/ir:hlo", "//xla/hlo/testlib:filecheck", "//xla/stream_executor:device_description", @@ -1171,6 +1172,7 @@ cc_library( "//xla/stream_executor/cuda:cuda_compute_capability", "//xla/tsl/platform:env", "//xla/tsl/platform:statusor", + "@com_google_absl//absl/algorithm:container", "@com_google_absl//absl/log", "@com_google_absl//absl/log:check", "@com_google_absl//absl/status", @@ -1712,6 +1714,8 @@ xla_test( "//xla/tests:literal_test_util", "//xla/tsl/platform:test", "//xla/tsl/testing:temporary_directory", + "@com_google_absl//absl/strings", + "@com_google_absl//absl/strings:str_format", "@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 0a118b107155db..60c7b25fb614ec 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 @@ -13,12 +13,15 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ +#include #include #include #include #include #include +#include "absl/strings/str_cat.h" +#include "absl/strings/str_format.h" #include "absl/strings/string_view.h" #include "xla/backends/gpu/tests/collective_ops_e2e_test_base.h" #include "xla/hlo/ir/hlo_computation.h" @@ -251,9 +254,10 @@ 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. +// Runs 10 consequent all-gathers and verifies that by default all of them use +// XLA's one-shot all-gather kernel and are placed into command buffers, and +// that excluding COLLECTIVES_KERNEL from --xla_gpu_enable_command_buffer +// disables command buffer capture for them. TEST_F(AllGatherTest, TenConsequentAllGathersWithCudaGraphs) { constexpr int32_t kNumReplicas = 2; constexpr int kNumAllGathers = 10; @@ -281,6 +285,11 @@ TEST_F(AllGatherTest, TenConsequentAllGathersWithCudaGraphs) { ASSERT_OK_AND_ASSIGN( tsl::testing::TemporaryDirectory dump_dir, tsl::testing::TemporaryDirectory::CreateForCurrentTestcase()); + const std::string default_dump_dir = + absl::StrCat(dump_dir.path(), "/default"); + const std::string no_graph_dump_dir = + absl::StrCat(dump_dir.path(), "/no_graphs"); + ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(kModuleStr, kNumReplicas)); @@ -288,15 +297,13 @@ TEST_F(AllGatherTest, TenConsequentAllGathersWithCudaGraphs) { 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); + // instead of NCCL. COLLECTIVES_KERNEL is enabled in + // --xla_gpu_enable_command_buffer by default. 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()); + debug_options.set_xla_dump_to(default_dump_dir); Literal input_r0 = LiteralUtil::CreateR1(std::vector(16, 1.0f)); Literal input_r1 = LiteralUtil::CreateR1(std::vector(16, 2.0f)); @@ -308,14 +315,15 @@ TEST_F(AllGatherTest, TenConsequentAllGathersWithCudaGraphs) { // 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")); + // By default, 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(default_dump_dir, "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")); + CountThunksInDump(default_dump_dir, "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 @@ -332,6 +340,173 @@ TEST_F(AllGatherTest, TenConsequentAllGathersWithCudaGraphs) { EXPECT_TRUE(LiteralTestUtil::Equal(expected, result.results[i])) << "Mismatch at replica " << i; } + + // Excluding COLLECTIVES_KERNEL from --xla_gpu_enable_command_buffer must + // disable command buffer capture for the collective kernels. + ASSERT_OK_AND_ASSIGN(auto no_graph_module, + ParseAndReturnVerifiedModule(kModuleStr, kNumReplicas)); + DebugOptions& no_graph_debug_options = + no_graph_module->mutable_config().mutable_debug_options(); + auto* enabled_commands = + no_graph_debug_options.mutable_xla_gpu_enable_command_buffer(); + enabled_commands->erase( + std::remove(enabled_commands->begin(), enabled_commands->end(), + DebugOptions::COLLECTIVES_KERNEL), + enabled_commands->end()); + no_graph_debug_options.set_xla_gpu_graph_min_graph_size(1); + no_graph_debug_options.set_xla_gpu_all_gather_combine_threshold_bytes(0); + no_graph_debug_options.set_xla_dump_to(no_graph_dump_dir); + + ASSERT_OK_AND_ASSIGN(ExecutionResult no_graph_result, + ExecuteReplicated(std::move(no_graph_module), args)); + VerifyOneShotAllGather(no_graph_result.optimized_module); + + ASSERT_OK_AND_ASSIGN( + CommandBufferThunkCounts one_shot_no_graph, + CountThunksInDump(no_graph_dump_dir, "kCollectiveKernel")); + EXPECT_EQ(one_shot_no_graph.in_command_buffer, 0); + EXPECT_EQ(one_shot_no_graph.outside_command_buffer, kNumAllGathers); + ASSERT_OK_AND_ASSIGN(CommandBufferThunkCounts nccl_no_graph, + CountThunksInDump(no_graph_dump_dir, "kAllGather")); + EXPECT_EQ( + nccl_no_graph.in_command_buffer + nccl_no_graph.outside_command_buffer, + 0); + ASSERT_EQ(no_graph_result.results.size(), kNumReplicas); + for (int i = 0; i < kNumReplicas; ++i) { + EXPECT_TRUE(LiteralTestUtil::Equal(expected, no_graph_result.results[i])) + << "Mismatch at replica " << i << " without graphs"; + } +} + +// Runs 20 one-shot all-gathers inside a 20-iteration while loop with WHILE +// command buffers enabled, and verifies that the while thunk and all 20 +// collective kernels are placed into command buffers and produce the expected +// result across repeated executions (exercising the persistent device-side +// invocation counter). +TEST_F(AllGatherTest, TwentyAllGathersInWhileLoopWithCudaGraphs) { + constexpr int32_t kNumReplicas = 2; + constexpr int kNumAllGathers = 20; + constexpr int kNumIterations = 20; + if (!CheckDeviceCount(kNumReplicas)) { + return; + } + + // Each step k in the loop body gathers in_k = 0.5 * prev + p0 (f32[128]) + // across the 2 replicas into ag_k (f32[256]) and sums the two 128-element + // slices: (0.5 * v + 1.0) + (0.5 * v + 2.0) = v + 3.0. Every all-gather + // therefore depends on both replicas' slices from the previous all-gather and + // adds 3.0 without growing the tensor shape or overflowing f32 across + // 20 * 20 = 400 collective launches. + std::string body_ops; + std::string prev = "acc"; + for (int k = 0; k < kNumAllGathers; ++k) { + absl::StrAppendFormat( + &body_ops, + " scaled%d = f32[128] multiply(%s, half)\n" + " in%d = f32[128] add(scaled%d, p0)\n" + " ag%d = f32[256] all-gather(in%d), dimensions={0}, " + "replica_groups={{0,1}}\n" + " s0_%d = f32[128] slice(ag%d), slice={[0:128]}\n" + " s1_%d = f32[128] slice(ag%d), slice={[128:256]}\n" + " sum%d = f32[128] add(s0_%d, s1_%d)\n", + k, prev, k, k, k, k, k, k, k, k, k, k, k); + prev = absl::StrCat("sum", k); + } + + const std::string module_str = absl::StrFormat( + R"( + HloModule test + + cond { + state = (s32[], f32[128], f32[128]) parameter(0) + i = s32[] get-tuple-element(state), index=0 + limit = s32[] constant(%d) + ROOT cmp = pred[] compare(i, limit), direction=LT + } + + body { + state = (s32[], f32[128], f32[128]) parameter(0) + i = s32[] get-tuple-element(state), index=0 + acc = f32[128] get-tuple-element(state), index=1 + p0 = f32[128] get-tuple-element(state), index=2 + one = s32[] constant(1) + next_i = s32[] add(i, one) + half_scalar = f32[] constant(0.5) + half = f32[128] broadcast(half_scalar), dimensions={} +%s ROOT next_state = (s32[], f32[128], f32[128]) tuple(next_i, %s, p0) + } + + ENTRY test_computation { + p0 = f32[128] parameter(0) + zero_i = s32[] constant(0) + zero_f = f32[] constant(0.0) + init_acc = f32[128] broadcast(zero_f), dimensions={} + init = (s32[], f32[128], f32[128]) tuple(zero_i, init_acc, p0) + loop = (s32[], f32[128], f32[128]) while(init), condition=cond, body=body + ROOT out = f32[128] get-tuple-element(loop), index=1 + } + )", + kNumIterations, body_ops, prev); + + ASSERT_OK_AND_ASSIGN( + tsl::testing::TemporaryDirectory dump_dir, + tsl::testing::TemporaryDirectory::CreateForCurrentTestcase()); + ASSERT_OK_AND_ASSIGN(auto module, + ParseAndReturnVerifiedModule(module_str, kNumReplicas)); + + DebugOptions& debug_options = + module->mutable_config().mutable_debug_options(); + debug_options.add_xla_gpu_enable_command_buffer(DebugOptions::WHILE); + debug_options.set_xla_gpu_graph_min_graph_size(1); + debug_options.set_xla_gpu_all_gather_combine_threshold_bytes(0); + debug_options.set_xla_gpu_enable_while_loop_unrolling( + DebugOptions::WHILE_LOOP_UNROLLING_NO_UNROLL); + debug_options.set_xla_dump_to(dump_dir.path()); + + Literal input_r0 = + LiteralUtil::CreateR1(std::vector(128, 1.0f)); + Literal input_r1 = + LiteralUtil::CreateR1(std::vector(128, 2.0f)); + std::vector> args = {{&input_r0}, {&input_r1}}; + ASSERT_OK_AND_ASSIGN(ExecutionResult result, + ExecuteReplicated(std::move(module), args)); + + VerifyOneShotAllGather(result.optimized_module); + + // Both the while thunk and all 20 one-shot all-gathers in its body must be + // recorded into command buffers, and none may fall back to NCCL. + ASSERT_OK_AND_ASSIGN(CommandBufferThunkCounts while_counts, + CountThunksInDump(dump_dir.path(), "kWhile")); + EXPECT_EQ(while_counts.in_command_buffer, 1); + EXPECT_EQ(while_counts.outside_command_buffer, 0); + 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); + + // 20 iterations * 20 all-gathers * 3.0 per all-gather = 1200.0. + const float expected_val = 3.0f * kNumAllGathers * kNumIterations; + Literal expected = + LiteralUtil::CreateR1(std::vector(128, expected_val)); + ASSERT_EQ(result.results.size(), kNumReplicas); + for (int i = 0; i < kNumReplicas; ++i) { + EXPECT_TRUE(LiteralTestUtil::Equal(expected, result.results[i])) + << "Mismatch at replica " << i << " on first execution"; + } + + // Re-run the same executable to verify that the device-side counters persist + // across executions and the replayed CUDA graph still produces the expected + // result. + ASSERT_OK_AND_ASSIGN(std::vector second_results, + ExecuteReplicated(result.executable.get(), args)); + ASSERT_EQ(second_results.size(), kNumReplicas); + for (int i = 0; i < kNumReplicas; ++i) { + EXPECT_TRUE(LiteralTestUtil::Equal(expected, second_results[i])) + << "Mismatch at replica " << i << " on second execution"; + } } } // namespace 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 fa2089ef918828..97780efd450c51 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 @@ -1101,9 +1101,10 @@ 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. +// Runs 10 consequent all-reduces and verifies that by default all of them use +// XLA's one-shot all-reduce kernel and are placed into command buffers, and +// that excluding COLLECTIVES_KERNEL from --xla_gpu_enable_command_buffer +// disables command buffer capture for them. TEST_F(AllReduceCollectiveKernelTest, TenConsequentAllReducesWithCudaGraphs) { constexpr int64_t kNumReplicas = 2; constexpr int kNumAllReduces = 10; @@ -1158,6 +1159,11 @@ TEST_F(AllReduceCollectiveKernelTest, TenConsequentAllReducesWithCudaGraphs) { ASSERT_OK_AND_ASSIGN( tsl::testing::TemporaryDirectory dump_dir, tsl::testing::TemporaryDirectory::CreateForCurrentTestcase()); + const std::string default_dump_dir = + absl::StrCat(dump_dir.path(), "/default"); + const std::string no_graph_dump_dir = + absl::StrCat(dump_dir.path(), "/no_graphs"); + ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(kModuleStr, kNumReplicas)); @@ -1165,15 +1171,13 @@ TEST_F(AllReduceCollectiveKernelTest, TenConsequentAllReducesWithCudaGraphs) { 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); + // instead of NCCL. COLLECTIVES_KERNEL is enabled in + // --xla_gpu_enable_command_buffer by default. 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()); + debug_options.set_xla_dump_to(default_dump_dir); Literal input_r0 = LiteralUtil::CreateR1(std::vector(1024, 1.0f)); @@ -1187,14 +1191,15 @@ TEST_F(AllReduceCollectiveKernelTest, TenConsequentAllReducesWithCudaGraphs) { // 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")); + // By default, 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(default_dump_dir, "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")); + CountThunksInDump(default_dump_dir, "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 @@ -1207,6 +1212,175 @@ TEST_F(AllReduceCollectiveKernelTest, TenConsequentAllReducesWithCudaGraphs) { EXPECT_TRUE(LiteralTestUtil::Equal(expected, result.results[i])) << "Mismatch at replica " << i; } + + // Excluding COLLECTIVES_KERNEL from --xla_gpu_enable_command_buffer must + // disable command buffer capture for the collective kernels. + ASSERT_OK_AND_ASSIGN(auto no_graph_module, + ParseAndReturnVerifiedModule(kModuleStr, kNumReplicas)); + DebugOptions& no_graph_debug_options = + no_graph_module->mutable_config().mutable_debug_options(); + auto* enabled_commands = + no_graph_debug_options.mutable_xla_gpu_enable_command_buffer(); + enabled_commands->erase( + std::remove(enabled_commands->begin(), enabled_commands->end(), + DebugOptions::COLLECTIVES_KERNEL), + enabled_commands->end()); + no_graph_debug_options.set_xla_gpu_graph_min_graph_size(1); + no_graph_debug_options.set_xla_gpu_all_reduce_combine_threshold_bytes(0); + no_graph_debug_options.set_xla_dump_to(no_graph_dump_dir); + + ASSERT_OK_AND_ASSIGN(ExecutionResult no_graph_result, + ExecuteReplicated(std::move(no_graph_module), args)); + VerifyOneShotAllReduce(no_graph_result.optimized_module); + + ASSERT_OK_AND_ASSIGN( + CommandBufferThunkCounts one_shot_no_graph, + CountThunksInDump(no_graph_dump_dir, "kCollectiveKernel")); + EXPECT_EQ(one_shot_no_graph.in_command_buffer, 0); + EXPECT_EQ(one_shot_no_graph.outside_command_buffer, kNumAllReduces); + ASSERT_OK_AND_ASSIGN(CommandBufferThunkCounts nccl_no_graph, + CountThunksInDump(no_graph_dump_dir, "kAllReduce")); + EXPECT_EQ( + nccl_no_graph.in_command_buffer + nccl_no_graph.outside_command_buffer, + 0); + ASSERT_EQ(no_graph_result.results.size(), kNumReplicas); + for (int i = 0; i < kNumReplicas; ++i) { + EXPECT_TRUE(LiteralTestUtil::Equal(expected, no_graph_result.results[i])) + << "Mismatch at replica " << i << " without graphs"; + } +} + +// Runs 20 one-shot all-reduces inside a 20-iteration while loop with WHILE +// command buffers enabled, and verifies that the while thunk and all 20 +// collective kernels are placed into command buffers and produce the expected +// result across repeated executions (exercising the persistent device-side +// invocation counter). +TEST_F(AllReduceCollectiveKernelTest, + TwentyAllReducesInWhileLoopWithCudaGraphs) { + constexpr int64_t kNumReplicas = 2; + constexpr int kNumAllReduces = 20; + constexpr int kNumIterations = 20; + if (!CheckDeviceCount(kNumReplicas)) { + return; + } + + // Each step k in the loop body computes ar_k = all-reduce(0.5 * prev + p0), + // where p0 is 1.0 on replica 0 and 2.0 on replica 1. When prev == v on both + // replicas, (0.5 * v + 1.0) + (0.5 * v + 2.0) = v + 3.0, so every all-reduce + // adds 3.0 without overflowing f32 across 20 * 20 = 400 collective launches. + std::string body_ops; + std::string prev = "acc"; + for (int k = 0; k < kNumAllReduces; ++k) { + absl::StrAppendFormat( + &body_ops, + " scaled%d = f32[1024] multiply(%s, half)\n" + " in%d = f32[1024] add(scaled%d, p0)\n" + " ar%d = f32[1024] all-reduce(in%d), to_apply=add, " + "replica_groups={{0,1}}\n", + k, prev, k, k, k, k); + prev = absl::StrCat("ar", k); + } + + const std::string module_str = absl::StrFormat( + R"( + HloModule test + + add { + x = f32[] parameter(0) + y = f32[] parameter(1) + ROOT sum = f32[] add(x, y) + } + + cond { + state = (s32[], f32[1024], f32[1024]) parameter(0) + i = s32[] get-tuple-element(state), index=0 + limit = s32[] constant(%d) + ROOT cmp = pred[] compare(i, limit), direction=LT + } + + body { + state = (s32[], f32[1024], f32[1024]) parameter(0) + i = s32[] get-tuple-element(state), index=0 + acc = f32[1024] get-tuple-element(state), index=1 + p0 = f32[1024] get-tuple-element(state), index=2 + one = s32[] constant(1) + next_i = s32[] add(i, one) + half_scalar = f32[] constant(0.5) + half = f32[1024] broadcast(half_scalar), dimensions={} +%s ROOT next_state = (s32[], f32[1024], f32[1024]) tuple(next_i, %s, p0) + } + + ENTRY test_computation { + p0 = f32[1024] parameter(0) + zero_i = s32[] constant(0) + zero_f = f32[] constant(0.0) + init_acc = f32[1024] broadcast(zero_f), dimensions={} + init = (s32[], f32[1024], f32[1024]) tuple(zero_i, init_acc, p0) + loop = (s32[], f32[1024], f32[1024]) while(init), condition=cond, body=body + ROOT out = f32[1024] get-tuple-element(loop), index=1 + } + )", + kNumIterations, body_ops, prev); + + ASSERT_OK_AND_ASSIGN( + tsl::testing::TemporaryDirectory dump_dir, + tsl::testing::TemporaryDirectory::CreateForCurrentTestcase()); + ASSERT_OK_AND_ASSIGN(auto module, + ParseAndReturnVerifiedModule(module_str, kNumReplicas)); + + DebugOptions& debug_options = + module->mutable_config().mutable_debug_options(); + debug_options.add_xla_gpu_enable_command_buffer(DebugOptions::WHILE); + debug_options.set_xla_gpu_graph_min_graph_size(1); + debug_options.set_xla_gpu_all_reduce_combine_threshold_bytes(0); + debug_options.set_xla_gpu_enable_while_loop_unrolling( + DebugOptions::WHILE_LOOP_UNROLLING_NO_UNROLL); + 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)); + + VerifyOneShotAllReduce(result.optimized_module); + + // Both the while thunk and all 20 one-shot all-reduces in its body must be + // recorded into command buffers, and none may fall back to NCCL. + ASSERT_OK_AND_ASSIGN(CommandBufferThunkCounts while_counts, + CountThunksInDump(dump_dir.path(), "kWhile")); + EXPECT_EQ(while_counts.in_command_buffer, 1); + EXPECT_EQ(while_counts.outside_command_buffer, 0); + 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); + + // 20 iterations * 20 all-reduces * 3.0 per all-reduce = 1200.0. + const float expected_val = 3.0f * kNumAllReduces * kNumIterations; + Literal expected = + LiteralUtil::CreateR1(std::vector(1024, expected_val)); + ASSERT_EQ(result.results.size(), kNumReplicas); + for (int i = 0; i < kNumReplicas; ++i) { + EXPECT_TRUE(LiteralTestUtil::Equal(expected, result.results[i])) + << "Mismatch at replica " << i << " on first execution"; + } + + // Re-run the same executable to verify that the device-side counters persist + // across executions and the replayed CUDA graph still produces the expected + // result. + ASSERT_OK_AND_ASSIGN(std::vector second_results, + ExecuteReplicated(result.executable.get(), args)); + ASSERT_EQ(second_results.size(), kNumReplicas); + for (int i = 0; i < kNumReplicas; ++i) { + EXPECT_TRUE(LiteralTestUtil::Equal(expected, second_results[i])) + << "Mismatch at replica " << i << " on second execution"; + } } } // namespace 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 ab11d842402e53..0a61389bad351c 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 @@ -24,6 +24,7 @@ limitations under the License. #include #include +#include "absl/algorithm/container.h" #include "absl/log/check.h" #include "absl/log/log.h" #include "absl/status/status.h" @@ -204,6 +205,12 @@ absl::StatusOr CountThunksInDump( ABSL_RETURN_IF_ERROR(tsl::Env::Default()->GetMatchingPaths( tsl::io::JoinPath(dump_dir, "*thunk_sequence_after_thunk_passes*.txt"), &dump_files)); + if (dump_files.empty()) { + // When thunk passes make no changes (e.g. no command buffers are formed), + // only the initial thunk_sequence.txt dump is written. + ABSL_RETURN_IF_ERROR(tsl::Env::Default()->GetMatchingPaths( + tsl::io::JoinPath(dump_dir, "*thunk_sequence.txt"), &dump_files)); + } if (dump_files.size() != 1) { return absl::FailedPreconditionError( absl::StrCat("Expected exactly one thunk sequence dump in ", dump_dir, @@ -213,19 +220,22 @@ absl::StatusOr CountThunksInDump( 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. + // Each thunk line in the dump has the form ": ...". + // Thunks nested inside a top-level kCommandBuffer thunk are indented under + // it; non-thunk lines (such as WhileThunk's "condition:" and "body:" headers) + // do not start with ": " and are ignored. 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()) { + absl::string_view thunk = absl::StripLeadingAsciiWhitespace(line); + size_t pos = thunk.find(": "); + if (pos == absl::string_view::npos || pos == 0 || + !absl::c_all_of(thunk.substr(0, pos), + [](char c) { return absl::ascii_isdigit(c); })) { 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); - } + const bool is_top_level = thunk.size() == line.size(); + thunk.remove_prefix(pos + 2); if (is_top_level) { in_command_buffer = absl::StartsWith(thunk, "kCommandBuffer"); } diff --git a/third_party/xla/xla/backends/gpu/tests/swap_conv_operands_test.cc b/third_party/xla/xla/backends/gpu/tests/swap_conv_operands_test.cc index 7ed04346490c4f..941aef2f58280f 100644 --- a/third_party/xla/xla/backends/gpu/tests/swap_conv_operands_test.cc +++ b/third_party/xla/xla/backends/gpu/tests/swap_conv_operands_test.cc @@ -26,6 +26,7 @@ limitations under the License. #include "xla/hlo/testlib/filecheck.h" #include "xla/stream_executor/device_description.h" #include "xla/tests/hlo_pjrt_interpreter_reference_mixin.h" +#include "xla/xla.pb.h" namespace xla::gpu { namespace { @@ -33,6 +34,13 @@ namespace { class SwapConvOperandsTest : public HloInterpreterReferenceMixin { public: + DebugOptions GetDebugOptionsForTest() const override { + DebugOptions debug_options = HloInterpreterReferenceMixin< + HloPjRtGpuTestBase>::GetDebugOptionsForTest(); + debug_options.set_xla_gpu_experimental_enable_conv_fusion(false); + return debug_options; + } + void MatchOptimizedHlo(absl::string_view hlo, absl::string_view expected_hlo) { ASSERT_OK_AND_ASSIGN(std::unique_ptr optimized_module, diff --git a/third_party/xla/xla/backends/gpu/transforms/BUILD b/third_party/xla/xla/backends/gpu/transforms/BUILD index c78569045a3131..09a84ce25588e7 100644 --- a/third_party/xla/xla/backends/gpu/transforms/BUILD +++ b/third_party/xla/xla/backends/gpu/transforms/BUILD @@ -2319,8 +2319,8 @@ cc_library( "//xla:xla_proto_cc", "//xla/hlo/ir:hlo", "//xla/service:computation_layout", + "//xla/service:hlo_value", "//xla/service:layout_assignment", - "//xla/service:logical_buffer", "//xla/service:matmul_indexing_utils", "//xla/service:memory_annotations_hdr", "//xla/service/gpu:conv_utils", @@ -3280,6 +3280,7 @@ cc_library( "//xla/service/gpu/model:triton_emitter_constraints", "//xla/stream_executor:device_description", "//xla/tools:hlo_decomposer_lib", + "//xla/tsl/platform:env", "@com_google_absl//absl/algorithm:container", "@com_google_absl//absl/container:flat_hash_map", "@com_google_absl//absl/container:flat_hash_set", @@ -3311,8 +3312,10 @@ xla_cc_test( "//xla/service/gpu:alias_info", "//xla/service/gpu:backend_configs_cc", "//xla/service/gpu:gpu_device_info_for_tests", + "//xla/service/gpu/model:gpu_indexing_performance_model", "//xla/stream_executor:device_description", "//xla/stream_executor/cuda:cuda_compute_capability", + "//xla/tsl/platform:env", "//xla/tsl/platform:errors", "@com_google_absl//absl/log", "@com_google_absl//absl/log:check", diff --git a/third_party/xla/xla/backends/gpu/transforms/fusion_block_level_rewriter.cc b/third_party/xla/xla/backends/gpu/transforms/fusion_block_level_rewriter.cc index 1304cdeb81d8a7..027abcc82d77a6 100644 --- a/third_party/xla/xla/backends/gpu/transforms/fusion_block_level_rewriter.cc +++ b/third_party/xla/xla/backends/gpu/transforms/fusion_block_level_rewriter.cc @@ -268,9 +268,10 @@ absl::StatusOr ProcessFusionInstruction( auto fusion_adaptor = HloFusionAdaptor::ForInstruction( Cast(fusion_instruction)); - ABSL_ASSIGN_OR_RETURN( - TiledRunTimeDataOrError tiled_runtime_data_or_error, - indexing_performance_model.TryFindBestTilingForFusion(*fusion_adaptor)); + ABSL_ASSIGN_OR_RETURN(TiledRunTimeDataOrError tiled_runtime_data_or_error, + indexing_performance_model + .TryFindBestTilingForFusionAsync(*fusion_adaptor) + .Await()); if (const auto* fusion_decision = std::get_if(&tiled_runtime_data_or_error)) { diff --git a/third_party/xla/xla/backends/gpu/transforms/layout_assignment.cc b/third_party/xla/xla/backends/gpu/transforms/layout_assignment.cc index a1f641becd85d1..ff0811a8a21a31 100644 --- a/third_party/xla/xla/backends/gpu/transforms/layout_assignment.cc +++ b/third_party/xla/xla/backends/gpu/transforms/layout_assignment.cc @@ -50,8 +50,8 @@ limitations under the License. #include "xla/service/gpu/matmul_utils.h" #include "xla/service/gpu/reduction_utils.h" #include "xla/service/gpu/stream_executor_util.h" +#include "xla/service/hlo_value.h" #include "xla/service/layout_assignment.h" -#include "xla/service/logical_buffer.h" #include "xla/service/matmul_indexing_utils.h" #include "xla/service/memory_annotations.h" #include "xla/shape.h" @@ -253,14 +253,13 @@ absl::Status GpuLayoutAssignment::AddBackendConstraintsToDnnConvCustomCall( // The custom call returns a tuple of (actual_result, scratch_buffer); // call_result_buf is the logical buffer for actual_result, the thing that // contains the result of the conv call. - ABSL_ASSIGN_OR_RETURN( - const LogicalBuffer* call_result_buf, - points_to_analysis_->GetBufferDefinedAt(instr, /*index=*/{0})); + const HloValue& call_result_buf = + dataflow_analysis_->GetValueDefinedAt(instr, /*index=*/{0}); // Set layouts of the instructions' shapes. ABSL_RETURN_IF_ERROR(SetOperandLayout(lhs_shape, instr, 0)); ABSL_RETURN_IF_ERROR(SetOperandLayout(rhs_shape, instr, 1)); - ABSL_RETURN_IF_ERROR(SetBufferLayout(result_shape.layout(), *call_result_buf)); + ABSL_RETURN_IF_ERROR(SetBufferLayout(result_shape.layout(), call_result_buf)); // For fused convolutions, instr->operand(2), if exists, is the bias buffer. // There is no need to assign layout to it, as it has only one dimension. // instr->operand(3), if exists, is the side input buffer. @@ -648,30 +647,22 @@ absl::Status GpuLayoutAssignment::AddBackendConstraints( Shape shape = instruction->operand(i)->shape(); *shape.mutable_layout() = keys_layout; ABSL_RETURN_IF_ERROR(SetOperandLayout(shape, instruction, i)); - const LogicalBuffer* output_buffer; - if (instruction->shape().IsArray()) { - ABSL_ASSIGN_OR_RETURN( - output_buffer, - points_to_analysis_->GetBufferDefinedAt(instruction, {})); - } else { - ABSL_ASSIGN_OR_RETURN( - output_buffer, - points_to_analysis_->GetBufferDefinedAt(instruction, {i})); - } - ABSL_RETURN_IF_ERROR(SetBufferLayout(keys_layout, *output_buffer)); + const HloValue& output_buffer = + instruction->shape().IsArray() + ? dataflow_analysis_->GetValueDefinedAt(instruction, {}) + : dataflow_analysis_->GetValueDefinedAt(instruction, {i}); + ABSL_RETURN_IF_ERROR(SetBufferLayout(keys_layout, output_buffer)); } } else if (IsCustomCallToTopK(*instruction)) { // The output of the TopK custom call needs to have default layout. Layout default_layout = LayoutUtil::GetDefaultLayoutForRank( instruction->operand(0)->shape().dimensions().size()); - ABSL_ASSIGN_OR_RETURN( - auto values_buffer, - points_to_analysis_->GetBufferDefinedAt(instruction, {0})); - ABSL_RETURN_IF_ERROR(SetBufferLayout(default_layout, *values_buffer)); - ABSL_ASSIGN_OR_RETURN( - auto indices_buffer, - points_to_analysis_->GetBufferDefinedAt(instruction, {1})); - ABSL_RETURN_IF_ERROR(SetBufferLayout(default_layout, *indices_buffer)); + const HloValue& values_buffer = + dataflow_analysis_->GetValueDefinedAt(instruction, {0}); + ABSL_RETURN_IF_ERROR(SetBufferLayout(default_layout, values_buffer)); + const HloValue& indices_buffer = + dataflow_analysis_->GetValueDefinedAt(instruction, {1}); + ABSL_RETURN_IF_ERROR(SetBufferLayout(default_layout, indices_buffer)); } else if (HloPredicateIsOp(instruction)) { Shape operand_shape = instruction->operand(0)->shape(); Shape output_shape = instruction->shape(); diff --git a/third_party/xla/xla/backends/gpu/transforms/layout_assignment_a100.hlo b/third_party/xla/xla/backends/gpu/transforms/layout_assignment_a100.hlo index 3b5c92a19388d2..cac1d16a5b4b3b 100644 --- a/third_party/xla/xla/backends/gpu/transforms/layout_assignment_a100.hlo +++ b/third_party/xla/xla/backends/gpu/transforms/layout_assignment_a100.hlo @@ -1,17 +1,18 @@ // RUN: hlo-opt %s --platform=gpu --stage=hlo --xla_gpu_target_config_filename=%S/../../../backends/gpu/target_config/specs/a100_pcie_80.txtpb --split-input-file | FileCheck %s - +// CHECK: %conv_fprop_fusion_comp +// CHECK: convolution +// CHECK-SAME: window={size=3x3 pad=1_1x1_1}, dim_labels=b01f_o01i->b01f, convolution_kind=fprop // CHECK: %wrapped_transpose_computation // CHECK-NEXT: bf16[3,3,16,32]{3,2,1,0} parameter(0) -// CHECK-NEXT: bf16[32,3,3,16]{3,2,1,0} transpose +// CHECK-NEXT: ROOT {{.*}} = bf16[32,3,3,16]{3,2,1,0} transpose // CHECK-SAME: dimensions={3,0,1,2} -// CHECK: (bf16[1,64,64,32]{3,2,1,0}, u8[0]{0}) custom-call -// CHECK-SAME: window={size=3x3 pad=1_1x1_1}, dim_labels=b01f_o01i->b01f, custom_call_target="__cudnn$convForward - +// CHECK: ENTRY %main +// CHECK: %[[WRAPPED_TRANSPOSE:.*]] = bf16[32,3,3,16]{3,2,1,0} fusion({{.*}}), {{.*}}calls=%wrapped_transpose_computation +// CHECK: ROOT {{.*}} = bf16[1,64,64,32]{3,2,1,0} fusion({{.*}}, %[[WRAPPED_TRANSPOSE]]), kind=kCustom, calls=%conv_fprop_fusion_comp HloModule ConvCuDNN - ENTRY main { Arg_0.1 = bf16[1,64,64,16]{3,2,1,0} parameter(0), sharding={replicated} Arg_1.2 = bf16[3,3,16,32]{3,2,1,0} parameter(1), sharding={replicated} ROOT convolution.3 = bf16[1,64,64,32]{3,2,1,0} convolution(Arg_0.1, Arg_1.2), window={size=3x3 pad=1_1x1_1}, dim_labels=b01f_01io->b01f -} +} \ No newline at end of file diff --git a/third_party/xla/xla/backends/gpu/transforms/layout_assignment_h100.hlo b/third_party/xla/xla/backends/gpu/transforms/layout_assignment_h100.hlo index 1b82bb55c80b2b..5d4b5daf721ee1 100644 --- a/third_party/xla/xla/backends/gpu/transforms/layout_assignment_h100.hlo +++ b/third_party/xla/xla/backends/gpu/transforms/layout_assignment_h100.hlo @@ -1,17 +1,18 @@ // RUN: hlo-opt %s --platform=gpu --stage=hlo --xla_gpu_target_config_filename=%S/../../../backends/gpu/target_config/specs/h100_sxm.txtpb --split-input-file | FileCheck %s - +// CHECK: %conv_fprop_fusion_comp +// CHECK: convolution +// CHECK-SAME: window={size=3x3 pad=1_1x1_1}, dim_labels=b01f_o01i->b01f, convolution_kind=fprop // CHECK: %wrapped_transpose_computation // CHECK-NEXT: f8e4m3fn[3,3,16,32]{3,2,1,0} parameter(0) -// CHECK-NEXT: f8e4m3fn[32,3,3,16]{3,2,1,0} transpose +// CHECK-NEXT: ROOT {{.*}} = f8e4m3fn[32,3,3,16]{3,2,1,0} transpose // CHECK-SAME: dimensions={3,0,1,2} -// CHECK: (f8e4m3fn[1,64,64,32]{3,2,1,0}, u8[0]{0}) custom-call -// CHECK-SAME: window={size=3x3 pad=1_1x1_1}, dim_labels=b01f_o01i->b01f, custom_call_target="__cudnn$convForward - +// CHECK: ENTRY %main +// CHECK: %[[WRAPPED_TRANSPOSE:.*]] = f8e4m3fn[32,3,3,16]{3,2,1,0} fusion({{.*}}), {{.*}}calls=%wrapped_transpose_computation +// CHECK: ROOT {{.*}} = f8e4m3fn[1,64,64,32]{3,2,1,0} fusion({{.*}}, %[[WRAPPED_TRANSPOSE]]), kind=kCustom, calls=%conv_fprop_fusion_comp HloModule ConvCuDNN - ENTRY main { Arg_0.1 = f8e4m3fn[1,64,64,16]{3,2,1,0} parameter(0), sharding={replicated} Arg_1.2 = f8e4m3fn[3,3,16,32]{3,2,1,0} parameter(1), sharding={replicated} ROOT convolution.3 = f8e4m3fn[1,64,64,32]{3,2,1,0} convolution(Arg_0.1, Arg_1.2), window={size=3x3 pad=1_1x1_1}, dim_labels=b01f_01io->b01f -} +} \ No newline at end of file diff --git a/third_party/xla/xla/backends/gpu/transforms/layout_assignment_v100.hlo b/third_party/xla/xla/backends/gpu/transforms/layout_assignment_v100.hlo index d5baeb8a42af7d..ec15e5da5f3e92 100644 --- a/third_party/xla/xla/backends/gpu/transforms/layout_assignment_v100.hlo +++ b/third_party/xla/xla/backends/gpu/transforms/layout_assignment_v100.hlo @@ -1,17 +1,18 @@ // RUN: hlo-opt %s --platform=gpu --stage=hlo --xla_gpu_target_config_filename=%S/../../../backends/gpu/target_config/specs/v100.txtpb --split-input-file | FileCheck %s - +// CHECK: %conv_fprop_fusion_comp +// CHECK: convolution +// CHECK-SAME: window={size=3x3 pad=1_1x1_1}, dim_labels=b01f_o01i->b01f, convolution_kind=fprop // CHECK: %wrapped_transpose_computation // CHECK-NEXT: f16[3,3,16,32]{3,2,1,0} parameter(0) -// CHECK-NEXT: f16[32,3,3,16]{3,2,1,0} transpose +// CHECK-NEXT: ROOT {{.*}} = f16[32,3,3,16]{3,2,1,0} transpose // CHECK-SAME: dimensions={3,0,1,2} -// CHECK: (f16[1,64,64,32]{3,2,1,0}, u8[0]{0}) custom-call -// CHECK-SAME: window={size=3x3 pad=1_1x1_1}, dim_labels=b01f_o01i->b01f, custom_call_target="__cudnn$convForward - +// CHECK: ENTRY %main +// CHECK: %[[WRAPPED_TRANSPOSE:.*]] = f16[32,3,3,16]{3,2,1,0} fusion({{.*}}), {{.*}}calls=%wrapped_transpose_computation +// CHECK: ROOT {{.*}} = f16[1,64,64,32]{3,2,1,0} fusion({{.*}}, %[[WRAPPED_TRANSPOSE]]), kind=kCustom, calls=%conv_fprop_fusion_comp HloModule ConvCuDNN - ENTRY main { Arg_0.1 = f16[1,64,64,16]{3,2,1,0} parameter(0), sharding={replicated} Arg_1.2 = f16[3,3,16,32]{3,2,1,0} parameter(1), sharding={replicated} ROOT convolution.3 = f16[1,64,64,32]{3,2,1,0} convolution(Arg_0.1, Arg_1.2), window={size=3x3 pad=1_1x1_1}, dim_labels=b01f_01io->b01f -} +} \ No newline at end of file diff --git a/third_party/xla/xla/backends/gpu/transforms/priority_fusion.cc b/third_party/xla/xla/backends/gpu/transforms/priority_fusion.cc index 909690650944b6..77fe80a539d3a0 100644 --- a/third_party/xla/xla/backends/gpu/transforms/priority_fusion.cc +++ b/third_party/xla/xla/backends/gpu/transforms/priority_fusion.cc @@ -723,7 +723,8 @@ class PriorityFusionQueue { producer, consumer, use_multi_output_fusion); absl::StatusOr result_or_status = - combined_gpu_performance_model_.TryFindBestTilingForFusion(*fusion); + combined_gpu_performance_model_.TryFindBestTilingForFusionAsync(*fusion) + .Await(); // Convert absl::Status into FusionDecision. We don't distinguish between // status and FusionDecision here, because both indicate that tile analysis diff --git a/third_party/xla/xla/backends/gpu/transforms/priority_fusion_test.cc b/third_party/xla/xla/backends/gpu/transforms/priority_fusion_test.cc index b3c092c3d9db6e..64e0197f544824 100644 --- a/third_party/xla/xla/backends/gpu/transforms/priority_fusion_test.cc +++ b/third_party/xla/xla/backends/gpu/transforms/priority_fusion_test.cc @@ -1329,7 +1329,7 @@ TEST_P(PriorityFusionTest, DoNotFuseInsideReducer) { } TEST_P(PriorityFusionTest, SkipsTilingsWithInfiniteRuntime) { - // This test verifies the fix in TryFindBestTilingForFusion that skips + // This test verifies the fix in TryFindBestTilingForFusionAsync that skips // tilings with infinite runtime estimates. // // The fix: After estimating runtime for each tiling candidate, check if diff --git a/third_party/xla/xla/backends/gpu/transforms/softmax_rewriter_triton.cc b/third_party/xla/xla/backends/gpu/transforms/softmax_rewriter_triton.cc index d863bea48b0f4b..d16ae5e238a76d 100644 --- a/third_party/xla/xla/backends/gpu/transforms/softmax_rewriter_triton.cc +++ b/third_party/xla/xla/backends/gpu/transforms/softmax_rewriter_triton.cc @@ -66,6 +66,7 @@ limitations under the License. #include "xla/shape.h" #include "xla/stream_executor/device_description.h" #include "xla/tools/hlo_decomposer.h" +#include "xla/tsl/platform/threadpool.h" #include "xla/util.h" #include "xla/xla_data.pb.h" @@ -353,12 +354,16 @@ DecideIfShouldFuseAndMaybeSetBlockLevelParameters( const se::DeviceDescription& device_info, const HloCostAnalysis::ShapeSizeFunction& shape_size, const GpuAliasInfo* alias_info, MLIRContext* mlir_context, - bool use_cost_model_to_evaluate_fusions) { + bool use_cost_model_to_evaluate_fusions, + tsl::thread::ThreadPool* thread_pool = nullptr) { auto fusion_adaptor = HloFusionAdaptor::ForInstruction(normalization_fusion); - ABSL_ASSIGN_OR_RETURN( - TiledRunTimeDataOrError tiled_runtime_data_or, - indexing_performance_model.TryFindBestTilingForFusion(*fusion_adaptor)); + ABSL_ASSIGN_OR_RETURN(TiledRunTimeDataOrError tiled_runtime_data_or, + indexing_performance_model + .TryFindBestTilingForFusionAsync( + *fusion_adaptor, + thread_pool ? thread_pool->AsExecutor() : nullptr) + .Await()); if (const auto* fusion_decision = std::get_if(&tiled_runtime_data_or)) { @@ -406,7 +411,8 @@ absl::StatusOr MaybeFuseDiamondImpl( const se::DeviceDescription& device_info, const HloCostAnalysis::ShapeSizeFunction& shape_size, const GpuAliasInfo* alias_info, MLIRContext* mlir_context, - bool use_cost_model_to_evaluate_fusions) { + bool use_cost_model_to_evaluate_fusions, + tsl::thread::ThreadPool* thread_pool = nullptr) { ABSL_ASSIGN_OR_RETURN(HloFusionInstruction * normalization_fusion, MakeFusionForDiamond(diamond)); HloInstruction* root = diamond.root; @@ -417,7 +423,7 @@ absl::StatusOr MaybeFuseDiamondImpl( DecideIfShouldFuseAndMaybeSetBlockLevelParameters( normalization_fusion, indexing_performance_model, device_info, shape_size, alias_info, mlir_context, - use_cost_model_to_evaluate_fusions)); + use_cost_model_to_evaluate_fusions, thread_pool)); if (fusion_decision.IsForbidden()) { VLOG(2) << "Not fusing: " << fusion_decision.Explain(); @@ -674,11 +680,13 @@ absl::StatusOr SoftmaxRewriterTriton::MaybeFuseNormalizationDiamond( diamond.root->GetModule() ->config() .debug_options() - .xla_gpu_experimental_enable_same_shape_multi_output_fusion()); + .xla_gpu_experimental_enable_same_shape_multi_output_fusion(), + mlir_context_pool_); return MaybeFuseDiamondImpl(diamond, indexing_performance_model, device_info_, shape_size_, alias_info_, mlir_context_, - use_cost_model_to_evaluate_fusions_); + use_cost_model_to_evaluate_fusions_, + thread_pool_); } absl::StatusOr SoftmaxRewriterTriton::RunImpl( diff --git a/third_party/xla/xla/backends/gpu/transforms/softmax_rewriter_triton.h b/third_party/xla/xla/backends/gpu/transforms/softmax_rewriter_triton.h index 426dc73c7ef912..19b442ab94efa9 100644 --- a/third_party/xla/xla/backends/gpu/transforms/softmax_rewriter_triton.h +++ b/third_party/xla/xla/backends/gpu/transforms/softmax_rewriter_triton.h @@ -27,10 +27,15 @@ limitations under the License. #include "xla/hlo/ir/hlo_module.h" #include "xla/hlo/pass/hlo_pass_interface.h" #include "xla/service/gpu/alias_info.h" +#include "xla/service/gpu/model/gpu_indexing_performance_model.h" #include "xla/service/hlo_cost_analysis.h" #include "xla/service/instruction_fusion.h" #include "xla/stream_executor/device_description.h" +namespace tsl::thread { +class ThreadPool; +} // namespace tsl::thread + namespace xla { namespace gpu { @@ -55,13 +60,17 @@ class SoftmaxRewriterTriton : public HloModulePass { const GpuAliasInfo* alias_info, mlir::MLIRContext* mlir_context, bool only_fuse_if_profitable, - bool use_experimental_tiling) + bool use_experimental_tiling, + tsl::thread::ThreadPool* thread_pool = nullptr, + MlirContextPool* mlir_context_pool = nullptr) : device_info_(device_info), shape_size_(shape_size), alias_info_(alias_info), use_cost_model_to_evaluate_fusions_(only_fuse_if_profitable), use_experimental_tiling_(use_experimental_tiling), - mlir_context_(mlir_context) {} + mlir_context_(mlir_context), + thread_pool_(thread_pool), + mlir_context_pool_(mlir_context_pool) {} absl::string_view name() const override { return "triton-softmax-rewriter"; } @@ -114,6 +123,8 @@ class SoftmaxRewriterTriton : public HloModulePass { bool use_cost_model_to_evaluate_fusions_; bool use_experimental_tiling_; mlir::MLIRContext* mlir_context_; + tsl::thread::ThreadPool* thread_pool_ = nullptr; + MlirContextPool* mlir_context_pool_ = nullptr; }; } // namespace gpu diff --git a/third_party/xla/xla/backends/gpu/transforms/softmax_rewriter_triton_test.cc b/third_party/xla/xla/backends/gpu/transforms/softmax_rewriter_triton_test.cc index 6af936e36b83f7..182225da938efb 100644 --- a/third_party/xla/xla/backends/gpu/transforms/softmax_rewriter_triton_test.cc +++ b/third_party/xla/xla/backends/gpu/transforms/softmax_rewriter_triton_test.cc @@ -35,12 +35,15 @@ limitations under the License. #include "xla/service/gpu/alias_info.h" #include "xla/service/gpu/backend_configs.pb.h" #include "xla/service/gpu/gpu_device_info_for_tests.h" +#include "xla/service/gpu/model/gpu_indexing_performance_model.h" #include "xla/service/hlo_cost_analysis.h" #include "xla/service/instruction_fusion.h" #include "xla/service/pattern_matcher.h" #include "xla/stream_executor/cuda/cuda_compute_capability.h" #include "xla/stream_executor/device_description.h" +#include "xla/tsl/platform/env.h" #include "xla/tsl/platform/errors.h" +#include "xla/tsl/platform/threadpool.h" #include "xla/xla.pb.h" namespace xla { @@ -1165,6 +1168,49 @@ ENTRY main { .WithPredicate(HasBlockLevelFusionConfig))); } +TEST_P(SoftmaxRewriterTritonTest, CanFuseWithParallelTilingSearch) { + const std::string hlo_string = R"( +HloModule softmax +max_computation { + arg_0 = f32[] parameter(0) + arg_1 = f32[] parameter(1) + ROOT maximum = f32[] maximum(arg_0, arg_1) +} +add_computation { + arg_0 = f32[] parameter(0) + arg_1 = f32[] parameter(1) + ROOT add = f32[] add(arg_0, arg_1) +} +ENTRY main { + param_0 = f32[127,125]{1,0} parameter(0) + constant_neg_inf = f32[] constant(-inf) + reduce = f32[127]{0} reduce(param_0, constant_neg_inf), dimensions={1}, to_apply=max_computation + broadcast = f32[127,125]{1,0} broadcast(reduce), dimensions={0} + ROOT subtract = f32[127,125]{1,0} subtract(param_0, broadcast) +})"; + ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); + tsl::thread::ThreadPool thread_pool(tsl::Env::Default(), "test_pool", 4); + // Mirrors the contexts GpuCompiler pools: multithreading is disabled, so the + // cost model must give each candidate its own context. + MlirContextPool mlir_context_pool( + [] { + return std::make_unique( + mlir::MLIRContext::Threading::DISABLED); + }, + /*preallocate=*/4); + SoftmaxRewriterTriton rewriter( + device_info_, HloCostAnalysis::DefaultShapeSize, &alias_info_, + &mlir_context_, + /*only_fuse_if_profitable=*/false, + /*use_experimental_tiling=*/GetParam(), &thread_pool, &mlir_context_pool); + EXPECT_THAT(rewriter.Run(module.get()), IsOkAndHolds(true)); + EXPECT_TRUE(verifier().Run(module.get()).status().ok()); + EXPECT_THAT( + module->entry_computation()->root_instruction(), + GmockMatch( + m::Fusion(m::Parameter()).WithPredicate(HasBlockLevelFusionConfig))); +} + INSTANTIATE_TEST_SUITE_P(SoftmaxRewriterTritonTestSuite, SoftmaxRewriterTritonTest, ::testing::Bool(), [](const ::testing::TestParamInfo& info) { diff --git a/third_party/xla/xla/backends/profiler/cpu/BUILD b/third_party/xla/xla/backends/profiler/cpu/BUILD index 9487b110986d65..1b5b07f98ff20a 100644 --- a/third_party/xla/xla/backends/profiler/cpu/BUILD +++ b/third_party/xla/xla/backends/profiler/cpu/BUILD @@ -79,15 +79,17 @@ cc_library( ]), deps = [ "//xla:status_macros", - "//xla/backends/profiler/util:metadata_registry", + "//xla/backends/profiler/util:metadata_registry", # buildcleaner: keep "//xla/python/profiler/internal:python_hooks", "//xla/tsl/platform:logging", + "//xla/tsl/profiler/utils:time_utils", "//xla/tsl/profiler/utils:xplane_builder", "//xla/tsl/profiler/utils:xplane_schema", "//xla/tsl/profiler/utils:xplane_utils", "@com_google_absl//absl/log", "@com_google_absl//absl/status", "@com_google_absl//absl/status:statusor", + "@rules_python//python/cc:current_py_cc_headers", # buildcleaner: keep "@tsl//tsl/profiler/lib:profiler_interface", "@tsl//tsl/profiler/protobuf:xplane_proto_cc", ], diff --git a/third_party/xla/xla/backends/profiler/cpu/python_tracer.cc b/third_party/xla/xla/backends/profiler/cpu/python_tracer.cc index f09bb2b1e33dd5..38dce9159357dd 100644 --- a/third_party/xla/xla/backends/profiler/cpu/python_tracer.cc +++ b/third_party/xla/xla/backends/profiler/cpu/python_tracer.cc @@ -14,8 +14,10 @@ limitations under the License. ==============================================================================*/ #include "xla/backends/profiler/cpu/python_tracer.h" +#include #include #include +#include #include #include @@ -25,6 +27,7 @@ limitations under the License. #include "xla/python/profiler/internal/python_hooks.h" #include "xla/status_macros.h" #include "xla/tsl/platform/logging.h" +#include "xla/tsl/profiler/utils/time_utils.h" #include "xla/tsl/profiler/utils/xplane_builder.h" #include "xla/tsl/profiler/utils/xplane_schema.h" #include "xla/tsl/profiler/utils/xplane_utils.h" @@ -57,6 +60,8 @@ class PythonTracer : public tsl::profiler::ProfilerInterface { private: bool recording_ = false; const PythonHooksOptions options_; + // Start timestamp (in nanoseconds) of the current recording/consume window. + uint64_t start_timestamp_ns_ = 0; std::unique_ptr context_; PythonTracer(const PythonTracer&) = delete; @@ -71,6 +76,7 @@ absl::Status PythonTracer::Start() { // TENSORFLOW_STATUS_OK } VLOG(1) << __FUNCTION__; recording_ = true; + start_timestamp_ns_ = tsl::profiler::GetCurrentTimeNanos(); PythonHooks::GetSingleton()->Start(options_); return absl::OkStatus(); } @@ -98,6 +104,8 @@ absl::Status PythonTracer::CollectData( // TENSORFLOW_STATUS_OK absl::StatusOr PythonTracer::Consume() { VLOG(1) << "PythonTracer::Consume called, recording=" << recording_; PythonTracerChunk chunk; + chunk.start_timestamp_ns = + std::exchange(start_timestamp_ns_, tsl::profiler::GetCurrentTimeNanos()); if (recording_) { chunk.consumed_data = PythonHooks::GetSingleton()->Consume(); } else if (context_) { @@ -110,7 +118,7 @@ absl::StatusOr PythonTracer::Consume() { size_t estimated_size = 0; size_t total_events = 0; - for (const auto& thread_data : chunk.consumed_data) { + for (const PerThreadConsumeData& thread_data : chunk.consumed_data) { estimated_size += sizeof(PerThreadConsumeData) + thread_data.events.size() * sizeof(TraceEventInfo); total_events += thread_data.events.size(); @@ -141,10 +149,36 @@ absl::Status PythonTracer::Serialize(std::any data, space, tsl::profiler::kPythonTracerPlaneName); tsl::profiler::XPlaneBuilder plane(raw_plane); - for (const auto& thread_data : chunk->consumed_data) { + for (const PerThreadConsumeData& thread_data : chunk->consumed_data) { + // Skip threads with no recorded events. + if (thread_data.events.empty()) { + continue; + } + + // Anchor the line to the chunk start timestamp, clamping down to the + // earliest event start timestamp if any event started before the chunk + // window (or if the chunk start timestamp was unset). This guarantees + // non-negative event offsets and prevents int64_t picosecond overflow. + uint64_t line_start_ns = chunk->start_timestamp_ns > 0 + ? chunk->start_timestamp_ns + : thread_data.events.front().start_time_ns; + for (const TraceEventInfo& event : thread_data.events) { + line_start_ns = std::min(line_start_ns, event.start_time_ns); + } + tsl::profiler::XLineBuilder line = plane.GetOrCreateLine(thread_data.thread_id); - for (const auto& event : thread_data.events) { + if (line.NumEvents() == 0 && line.TimestampNs() == 0) { + // Newly created lines have no events yet (since empty event chunks are + // skipped above), so we can set the initial baseline timestamp directly. + line.SetTimestampNs(line_start_ns); + } else if (line_start_ns < static_cast(line.TimestampNs())) { + // If this is a subsequent chunk, with event timestamps earlier than the + // existing line start timestamp, we adjust the line start timestamp and + // event offsets. + line.SetTimestampNsAndAdjustEventOffsets(line_start_ns); + } + for (const TraceEventInfo& event : thread_data.events) { tsl::profiler::XEventBuilder xevent = line.AddEvent(*plane.GetOrCreateEventMetadata(event.name)); xevent.SetTimestampNs(event.start_time_ns); diff --git a/third_party/xla/xla/backends/profiler/cpu/python_tracer.h b/third_party/xla/xla/backends/profiler/cpu/python_tracer.h index 76c4d0d498ab08..5d9b5fc8fef1be 100644 --- a/third_party/xla/xla/backends/profiler/cpu/python_tracer.h +++ b/third_party/xla/xla/backends/profiler/cpu/python_tracer.h @@ -15,17 +15,22 @@ limitations under the License. #ifndef XLA_BACKENDS_PROFILER_CPU_PYTHON_TRACER_H_ #define XLA_BACKENDS_PROFILER_CPU_PYTHON_TRACER_H_ +#include #include #include #include "xla/python/profiler/internal/python_hooks.h" #include "tsl/profiler/lib/profiler_interface.h" -#include "tsl/profiler/protobuf/xplane.pb.h" namespace xla { namespace profiler { +// A chunk of Python trace data collected during a streaming profiling interval. struct PythonTracerChunk { + // Timestamp (in nanoseconds) when this chunk's profiling interval began + // (either at Start() or at the preceding Consume() call). Used as the + // initial baseline timestamp for thread XLines during Serialize(). + uint64_t start_timestamp_ns = 0; std::vector consumed_data; }; diff --git a/third_party/xla/xla/benchmarks/README.md b/third_party/xla/xla/benchmarks/README.md index 484da58bc44987..e3efcc33d4cdbc 100644 --- a/third_party/xla/xla/benchmarks/README.md +++ b/third_party/xla/xla/benchmarks/README.md @@ -24,6 +24,9 @@ python3 xla/benchmarks/run_benchmarks.py --benchmarks=dense_matmul --csv_path=

+# Run collectives benchmark suite and write results to a CSV file +python3 xla/benchmarks/run_benchmarks.py --benchmarks=collectives --csv_path= + # Run full benchmark suite and write results to multiple CSV files python3 xla/benchmarks/run_benchmarks.py --csv_path= ``` diff --git a/third_party/xla/xla/benchmarks/benchmark_configs.py b/third_party/xla/xla/benchmarks/benchmark_configs.py index d13b0155b4acdf..fe46e6393cc326 100644 --- a/third_party/xla/xla/benchmarks/benchmark_configs.py +++ b/third_party/xla/xla/benchmarks/benchmark_configs.py @@ -21,6 +21,8 @@ from jax.experimental.pallas import tpu as pltpu import jax.numpy as jnp +from xla.benchmarks.collectives import all_gather_benchmark +from xla.benchmarks.collectives import collectives_benchmark from xla.benchmarks.core import platform_info from xla.benchmarks.dma_microbenchmarks import chip_to_chip_dma_benchmark from xla.benchmarks.dma_microbenchmarks import chiplet_to_chiplet_dma_benchmark @@ -232,6 +234,28 @@ def get_dma_configs( return configs +_COLLECTIVE_SUITES: tuple[ + tuple[str, type[collectives_benchmark.CollectivesBenchmarks]], ... +] = (("all_gather", all_gather_benchmark.AllGatherBenchmarks),) + + +def get_collectives_configs( + chip_version: pltpu.ChipVersion | None = None, +) -> list[collectives_benchmark.CollectiveBenchmarkConfig]: + """Generates configs for all collective microbenchmarks.""" + del chip_version # Unused. + configs = [] + for suite, test_class in _COLLECTIVE_SUITES: + for name in sorted(dir(test_class)): + if name.startswith("test_"): + configs.append( + collectives_benchmark.CollectiveBenchmarkConfig( + suite=suite, test_class=test_class, test_name=name + ) + ) + return configs + + BENCHMARK_FACTORIES: Mapping[ str, Callable[[pltpu.ChipVersion | None], list[Any]] ] = immutabledict.immutabledict({ @@ -239,4 +263,5 @@ def get_dma_configs( "subchannel_matmul": get_subchannel_matmul_configs, "jax_matmul": get_jax_matmul_configs, "dma": get_dma_configs, + "collectives": get_collectives_configs, }) diff --git a/third_party/xla/xla/benchmarks/benchmark_configs_test.py b/third_party/xla/xla/benchmarks/benchmark_configs_test.py index acfa18f988de10..638ee240c8205f 100644 --- a/third_party/xla/xla/benchmarks/benchmark_configs_test.py +++ b/third_party/xla/xla/benchmarks/benchmark_configs_test.py @@ -23,6 +23,7 @@ from xla.benchmarks import benchmark_configs from xla.benchmarks import results_utils +from xla.benchmarks.collectives import collectives_benchmark from xla.benchmarks.dma_microbenchmarks import memory_base from xla.benchmarks.jax_microbenchmarks import matmul_lib from xla.benchmarks.pallas_microbenchmarks import dense_matmul_lib @@ -214,6 +215,78 @@ def test_bw(self): self.assertAlmostEqual(df["latency_us"].iloc[0], 15.0) self.assertGreater(df["bandwidth_gbps"].iloc[0], 0.0) + def test_collectives_configs(self): + configs = benchmark_configs.get_collectives_configs() + self.assertLen(configs, 2) + for cfg in configs: + self.assertIsInstance( + cfg, collectives_benchmark.CollectiveBenchmarkConfig + ) + self.assertEqual(cfg.suite, "all_gather") + self.assertStartsWith(cfg.test_name, "test_all_gather_") + self.assertEqual( + cfg.as_dict(), {"suite": cfg.suite, "test": cfg.test_name} + ) + + def test_collectives_benchmark_end_to_end_and_skip(self): + # 1. On CPU without TPU, CollectiveBenchmark.run() catches SkipTest. + cfg_skip = benchmark_configs.get_collectives_configs()[0] + self.assertEqual(cfg_skip.get_benchmark().run(), [None]) + + # 2. With TPU info patched using spec_set=True, run() executes the test + # method, records bandwidth/latency metrics, and populates the table. + class _FakeCollectiveSuite(collectives_benchmark.CollectivesBenchmarks): + + def setUp(self): + absltest.TestCase.setUp(self) + self.latencies_us = [] + self.metrics = {} + + def test_ag(self): + self._print_collective_bandwidth_statistics( + op_name="all-gather-major-dim", + per_device_size_bytes=256 * 1024 * 1024, + total_size_bytes=1024 * 1024 * 1024, + bus_data_bytes=768 * 1024 * 1024, + num_devices=4, + latencies_us=[100.0, 200.0], + ) + + cfg = collectives_benchmark.CollectiveBenchmarkConfig( + suite="all_gather", test_class=_FakeCollectiveSuite, test_name="test_ag" + ) + with unittest.mock.patch.object( + pltpu, + "get_tpu_info", + spec_set=True, + return_value=pltpu.get_tpu_info_for_chip(pltpu.ChipVersion.TPU_V5E, 1), + ): + prof_results = cfg.get_benchmark().run() + + df = results_utils.create_results_table([(cfg, prof_results)]) + self.assertEqual( + list(df.columns), + [ + "suite", + "test", + "per_device_size_mib", + "total_size_mib", + "num_devices", + "bus_bandwidth_gbps", + "effective_bandwidth_gbps", + "latency_us", + "flops", + ], + ) + self.assertEqual(df["suite"].iloc[0], "all_gather") + self.assertEqual(df["test"].iloc[0], "test_ag") + self.assertEqual(df["per_device_size_mib"].iloc[0], 256.0) + self.assertEqual(df["total_size_mib"].iloc[0], 1024.0) + self.assertEqual(df["num_devices"].iloc[0], 4) + self.assertAlmostEqual(df["latency_us"].iloc[0], 150.0) + self.assertGreater(df["bus_bandwidth_gbps"].iloc[0], 0.0) + self.assertGreater(df["effective_bandwidth_gbps"].iloc[0], 0.0) + if __name__ == "__main__": absltest.main() diff --git a/third_party/xla/xla/benchmarks/collectives/all_gather_benchmark.py b/third_party/xla/xla/benchmarks/collectives/all_gather_benchmark.py new file mode 100644 index 00000000000000..5eccc95260ae88 --- /dev/null +++ b/third_party/xla/xla/benchmarks/collectives/all_gather_benchmark.py @@ -0,0 +1,160 @@ +# Copyright 2026 The OpenXLA Authors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""JAX microbenchmark for TPU all-gather collective bandwidth and latency.""" + +from absl.testing import absltest +import jax +from jax import numpy as jnp +from jax.experimental import shard_map +import numpy as np + +from xla.benchmarks.collectives import collectives_benchmark + + +def zero_crop(x): + """Prevents x from being a live-out buffer by passing through a ZeroCrop FFI call.""" + if jax.default_backend() == "cpu": + return x + return jax.ffi.ffi_call( + "ZeroCrop", + result_shape_dtypes=jax.ShapeDtypeStruct(x.shape, x.dtype), + has_side_effect=True, + )(x) + + +class AllGatherBenchmarks(collectives_benchmark.CollectivesBenchmarks): + """JAX benchmarks for measuring single-tray all-gather bandwidth and latency.""" + + def _setup_mesh_and_axis(self): + """Creates the single-tray device mesh and returns (mesh, ag_axis, num_devices).""" + devices = self.get_single_tray_devices() + num_devices = len(devices) + + num_chips = len({tuple(d.coords) for d in devices}) + cores_per_chip = num_devices // num_chips + + if cores_per_chip > 1: + mesh = jax.sharding.Mesh( + np.array(devices).reshape((num_chips, cores_per_chip)), + ("chips", "cores"), + ) + ag_axis = ("chips", "cores") + else: + mesh = jax.sharding.Mesh(np.array(devices), ("chips",)) + ag_axis = "chips" + + return mesh, ag_axis, num_devices + + def _run_all_gather_benchmark( + self, + global_shape, + in_partition, + gather_axis, + op_name, + ): + """Runs a single-HLO all-gather benchmark and logs bandwidth/latency stats.""" + mesh, ag_axis, num_devices = self._setup_mesh_and_axis() + out_partition = jax.sharding.PartitionSpec(*([None] * len(global_shape))) + + sharding = jax.sharding.NamedSharding(mesh, in_partition) + arr = jax.random.uniform( + jax.random.key(0), + global_shape, + dtype=jnp.bfloat16, + ) + array_on_devices = jax.device_put(arr, sharding) + per_device_bytes = arr.nbytes // num_devices + + @jax.jit + def all_gather_fn(x): + return shard_map.shard_map( + lambda x_shard: zero_crop( + jax.lax.all_gather(x_shard, ag_axis, axis=gather_axis) + ), + mesh=mesh, + in_specs=in_partition, + out_specs=out_partition, + check_rep=False, + )(x) + + # Warmup run to compile and initialize buffers before tracing. + warmup_res = all_gather_fn(array_on_devices) + jax.block_until_ready(warmup_res) + + execution_times_us = self._measure_execution_times( + all_gather_fn, + ["all-gather"], + array_on_devices, + ) + self._print_all_gather_bandwidth_statistics( + per_device_size_bytes=per_device_bytes, + num_devices=num_devices, + latencies_us=execution_times_us, + op_name=op_name, + ) + + def _print_all_gather_bandwidth_statistics( + self, + per_device_size_bytes, + num_devices, + latencies_us, + op_name="all-gather", + ): + """Computes all-gather data sizes and prints bandwidth/latency statistics.""" + total_gathered_bytes = per_device_size_bytes * num_devices + # In ring-based all-gather, each device transmits (N-1)/N of the + # gathered data. + bus_data_bytes = per_device_size_bytes * (num_devices - 1) + self._print_collective_bandwidth_statistics( + op_name=op_name, + per_device_size_bytes=per_device_size_bytes, + total_size_bytes=total_gathered_bytes, + bus_data_bytes=bus_data_bytes, + num_devices=num_devices, + latencies_us=latencies_us, + ) + + def _get_per_device_chunk_mib(self, num_devices): + """Returns per-device chunk size in MiB from collective_size_mib and num_devices.""" + if self.collective_size_mib % num_devices != 0: + raise ValueError( + f"collective_size_mib ({self.collective_size_mib}) must be divisible" + f" by num_devices ({num_devices})." + ) + return self.collective_size_mib // num_devices + + def test_all_gather_major_dim_bandwidth(self): + _, ag_axis, num_devices = self._setup_mesh_and_axis() + chunk_size_mib = self._get_per_device_chunk_mib(num_devices) + self._run_all_gather_benchmark( + global_shape=(num_devices, chunk_size_mib * 512, 1024), + in_partition=jax.sharding.PartitionSpec(ag_axis, None, None), + gather_axis=0, + op_name="all-gather-major-dim", + ) + + def test_all_gather_minor_dim_bandwidth(self): + _, ag_axis, num_devices = self._setup_mesh_and_axis() + chunk_size_mib = self._get_per_device_chunk_mib(num_devices) + self._run_all_gather_benchmark( + global_shape=(chunk_size_mib, num_devices * 512, 1024), + in_partition=jax.sharding.PartitionSpec(None, ag_axis, None), + gather_axis=1, + op_name="all-gather-minor-dim", + ) + + +if __name__ == "__main__": + jax.config.config_with_absl() + absltest.main() diff --git a/third_party/xla/xla/benchmarks/collectives/collectives_benchmark.py b/third_party/xla/xla/benchmarks/collectives/collectives_benchmark.py new file mode 100644 index 00000000000000..23b216829a7d86 --- /dev/null +++ b/third_party/xla/xla/benchmarks/collectives/collectives_benchmark.py @@ -0,0 +1,452 @@ +# Copyright 2026 The OpenXLA Authors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Base class and utilities for JAX collective microbenchmarks.""" + +import gzip +import json +import os +import pathlib +import shutil +import sys +import tempfile + +from absl import flags +from absl import logging +from absl.testing import absltest +import immutabledict +import jax +import jax.experimental.pallas.tpu as pltpu +import numpy as np + +from xla.benchmarks.core import benchmark +from xla.benchmarks.jax_microbenchmarks import jax_profiler_utils + +_NUMBER_OF_MEASUREMENTS = flags.DEFINE_integer( + "number_of_measurements", + default=5, + help="Number of measurements to take. Default: 5", + allow_override=True, +) + +_COLLECTIVE_SIZE_MIB = flags.DEFINE_integer( + "collective_size_mib", + default=1024, + help="Total collective size in MiB across all devices. Default: 1024", +) + +EXPECTED_CHIPS_PER_HOST = immutabledict.immutabledict({ + pltpu.ChipVersion.TPU_V4: 4, + pltpu.ChipVersion.TPU_V4I: 4, + pltpu.ChipVersion.TPU_V5E: 4, + pltpu.ChipVersion.TPU_V5P: 4, + pltpu.ChipVersion.TPU_V6E: 4, + pltpu.ChipVersion.TPU_7: 4, + pltpu.ChipVersion.TPU_7X: 4, + pltpu.ChipVersion.TPU_8I: 2, + pltpu.ChipVersion.TPU_8T: 2, +}) + + +class CollectivesBenchmarks(absltest.TestCase): + """Base benchmark class for measuring JAX collective operations.""" + + def setUp(self): + super().setUp() + if not any(device.platform == "tpu" for device in jax.devices()): + self.skipTest("This test requires TPU hardware.") + self.number_of_measurements = _NUMBER_OF_MEASUREMENTS.value + self.collective_size_mib = _COLLECTIVE_SIZE_MIB.value + self.latencies_us = [] + self.metrics = {} + + def verify_single_tray_setup(self): + """Verifies that the test is running on exactly one full TPU tray.""" + if not any(device.platform == "tpu" for device in jax.devices()): + self.skipTest("This test requires TPU hardware.") + + if jax.process_count() > 1: + self.skipTest( + "Single-host test expected process_count=1, got" + f" {jax.process_count()}." + ) + + local_devices = jax.local_devices() + all_devices = jax.devices() + if len(local_devices) != len(all_devices): + self.skipTest( + f"Not all devices are local to this host: local={len(local_devices)}," + f" all={len(all_devices)}." + ) + + tpu_info = pltpu.get_tpu_info() + chip_version = tpu_info.chip_version + if not isinstance(chip_version, pltpu.ChipVersion): + logging.warning( + "Unknown TPU generation %s for chip count verification.", chip_version + ) + return + expected_chips = EXPECTED_CHIPS_PER_HOST.get(chip_version) + if expected_chips is None: + logging.warning( + "Unknown TPU generation %s for chip count verification.", chip_version + ) + return + + # Check unique chip coordinates across local TPU devices. + coords = [tuple(d.coords) for d in local_devices] + num_unique_chips = len(set(coords)) + if num_unique_chips != expected_chips: + self.skipTest( + f"Expected {expected_chips} TPU chips for a single {chip_version}" + f" tray, but found {num_unique_chips}." + ) + + # Check device count based on cores per chip. + num_cores_per_chip = chip_version.num_physical_tensor_cores_per_chip + expected_devices_split = expected_chips * num_cores_per_chip + expected_devices_megacore = expected_chips + + if len(local_devices) not in ( + expected_devices_split, + expected_devices_megacore, + ): + self.skipTest( + f"Expected {expected_devices_split} (or {expected_devices_megacore}) " + f"devices for a single {chip_version} tray, got " + f"{len(local_devices)}." + ) + + def get_single_tray_devices(self): + """Verifies the single-tray setup and returns the list of local devices.""" + self.verify_single_tray_setup() + return jax.local_devices() + + def _normalize_kernel_names(self, kernel_names): + """Normalizes kernel names into a list of strings.""" + if isinstance(kernel_names, str): + return [kernel_names] + elif isinstance(kernel_names, (tuple, list)): + return list(kernel_names) + raise TypeError( + "`kernel_names` must be a string, tuple, or list of strings. " + f"Got: {type(kernel_names)}" + ) + + def _extract_execution_times(self, profiler_dir, kernel_names): + """Extracts execution times from profiler data for matching kernel names.""" + undeclared_outputs_dir = os.environ.get("TEST_UNDECLARED_OUTPUTS_DIR") + if undeclared_outputs_dir: + test_name = self._testMethodName + for profile_file in pathlib.Path(profiler_dir).rglob("*"): + if profile_file.is_file() and ( + profile_file.name.endswith(".xplane.pb") + or profile_file.name.endswith(".trace.json.gz") + ): + dest = ( + pathlib.Path(undeclared_outputs_dir) + / f"{test_name}_{profile_file.name}" + ) + shutil.copy2(profile_file, dest) + logging.info( + "Saved XProf profile artifact to undeclared outputs: %s", dest + ) + + trace_files = list(pathlib.Path(profiler_dir).glob("**/*.trace.json.gz")) + if not trace_files: + raise FileNotFoundError(f"Could not find trace.json.gz in {profiler_dir}") + + kernel_durs_per_core = {} + sc_kernel_durs_per_core = {} + for trace_file in trace_files: + with gzip.open(trace_file, "rt") as f: + trace_events = json.load(f).get("traceEvents", []) + + has_sc_offload = any( + e.get("args", {}).get("hlo_category") == "async-done" + and e.get("args", {}).get("offload_type") == "OFFLOAD_COLLECTIVE" + for e in trace_events + ) + matched_count = 0 + sample_events = [] + for event in trace_events: + event_name = event.get("name", "") + args = event.get("args", {}) + tf_op = args.get("tf_op", "") + hlo_op = args.get("hlo_op", "") + if "dur" in event: + is_host_wrapper = event_name.startswith(( + "PjitFunction", + "CommonPjRt", + "PJRT_", + "PjRt", + "$", + "Mutex", + "PythonRefManager", + "ParseArguments", + "Wait for ", + )) + if len(sample_events) < 30 and event_name and not is_host_wrapper: + sample_events.append( + f"pid={event.get('pid')} name='{event_name}'" + f" tf_op='{tf_op}' hlo_op='{hlo_op}'" + ) + if is_host_wrapper: + continue + hlo_category = args.get("hlo_category", "") + if hlo_category in ("async-start", "async-update"): + continue + offload_type = args.get("offload_type", "") + is_sc_offload = ( + hlo_category == "async-done" + and offload_type == "OFFLOAD_COLLECTIVE" + ) + for kernel_name in kernel_names: + k_hyphen = kernel_name.replace("_", "-") + k_under = kernel_name.replace("-", "_") + is_name_match = event_name.startswith(( + k_hyphen, + k_under, + "ragged-all-to-all", + )) + if has_sc_offload: + if is_sc_offload: + pid = event.get("pid") + if pid not in kernel_durs_per_core: + kernel_durs_per_core[pid] = [] + kernel_durs_per_core[pid].append(event["dur"]) + matched_count += 1 + logging.info( + "[Trace Event] pid=%s, dur=%.2f us, name='%s', tf_op='%s'," + " hlo_op='%s', hlo_category='%s', offload_type='%s'", + pid, + event["dur"], + event_name, + tf_op, + hlo_op, + hlo_category, + offload_type, + ) + break + elif is_name_match and not offload_type: + pid = event.get("pid") + if pid not in sc_kernel_durs_per_core: + sc_kernel_durs_per_core[pid] = [] + sc_kernel_durs_per_core[pid].append(event["dur"]) + break + elif is_name_match: + pid = event.get("pid") + if pid not in kernel_durs_per_core: + kernel_durs_per_core[pid] = [] + kernel_durs_per_core[pid].append(event["dur"]) + matched_count += 1 + logging.info( + "[Trace Event] pid=%s, dur=%.2f us, name='%s', tf_op='%s'," + " hlo_op='%s', hlo_category='%s', offload_type='%s'", + pid, + event["dur"], + event_name, + tf_op, + hlo_op, + hlo_category, + offload_type, + ) + break + if matched_count == 0: + logging.warning( + "No events matched %s in %s. Non-host sample trace events with" + " dur: %s", + kernel_names, + trace_file.name, + sample_events, + ) + logging.info( + "Processed trace file %s: found %d matching events across %d" + " pids/cores.", + trace_file.name, + matched_count, + len(kernel_durs_per_core), + ) + + if sc_kernel_durs_per_core: + sc_durs = [] + for pid, durs in sc_kernel_durs_per_core.items(): + sc_durs.extend(durs) + if sc_durs: + logging.info( + "SparseCore pure kernel execution durations across %d PIDs (%d" + " samples): %.2f +/- %.2f us", + len(sc_kernel_durs_per_core), + len(sc_durs), + np.mean(sc_durs), + np.std(sc_durs), + ) + + execution_times_us = [] + for pid, durs in kernel_durs_per_core.items(): + if durs: + logging.info("PID %s event durations: %s", pid, durs) + num_measurements = self.number_of_measurements + if len(durs) >= num_measurements and len(durs) % num_measurements == 0: + k = len(durs) // num_measurements + for i in range(0, len(durs), k): + execution_times_us.append(sum(durs[i : i + k])) + else: + execution_times_us.extend(durs) + + if not execution_times_us: + raise ValueError( + f"Could not find execution times matching {kernel_names}." + ) + return execution_times_us + + def _measure_execution_times( + self, + kernel, + kernel_names, + *kernel_args, + **kernel_kwargs, + ): + """Collect XProf measurements using the JAX Profiler API.""" + normalized_names = self._normalize_kernel_names(kernel_names) + + # Warmup run to compile and initialize buffers before tracing. + warmup_res = kernel(*kernel_args, **kernel_kwargs) + jax.block_until_ready(warmup_res) + + with tempfile.TemporaryDirectory() as tmpdir: + with jax.profiler.trace(tmpdir): + for _ in range(self.number_of_measurements): + result = kernel(*kernel_args, **kernel_kwargs) + jax.block_until_ready(result) + return self._extract_execution_times(tmpdir, normalized_names) + + def _print_collective_bandwidth_statistics( + self, + op_name, + per_device_size_bytes, + total_size_bytes, + bus_data_bytes, + num_devices, + latencies_us, + ): + """Prints bandwidth and latency statistics for a collective operation.""" + latencies_ns = np.array(latencies_us) * 1e3 + bus_bandwidth_gbps = bus_data_bytes / latencies_ns + effective_bandwidth_gbps = total_size_bytes / latencies_ns + + avg_bus_bw = float(np.mean(bus_bandwidth_gbps)) + std_bus_bw = float(np.std(bus_bandwidth_gbps)) + avg_eff_bw = float(np.mean(effective_bandwidth_gbps)) + std_eff_bw = float(np.std(effective_bandwidth_gbps)) + avg_latency_us = float(np.mean(latencies_us)) + std_latency_us = float(np.std(latencies_us)) + + tpu_info = pltpu.get_tpu_info() + + coords = [tuple(getattr(d, "coords", (0,))) for d in jax.local_devices()] + mesh_dims = "x".join(str(max(dim) + 1) for dim in zip(*coords)) + num_chips = len(set(coords)) + cores_per_chip = num_devices // num_chips if num_chips else 1 + topology_str = ( + f"{mesh_dims} mesh ({num_chips} chips x {cores_per_chip} cores)" + ) + + flags_list = [arg for arg in sys.argv[1:] if arg.startswith("--xla_")] + flags_str = " ".join(flags_list) if flags_list else "default" + + logging.info("==================================================") + logging.info("Test: %s", self._testMethodName) + logging.info("\tCollective Op: %s", op_name) + logging.info("\tTPU generation: %s", tpu_info.chip_version) + logging.info("\tNumber of devices: %d", num_devices) + logging.info("\tTopology: %s", topology_str) + logging.info("\tFlags: %s", flags_str) + if per_device_size_bytes >= 1024**2: + logging.info( + "\tPer-device chunk size: %.2f MiB", + per_device_size_bytes / (1024**2), + ) + else: + logging.info( + "\tPer-device chunk size: %.2f KiB", per_device_size_bytes / 1024 + ) + if total_size_bytes >= 1024**2: + logging.info( + "\tTotal collective size: %.2f MiB", total_size_bytes / (1024**2) + ) + else: + logging.info( + "\tTotal collective size: %.2f KiB", total_size_bytes / 1024 + ) + logging.info("\tLatency: %.2f +/- %.2f us", avg_latency_us, std_latency_us) + logging.info("\tBus Bandwidth: %.2f +/- %.2f GB/s", avg_bus_bw, std_bus_bw) + logging.info( + "\tEffective Bandwidth: %.2f +/- %.2f GB/s", avg_eff_bw, std_eff_bw + ) + logging.info("==================================================") + + self.latencies_us = [float(x) for x in latencies_us] + self.metrics = { + "per_device_size_mib": float(per_device_size_bytes / (1024**2)), + "total_size_mib": float(total_size_bytes / (1024**2)), + "num_devices": num_devices, + "bus_bandwidth_gbps": avg_bus_bw, + "effective_bandwidth_gbps": avg_eff_bw, + } + + +class CollectiveBenchmarkConfig(benchmark.BenchmarkConfig): + """Config identifying a single collective test method for the benchmark driver.""" + + def __init__(self, suite, test_class, test_name): + self.suite = suite + self.test_class = test_class + self.test_name = test_name + self.metrics = {} + + def as_dict(self): + return {"suite": self.suite, "test": self.test_name, **self.metrics} + + def get_benchmark(self): + return CollectiveBenchmark(self) + + +class CollectiveBenchmark(benchmark.Benchmark): + """Runs a collective test method and reports its latencies to the driver.""" + + def __init__(self, config): + self._config = config + + # Required by benchmark.Benchmark ABC; unused since run() delegates to test. + def get_input_shapes_and_dtypes(self): + return [] + + def target_fn(self): + return lambda: None + + def kernel_name(self): + return self._config.test_name + + def run(self, **kwargs): + del kwargs + test = self._config.test_class(self._config.test_name) + try: + test.setUp() + getattr(test, self._config.test_name)() + except absltest.SkipTest as e: + logging.warning("Skipping %s: %s", self._config.test_name, e) + return [None] + self._config.metrics.update(test.metrics) + return [jax_profiler_utils.JaxProfilerResult(runtimes_us=test.latencies_us)] + diff --git a/third_party/xla/xla/bit_set.h b/third_party/xla/xla/bit_set.h new file mode 100644 index 00000000000000..875d5cabc22a27 --- /dev/null +++ b/third_party/xla/xla/bit_set.h @@ -0,0 +1,180 @@ +/* Copyright 2026 The OpenXLA Authors. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#ifndef XLA_BIT_SET_H_ +#define XLA_BIT_SET_H_ + +#include +#include +#include +#include + +#include "absl/container/inlined_vector.h" +#include "absl/log/check.h" + +namespace xla { + +// A fixed-size bit set (or bitmap) that stores up to `InlinedWords * 64` bits +// inline without heap allocation, spilling to the heap only when `num_bits` +// exceeds this capacity. Useful for tracking sets of dense integer indices +// over a typically small range `[0, num_bits)`. +template +class InlinedBitSet { + public: + using Word = uint64_t; + static constexpr int kWordSize = sizeof(Word) * 8; + + InlinedBitSet() = default; + explicit InlinedBitSet(int num_bits) : words_(NumWords(num_bits), 0) {} + + void Set(int index) { + DCHECK_GE(index, 0); + const size_t word_idx = IndexToWord(index); + words_[word_idx] |= (Word{1} << (index % kWordSize)); + } + + // Sets all bits in [0, num_bits) to 1 + void SetAll(int num_bits) { + DCHECK_GE(num_bits, 0); + DCHECK_LE(static_cast(num_bits), words_.size() * kWordSize); + const size_t full_words = num_bits / kWordSize; + std::fill_n(words_.begin(), full_words, ~Word{0}); + const size_t remainder = num_bits % kWordSize; + if (remainder != 0) { + words_[full_words] |= (Word{1} << remainder) - 1; + } + } + + void Clear(int index) { + const size_t word_idx = IndexToWord(index); + words_[word_idx] &= ~(Word{1} << (index % kWordSize)); + } + + bool Test(int index) const { + const size_t word_idx = IndexToWord(index); + return (words_[word_idx] & (Word{1} << (index % kWordSize))) != 0; + } + + InlinedBitSet& operator|=(const InlinedBitSet& other) { + DCHECK_EQ(words_.size(), other.words_.size()); + for (size_t i = 0; i < words_.size(); ++i) { + words_[i] |= other.words_[i]; + } + return *this; + } + + private: + static size_t NumWords(int num_bits) { + DCHECK_GE(num_bits, 0); + return (num_bits + kWordSize - 1) / kWordSize; + } + + size_t IndexToWord(int index) const { + DCHECK_GE(index, 0); + const size_t word_idx = index / kWordSize; + DCHECK_LT(word_idx, words_.size()); + return word_idx; + } + + absl::InlinedVector words_; +}; + +// Holds num_rows bit sets of num_bits bits each in a single heap allocation. +// Useful when an analysis keeps one bit set (of the same size) per element of +// a large dense index space, for example per instruction: one allocation +// instead of one per element, and the rows sit next to each other in memory, +// so an OR of one row into another streams over two contiguous word ranges. +// Rows are addressed by index. +// +// Supported ops: +// - Set: sets one bit of a row. +// - Test: reads one bit of a row. +// - Or: ORs one row into another. +// - OrIgnoringBit: ORs one row into another, leaving one bit out. +class DenseBitSets { + public: + using Word = uint64_t; + static constexpr int kWordSize = sizeof(Word) * 8; + + DenseBitSets(int num_rows, int num_bits) + : num_rows_(CheckNonNegative(num_rows)), + num_bits_(CheckNonNegative(num_bits)), + words_per_row_(num_bits_ / kWordSize + (num_bits_ % kWordSize != 0)), + words_(static_cast(num_rows_) * words_per_row_, 0) {} + + int num_rows() const { return num_rows_; } + int num_bits() const { return num_bits_; } + + // Sets bit 'bit' of row 'row'. + void Set(int row, int bit) { Row(row)[WordIndex(bit)] |= Mask(bit); } + + // Returns whether bit 'bit' of row 'row' is set. + bool Test(int row, int bit) const { + return (Row(row)[WordIndex(bit)] & Mask(bit)) != 0; + } + + // row |= other. + void Or(int row, int other) { + Word* dst = Row(row); + const Word* src = Row(other); + for (int w = 0; w < words_per_row_; ++w) { + dst[w] |= src[w]; + } + } + + // row |= other, with ignored_bit of other left out: a plain Or, which the + // compiler vectorizes, then the destination's own value of the bit is put + // back. + void OrIgnoringBit(int row, int other, int ignored_bit) { + Word& word = Row(row)[WordIndex(ignored_bit)]; + const Word own_bit = word & Mask(ignored_bit); + Or(row, other); + word = (word & ~Mask(ignored_bit)) | own_bit; + } + + private: + static int CheckNonNegative(int count) { + DCHECK_GE(count, 0); + return count; + } + + static constexpr Word Mask(int bit) { return Word{1} << (bit % kWordSize); } + + int WordIndex(int bit) const { + DCHECK_GE(bit, 0); + DCHECK_LT(bit, num_bits_); + return bit / kWordSize; + } + + // The word offset of a row; computed in size_t because the whole block can + // hold more than 2^31 words. + size_t RowOffset(int row) const { + DCHECK_GE(row, 0); + DCHECK_LT(row, num_rows_); + return static_cast(row) * words_per_row_; + } + + Word* Row(int row) { return words_.data() + RowOffset(row); } + const Word* Row(int row) const { return words_.data() + RowOffset(row); } + + int num_rows_; + int num_bits_; + int words_per_row_; + std::vector words_; +}; + +} // namespace xla + +#endif // XLA_BIT_SET_H_ diff --git a/third_party/xla/xla/bit_set_test.cc b/third_party/xla/xla/bit_set_test.cc new file mode 100644 index 00000000000000..e9a5a634e034bf --- /dev/null +++ b/third_party/xla/xla/bit_set_test.cc @@ -0,0 +1,404 @@ +/* Copyright 2026 The OpenXLA Authors. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#include "xla/bit_set.h" + +#include +#include + +#include "absl/container/flat_hash_set.h" +#include "xla/tsl/platform/test.h" +#include "xla/tsl/platform/test_benchmark.h" + +namespace xla { +namespace { + +using ::benchmark::DoNotOptimize; +using ::testing::benchmark::State; + +class BruteForceBitSet { + public: + void Set(int index) { bits_.insert(index); } + + void SetAll(int num_bits) { + for (int i = 0; i < num_bits; ++i) { + bits_.insert(i); + } + } + + void Clear(int index) { bits_.erase(index); } + + bool Test(int index) const { return bits_.contains(index); } + + BruteForceBitSet& operator|=(const BruteForceBitSet& other) { + bits_.insert(other.bits_.begin(), other.bits_.end()); + return *this; + } + + private: + absl::flat_hash_set bits_; +}; + +TEST(InlinedBitSetTest, DefaultInitializedAllFalse) { + InlinedBitSet<> bits(130); + + EXPECT_FALSE(bits.Test(0)); + EXPECT_FALSE(bits.Test(63)); + EXPECT_FALSE(bits.Test(64)); + EXPECT_FALSE(bits.Test(127)); + EXPECT_FALSE(bits.Test(128)); +} + +TEST(InlinedBitSetTest, Set) { + InlinedBitSet<> bits(130); + + bits.Set(0); + bits.Set(63); + bits.Set(64); + bits.Set(128); + + EXPECT_TRUE(bits.Test(0)); + EXPECT_TRUE(bits.Test(63)); + EXPECT_TRUE(bits.Test(64)); + EXPECT_FALSE(bits.Test(65)); + EXPECT_FALSE(bits.Test(127)); + EXPECT_TRUE(bits.Test(128)); +} + +TEST(InlinedBitSetTest, Clear) { + InlinedBitSet<> bits(130); + bits.Set(63); + bits.Set(64); + + bits.Clear(63); + + EXPECT_FALSE(bits.Test(63)); + EXPECT_TRUE(bits.Test(64)); +} + +TEST(InlinedBitSetTest, SetAllZeroBits) { + InlinedBitSet<> bits(0); + + bits.SetAll(0); +} + +TEST(InlinedBitSetTest, SetAllExactWordBoundary) { + InlinedBitSet<> bits(64); + + bits.SetAll(64); + + for (int i = 0; i < 64; ++i) { + EXPECT_TRUE(bits.Test(i)); + } +} + +TEST(InlinedBitSetTest, SetAllNonExactBoundary) { + InlinedBitSet<> bits(70); + + bits.SetAll(70); + + for (int i = 0; i < 70; ++i) { + EXPECT_TRUE(bits.Test(i)); + } +} + +TEST(InlinedBitSetTest, SetAllLessThanCapacity) { + InlinedBitSet<> bits(128); + + bits.SetAll(50); + + for (int i = 0; i < 50; ++i) { + EXPECT_TRUE(bits.Test(i)); + } + for (int i = 50; i < 128; ++i) { + EXPECT_FALSE(bits.Test(i)); + } +} + +TEST(InlinedBitSetTest, SetAllZeroOnNonEmpty) { + InlinedBitSet<> bits(64); + + bits.SetAll(0); + + for (int i = 0; i < 64; ++i) { + EXPECT_FALSE(bits.Test(i)); + } +} + +TEST(InlinedBitSetTest, BitwiseOr) { + InlinedBitSet<> a(130); + InlinedBitSet<> b(130); + a.Set(10); + a.Set(100); + b.Set(70); + b.Set(100); + + a |= b; + + EXPECT_TRUE(a.Test(10)); + EXPECT_TRUE(a.Test(70)); + EXPECT_TRUE(a.Test(100)); + EXPECT_FALSE(a.Test(0)); + EXPECT_FALSE(a.Test(129)); +} + +TEST(InlinedBitSetTest, MatchesBruteForceOnRandomOperations) { + constexpr int kNumBits = 200; + constexpr int kNumOps = 500; + std::mt19937 rng(42); + std::uniform_int_distribution op_dist(0, 3); + std::uniform_int_distribution bit_dist(0, kNumBits - 1); + std::uniform_int_distribution count_dist(0, kNumBits); + InlinedBitSet<> bits(kNumBits); + InlinedBitSet<> other_bits(kNumBits); + BruteForceBitSet bf_bits; + BruteForceBitSet bf_other_bits; + + for (int step = 0; step < kNumOps; ++step) { + const int idx = bit_dist(rng); + other_bits.Set(idx); + bf_other_bits.Set(idx); + switch (op_dist(rng)) { + case 0: + bits.Set(idx); + bf_bits.Set(idx); + break; + case 1: + bits.Clear(idx); + bf_bits.Clear(idx); + break; + case 2: { + const int count = count_dist(rng); + bits.SetAll(count); + bf_bits.SetAll(count); + break; + } + case 3: + bits |= other_bits; + bf_bits |= bf_other_bits; + break; + } + } + + for (int i = 0; i < kNumBits; ++i) { + EXPECT_EQ(bits.Test(i), bf_bits.Test(i)) << "Mismatch at index " << i; + } +} + +// Reference implementation of DenseBitSets: rows of bits as a vector of +// vectors. +class BruteForceBitSetRows { + public: + BruteForceBitSetRows(int num_rows, int num_bits) + : rows_(num_rows, std::vector(num_bits, false)) {} + + void Set(int row, int bit) { rows_[row][bit] = true; } + bool Test(int row, int bit) const { return rows_[row][bit]; } + + void Or(int row, int other) { + for (int bit = 0; bit < rows_[row].size(); ++bit) { + rows_[row][bit] = rows_[row][bit] || rows_[other][bit]; + } + } + + void OrIgnoringBit(int row, int other, int ignored_bit) { + for (int bit = 0; bit < rows_[row].size(); ++bit) { + if (bit != ignored_bit) { + rows_[row][bit] = rows_[row][bit] || rows_[other][bit]; + } + } + } + + private: + std::vector> rows_; +}; + +TEST(DenseBitSetsTest, EmptyShapesAreAllowed) { + DenseBitSets no_rows(0, 130); + EXPECT_EQ(no_rows.num_rows(), 0); + EXPECT_EQ(no_rows.num_bits(), 130); + + DenseBitSets no_bits(3, 0); + EXPECT_EQ(no_bits.num_rows(), 3); + EXPECT_EQ(no_bits.num_bits(), 0); + no_bits.Or(0, 1); // Rows of zero words. + + DenseBitSets empty(0, 0); + EXPECT_EQ(empty.num_rows(), 0); + EXPECT_EQ(empty.num_bits(), 0); +} + +TEST(DenseBitSetsTest, DefaultInitializedAllFalse) { + DenseBitSets bits(3, 130); + + EXPECT_EQ(bits.num_rows(), 3); + EXPECT_EQ(bits.num_bits(), 130); + for (int row = 0; row < 3; ++row) { + EXPECT_FALSE(bits.Test(row, 0)); + EXPECT_FALSE(bits.Test(row, 63)); + EXPECT_FALSE(bits.Test(row, 64)); + EXPECT_FALSE(bits.Test(row, 129)); + } +} + +TEST(DenseBitSetsTest, SetIsConfinedToItsRow) { + DenseBitSets bits(3, 130); + + bits.Set(1, 0); + bits.Set(1, 63); + bits.Set(1, 64); + bits.Set(1, 129); + + EXPECT_TRUE(bits.Test(1, 0)); + EXPECT_TRUE(bits.Test(1, 63)); + EXPECT_TRUE(bits.Test(1, 64)); + EXPECT_TRUE(bits.Test(1, 129)); + EXPECT_FALSE(bits.Test(1, 1)); + EXPECT_FALSE(bits.Test(1, 128)); + for (int bit : {0, 63, 64, 129}) { + EXPECT_FALSE(bits.Test(0, bit)); + EXPECT_FALSE(bits.Test(2, bit)); + } +} + +TEST(DenseBitSetsTest, OrAndOrIgnoringBit) { + DenseBitSets bits(2, 130); + bits.Set(0, 5); + bits.Set(1, 5); + bits.Set(1, 64); + bits.Set(1, 129); + + bits.Or(0, 1); + + EXPECT_TRUE(bits.Test(0, 5)); + EXPECT_TRUE(bits.Test(0, 64)); + EXPECT_TRUE(bits.Test(0, 129)); + EXPECT_FALSE(bits.Test(0, 6)); + + DenseBitSets more(2, 130); + more.Set(0, 64); + more.Set(1, 64); + more.Set(1, 65); + more.Set(1, 129); + + more.OrIgnoringBit(0, 1, 129); + more.OrIgnoringBit(0, 1, 64); + + EXPECT_TRUE(more.Test(0, 64)); // Already set in the row; never cleared. + EXPECT_TRUE(more.Test(0, 65)); + EXPECT_TRUE(more.Test(0, 129)); // Copied by the second call. + EXPECT_TRUE(more.Test(1, 64)); +} + +TEST(DenseBitSetsTest, MatchesBruteForceOnRandomOperations) { + constexpr int kNumInstances = 200; + constexpr int kNumOps = 300; + std::mt19937 rng(7); + std::uniform_int_distribution rows_dist(1, 12); + std::uniform_int_distribution bits_dist(1, 300); + std::uniform_int_distribution op_dist(0, 2); + + for (int instance = 0; instance < kNumInstances; ++instance) { + const int num_rows = rows_dist(rng); + // Every fourth instance sits on a word boundary. + const int num_bits = + instance % 4 == 0 ? 64 * (1 + instance % 3) : bits_dist(rng); + std::uniform_int_distribution row_dist(0, num_rows - 1); + std::uniform_int_distribution bit_dist(0, num_bits - 1); + DenseBitSets bits(num_rows, num_bits); + BruteForceBitSetRows bf_bits(num_rows, num_bits); + + for (int step = 0; step < kNumOps; ++step) { + const int row = row_dist(rng); + const int other = row_dist(rng); + const int bit = bit_dist(rng); + switch (op_dist(rng)) { + case 0: + bits.Set(row, bit); + bf_bits.Set(row, bit); + break; + case 1: + bits.Or(row, other); + bf_bits.Or(row, other); + break; + case 2: + bits.OrIgnoringBit(row, other, bit); + bf_bits.OrIgnoringBit(row, other, bit); + break; + } + ASSERT_EQ(bits.Test(row, bit), bf_bits.Test(row, bit)) + << "instance " << instance << " step " << step << " row " << row + << " bit " << bit; + } + for (int row = 0; row < num_rows; ++row) { + for (int bit = 0; bit < num_bits; ++bit) { + ASSERT_EQ(bits.Test(row, bit), bf_bits.Test(row, bit)) + << "instance " << instance << " row " << row << " bit " << bit; + } + } + } +} + +// Benchmarks: the operations TryRemoveDeadWhileParams issues, for rows of 64 +// to 65,536 bits (more than twice the largest while tuple seen so far): +// - Or, once per operand edge, +// - OrIgnoringBit, once per body output, +// - Set and Test, once per get-tuple-element. +void BM_DenseBitSetsOr(State& state) { + const int num_bits = state.range(0); + DenseBitSets bits(2, num_bits); + bits.Set(1, num_bits - 1); + for (auto s : state) { + bits.Or(0, 1); + DoNotOptimize(bits); + } + state.SetBytesProcessed(state.iterations() * ((num_bits + 63) / 64) * 8); +} +BENCHMARK(BM_DenseBitSetsOr)->Arg(64)->Arg(1024)->Arg(16384)->Arg(65536); + +void BM_DenseBitSetsOrIgnoringBit(State& state) { + const int num_bits = state.range(0); + DenseBitSets bits(2, num_bits); + bits.Set(1, num_bits - 1); + for (auto s : state) { + bits.OrIgnoringBit(0, 1, num_bits / 2); + DoNotOptimize(bits); + } + state.SetBytesProcessed(state.iterations() * ((num_bits + 63) / 64) * 8); +} +BENCHMARK(BM_DenseBitSetsOrIgnoringBit) + ->Arg(64) + ->Arg(1024) + ->Arg(16384) + ->Arg(65536); + +void BM_DenseBitSetsSetTest(State& state) { + const int num_bits = state.range(0); + DenseBitSets bits(1024, num_bits); + int row = 0; + int bit = 0; + bool sink = false; + for (auto s : state) { + bits.Set(row, bit); + sink ^= bits.Test((row + 1) % 1024, bit); + row = (row + 7) % 1024; + bit = (bit + 13) % num_bits; + } + DoNotOptimize(sink); +} +BENCHMARK(BM_DenseBitSetsSetTest)->Arg(64)->Arg(1024)->Arg(16384)->Arg(65536); + +} // namespace +} // namespace xla diff --git a/third_party/xla/xla/codegen/emitters/transforms/simplify_affine.cc b/third_party/xla/xla/codegen/emitters/transforms/simplify_affine.cc index 15f442ed6544e5..c3cd5b545fb632 100644 --- a/third_party/xla/xla/codegen/emitters/transforms/simplify_affine.cc +++ b/third_party/xla/xla/codegen/emitters/transforms/simplify_affine.cc @@ -90,9 +90,45 @@ void CollectArgs(SymbolicExpr expr, SymbolicExprType type, ret.push_back(expr); } +// Returns true if the top-level operation of `expr` can be lowered to an +// unsigned arith op. Operands of `expr` are not checked. +bool IsUnsignedLoweringSupported(SymbolicExpr expr, + RangeEvaluator& range_evaluator) { + if (!expr.IsBinaryOp()) { + return true; + } + // Mod and div can be lowered if their LHS is >= 0 and their RHS is a + // constant positive number. + if (expr.GetType() == SymbolicExprType::kMod || + expr.GetType() == SymbolicExprType::kFloorDiv || + expr.GetType() == SymbolicExprType::kCeilDiv) { + if (!range_evaluator.IsAlwaysPositiveOrZero(expr.GetLHS())) { + return false; + } + auto rhs_range = range_evaluator.ComputeExpressionRange(expr.GetRHS()); + if (!rhs_range.IsPoint() || rhs_range.lower <= 0) { + return false; + } + } + // Max and min can be lowered if both operands are >= 0. + if (expr.GetType() == SymbolicExprType::kMax || + expr.GetType() == SymbolicExprType::kMin) { + if (!range_evaluator.IsAlwaysPositiveOrZero(expr.GetLHS()) || + !range_evaluator.IsAlwaysPositiveOrZero(expr.GetRHS())) { + return false; + } + } + return true; +} + struct ExpressionEvaluator { - ExpressionEvaluator(ImplicitLocOpBuilder& builder, ValueRange operands) - : builder(builder), operands(operands) { + // If `range_evaluator` is null, all mod, div, min and max operations are + // lowered to unsigned arith ops. The caller must ensure that this is correct, + // e.g. by checking `IsLoweringSupported`. Otherwise, signed arith ops are + // emitted for the operations whose unsigned lowering is not provably correct. + ExpressionEvaluator(ImplicitLocOpBuilder& builder, ValueRange operands, + RangeEvaluator* range_evaluator = nullptr) + : builder(builder), operands(operands), range_evaluator(range_evaluator) { for (int i = 0; i < operands.size(); ++i) { variable_distances.push_back(Distance(builder, operands[i])); } @@ -123,6 +159,7 @@ struct ExpressionEvaluator { ImplicitLocOpBuilder& builder; ValueRange operands; + RangeEvaluator* range_evaluator; SmallVector variable_distances; }; @@ -153,26 +190,56 @@ Value ExpressionEvaluator::EvaluateAddMul(SymbolicExpr expr) { } Value ExpressionEvaluator::EvaluateExpression(SymbolicExpr expr) { + bool use_signed = range_evaluator != nullptr && + !IsUnsignedLoweringSupported(expr, *range_evaluator); switch (expr.GetType()) { case SymbolicExprType::kAdd: return EvaluateAddMul(expr); case SymbolicExprType::kMul: return EvaluateAddMul(expr); - case SymbolicExprType::kMod: - return builder.create(EvaluateExpression(expr.GetLHS()), - EvaluateExpression(expr.GetRHS())); - case SymbolicExprType::kFloorDiv: - return builder.create(EvaluateExpression(expr.GetLHS()), - EvaluateExpression(expr.GetRHS())); - case SymbolicExprType::kCeilDiv: - return builder.create( - EvaluateExpression(expr.GetLHS()), EvaluateExpression(expr.GetRHS())); - case SymbolicExprType::kMax: - return builder.create(EvaluateExpression(expr.GetLHS()), - EvaluateExpression(expr.GetRHS())); - case SymbolicExprType::kMin: - return builder.create(EvaluateExpression(expr.GetLHS()), - EvaluateExpression(expr.GetRHS())); + case SymbolicExprType::kMod: { + Value lhs = EvaluateExpression(expr.GetLHS()); + Value rhs = EvaluateExpression(expr.GetRHS()); + if (!use_signed) { + return builder.create(lhs, rhs); + } + // SymbolicExpr uses floor semantics: lhs mod rhs has the sign of rhs. + Value quotient = builder.create(lhs, rhs); + return builder.create( + lhs, builder.create(quotient, rhs)); + } + case SymbolicExprType::kFloorDiv: { + Value lhs = EvaluateExpression(expr.GetLHS()); + Value rhs = EvaluateExpression(expr.GetRHS()); + if (use_signed) { + return builder.create(lhs, rhs); + } + return builder.create(lhs, rhs); + } + case SymbolicExprType::kCeilDiv: { + Value lhs = EvaluateExpression(expr.GetLHS()); + Value rhs = EvaluateExpression(expr.GetRHS()); + if (use_signed) { + return builder.create(lhs, rhs); + } + return builder.create(lhs, rhs); + } + case SymbolicExprType::kMax: { + Value lhs = EvaluateExpression(expr.GetLHS()); + Value rhs = EvaluateExpression(expr.GetRHS()); + if (use_signed) { + return builder.create(lhs, rhs); + } + return builder.create(lhs, rhs); + } + case SymbolicExprType::kMin: { + Value lhs = EvaluateExpression(expr.GetLHS()); + Value rhs = EvaluateExpression(expr.GetRHS()); + if (use_signed) { + return builder.create(lhs, rhs); + } + return builder.create(lhs, rhs); + } case SymbolicExprType::kConstant: return builder.create(expr.GetValue()); case SymbolicExprType::kVariable: @@ -182,32 +249,14 @@ Value ExpressionEvaluator::EvaluateExpression(SymbolicExpr expr) { } } +// Returns true if all operations in `expr` can be lowered to unsigned arith +// ops. bool IsLoweringSupported(SymbolicExpr expr, RangeEvaluator& range_evaluator) { if (!expr.IsBinaryOp()) { return true; } - // Mod and div can be lowered if their LHS is >= 0 and their RHS is a - // constant positive number. - if (expr.GetType() == SymbolicExprType::kMod || - expr.GetType() == SymbolicExprType::kFloorDiv || - expr.GetType() == SymbolicExprType::kCeilDiv) { - if (!range_evaluator.IsAlwaysPositiveOrZero(expr.GetLHS())) { - return false; - } - auto rhs_range = range_evaluator.ComputeExpressionRange(expr.GetRHS()); - if (!rhs_range.IsPoint() || rhs_range.lower <= 0) { - return false; - } - } - // Max and min can be lowered if both operands are >= 0. - if (expr.GetType() == SymbolicExprType::kMax || - expr.GetType() == SymbolicExprType::kMin) { - if (!range_evaluator.IsAlwaysPositiveOrZero(expr.GetLHS()) || - !range_evaluator.IsAlwaysPositiveOrZero(expr.GetRHS())) { - return false; - } - } - return IsLoweringSupported(expr.GetLHS(), range_evaluator) && + return IsUnsignedLoweringSupported(expr, range_evaluator) && + IsLoweringSupported(expr.GetLHS(), range_evaluator) && IsLoweringSupported(expr.GetRHS(), range_evaluator); } @@ -274,23 +323,27 @@ struct RewriteApplyIndexingOp : OpRewritePattern { results.reserve(symbolic_map.GetNumResults()); for (unsigned i = 0; i < symbolic_map.GetNumResults(); ++i) { SymbolicExpr result_expr = symbolic_map.GetResult(i); - // If the expression cannot be lowered, we convert it to affine.apply, - // since it supports more expression types. if (IsLoweringSupported(result_expr, range_evaluator)) { results.push_back( ExpressionEvaluator(b, operands).EvaluateExpression(result_expr)); - } else { - // TODO: b/446856305 - Create a SymbolicApplyOp. For now, we convert the - // SymbolicMap back to AffineMap and fall back to AffineApplyOp. - AffineMap sub_map = SymbolicMapToAffineMap(symbolic_map.GetSubMap({i})); - if (!sub_map) { - return rewriter.notifyMatchFailure(op, - "cannot fallback to affine.apply"); - } + continue; + } + // If the expression cannot be lowered to unsigned ops, we convert it to + // affine.apply, since it supports more expression types. + // TODO: b/446856305 - Create a SymbolicApplyOp. For now, we convert the + // SymbolicMap back to AffineMap and fall back to AffineApplyOp. + if (AffineMap sub_map = + SymbolicMapToAffineMap(symbolic_map.GetSubMap({i}))) { results.push_back(b.create( sub_map, operands.take_front(symbolic_map.GetNumDims() + symbolic_map.GetNumSymbols()))); + continue; } + // affine.apply does not support min and max. Lower the expression to + // arith ops, using signed ops wherever unsigned ones are not provably + // correct. + results.push_back(ExpressionEvaluator(b, operands, &range_evaluator) + .EvaluateExpression(result_expr)); } rewriter.replaceOp(op, results); return mlir::success(); diff --git a/third_party/xla/xla/codegen/emitters/transforms/tests/simplify_affine.mlir b/third_party/xla/xla/codegen/emitters/transforms/tests/simplify_affine.mlir index e6e3ffe69d62ab..8d781e6977d758 100644 --- a/third_party/xla/xla/codegen/emitters/transforms/tests/simplify_affine.mlir +++ b/third_party/xla/xla/codegen/emitters/transforms/tests/simplify_affine.mlir @@ -260,15 +260,38 @@ func.func @lower_ceildiv_constant(%arg0: index) -> index { // ----- -func.func @unsafe_max(%arg0: index {xla.range = [-10 : index, 10 : index]}, %arg1: index) -> index { +func.func @signed_max(%arg0: index {xla.range = [-10 : index, 10 : index]}, %arg1: index) -> index { %0 = xla.apply_indexing #xla.indexing_map< "()[s0, s1] -> (max(s0, s1))," "domain: s0 in [-10, 10], s1 in [0, 5]">[%arg0, %arg1] return %0 : index } -// CHECK-LABEL: @unsafe_max -// CHECK: xla.apply_indexing +// CHECK-LABEL: @signed_max +// CHECK-SAME: (%[[ARG0:.*]]: index {{.*}}, %[[ARG1:.*]]: index) +// CHECK-NOT: xla.apply_indexing +// CHECK: %[[RET:.*]] = arith.maxsi %[[ARG0]], %[[ARG1]] +// CHECK-NEXT: return %[[RET]] + +// ----- + +// min/max can't be expressed with affine.apply, so possibly negative operands +// are lowered to signed arith ops. Operations with provably non-negative +// operands still use unsigned ops. +func.func @signed_clamp_of_floordiv(%arg0: index, %arg1: index) -> index { + %0 = xla.apply_indexing + #xla.indexing_map< + "(d0, d1) -> (min(max((-(d1 mod 3) - (d1 floordiv 3) * 24 - d0 * 13920 + 111317) floordiv 3, 0), 8))," + "domain: d0 in [0, 7], d1 in [0, 1739]">(%arg0, %arg1) + return %0 : index +} +// CHECK-LABEL: @signed_clamp_of_floordiv +// CHECK-NOT: xla.apply_indexing +// CHECK-NOT: affine.apply +// CHECK: %[[DIV:.*]] = arith.floordivsi +// CHECK: %[[MAX:.*]] = arith.maxsi %[[DIV]] +// CHECK: %[[MIN:.*]] = arith.minui %[[MAX]] +// CHECK-NEXT: return %[[MIN]] // ----- diff --git a/third_party/xla/xla/codegen/tiling/experimental/tile_propagation.cc b/third_party/xla/xla/codegen/tiling/experimental/tile_propagation.cc index 8b2abc3b0c8219..175997a73cfe58 100644 --- a/third_party/xla/xla/codegen/tiling/experimental/tile_propagation.cc +++ b/third_party/xla/xla/codegen/tiling/experimental/tile_propagation.cc @@ -1515,6 +1515,65 @@ absl::StatusOr PropagateTileToOutputForBitcastOp( return Tiles{std::move(output_tile)}; } +// Calculates the required input tile for a ReduceScatter operation. +// +// In a ReduceScatter, the input tensor is divided into "shards" along the +// scatter dimension and distributed across R replicas. For example: +// * Replicas (R) = 2 +// * Input shape = [128, 256] +// * Output shape = [64, 256] (Each replica gets a 64-row shard) +// +// This method does two things to the tile along the scattered dimension: +// +// 1. Scales it up by R: +// Because the data will be scattered, a thread block needs to process R +// times as much input data. If the block is assigned an output tile of +// 8 rows, it must actually compute 16 rows of the input. We multiply +// the offset, size, and upper bounds by R to reflect this. +// +// 2. Enforces single-destination routing: +// We add a divisibility constraint to ensure that our scaled-up input +// tile (e.g., 16 rows) fits perfectly inside a single replica's shard +// (64 rows). This guarantees that the tile never crosses the boundary +// between shards (e.g., crossing row 64), meaning this thread block +// only ever has to send its data to exactly *one* target GPU. +absl::StatusOr PropagateTileToInputForReduceScatterOp( + const TilingSpace& tiling_space, + const HloReduceScatterInstruction& reduce_scatter, + const Tile& output_tile) { + // Crash OK. Must be checked while forming the fusion. + CHECK_EQ(reduce_scatter.operand_count(), 1) + << "Multi-operand ReduceScatter is not yet supported."; + const int64_t scatter_dim = reduce_scatter.scatter_dimension(); + const int64_t shard_size = reduce_scatter.shape().dimensions(scatter_dim); + const int64_t num_replicas = + reduce_scatter.operand(0)->shape().dimensions(scatter_dim) / shard_size; + + SmallVector dim_tiles(output_tile.dim_tiles()); + DimTile& scatter_tile = dim_tiles[scatter_dim]; + if (!IsConstantValue(scatter_tile.stride, 1)) { + return absl::UnimplementedError( + absl::StrCat("Non-unit stride on the scatter dimension of ", + reduce_scatter.ToString())); + } + scatter_tile.offset = scatter_tile.offset * num_replicas; + scatter_tile.size = scatter_tile.size * num_replicas; + scatter_tile.upper_bound = scatter_tile.upper_bound * num_replicas; + + if (!tiling_space.IsSymbolic()) { + std::optional size = TryGetConstantValue(scatter_tile.size); + if (!size.has_value() || shard_size % *size != 0 || + !scatter_tile.offset.Canonicalize().IsMultipleOf(*size)) { + return absl::FailedPreconditionError(absl::StrCat( + "Tiling propagation rejected for ", reduce_scatter.ToString(), + ": operand tile ", scatter_tile.size.ToString(), " at ", + scatter_tile.offset.ToString(), + " crosses shard boundary (shard_size=", shard_size, ")")); + } + } + return Tiles{output_tile.CloneWithNewDims(std::move(dim_tiles))}; +} + } // namespace std::string ToString(const Tiles& tiles) { @@ -1589,6 +1648,10 @@ absl::StatusOr PropagateTileToInput(TilingSpace& tiling_space, if (hlo.opcode() == HloOpcode::kAllGather) { return PropagateTileToInputForAllGatherOp(tiling_space, hlo, output_tile); } + if (hlo.opcode() == HloOpcode::kReduceScatter) { + return PropagateTileToInputForReduceScatterOp( + tiling_space, *Cast(&hlo), output_tile); + } if (hlo.opcode() == HloOpcode::kBitcast) { return PropagateTileToInputForBitcastOp(hlo, output_tile); } diff --git a/third_party/xla/xla/codegen/tiling/experimental/tile_propagation_test.cc b/third_party/xla/xla/codegen/tiling/experimental/tile_propagation_test.cc index 070653530648b2..3182cd63ce52f8 100644 --- a/third_party/xla/xla/codegen/tiling/experimental/tile_propagation_test.cc +++ b/third_party/xla/xla/codegen/tiling/experimental/tile_propagation_test.cc @@ -198,6 +198,161 @@ TEST_F(TilePropagationTest, CanPropagateToInputsOfAllGatherOp) { )")); } +TEST_F(TilePropagationTest, CanPropagateToInputsOfReduceScatterOp) { + HloInstruction* root = ParseAndGetRoot(R"( + HloModule m + add { + lhs = f32[] parameter(0) + rhs = f32[] parameter(1) + ROOT add = f32[] add(lhs, rhs) + } + ENTRY e { + p0 = f32[128,256] parameter(0) + ROOT reduce_scatter = f32[64,256] reduce-scatter(p0), replica_groups={{0,1}}, dimensions={0}, to_apply=add + } + )"); + ASSERT_OK_AND_ASSIGN( + std::unique_ptr tiling_space, + TilingSpace::Create(*HloFusionAdaptor::ForInstruction(root), + &mlir_context_)); + ASSERT_OK_AND_ASSIGN( + auto tiled_operands, + PropagateTileToInput( + *tiling_space, *root, + GetTestTile(*tiling_space, root->shape().dimensions()), 0)); + EXPECT_THAT(tiled_operands, MatchToString(R"( + 0) (tid_0, tid_1) + -> offsets [tid_0 * ts_0 * 2, tid_1 * ts_1] + sizes [ts_0 * 2, ts_1] + strides [1, 2] + upper bounds [128, 256] + )")); +} + +TEST_F(TilePropagationTest, CanPropagateFromReduceScatterToDotOp) { + HloInstruction* fusion = ParseAndGetRoot(R"( + HloModule m + add { + lhs = f32[] parameter(0) + rhs = f32[] parameter(1) + ROOT add = f32[] add(lhs, rhs) + } + fusion { + p0 = f32[128,256] parameter(0) + p1 = f32[256,512] parameter(1) + dot = f32[128,512] dot(p0, p1), lhs_contracting_dims={1}, rhs_contracting_dims={0} + ROOT reduce_scatter = f32[64,512] reduce-scatter(dot), replica_groups={{0,1}}, dimensions={0}, to_apply=add + } + ENTRY e { + p0 = f32[128,256] parameter(0) + p1 = f32[256,512] parameter(1) + ROOT fusion = f32[64,512] fusion(p0, p1), kind=kCustom, calls=fusion + } + )"); + const HloInstruction* root = fusion->fused_expression_root(); + ASSERT_OK_AND_ASSIGN( + std::unique_ptr tiling_space, + TilingSpace::Create(*HloFusionAdaptor::ForInstruction(fusion), + &mlir_context_)); + ASSERT_OK_AND_ASSIGN( + auto rs_operands, + PropagateTileToInput( + *tiling_space, *root, + GetTestTile(*tiling_space, root->shape().dimensions()), 0)); + ASSERT_OK_AND_ASSIGN(auto dot_operands, + PropagateTileToInput(*tiling_space, *root->operand(0), + rs_operands[0], 0)); + EXPECT_THAT(dot_operands, MatchToString(R"( + 0) (tid_0, tid_1, tid_2) + -> offsets [tid_0 * ts_0 * 2, tid_2 * ts_2] + sizes [ts_0 * 2, ts_2] + strides [1, 1] + upper bounds [128, 256] + 1) (tid_0, tid_1, tid_2) + -> offsets [tid_2 * ts_2, tid_1 * ts_1] + sizes [ts_2, ts_1] + strides [1, 2] + upper bounds [256, 512] + )")); +} + +TEST_F(TilePropagationTest, CanPropagateToInputsOfReduceScatterOpScatterDim1) { + HloInstruction* root = ParseAndGetRoot(R"( + HloModule m + add { + lhs = f32[] parameter(0) + rhs = f32[] parameter(1) + ROOT add = f32[] add(lhs, rhs) + } + ENTRY e { + p0 = f32[128,256] parameter(0) + ROOT reduce_scatter = f32[128,128] reduce-scatter(p0), replica_groups={{0,1}}, dimensions={1}, to_apply=add + } + )"); + ASSERT_OK_AND_ASSIGN( + std::unique_ptr tiling_space, + TilingSpace::Create(*HloFusionAdaptor::ForInstruction(root), + &mlir_context_)); + ASSERT_OK_AND_ASSIGN(auto tiled_operands, + PropagateTileToInput(*tiling_space, *root, + tiling_space->tiled_roots()[0], 0)); + EXPECT_THAT(tiled_operands, MatchToString(R"( + 0) (tid_0, tid_1) + -> offsets [tid_0 * ts_0, tid_1 * ts_1 * 2] + sizes [ts_0, ts_1 * 2] + strides [1, 1] + upper bounds [128, 256] + )")); +} + +TEST_F(TilePropagationTest, ReduceScatterRejectsNonUnitScatterStride) { + HloInstruction* root = ParseAndGetRoot(R"( + HloModule m + add { + lhs = f32[] parameter(0) + rhs = f32[] parameter(1) + ROOT add = f32[] add(lhs, rhs) + } + ENTRY e { + p0 = f32[128,256] parameter(0) + ROOT reduce_scatter = f32[128,128] reduce-scatter(p0), replica_groups={{0,1}}, dimensions={1}, to_apply=add + } + )"); + ASSERT_OK_AND_ASSIGN( + std::unique_ptr tiling_space, + TilingSpace::Create(*HloFusionAdaptor::ForInstruction(root), + &mlir_context_)); + // GetTestTile assigns stride = index + 1 = 2 on dimension 1. + EXPECT_THAT(PropagateTileToInput( + *tiling_space, *root, + GetTestTile(*tiling_space, root->shape().dimensions()), 0), + StatusIs(absl::StatusCode::kUnimplemented)); +} + +TEST_F(TilePropagationTest, ReduceScatterRejectsTileCrossingShards) { + HloInstruction* root = ParseAndGetRoot(R"( + HloModule m + add { + lhs = f32[] parameter(0) + rhs = f32[] parameter(1) + ROOT add = f32[] add(lhs, rhs) + } + ENTRY e { + p0 = f32[192,256] parameter(0) + ROOT reduce_scatter = f32[96,256] reduce-scatter(p0), replica_groups={{0,1}}, dimensions={0}, to_apply=add + } + )"); + ASSERT_OK_AND_ASSIGN( + std::unique_ptr tiling_space, + TilingSpace::Create(*HloFusionAdaptor::ForInstruction(root), + &mlir_context_)); + // Operand tile 32 * 2 = 64: block 1 reads rows [64, 128), crossing 96. + ASSERT_OK(tiling_space->AssignTileSizes({32, 256})); + EXPECT_THAT(PropagateTileToInput(*tiling_space, *root, + tiling_space->tiled_roots()[0], 0), + StatusIs(absl::StatusCode::kFailedPrecondition)); +} + TEST_F(TilePropagationTest, CanPropagateToInputOfBroadcastOp) { HloInstruction* root = ParseAndGetRoot(R"( HloModule m diff --git a/third_party/xla/xla/codegen/tiling/experimental/tiled_hlo_test.cc b/third_party/xla/xla/codegen/tiling/experimental/tiled_hlo_test.cc index 8e1cb53ad9576e..e9fb30a0260fce 100644 --- a/third_party/xla/xla/codegen/tiling/experimental/tiled_hlo_test.cc +++ b/third_party/xla/xla/codegen/tiling/experimental/tiled_hlo_test.cc @@ -700,6 +700,47 @@ TEST_P(TileAnalysisTest, CollectiveDotBasic) { )")); } +TEST_P(TileAnalysisTest, CollectiveDotReduceScatter) { + ASSERT_OK_AND_ASSIGN(const TiledHloComputation tiled_computation, + ParseAndTile(R"hlo( + add { + lhs = f32[] parameter(0) + rhs = f32[] parameter(1) + ROOT add = f32[] add(lhs, rhs) + } + + fusion { + p0 = f32[128,256] parameter(0) + p1 = f32[256,512] parameter(1) + dot = f32[128,512] dot(p0, p1), lhs_contracting_dims={1}, rhs_contracting_dims={0} + ROOT rs = f32[64,512] reduce-scatter(dot), replica_groups={{0,1}}, dimensions={0}, to_apply=add + } + + ENTRY main { + p0 = f32[128,256] parameter(0) + p1 = f32[256,512] parameter(1) + ROOT fusion = f32[64,512] fusion(p0, p1), kind=kCustom, calls=fusion + })hlo", + {8, 32, 32})); + + EXPECT_THAT(tiled_computation, MatchString(R"( + Dimensions: + 0 type: parallel size: 64 tile size: 8 dim ID:0 hlo: %rs = f32[64,512]{1,0} reduce-scatter(%dot), replica_groups={{0,1}}, dimensions={0}, to_apply=%add + 1 type: parallel size: 512 tile size: 32 dim ID:1 hlo: %rs = f32[64,512]{1,0} reduce-scatter(%dot), replica_groups={{0,1}}, dimensions={0}, to_apply=%add + 2 type: sequential size: 256 tile size: 32 dim ID:2 hlo: %dot = f32[128,512]{1,0} dot(%p0, %p1), lhs_contracting_dims={1}, rhs_contracting_dims={0} + Root tiles: + 0 root tile: offsets [tid_0 * 8, tid_1 * 32] sizes [8, 32] strides [1, 1] upper bounds [64, 512] + + Tiled HLO: + dot.tile_0 = dot(p0.1.tile_0, p1.1.tile_0) offsets [tid_0 * 16, tid_1 * 32] sizes [16, 32] strides [1, 1] upper bounds [128, 512] + region #0 { + p0.1.tile_0 = parameter(0) offsets [tid_0 * 16, tid_2 * 32] sizes [16, 32] strides [1, 1] upper bounds [128, 256] + p1.1.tile_0 = parameter(1) offsets [tid_2 * 32, tid_1 * 32] sizes [32, 32] strides [1, 1] upper bounds [256, 512] + } + rs.tile_0 = reduce-scatter(dot.tile_0) offsets [tid_0 * 8, tid_1 * 32] sizes [8, 32] strides [1, 1] upper bounds [64, 512] + )")); +} + TEST_F(SingleOutputFusionTileAnalysisTest, SimpleFusion) { absl::StatusOr parse_result = ParseAndTile(R"hlo( fusion { diff --git a/third_party/xla/xla/codegen/xtile/codegen/experimental_fusion_emitter.cc b/third_party/xla/xla/codegen/xtile/codegen/experimental_fusion_emitter.cc index f3494bbe77856e..f161de27bae98a 100644 --- a/third_party/xla/xla/codegen/xtile/codegen/experimental_fusion_emitter.cc +++ b/third_party/xla/xla/codegen/xtile/codegen/experimental_fusion_emitter.cc @@ -133,6 +133,52 @@ ArrayRef MakeArrayRef(const absl::Span span) { return ArrayRef(span.data(), span.size()); } +absl::StatusOr EmitAllGather( + EmitterContext& emitter_ctx, const HloAllGatherInstruction* all_gather, + const ge::TiledHloInstruction& tiled_all_gather, ValueRange operands) { + if (all_gather->device_list()->replica_groups().empty()) { + return Internal( + "Triton emitting AllGather without replica groups is not supported."); + } + + llvm::SmallVector flattened_replica_group_ids; + for (const auto& replica_group : all_gather->replica_groups()) { + for (const auto& replica_id : replica_group.replica_ids()) { + flattened_replica_group_ids.push_back(replica_id); + } + } + + std::optional channel_handle = all_gather->channel_id(); + bool use_global_device_ids = all_gather->use_global_device_ids(); + + ImplicitLocOpBuilder& b = emitter_ctx.b(); + ABSL_ASSIGN_OR_RETURN( + auto output_element_type, + xtile::PrimitiveTypeToMlirType(b, all_gather->shape().element_type())); + ABSL_ASSIGN_OR_RETURN(SmallVector tile_sizes, + tiled_all_gather.tile().GetStaticTileSizes()); + auto output_type = + mlir::RankedTensorType::get(tile_sizes, output_element_type); + + auto replica_groups_type = mlir::RankedTensorType::get( + {static_cast(all_gather->replica_groups().size()), + static_cast( + all_gather->replica_groups()[0].replica_ids_size())}, + b.getI64Type()); + auto replica_groups_attr = mlir::DenseIntElementsAttr::get( + replica_groups_type, flattened_replica_group_ids); + auto channel_handle_attr = + channel_handle ? mlir::stablehlo::ChannelHandleAttr::get(b.getContext(), + *channel_handle, + /*type=*/0) + : nullptr; + auto all_gather_op = mlir::stablehlo::AllGatherOp::create( + b, mlir::TypeRange{output_type}, operands, + all_gather->all_gather_dimension(), replica_groups_attr, + channel_handle_attr, use_global_device_ids); + return mlir::cast(all_gather_op.getResult(0)); +} + absl::StatusOr EmitAllReduce( EmitterContext& emitter_ctx, const HloAllReduceInstruction* all_reduce, const ge::TiledHloInstruction& tiled_all_reduce, ValueRange operands) { @@ -1884,8 +1930,8 @@ absl::StatusOr EmitTiledHloInstruction( // Please keep the cases in alphabetical order. switch (hlo->opcode()) { case HloOpcode::kAllGather: { - // AllGather is a no-op. Tile extraction handles the data movement. - return emitter_ctx.TiledHloToTensorValue(*tiled_hlo.operand(0)); + return EmitAllGather(emitter_ctx, xla::Cast(hlo), + tiled_hlo, operands); } case HloOpcode::kAllReduce: { const HloComputation* computation = diff --git a/third_party/xla/xla/codegen/xtile/ir/tests/ops.mlir b/third_party/xla/xla/codegen/xtile/ir/tests/ops.mlir index d05c7e37e3741d..2f9e60aacc1f54 100644 --- a/third_party/xla/xla/codegen/xtile/ir/tests/ops.mlir +++ b/third_party/xla/xla/codegen/xtile/ir/tests/ops.mlir @@ -133,7 +133,7 @@ func.func @type_mismatch_insert(%src: tensor<24xf64>, %dst: memref<1024xf32>) { // ----- func.func @dot_scaled(%lhs: tensor<128x128xf32>, %lhs_scale: tensor<128x4xi8>, %rhs: tensor<128x256xf32>, %rhs_scale: tensor<256x4xi8>, %acc: tensor<128x256xf32>) -> tensor<128x256xf32> { - %0 = xtile.dot_scaled %lhs scale %lhs_scale, %rhs scale %rhs_scale {fastMath = true, lhs_elem_type = f32, rhs_elem_type = f32} : tensor<128x128xf32>, tensor<128x4xi8> * tensor<128x256xf32>, tensor<256x4xi8> -> tensor<128x256xf32> + %0 = xtile.dot_scaled %lhs scale %lhs_scale, %rhs scale %rhs_scale : tensor<128x128xf32>, tensor<128x4xi8> * tensor<128x256xf32>, tensor<256x4xi8> -> tensor<128x256xf32> return %0 : tensor<128x256xf32> } @@ -141,7 +141,7 @@ func.func @dot_scaled(%lhs: tensor<128x128xf32>, %lhs_scale: tensor<128x4xi8>, % func.func @invalid_dot_scaled_scale_type(%lhs: tensor<128x128xf32>, %lhs_scale: tensor<128x4xbf16>, %rhs: tensor<128x256xf32>, %rhs_scale: tensor<256x4xi8>) -> tensor<128x256xf32> { // expected-error@+1 {{LHS scale element type 'bf16' is not supported}} - %0 = xtile.dot_scaled %lhs scale %lhs_scale, %rhs scale %rhs_scale {fastMath = true, lhs_elem_type = f32, rhs_elem_type = f32} : tensor<128x128xf32>, tensor<128x4xbf16> * tensor<128x256xf32>, tensor<256x4xi8> -> tensor<128x256xf32> + %0 = xtile.dot_scaled %lhs scale %lhs_scale, %rhs scale %rhs_scale : tensor<128x128xf32>, tensor<128x4xbf16> * tensor<128x256xf32>, tensor<256x4xi8> -> tensor<128x256xf32> return %0 : tensor<128x256xf32> } @@ -149,7 +149,7 @@ func.func @invalid_dot_scaled_scale_type(%lhs: tensor<128x128xf32>, %lhs_scale: func.func @invalid_dot_scaled_operand_type(%lhs: tensor<128x128xi32>, %lhs_scale: tensor<128x4xi8>, %rhs: tensor<128x256xf32>, %rhs_scale: tensor<256x4xi8>) -> tensor<128x256xf32> { // expected-error@+1 {{'xtile.dot_scaled' op operand #0 must be ranked tensor of floating-point or 8-bit signless integer values, but got 'tensor<128x128xi32>'}} - %0 = xtile.dot_scaled %lhs scale %lhs_scale, %rhs scale %rhs_scale {fastMath = true, lhs_elem_type = f32, rhs_elem_type = f32} : tensor<128x128xi32>, tensor<128x4xi8> * tensor<128x256xf32>, tensor<256x4xi8> -> tensor<128x256xf32> + %0 = xtile.dot_scaled %lhs scale %lhs_scale, %rhs scale %rhs_scale : tensor<128x128xi32>, tensor<128x4xi8> * tensor<128x256xf32>, tensor<256x4xi8> -> tensor<128x256xf32> return %0 : tensor<128x256xf32> } @@ -157,7 +157,7 @@ func.func @invalid_dot_scaled_operand_type(%lhs: tensor<128x128xi32>, %lhs_scale func.func @invalid_dot_scaled_logical_elem_type(%lhs: tensor<128x128xf32>, %lhs_scale: tensor<128x4xi8>, %rhs: tensor<128x256xf32>, %rhs_scale: tensor<256x4xi8>) -> tensor<128x256xf32> { // expected-error@+1 {{LHS logical element type 'i32' is not supported}} - %0 = xtile.dot_scaled %lhs scale %lhs_scale, %rhs scale %rhs_scale {fastMath = true, lhs_elem_type = i32, rhs_elem_type = f32} : tensor<128x128xf32>, tensor<128x4xi8> * tensor<128x256xf32>, tensor<256x4xi8> -> tensor<128x256xf32> + %0 = xtile.dot_scaled %lhs scale %lhs_scale, %rhs scale %rhs_scale : tensor<128x128xf32>, tensor<128x4xi8> * tensor<128x256xf32>, tensor<256x4xi8> -> tensor<128x256xf32> return %0 : tensor<128x256xf32> } @@ -189,7 +189,7 @@ func.func @illegal_mask_out_of_bounds(%src: tensor<32xf64>, %mask: f64) -> tenso // ----- func.func @scan_op(%input: tensor<10x20x30xf32>, %init: tensor<10x20x1xf32>) -> (tensor<10x20x30xf32>, tensor<10x20x1xf32>) { - %output, %carry = xtile.scan(%input) inits(%init) dimension = 2 {scan_dim_size = 30 : i64} : (tensor<10x20x30xf32>), (tensor<10x20x1xf32>) -> (tensor<10x20x30xf32>), (tensor<10x20x1xf32>) { + %output, %carry = xtile.scan(%input) inits(%init) dimension = 2 : (tensor<10x20x30xf32>), (tensor<10x20x1xf32>) -> (tensor<10x20x30xf32>), (tensor<10x20x1xf32>) { ^bb0(%arg0: tensor<10x20xf32>, %arg1: tensor<10x20xf32>): %add = stablehlo.add %arg0, %arg1 : tensor<10x20xf32> stablehlo.return %add : tensor<10x20xf32> @@ -201,7 +201,7 @@ func.func @scan_op(%input: tensor<10x20x30xf32>, %init: tensor<10x20x1xf32>) -> func.func @scan_op_rank_reduced_init(%input: tensor<10x20x30xf32>, %init: tensor<10x20xf32>) -> (tensor<10x20x30xf32>, tensor<10x20xf32>) { // expected-error @+1 {{init #0 type 'tensor<10x20xf32>' does not match expected type 'tensor<10x20x1xf32>'}} - %output, %carry = xtile.scan(%input) inits(%init) dimension = 2 {scan_dim_size = 30 : i64} : (tensor<10x20x30xf32>), (tensor<10x20xf32>) -> (tensor<10x20x30xf32>), (tensor<10x20xf32>) { + %output, %carry = xtile.scan(%input) inits(%init) dimension = 2 : (tensor<10x20x30xf32>), (tensor<10x20xf32>) -> (tensor<10x20x30xf32>), (tensor<10x20xf32>) { ^bb0(%arg0: tensor<10x20xf32>, %arg1: tensor<10x20xf32>): %add = stablehlo.add %arg0, %arg1 : tensor<10x20xf32> stablehlo.return %add : tensor<10x20xf32> @@ -213,7 +213,7 @@ func.func @scan_op_rank_reduced_init(%input: tensor<10x20x30xf32>, %init: tensor func.func @scan_op_dimension_out_of_range(%input: tensor<10x20x30xf32>, %init: tensor<10x20x1xf32>) -> (tensor<10x20x30xf32>, tensor<10x20x1xf32>) { // expected-error @+1 {{scan dimension 3 is out of range for input #0 of rank 3}} - %output, %carry = xtile.scan(%input) inits(%init) dimension = 3 {scan_dim_size = 30 : i64} : (tensor<10x20x30xf32>), (tensor<10x20x1xf32>) -> (tensor<10x20x30xf32>), (tensor<10x20x1xf32>) { + %output, %carry = xtile.scan(%input) inits(%init) dimension = 3 : (tensor<10x20x30xf32>), (tensor<10x20x1xf32>) -> (tensor<10x20x30xf32>), (tensor<10x20x1xf32>) { ^bb0(%arg0: tensor<10x20xf32>, %arg1: tensor<10x20xf32>): %add = stablehlo.add %arg0, %arg1 : tensor<10x20xf32> stablehlo.return %add : tensor<10x20xf32> @@ -225,7 +225,7 @@ func.func @scan_op_dimension_out_of_range(%input: tensor<10x20x30xf32>, %init: t func.func @scan_op_init_count_mismatch(%input: tensor<10x20x30xf32>, %init: tensor<10x20x1xf32>) -> (tensor<10x20x30xf32>, tensor<10x20x1xf32>) { // expected-error @+1 {{expects one init, output and carry per input, but got 1 inputs, 2 inits, 1 outputs and 1 carries}} - %output, %carry = xtile.scan(%input) inits(%init, %init) dimension = 2 {scan_dim_size = 30 : i64} : (tensor<10x20x30xf32>), (tensor<10x20x1xf32>, tensor<10x20x1xf32>) -> (tensor<10x20x30xf32>), (tensor<10x20x1xf32>) { + %output, %carry = xtile.scan(%input) inits(%init, %init) dimension = 2 : (tensor<10x20x30xf32>), (tensor<10x20x1xf32>, tensor<10x20x1xf32>) -> (tensor<10x20x30xf32>), (tensor<10x20x1xf32>) { ^bb0(%arg0: tensor<10x20xf32>, %arg1: tensor<10x20xf32>): %add = stablehlo.add %arg0, %arg1 : tensor<10x20xf32> stablehlo.return %add : tensor<10x20xf32> @@ -237,7 +237,7 @@ func.func @scan_op_init_count_mismatch(%input: tensor<10x20x30xf32>, %init: tens func.func @scan_op_negative_dimension(%input: tensor<10x20x30xf32>, %init: tensor<10x20x1xf32>) -> (tensor<10x20x30xf32>, tensor<10x20x1xf32>) { // expected-error @+1 {{scan dimension -1 is out of range for input #0 of rank 3}} - %output, %carry = xtile.scan(%input) inits(%init) dimension = -1 {scan_dim_size = 30 : i64} : (tensor<10x20x30xf32>), (tensor<10x20x1xf32>) -> (tensor<10x20x30xf32>), (tensor<10x20x1xf32>) { + %output, %carry = xtile.scan(%input) inits(%init) dimension = -1 : (tensor<10x20x30xf32>), (tensor<10x20x1xf32>) -> (tensor<10x20x30xf32>), (tensor<10x20x1xf32>) { ^bb0(%arg0: tensor<10x20xf32>, %arg1: tensor<10x20xf32>): %add = stablehlo.add %arg0, %arg1 : tensor<10x20xf32> stablehlo.return %add : tensor<10x20xf32> @@ -249,7 +249,7 @@ func.func @scan_op_negative_dimension(%input: tensor<10x20x30xf32>, %init: tenso func.func @scan_op_output_type_mismatch(%input: tensor<10x20x30xf32>, %init: tensor<10x20x1xf32>) -> (tensor<10x20x30xf64>, tensor<10x20x1xf32>) { // expected-error @+1 {{output #0 type 'tensor<10x20x30xf64>' does not match input type 'tensor<10x20x30xf32>'}} - %output, %carry = xtile.scan(%input) inits(%init) dimension = 2 {scan_dim_size = 30 : i64} : (tensor<10x20x30xf32>), (tensor<10x20x1xf32>) -> (tensor<10x20x30xf64>), (tensor<10x20x1xf32>) { + %output, %carry = xtile.scan(%input) inits(%init) dimension = 2 : (tensor<10x20x30xf32>), (tensor<10x20x1xf32>) -> (tensor<10x20x30xf64>), (tensor<10x20x1xf32>) { ^bb0(%arg0: tensor<10x20xf32>, %arg1: tensor<10x20xf32>): %add = stablehlo.add %arg0, %arg1 : tensor<10x20xf32> stablehlo.return %add : tensor<10x20xf32> @@ -261,7 +261,7 @@ func.func @scan_op_output_type_mismatch(%input: tensor<10x20x30xf32>, %init: ten func.func @scan_op_carry_type_mismatch(%input: tensor<10x20x30xf32>, %init: tensor<10x20x1xf32>) -> (tensor<10x20x30xf32>, tensor<10x20x30xf32>) { // expected-error @+1 {{carry #0 type 'tensor<10x20x30xf32>' does not match expected type 'tensor<10x20x1xf32>'}} - %output, %carry = xtile.scan(%input) inits(%init) dimension = 2 {scan_dim_size = 30 : i64} : (tensor<10x20x30xf32>), (tensor<10x20x1xf32>) -> (tensor<10x20x30xf32>), (tensor<10x20x30xf32>) { + %output, %carry = xtile.scan(%input) inits(%init) dimension = 2 : (tensor<10x20x30xf32>), (tensor<10x20x1xf32>) -> (tensor<10x20x30xf32>), (tensor<10x20x30xf32>) { ^bb0(%arg0: tensor<10x20xf32>, %arg1: tensor<10x20xf32>): %add = stablehlo.add %arg0, %arg1 : tensor<10x20xf32> stablehlo.return %add : tensor<10x20xf32> diff --git a/third_party/xla/xla/codegen/xtile/ir/transforms/verify_legal_xtile_ops.cc b/third_party/xla/xla/codegen/xtile/ir/transforms/verify_legal_xtile_ops.cc index 331bb845448054..91a224866506bb 100644 --- a/third_party/xla/xla/codegen/xtile/ir/transforms/verify_legal_xtile_ops.cc +++ b/third_party/xla/xla/codegen/xtile/ir/transforms/verify_legal_xtile_ops.cc @@ -84,18 +84,18 @@ std::optional IsLegalStablehloOp(mlir::Operation* op) { if (mlir::isa< // go/keep-sorted start mlir::stablehlo::AbsOp, mlir::stablehlo::AddOp, - mlir::stablehlo::AllReduceOp, mlir::stablehlo::AndOp, - mlir::stablehlo::BroadcastInDimOp, mlir::stablehlo::CompareOp, - mlir::stablehlo::ComplexOp, mlir::stablehlo::ConstantOp, - mlir::stablehlo::ConvertOp, mlir::stablehlo::DivOp, - mlir::stablehlo::DotGeneralOp, mlir::stablehlo::ImagOp, - mlir::stablehlo::MaxOp, mlir::stablehlo::MinOp, - mlir::stablehlo::MulOp, mlir::stablehlo::NegOp, mlir::stablehlo::OrOp, - mlir::stablehlo::PowOp, mlir::stablehlo::RealOp, - mlir::stablehlo::ReduceOp, mlir::stablehlo::RemOp, - mlir::stablehlo::ReshapeOp, mlir::stablehlo::ReturnOp, - mlir::stablehlo::SubtractOp, mlir::stablehlo::TransposeOp, - mlir::stablehlo::XorOp + mlir::stablehlo::AllGatherOp, mlir::stablehlo::AllReduceOp, + mlir::stablehlo::AndOp, mlir::stablehlo::BroadcastInDimOp, + mlir::stablehlo::CompareOp, mlir::stablehlo::ComplexOp, + mlir::stablehlo::ConstantOp, mlir::stablehlo::ConvertOp, + mlir::stablehlo::DivOp, mlir::stablehlo::DotGeneralOp, + mlir::stablehlo::ImagOp, mlir::stablehlo::MaxOp, + mlir::stablehlo::MinOp, mlir::stablehlo::MulOp, + mlir::stablehlo::NegOp, mlir::stablehlo::OrOp, mlir::stablehlo::PowOp, + mlir::stablehlo::RealOp, mlir::stablehlo::ReduceOp, + mlir::stablehlo::RemOp, mlir::stablehlo::ReshapeOp, + mlir::stablehlo::ReturnOp, mlir::stablehlo::SubtractOp, + mlir::stablehlo::TransposeOp, mlir::stablehlo::XorOp // go/keep-sorted end >(op)) { return std::nullopt; diff --git a/third_party/xla/xla/codegen/xtile/ir/xtile_dialect.td b/third_party/xla/xla/codegen/xtile/ir/xtile_dialect.td index 3f49d7032eba1f..38c4bbf47a1ce8 100644 --- a/third_party/xla/xla/codegen/xtile/ir/xtile_dialect.td +++ b/third_party/xla/xla/codegen/xtile/ir/xtile_dialect.td @@ -34,6 +34,7 @@ def XTileDialect : Dialect { ]; let useDefaultAttributePrinterParser = 1; + let useStrictPropertiesInAssemblyFormat = 1; } #endif // XLA_CODEGEN_XTILE_IR_XTILE_DIALECT diff --git a/third_party/xla/xla/codegen/xtile/ir/xtile_ops.td b/third_party/xla/xla/codegen/xtile/ir/xtile_ops.td index 1cbaab2442e73b..468a9b05823561 100644 --- a/third_party/xla/xla/codegen/xtile/ir/xtile_ops.td +++ b/third_party/xla/xla/codegen/xtile/ir/xtile_ops.td @@ -350,7 +350,7 @@ def ScanOp : XTile_Op<"scan", [Pure, AttrSizedOperandSegments, AttrSizedResultSe let assemblyFormat = [{ `(` $inputs `)` `inits` `(` $inits `)` `dimension` `=` $dimension - attr-dict + prop-dict attr-dict `:` `(` type($inputs) `)` `,` `(` type($inits) `)` `->` `(` type($outputs) `)` `,` `(` type($carries) `)` $body }]; @@ -391,7 +391,7 @@ def DotScaledOp : XTile_Op<"dot_scaled", [Pure, AttrSizedOperandSegments]> { let results = (outs RankedTensorOf<[AnyFloat]>:$result); let assemblyFormat = [{ - $lhs (`scale` $lhs_scale^)? `,` $rhs (`scale` $rhs_scale^)? attr-dict + $lhs (`scale` $lhs_scale^)? `,` $rhs (`scale` $rhs_scale^)? prop-dict attr-dict `:` type($lhs) (`,` type($lhs_scale)^)? `*` type($rhs) (`,` type($rhs_scale)^)? `->` type($result) }]; diff --git a/third_party/xla/xla/debug_options_flags.cc b/third_party/xla/xla/debug_options_flags.cc index 4a5a969daac4eb..4ff61a22dd4f24 100644 --- a/third_party/xla/xla/debug_options_flags.cc +++ b/third_party/xla/xla/debug_options_flags.cc @@ -54,7 +54,6 @@ limitations under the License. #include "xla/hlo/pass/hlo_pass_filter.h" #include "xla/parse_flags_from_env.h" #include "xla/service/collective_utils.h" -#include "xla/stream_executor/cuda/nvjitlink_support.h" #include "xla/stream_executor/cuda/ptx_compiler_support.h" #include "xla/tsl/platform/logging.h" // IWYU pragma: keep #include "xla/tsl/util/command_line_flags.h" @@ -303,6 +302,7 @@ DebugOptions DefaultDebugOptionsIgnoringFlags() { opts.set_xla_gpu_trace_annotation_level(0); opts.set_xla_gpu_enable_cupti_multi_subscriber(true); + opts.add_xla_gpu_enable_command_buffer(DebugOptions::COLLECTIVES_KERNEL); opts.add_xla_gpu_enable_command_buffer(DebugOptions::CONDITIONAL); opts.add_xla_gpu_enable_command_buffer(DebugOptions::CUBLAS); opts.add_xla_gpu_enable_command_buffer(DebugOptions::CUBLASLT); @@ -522,7 +522,7 @@ DebugOptions DefaultDebugOptionsIgnoringFlags() { opts.set_xla_gpu_experimental_matmul_perf_table_path(""); // TODO(b/366475196): Create XLA GPU without cuDNN, cuBLAS. opts.set_xla_gpu_experimental_disable_binary_libraries(false); - opts.set_xla_gpu_experimental_enable_conv_fusion(false); + opts.set_xla_gpu_experimental_enable_conv_fusion(true); opts.set_xla_gpu_dot_merger_threshold_mb(64); opts.set_xla_enable_fast_math(false); opts.set_xla_gpu_experimental_parallel_collective_overlap_limit(1); @@ -2958,7 +2958,7 @@ void MakeDebugOptionsFlags(std::vector* flag_list, : DebugOptions::LIB_NV_JIT_LINK_MODE_DISABLED); return true; }, - stream_executor::IsLibNvJitLinkSupported(), + /*default_value_for_display=*/true, "Use libnvjitlink for PTX-to-GPU-assembly compilation instead of " "calling ptxas.")); flag_list->push_back(tsl::Flag( diff --git a/third_party/xla/xla/debug_options_parsers_test.cc b/third_party/xla/xla/debug_options_parsers_test.cc index 78d35a6bd31642..8c08a475fb0dff 100644 --- a/third_party/xla/xla/debug_options_parsers_test.cc +++ b/third_party/xla/xla/debug_options_parsers_test.cc @@ -442,15 +442,15 @@ TEST(ParseRepeatedEnumModifiersTest, Invalid) { TEST(ParseRepeatedEnumFlagsTest, CommandBufferCmdType) { DebugOptions debug_options = DefaultDebugOptionsIgnoringFlags(); - // Check that the default setting has 6 types. + // Check that the default setting has 8 types. const auto& enabled_types = debug_options.xla_gpu_enable_command_buffer(); - ASSERT_EQ(enabled_types.size(), 7); + ASSERT_EQ(enabled_types.size(), 8); ASSERT_THAT( enabled_types, - ElementsAre(DebugOptions::CONDITIONAL, DebugOptions::CUBLAS, - DebugOptions::CUBLASLT, DebugOptions::CUDNN, - DebugOptions::CUSTOM_CALL, DebugOptions::DYNAMIC_SLICE_FUSION, - DebugOptions::FUSION)); + ElementsAre(DebugOptions::COLLECTIVES_KERNEL, DebugOptions::CONDITIONAL, + DebugOptions::CUBLAS, DebugOptions::CUBLASLT, + DebugOptions::CUDNN, DebugOptions::CUSTOM_CALL, + DebugOptions::DYNAMIC_SLICE_FUSION, DebugOptions::FUSION)); // Initialize the flag objects. std::vector flag_objects; @@ -459,30 +459,33 @@ TEST(ParseRepeatedEnumFlagsTest, CommandBufferCmdType) { // Removing options from the existing setting. SetXlaFlagsEnvVar("--xla_gpu_enable_command_buffer=-fusion,-cublas"); ParseFlagsFromEnvAndDieIfUnknown("XLA_FLAGS", flag_objects); - EXPECT_EQ(enabled_types.size(), 5); + EXPECT_EQ(enabled_types.size(), 6); EXPECT_THAT(enabled_types, - ElementsAre(DebugOptions::CONDITIONAL, DebugOptions::CUBLASLT, + ElementsAre(DebugOptions::COLLECTIVES_KERNEL, + DebugOptions::CONDITIONAL, DebugOptions::CUBLASLT, DebugOptions::CUDNN, DebugOptions::CUSTOM_CALL, DebugOptions::DYNAMIC_SLICE_FUSION)); // Removing an option that isn't there and adding a duplicate. SetXlaFlagsEnvVar("--xla_gpu_enable_command_buffer=+cublaslt,-fusion"); ParseFlagsFromEnvAndDieIfUnknown("XLA_FLAGS", flag_objects); - EXPECT_EQ(enabled_types.size(), 5); + EXPECT_EQ(enabled_types.size(), 6); EXPECT_THAT(enabled_types, - ElementsAre(DebugOptions::CONDITIONAL, DebugOptions::CUBLASLT, + ElementsAre(DebugOptions::COLLECTIVES_KERNEL, + DebugOptions::CONDITIONAL, DebugOptions::CUBLASLT, DebugOptions::CUDNN, DebugOptions::CUSTOM_CALL, DebugOptions::DYNAMIC_SLICE_FUSION)); // Adding an option. SetXlaFlagsEnvVar("--xla_gpu_enable_command_buffer=+cublas"); ParseFlagsFromEnvAndDieIfUnknown("XLA_FLAGS", flag_objects); - EXPECT_EQ(enabled_types.size(), 6); + EXPECT_EQ(enabled_types.size(), 7); EXPECT_THAT( enabled_types, - ElementsAre(DebugOptions::CONDITIONAL, DebugOptions::CUBLASLT, - DebugOptions::CUDNN, DebugOptions::CUSTOM_CALL, - DebugOptions::DYNAMIC_SLICE_FUSION, DebugOptions::CUBLAS)); + ElementsAre(DebugOptions::COLLECTIVES_KERNEL, DebugOptions::CONDITIONAL, + DebugOptions::CUBLASLT, DebugOptions::CUDNN, + DebugOptions::CUSTOM_CALL, DebugOptions::DYNAMIC_SLICE_FUSION, + DebugOptions::CUBLAS)); // Overwriting the default setting. SetXlaFlagsEnvVar("--xla_gpu_enable_command_buffer=custom_call,fusion"); diff --git a/third_party/xla/xla/hlo/analysis/BUILD b/third_party/xla/xla/hlo/analysis/BUILD index d0b2e163a1da12..952d475c49271a 100644 --- a/third_party/xla/xla/hlo/analysis/BUILD +++ b/third_party/xla/xla/hlo/analysis/BUILD @@ -381,12 +381,14 @@ cc_library( deps = [ "//xla:shape_tree", "//xla:shape_util", + "//xla:tuple_tree", "//xla:util", "//xla:xla_data_proto_cc", "//xla/hlo/ir:hlo", "@com_google_absl//absl/algorithm:container", "@com_google_absl//absl/container:flat_hash_map", "@com_google_absl//absl/container:flat_hash_set", + "@com_google_absl//absl/container:node_hash_map", "@com_google_absl//absl/log", "@com_google_absl//absl/log:check", "@com_google_absl//absl/memory", diff --git a/third_party/xla/xla/hlo/analysis/hlo_dataflow_analysis.cc b/third_party/xla/xla/hlo/analysis/hlo_dataflow_analysis.cc index 83ca7aa091d94c..76d6ecf2ab1a3c 100644 --- a/third_party/xla/xla/hlo/analysis/hlo_dataflow_analysis.cc +++ b/third_party/xla/xla/hlo/analysis/hlo_dataflow_analysis.cc @@ -107,12 +107,13 @@ using absl::StrCat; HloDataflowAnalysis::HloDataflowAnalysis( const HloModule& module, bool ssa_form, bool bitcast_defines_value, absl::flat_hash_set execution_threads, - bool propagate_through_calls) + bool propagate_through_calls, bool propagate_through_control_flow) : module_(module), execution_threads_(std::move(execution_threads)), ssa_form_(ssa_form), bitcast_defines_value_(bitcast_defines_value), propagate_through_calls_(propagate_through_calls), + propagate_through_control_flow_(propagate_through_control_flow), call_graph_(CallGraph::Build(&module)) {} bool HloDataflowAnalysis::AreTransitiveUsesElementwiseOrTuple( @@ -477,8 +478,10 @@ bool HloDataflowAnalysis::UpdateAsyncChainOutputValueSet( HloInstruction* async_op) { CHECK(async_op->IsAsynchronous()); bool changed = false; - bool is_thread_included = HloInstruction::IsThreadIncluded( - async_op->async_execution_thread(), execution_threads_); + bool is_thread_included = + propagate_through_control_flow_ && + HloInstruction::IsThreadIncluded(async_op->async_execution_thread(), + execution_threads_); if (!is_thread_included && async_op->opcode() == HloOpcode::kAsyncStart) { // AsyncStart in a non-included thread has the output values defined, no @@ -563,6 +566,9 @@ bool HloDataflowAnalysis::UpdateAsyncStartValueSet( changed |= UpdateAsyncChainOperandValueSet(async_start, i, async_start->operand(i)); } + if (!propagate_through_control_flow_) { + return changed; + } bool is_dus = async_start->async_wrapped_opcode() == HloOpcode::kDynamicUpdateSlice; @@ -745,6 +751,9 @@ bool HloDataflowAnalysis::UpdateCallValueSet(HloInstruction* call) { bool HloDataflowAnalysis::UpdateConditionalValueSet( HloInstruction* conditional) { + if (!propagate_through_control_flow_) { + return false; + } CHECK_EQ(conditional->opcode(), HloOpcode::kConditional); std::vector inputs(conditional->branch_count()); for (int j = 0; j < conditional->branch_count(); ++j) { @@ -784,7 +793,8 @@ bool HloDataflowAnalysis::UpdateOptimizationBarrierValueSet( // Optimization Barriers just forward their operand. Given that barriers can // have a tuple operand, we iterate through its indexes, like for copies. // Unlike copies though we also propagate the top-level value. - CHECK_EQ(barrier->opcode(), HloOpcode::kOptimizationBarrier); + CHECK(barrier->opcode() == HloOpcode::kOptimizationBarrier || + barrier->IsCustomCall(kCallMarkerAfterTarget)); bool changed = false; for (auto& pair : GetInstructionValueSet(barrier)) { const ShapeIndex& index = pair.first; @@ -884,6 +894,9 @@ bool HloDataflowAnalysis::UpdateParameterValueSet(HloInstruction* parameter) { callsite.instruction()->operand(parameter->parameter_number()))); } } else if (opcode == HloOpcode::kWhile) { + if (!propagate_through_control_flow_) { + continue; + } // In a while instruction, the while operand (ie, the init value) and the // backedge are dataflow inputs to the parameter instruction. This is the // case for parameters of both the body and condition computations. @@ -900,6 +913,9 @@ bool HloDataflowAnalysis::UpdateParameterValueSet(HloInstruction* parameter) { } need_phi = true; } else if (opcode == HloOpcode::kConditional) { + if (!propagate_through_control_flow_) { + continue; + } CHECK_EQ(parameter->parameter_number(), 0); auto conditional = callsite.instruction(); // Conditional has branch_count+1 operands. Operand 0 is the branch_index, @@ -921,6 +937,9 @@ bool HloDataflowAnalysis::UpdateParameterValueSet(HloInstruction* parameter) { CHECK(found_parent); need_phi = true; } else if (opcode == HloOpcode::kAsyncStart) { + if (!propagate_through_control_flow_) { + continue; + } const HloInstruction* async_done = callsite.instruction()->async_chain_done(); // When an async chain crosses a while loop boundary, async_chain_done() @@ -977,7 +996,8 @@ bool HloDataflowAnalysis::UpdateParameterValueSet(HloInstruction* parameter) { } bool HloDataflowAnalysis::UpdateTupleValueSet(HloInstruction* tuple) { - CHECK_EQ(tuple->opcode(), HloOpcode::kTuple); + CHECK(tuple->opcode() == HloOpcode::kTuple || + tuple->IsCustomCall(kCallMarkerBeforeTarget)); bool changed = false; for (int64_t i = 0; i < tuple->operands().size(); ++i) { // Copy the value set(s) of each operand into the respective position in the @@ -1002,6 +1022,9 @@ bool HloDataflowAnalysis::UpdateTupleValueSet(HloInstruction* tuple) { } bool HloDataflowAnalysis::UpdateWhileValueSet(HloInstruction* xla_while) { + if (!propagate_through_control_flow_) { + return false; + } CHECK_EQ(xla_while->opcode(), HloOpcode::kWhile); const InstructionValueSet* const inputs[] = { &GetInstructionValueSet(xla_while->while_body()->root_instruction()), @@ -1220,6 +1243,14 @@ bool HloDataflowAnalysis::UpdateInstructionValueSet( return UpdateCollectivePermuteDoneValueSet(instruction); case HloOpcode::kOptimizationBarrier: return UpdateOptimizationBarrierValueSet(instruction); + case HloOpcode::kCustomCall: + if (instruction->custom_call_target() == kCallMarkerBeforeTarget) { + return UpdateTupleValueSet(instruction); + } + if (instruction->custom_call_target() == kCallMarkerAfterTarget) { + return UpdateOptimizationBarrierValueSet(instruction); + } + return false; default: break; } @@ -1340,6 +1371,10 @@ void HloDataflowAnalysis::Propagate() { callsite.instruction()->opcode() == HloOpcode::kCall) { continue; } + if (!propagate_through_control_flow_ && + callsite.instruction()->opcode() != HloOpcode::kCall) { + continue; + } if (callsite.instruction()->opcode() == HloOpcode::kWhile) { // Add the while itself, and the body and condition parameters. add_to_worklist(callsite.instruction()); @@ -1395,6 +1430,9 @@ absl::Status HloDataflowAnalysis::InitializeInstructionValueSets() { const CallGraphNode& call_graph_node = call_graph_->GetNode(computation); const bool is_regular_call_computation = IsRegularCallComputation(call_graph_node); + const bool is_control_flow_computation = + !call_graph_node.caller_callsites().empty() && + !is_regular_call_computation; for (HloInstruction* instruction : computation->MakeInstructionPostOrder()) { // Create an empty shape tree. @@ -1436,9 +1474,13 @@ absl::Status HloDataflowAnalysis::InitializeInstructionValueSets() { define_all_values(); } break; - case HloOpcode::kAddDependency: case HloOpcode::kWhile: case HloOpcode::kConditional: + if (!propagate_through_control_flow_) { + define_all_values(); + } + break; + case HloOpcode::kAddDependency: case HloOpcode::kGetTupleElement: case HloOpcode::kDomain: case HloOpcode::kOptimizationBarrier: @@ -1459,7 +1501,9 @@ absl::Status HloDataflowAnalysis::InitializeInstructionValueSets() { } if (call_graph_node.caller_callsites().empty() || call_graph_node.context() == CallContext::kEmbedded || - (!propagate_through_calls_ && is_regular_call_computation)) { + (!propagate_through_calls_ && is_regular_call_computation) || + (!propagate_through_control_flow_ && + is_control_flow_computation)) { // Parameters of computations called in a parallel context (eg, map // and reduce) as well as parameters of dead computations define all // values in their output. Otherwise the values of the parameter @@ -1474,6 +1518,14 @@ absl::Status HloDataflowAnalysis::InitializeInstructionValueSets() { // values flow from their operands. define_value_at(/*index=*/{}); break; + case HloOpcode::kCustomCall: + if (instruction->custom_call_target() == kCallMarkerBeforeTarget) { + define_value_at(/*index=*/{}); + } else if (instruction->custom_call_target() != + kCallMarkerAfterTarget) { + define_all_values(); + } + break; case HloOpcode::kAsyncStart: { // AsyncStart produces a tuple of {{aliased operands}, {destination}, // contexts}. It defines all of the tuple-shaped values and the @@ -1481,8 +1533,10 @@ absl::Status HloDataflowAnalysis::InitializeInstructionValueSets() { // // If the thread is excluded, then we don't track the contained // dataflow, and define the destination values too. - bool thread_included = HloInstruction::IsThreadIncluded( - instruction->async_execution_thread(), execution_threads_); + bool thread_included = + propagate_through_control_flow_ && + HloInstruction::IsThreadIncluded( + instruction->async_execution_thread(), execution_threads_); define_all_values([&](const ShapeIndex& index) { return ShapeUtil::GetSubshape(instruction->shape(), index) .IsTuple() || @@ -1493,23 +1547,55 @@ absl::Status HloDataflowAnalysis::InitializeInstructionValueSets() { } case HloOpcode::kAsyncUpdate: { // AsyncUpdate produces a tuple of {{aliased operands}, {destination}, - // contexts}. It defines all of the tuple-shaped values and the - // contexts. + // contexts}. It defines all of the tuple values and the contexts. + // When the async thread is excluded (or control flow propagation is + // disabled) and the output at {1, ...} is late bound, AsyncUpdate + // also defines the newly bound output values. + bool thread_included = + propagate_through_control_flow_ && + HloInstruction::IsThreadIncluded( + instruction->async_execution_thread(), execution_threads_); + const Shape& prev_shape = instruction->operand(0)->shape(); define_all_values([&](const ShapeIndex& index) { - return ShapeUtil::GetSubshape(instruction->shape(), index) - .IsTuple() || - index.front() > 1; + if (ShapeUtil::GetSubshape(instruction->shape(), index).IsTuple() || + index.front() > 1) { + return true; + } + if (!thread_included && index.front() == 1) { + return !ShapeUtil::IndexIsValid(prev_shape, index) || + !ShapeUtil::Compatible( + ShapeUtil::GetSubshape(instruction->shape(), index), + ShapeUtil::GetSubshape(prev_shape, index)); + } + return false; }); break; } - case HloOpcode::kAsyncDone: + case HloOpcode::kAsyncDone: { // AsyncDone's output aliases its output. It defines all remaining - // tuple-shaped values. + // tuple values, plus any late bound output values when the async + // thread is excluded (or control flow propagation is disabled). + bool thread_included = + propagate_through_control_flow_ && + HloInstruction::IsThreadIncluded( + instruction->async_execution_thread(), execution_threads_); + const Shape& prev_shape = instruction->operand(0)->shape(); define_all_values([&](const ShapeIndex& index) { - return ShapeUtil::GetSubshape(instruction->shape(), index) - .IsTuple(); + if (ShapeUtil::GetSubshape(instruction->shape(), index).IsTuple()) { + return true; + } + if (!thread_included) { + ShapeIndex src_index = {1}; + src_index.insert(src_index.end(), index.begin(), index.end()); + return !ShapeUtil::IndexIsValid(prev_shape, src_index) || + !ShapeUtil::Compatible( + ShapeUtil::GetSubshape(instruction->shape(), index), + ShapeUtil::GetSubshape(prev_shape, src_index)); + } + return false; }); break; + } case HloOpcode::kCopyStart: // CopyStart produces a tuple of {destination buffer, aliased operand, // U32 context}. @@ -1655,13 +1741,14 @@ absl::StatusOr> HloDataflowAnalysis::Run( const HloModule& module, bool ssa_form, bool bitcast_defines_value, absl::flat_hash_set execution_threads, bool propagate_through_calls, - std::optional> precompute_uses) { + std::optional> precompute_uses, + bool propagate_through_control_flow) { VLOG(1) << "HloDataflowAnalysis::Run on module " << module.name(); XLA_VLOG_LINES(2, module.ToString()); - auto dataflow_analysis = absl::WrapUnique( - new HloDataflowAnalysis(module, ssa_form, bitcast_defines_value, - execution_threads, propagate_through_calls)); + auto dataflow_analysis = absl::WrapUnique(new HloDataflowAnalysis( + module, ssa_form, bitcast_defines_value, execution_threads, + propagate_through_calls, propagate_through_control_flow)); ABSL_RETURN_IF_ERROR(dataflow_analysis->RunImpl()); if (precompute_uses.has_value()) { std::vector values; diff --git a/third_party/xla/xla/hlo/analysis/hlo_dataflow_analysis.h b/third_party/xla/xla/hlo/analysis/hlo_dataflow_analysis.h index 202a9c6cb60cd9..b05ca2ca8d0447 100644 --- a/third_party/xla/xla/hlo/analysis/hlo_dataflow_analysis.h +++ b/third_party/xla/xla/hlo/analysis/hlo_dataflow_analysis.h @@ -88,7 +88,8 @@ class HloDataflowAnalysis { absl::flat_hash_set execution_threads = {}, bool propagate_through_calls = true, std::optional> precompute_uses = - std::nullopt); + std::nullopt, + bool propagate_through_control_flow = true); // Returns true if 'instruction' defines an HLO value at the given shape index // of its output. @@ -219,7 +220,8 @@ class HloDataflowAnalysis { HloDataflowAnalysis(const HloModule& module, bool ssa_form, bool bitcast_defines_value, absl::flat_hash_set execution_threads, - bool propagate_through_calls = true); + bool propagate_through_calls = true, + bool propagate_through_control_flow = true); // Runs dataflow analysis on the module attached to this HloDataflowAnalysis. absl::Status RunImpl(); @@ -335,6 +337,7 @@ class HloDataflowAnalysis { const bool ssa_form_; const bool bitcast_defines_value_; bool propagate_through_calls_ = true; + bool propagate_through_control_flow_ = true; std::unique_ptr call_graph_; diff --git a/third_party/xla/xla/hlo/analysis/hlo_dataflow_analysis_test.cc b/third_party/xla/xla/hlo/analysis/hlo_dataflow_analysis_test.cc index 9265fb58d1f346..fbd952fbe8fe62 100644 --- a/third_party/xla/xla/hlo/analysis/hlo_dataflow_analysis_test.cc +++ b/third_party/xla/xla/hlo/analysis/hlo_dataflow_analysis_test.cc @@ -4332,5 +4332,109 @@ ENTRY main { } } +TEST_F(HloDataflowAnalysisTest, DisablePropagateThroughControlFlow) { + const char* hlo_text = R"hlo( +HloModule module + +callee { + ROOT call_param = f32[] parameter(0) +} + +while_cond { + param = f32[] parameter(0) + ROOT cond = pred[] constant(false) +} + +while_body { + param = f32[] parameter(0) + ROOT add = f32[] add(param, param) +} + +ENTRY main { + const0 = f32[] constant(1.0) + call_inst = f32[] call(const0), to_apply=callee + ROOT while_inst = f32[] while(call_inst), condition=while_cond, body=while_body +} +)hlo"; + ASSERT_OK_AND_ASSIGN(module_, ParseAndReturnVerifiedModule( + hlo_text, GetModuleConfigForTest())); + HloInstruction* const0 = FindInstruction(module_.get(), "const0"); + HloInstruction* call_inst = FindInstruction(module_.get(), "call_inst"); + HloInstruction* call_param = + module_->GetComputationWithName("callee")->parameter_instruction(0); + HloInstruction* while_inst = FindInstruction(module_.get(), "while_inst"); + HloInstruction* body_param = + module_->GetComputationWithName("while_body")->parameter_instruction(0); + + { + ASSERT_OK_AND_ASSIGN( + auto analysis, + HloDataflowAnalysis::Run( + *module_, /*ssa_form=*/false, /*bitcast_defines_value=*/false, + /*execution_threads=*/{}, /*propagate_through_calls=*/false, + /*precompute_uses=*/std::nullopt, + /*propagate_through_control_flow=*/false)); + EXPECT_TRUE(analysis->ValueIsDefinedAt(call_inst)); + EXPECT_TRUE(analysis->ValueIsDefinedAt(call_param)); + EXPECT_TRUE(analysis->ValueIsDefinedAt(while_inst)); + EXPECT_TRUE(analysis->ValueIsDefinedAt(body_param)); + EXPECT_THAT(analysis->GetValueSet(body_param).values(), + ::testing::Not( + ::testing::Contains(&analysis->GetValueDefinedAt(const0)))); + } + { + ASSERT_OK_AND_ASSIGN( + auto analysis, + HloDataflowAnalysis::Run( + *module_, /*ssa_form=*/false, /*bitcast_defines_value=*/false, + /*execution_threads=*/{}, /*propagate_through_calls=*/true, + /*precompute_uses=*/std::nullopt, + /*propagate_through_control_flow=*/false)); + EXPECT_FALSE(analysis->ValueIsDefinedAt(call_inst)); + EXPECT_FALSE(analysis->ValueIsDefinedAt(call_param)); + EXPECT_THAT(analysis->GetValueSet(call_inst).values(), + ::testing::Contains(&analysis->GetValueDefinedAt(const0))); + EXPECT_TRUE(analysis->ValueIsDefinedAt(while_inst)); + EXPECT_TRUE(analysis->ValueIsDefinedAt(body_param)); + EXPECT_THAT(analysis->GetValueSet(body_param).values(), + ::testing::Not( + ::testing::Contains(&analysis->GetValueDefinedAt(const0)))); + } +} + +TEST_F(HloDataflowAnalysisTest, CallMarkerCustomCalls) { + const char* hlo_text = R"hlo( +HloModule module + +ENTRY main { + p0 = f32[4] parameter(0) + p1 = f32[4] parameter(1) + before = (f32[4], f32[4]) custom-call(p0, p1), custom_call_target="__xla_internal_call_marker_before" + gte0 = f32[4] get-tuple-element(before), index=0 + gte1 = f32[4] get-tuple-element(before), index=1 + add = f32[4] add(gte0, gte1) + ROOT after = f32[4] custom-call(add), custom_call_target="__xla_internal_call_marker_after" +} +)hlo"; + ASSERT_OK_AND_ASSIGN(module_, ParseAndReturnVerifiedModule( + hlo_text, GetModuleConfigForTest())); + HloInstruction* p0 = FindInstruction(module_.get(), "p0"); + HloInstruction* p1 = FindInstruction(module_.get(), "p1"); + HloInstruction* before = FindInstruction(module_.get(), "before"); + HloInstruction* add = FindInstruction(module_.get(), "add"); + HloInstruction* after = FindInstruction(module_.get(), "after"); + + ASSERT_OK_AND_ASSIGN(auto analysis, HloDataflowAnalysis::Run(*module_)); + EXPECT_TRUE(analysis->ValueIsDefinedAt(before, {})); + EXPECT_FALSE(analysis->ValueIsDefinedAt(before, {0})); + EXPECT_EQ(&analysis->GetUniqueValueAt(before, {0}), + &analysis->GetValueDefinedAt(p0)); + EXPECT_EQ(&analysis->GetUniqueValueAt(before, {1}), + &analysis->GetValueDefinedAt(p1)); + EXPECT_FALSE(analysis->ValueIsDefinedAt(after)); + EXPECT_EQ(&analysis->GetUniqueValueAt(after), + &analysis->GetValueDefinedAt(add)); +} + } // namespace } // namespace xla diff --git a/third_party/xla/xla/hlo/analysis/hlo_replication_analysis.cc b/third_party/xla/xla/hlo/analysis/hlo_replication_analysis.cc index 5443f988a4991e..6045dec29aa0ce 100644 --- a/third_party/xla/xla/hlo/analysis/hlo_replication_analysis.cc +++ b/third_party/xla/xla/hlo/analysis/hlo_replication_analysis.cc @@ -46,6 +46,7 @@ limitations under the License. #include "xla/map_util.h" #include "xla/shape_tree.h" #include "xla/shape_util.h" +#include "xla/tuple_tree.h" #include "xla/xla_data.pb.h" namespace xla { @@ -310,70 +311,212 @@ HloReplicationAnalysis::DetermineHloInstructionIsReplicated( return HloReplication::UniqueOnAllDevices(); } +const TupleTree& +HloReplicationAnalysis::GetReplication(const HloInstruction* inst) const { + auto it = hlo_replication_.find(inst); + CHECK(it != hlo_replication_.end()) + << inst->name() << " has no replication yet"; + return it->second; +} + +bool HloReplicationAnalysis::CombineReplication( + const TupleTree& source, TupleTree* dest) { + bool updated = false; + // Both trees are built from compatible shapes, so their pre order nodes line + // up one to one. + auto source_node = source.begin(); + for (TupleTree::NodePair& node : *dest) { + CHECK(source_node != source.end()) + << "source tree ends before " << node.first; + CHECK(source_node->first == node.first) + << source_node->first << " vs " << node.first; + HloReplication new_replication = + MergeReplications(node.second, source_node->second); + if (!node.second.Equal(new_replication)) { + node.second = std::move(new_replication); + updated = true; + } + ++source_node; + } + return updated; +} + +bool HloReplicationAnalysis::AssignOrCombineReplication( + TupleTree replication, const HloInstruction* dest) { + auto [it, inserted] = + hlo_replication_.try_emplace(dest, std::move(replication)); + // try_emplace leaves `replication` untouched when the key exists. + if (inserted || CombineReplication(replication, &it->second)) { + OnReplicationChanged(dest); + return true; + } + return false; +} + +bool HloReplicationAnalysis::PropagateReplication(const HloInstruction* source, + const HloInstruction* dest) { + auto source_it = hlo_replication_.find(source); + if (source_it == hlo_replication_.end()) { + return false; + } + auto dest_it = hlo_replication_.find(dest); + if (dest_it != hlo_replication_.end()) { + if (!CombineReplication(source_it->second, &dest_it->second)) { + return false; + } + } else { + // Copy before inserting: the insertion may rehash the map. + TupleTree copy(source_it->second); + hlo_replication_.try_emplace(dest, std::move(copy)); + } + OnReplicationChanged(dest); + return true; +} + +bool HloReplicationAnalysis::MarkNotReplicated(const HloInstruction* inst) { + auto [it, inserted] = hlo_replication_.try_emplace( + inst, inst->shape(), HloReplication::UniqueOnAllDevices()); + if (inserted) { + OnReplicationChanged(inst); + return true; + } + // Merging any value with UniqueOnAllDevices() gives UniqueOnAllDevices(), so + // the existing tree is overwritten in place; only nodes that were not unique + // yet count as a change for the fixed point. + bool updated = false; + for (TupleTree::NodePair& node : it->second) { + if (!node.second.IsUniqueOnAllDevices()) { + node.second = HloReplication::UniqueOnAllDevices(); + updated = true; + } + } + if (updated) { + OnReplicationChanged(inst); + } + return updated; +} + +void HloReplicationAnalysis::OnReplicationChanged(const HloInstruction* inst) { + const HloComputation* computation = inst->parent(); + // With partial replication every visit evaluates every instruction. Users + // come later in post order, so a visit that evaluates every instruction of + // the computation reaches them, and only a change of its root matters. + if (support_partial_replication_ || + (computation == visiting_computation_ && visiting_all_instructions_ && + inst != computation->root_instruction())) { + return; + } + MarkDependentsDirty(inst); +} + +void HloReplicationAnalysis::MarkDependentsDirty(const HloInstruction* inst) { + const HloComputation* computation = inst->parent(); + if (inst == computation->root_instruction()) { + for (const HloInstruction* caller : computation->caller_instructions()) { + MarkDirty(caller); + } + } + // The first visit of a computation evaluates every instruction too. + if (computation == visiting_computation_ && visiting_all_instructions_) { + return; + } + auto it = computation_states_.find(computation); + if (it == computation_states_.end() || !it->second.visited) { + return; + } + for (const HloInstruction* user : inst->users()) { + MarkDirty(user); + } +} + +void HloReplicationAnalysis::MarkDirty(const HloInstruction* inst) { + // An instruction without a replication has not been evaluated yet. A visit + // that evaluates every instruction reaches it later, or is evaluating it now + // and MarkEvaluated then checks the computations it calls. + if (!hlo_replication_.contains(inst) || !dirty_.insert(inst).second) { + return; + } + const HloComputation* computation = inst->parent(); + if (computation_states_[computation].num_dirty++ == 0) { + for (const HloInstruction* caller : computation->caller_instructions()) { + MarkDirty(caller); + } + } +} + +void HloReplicationAnalysis::MarkEvaluated(const HloInstruction* inst) { + for (const HloComputation* callee : inst->called_computations()) { + auto it = computation_states_.find(callee); + if (it != computation_states_.end() && it->second.num_dirty > 0) { + MarkDirty(inst); + return; + } + } + if (dirty_.erase(inst) > 0) { + --computation_states_[inst->parent()].num_dirty; + } +} + bool HloReplicationAnalysis::ComputeHloReplicationOnComputation( const HloComputation* computation, bool mark_everything_not_replicated) { + // A repeat visit with everything marked unique cannot change anything. + if (mark_everything_not_replicated && + !computations_marked_not_replicated_.insert(computation).second) { + return false; + } + ComputationState& state = computation_states_[computation]; + if (state.post_order.empty()) { + state.post_order = computation->MakeInstructionPostOrder(); + } + // A repeat visit that does not mark evaluates only the dirty instructions: + // merging a value that is replicated or unique on all devices with itself + // changes nothing. A partially replicated value whose device sets are all + // singletons merges with itself to unique, so with partial replication every + // visit evaluates every instruction. + const bool evaluate_all = !state.visited || mark_everything_not_replicated || + support_partial_replication_; + if (!evaluate_all && state.num_dirty == 0) { + return false; + } + const HloComputation* outer_computation = visiting_computation_; + const bool outer_evaluates_all = visiting_all_instructions_; + visiting_computation_ = computation; + visiting_all_instructions_ = evaluate_all; bool changed = false; - for (const HloInstruction* inst : computation->MakeInstructionPostOrder()) { - // Assigns the shape tree to dest if dest doesn't have one yet, or combines - // it with the existing one by and'ing them. Returns if anything is updated. - auto assign_or_combine_shapetree = - [&](ShapeTree&& to_combine, - const HloInstruction* dest) { - auto it = hlo_replication_.find(dest); - if (it == hlo_replication_.end()) { - hlo_replication_[dest] = std::move(to_combine); - return true; - } - bool updated = false; - it->second.ForEachMutableElement( - [&](const ShapeIndex& index, HloReplication* element) { - HloReplication new_replication = - MergeReplications(*element, to_combine.element(index)); - if (!element->Equal(new_replication)) { - *element = std::move(new_replication); - updated = true; - } - }); - return updated; - }; - // Assigns or combines source's shape tree to dest. Returns if anything is - // updated. - auto propagate_shapetree = [&](const HloInstruction* source, - const HloInstruction* dest) { - auto source_it = hlo_replication_.find(source); - if (source_it == hlo_replication_.end()) { - return false; - } - return assign_or_combine_shapetree( - ShapeTree(source_it->second), dest); - }; + for (const HloInstruction* inst : state.post_order) { + if (!evaluate_all && !dirty_.contains(inst)) { + continue; + } + // Whether the evaluation visits the computations that `inst` calls. + bool visits_computations = false; // For the opcodes below that we do special handling, we don't need to // explicitly check mark_everything_not_replicated because if it is set, the // operands should already be marked as not replicated. if (inst->opcode() == HloOpcode::kWhile) { + visits_computations = true; // Since while body's input and output alias each other, we need to run it // multiple times until a fixed point is reached. while (true) { // First, propagate the input's and body root's shape trees to the // parameters of the body and condition. - bool updated = propagate_shapetree( + bool updated = PropagateReplication( inst->operand(0), inst->while_condition()->parameter_instruction(0)); - updated |= propagate_shapetree( + updated |= PropagateReplication( inst->while_body()->root_instruction(), inst->while_condition()->parameter_instruction(0)); - updated |= propagate_shapetree( + updated |= PropagateReplication( inst->operand(0), inst->while_body()->parameter_instruction(0)); updated |= - propagate_shapetree(inst->while_body()->root_instruction(), - inst->while_body()->parameter_instruction(0)); + PropagateReplication(inst->while_body()->root_instruction(), + inst->while_body()->parameter_instruction(0)); // Compute the condition. updated |= ComputeHloReplicationOnComputation( inst->while_condition(), mark_everything_not_replicated); // Compute the body. If the condition is not replicated, the while body // should be different across replicas. if (!ContainsKey(loops_known_with_same_iterations_, inst) && - !hlo_replication_[inst->while_condition()->root_instruction()] + !GetReplication(inst->while_condition()->root_instruction()) .element({}) .IsReplicatedOnAllDevices()) { updated |= ComputeHloReplicationOnComputation( @@ -388,11 +531,12 @@ bool HloReplicationAnalysis::ComputeHloReplicationOnComputation( changed = true; } // Propagate the input's and body root's shape trees to the while HLO. - changed |= propagate_shapetree(inst->operand(0), inst); + changed |= PropagateReplication(inst->operand(0), inst); changed |= - propagate_shapetree(inst->while_body()->root_instruction(), inst); + PropagateReplication(inst->while_body()->root_instruction(), inst); } else if (inst->opcode() == HloOpcode::kCall || inst->opcode() == HloOpcode::kFusion) { + visits_computations = true; auto called = inst->called_computations().front(); // If the called computation has several call sites, conservatively mark // everything inside as not-replicated, since we can't reason about @@ -403,23 +547,24 @@ bool HloReplicationAnalysis::ComputeHloReplicationOnComputation( /*mark_everything_not_replicated=*/true); } else { for (int64_t i = 0; i < inst->operand_count(); ++i) { - changed |= propagate_shapetree(inst->operand(i), - called->parameter_instruction(i)); + changed |= PropagateReplication(inst->operand(i), + called->parameter_instruction(i)); } changed |= ComputeHloReplicationOnComputation( called, mark_everything_not_replicated); } - changed |= propagate_shapetree(called->root_instruction(), inst); + changed |= PropagateReplication(called->root_instruction(), inst); } else if (inst->opcode() == HloOpcode::kConditional) { + visits_computations = true; // Propagate inputs' shape trees to the called computations' parameters. for (int64_t i = 0; i < inst->called_computations().size(); ++i) { - changed |= propagate_shapetree( + changed |= PropagateReplication( inst->operand(i + 1), inst->called_computations()[i]->parameter_instruction(0)); } // If the condition is not replicated, the conditional result should be // different across replicas. - if (!hlo_replication_[inst->operand(0)] + if (!GetReplication(inst->operand(0)) .element({}) .IsReplicatedOnAllDevices()) { for (auto called : inst->called_computations()) { @@ -427,64 +572,67 @@ bool HloReplicationAnalysis::ComputeHloReplicationOnComputation( called, /*mark_everything_not_replicated=*/true); } - changed |= assign_or_combine_shapetree( - ShapeTree(inst->shape(), - HloReplication::UniqueOnAllDevices()), - inst); + changed |= MarkNotReplicated(inst); } else { for (auto called : inst->called_computations()) { changed |= ComputeHloReplicationOnComputation( called, mark_everything_not_replicated); - changed |= propagate_shapetree(called->root_instruction(), inst); + changed |= PropagateReplication(called->root_instruction(), inst); } } } else if (inst->opcode() == HloOpcode::kTuple) { - ShapeTree shape_tree( - inst->shape(), HloReplication::ReplicatedOnAllDevices()); + // When marking, the tuple is unique at its own index like every other + // value of the computation; its elements come from marked operands. + TupleTree tree( + inst->shape(), mark_everything_not_replicated + ? HloReplication::UniqueOnAllDevices() + : HloReplication::ReplicatedOnAllDevices()); for (int64_t i = 0; i < inst->operand_count(); ++i) { - shape_tree.CopySubtreeFrom(hlo_replication_[inst->operand(i)], {}, {i}); + CHECK_OK(tree.CopyCompatibleSubtreeFrom( + GetReplication(inst->operand(i)), {}, {i})); } - changed |= assign_or_combine_shapetree(std::move(shape_tree), inst); + changed |= AssignOrCombineReplication(std::move(tree), inst); } else if (inst->opcode() == HloOpcode::kOptimizationBarrier) { - ShapeTree shape_tree = hlo_replication_[inst->operand(0)]; - changed |= assign_or_combine_shapetree(std::move(shape_tree), inst); + changed |= PropagateReplication(inst->operand(0), inst); } else if (inst->opcode() == HloOpcode::kGetTupleElement) { - ShapeTree shape_tree( - inst->shape(), HloReplication::ReplicatedOnAllDevices()); - shape_tree.CopySubtreeFrom(hlo_replication_[inst->operand(0)], - {inst->tuple_index()}, {}); - changed |= assign_or_combine_shapetree(std::move(shape_tree), inst); - } else if (inst->opcode() == HloOpcode::kInfeed && cross_partition_spmd_) { - ShapeTree shape_tree( - inst->shape(), HloReplication::UniqueOnAllDevices()); + TupleTree tree(inst->shape(), + HloReplication::ReplicatedOnAllDevices()); + CHECK_OK(tree.CopyCompatibleSubtreeFrom(GetReplication(inst->operand(0)), + {inst->tuple_index()}, {})); + changed |= AssignOrCombineReplication(std::move(tree), inst); + } else if (inst->opcode() == HloOpcode::kInfeed && cross_partition_spmd_ && + !mark_everything_not_replicated) { + // Replicated elements are equal on every partition only if all of them + // run the infeed equally often, so marked computations skip this rule. + TupleTree tree(inst->shape(), + HloReplication::UniqueOnAllDevices()); if (inst->has_sharding()) { auto sharding = inst->sharding().GetAsShapeTree(inst->shape()); - shape_tree.ForEachMutableElement( + tree.ForEachMutableElement( [&sharding](const ShapeIndex& index, HloReplication* data) { *data = sharding.element(index).IsReplicated() ? HloReplication::ReplicatedOnAllDevices() : HloReplication::UniqueOnAllDevices(); }); } - changed |= assign_or_combine_shapetree(std::move(shape_tree), inst); + changed |= AssignOrCombineReplication(std::move(tree), inst); + } else if (mark_everything_not_replicated) { + changed |= MarkNotReplicated(inst); } else { - if (mark_everything_not_replicated) { - changed |= assign_or_combine_shapetree( - ShapeTree(inst->shape(), - HloReplication::UniqueOnAllDevices()), - inst); - } else { - ShapeTree shape_tree( - inst->shape(), HloReplication::ReplicatedOnAllDevices()); - ShapeUtil::ForEachSubshape( - inst->shape(), [&](const Shape& subshape, const ShapeIndex& index) { - *shape_tree.mutable_element(index) = - DetermineHloInstructionIsReplicated(inst, index); - }); - changed |= assign_or_combine_shapetree(std::move(shape_tree), inst); + TupleTree tree(inst->shape(), + HloReplication::ReplicatedOnAllDevices()); + for (TupleTree::NodePair& node : tree) { + node.second = DetermineHloInstructionIsReplicated(inst, node.first); } + changed |= AssignOrCombineReplication(std::move(tree), inst); + } + if (state.num_dirty > 0 || visits_computations) { + MarkEvaluated(inst); } } + state.visited = true; + visiting_computation_ = outer_computation; + visiting_all_instructions_ = outer_evaluates_all; return changed; } @@ -495,8 +643,8 @@ absl::Status HloReplicationAnalysis::ComputeHloReplication() { auto entry = module_->entry_computation(); for (int i = 0; i < entry->num_parameters(); ++i) { auto param = entry->parameter_instruction(i); - ShapeTree shape_tree(param->shape(), - HloReplication::UniqueOnAllDevices()); + TupleTree tree(param->shape(), + HloReplication::UniqueOnAllDevices()); std::unique_ptr> sharding_tree = nullptr; if (cross_partition_spmd_) { @@ -524,7 +672,7 @@ absl::Status HloReplicationAnalysis::ComputeHloReplication() { if (sharding_tree != nullptr) { // In cross-partition spmd mode, set parameter replication status // based on the parameter's sharding. - *shape_tree.mutable_element(index) = + *tree.mutable_element(index) = sharding_tree->element(index).IsReplicated() ? HloReplication::ReplicatedOnAllDevices() : HloReplication::UniqueOnAllDevices(); @@ -535,13 +683,13 @@ absl::Status HloReplicationAnalysis::ComputeHloReplication() { if (!cross_partition_spmd_ && (*replication)[leaf_index]) { // Setting parameter replication status for replicas in // non cross-partition spmd mode. - *shape_tree.mutable_element(index) = + *tree.mutable_element(index) = HloReplication::ReplicatedOnAllDevices(); } if (cross_partition_spmd_ && !(*replication)[leaf_index]) { // Setting paramemter replication status for partitions in // cross-partition spmd mode. - *shape_tree.mutable_element(index) = + *tree.mutable_element(index) = HloReplication::UniqueOnAllDevices(); } ++leaf_index; @@ -549,10 +697,14 @@ absl::Status HloReplicationAnalysis::ComputeHloReplication() { return absl::OkStatus(); }); ABSL_RETURN_IF_ERROR(status); - hlo_replication_[param] = std::move(shape_tree); + hlo_replication_[param] = std::move(tree); } ComputeHloReplicationOnComputation(entry, /*mark_everything_not_replicated=*/false); + // Scratch state of the walk. + computation_states_.clear(); + dirty_.clear(); + computations_marked_not_replicated_.clear(); return absl::OkStatus(); } @@ -679,25 +831,27 @@ HloReplicationAnalysis::HloReplication::HloReplication() : state_(State::kReplicatedOnAllDevices) {} HloReplicationAnalysis::HloReplication::HloReplication( - HloReplicationAnalysis::HloReplication::State state, + HloReplicationAnalysis::HloReplication::State state) + : state_(state) { + DCHECK(state != State::kPartiallyReplicated); +} + +HloReplicationAnalysis::HloReplication::HloReplication( absl::Span> device_set_root_per_replica) - : state_(state), + : state_(State::kPartiallyReplicated), device_set_root_per_replica_( std::make_shared< HashOnConstruction>>>( - device_set_root_per_replica)) { - CHECK(state == State::kPartiallyReplicated || - device_set_root_per_replica_->empty()); -} + device_set_root_per_replica)) {} HloReplicationAnalysis::HloReplication HloReplicationAnalysis::HloReplication::ReplicatedOnAllDevices() { - return HloReplication(State::kReplicatedOnAllDevices, {}); + return HloReplication(State::kReplicatedOnAllDevices); } HloReplicationAnalysis::HloReplication HloReplicationAnalysis::HloReplication::UniqueOnAllDevices() { - return HloReplication(State::kUniqueOnAllDevices, {}); + return HloReplication(State::kUniqueOnAllDevices); } HloReplicationAnalysis::HloReplication @@ -724,8 +878,7 @@ HloReplicationAnalysis::HloReplication::PartiallyReplicated( } device_set_root_per_replica.push_back(std::move(device_set_root)); } - return HloReplication(State::kPartiallyReplicated, - device_set_root_per_replica); + return HloReplication(device_set_root_per_replica); } HloReplicationAnalysis::HloReplication @@ -781,9 +934,14 @@ HloReplicationAnalysis::HloReplication::Merge( bool HloReplicationAnalysis::HloReplication::Equal( const HloReplication& other) const { - return state_ == other.state_ && - device_set_root_per_replica_->hash_ == - other.device_set_root_per_replica_->hash_; + if (state_ != other.state_) { + return false; + } + if (state_ != State::kPartiallyReplicated) { + return true; + } + return device_set_root_per_replica_->hash_ == + other.device_set_root_per_replica_->hash_; } bool HloReplicationAnalysis::HloReplication::operator==( @@ -799,9 +957,16 @@ bool HloReplicationAnalysis::HloReplication::IsUniqueOnAllDevices() const { return state_ == State::kUniqueOnAllDevices; } +bool HloReplicationAnalysis::HloReplication::IsPartiallyReplicated() const { + return state_ == State::kPartiallyReplicated; +} + bool HloReplicationAnalysis::HloReplication::IsReplicatedWithinSubgroup( absl::Span device_ids) const { if (device_ids.empty()) return true; + if (state_ != State::kPartiallyReplicated) { + return state_ == State::kReplicatedOnAllDevices; + } for (const std::vector& device_set_roots : *device_set_root_per_replica_) { if (!absl::c_all_of(device_ids, diff --git a/third_party/xla/xla/hlo/analysis/hlo_replication_analysis.h b/third_party/xla/xla/hlo/analysis/hlo_replication_analysis.h index d050d44e19c9f5..3c4f8eb6844ae4 100644 --- a/third_party/xla/xla/hlo/analysis/hlo_replication_analysis.h +++ b/third_party/xla/xla/hlo/analysis/hlo_replication_analysis.h @@ -26,11 +26,15 @@ limitations under the License. #include "absl/container/flat_hash_map.h" #include "absl/container/flat_hash_set.h" +#include "absl/container/node_hash_map.h" #include "absl/status/status.h" #include "absl/status/statusor.h" #include "absl/types/span.h" +#include "xla/hlo/ir/hlo_computation.h" #include "xla/hlo/ir/hlo_instruction.h" #include "xla/hlo/ir/hlo_module.h" +#include "xla/shape_util.h" +#include "xla/tuple_tree.h" #include "xla/xla_data.pb.h" namespace xla { @@ -123,13 +127,17 @@ class HloReplicationAnalysis { bool operator==(const HloReplication& rhs) const; bool IsReplicatedOnAllDevices() const; bool IsUniqueOnAllDevices() const; + bool IsPartiallyReplicated() const; bool IsReplicatedWithinSubgroup(absl::Span device_ids) const; std::string ToString() const; template friend H AbslHashValue(H h, const HloReplication& r) { - return H::combine(std::move(h), r.state_, - *r.device_set_root_per_replica_); + h = H::combine(std::move(h), r.state_); + if (r.device_set_root_per_replica_ != nullptr) { + h = H::combine(std::move(h), *r.device_set_root_per_replica_); + } + return h; } private: @@ -138,8 +146,10 @@ class HloReplicationAnalysis { kUniqueOnAllDevices = 1, kPartiallyReplicated = 2, }; + // Only partially replicated values carry device sets; the other two states + // share a null pointer so that copying them is a plain state copy. + explicit HloReplication(State state); explicit HloReplication( - State state, absl::Span> device_set_root_per_replica); State state_; // Helper class that subclasses T, and computes the hash once on @@ -161,15 +171,17 @@ class HloReplicationAnalysis { return H::combine(std::move(h), r.hash_); } }; - // Empty if state_ is kReplicatedOnAllDevices or kUniqueOnAllDevices. - - // If cross_partition_spmd is true, groups_for_replicas_[k]'s size equals - // the number of partitions, and within replica k, groups_for_replicas_[k] - // maps each partition ID to the smallest partition ID in the set. + // Null if state_ is kReplicatedOnAllDevices or kUniqueOnAllDevices. // - // If cross_partition_spmd is false, groups_for_replicas_[k]'s size equals - // the number of replicas, and within partition k, groups_for_replicas_[k] - // maps each replica to the smallest replica ID in the set. + // If cross_partition_spmd is true, device_set_root_per_replica_[k]'s size + // equals the number of partitions, and within replica k, + // device_set_root_per_replica_[k] maps each partition ID to the smallest + // partition ID in the set. + // + // If cross_partition_spmd is false, device_set_root_per_replica_[k]'s size + // equals the number of replicas, and within partition k, + // device_set_root_per_replica_[k] maps each replica to the smallest replica + // ID in the set. std::shared_ptr>>> device_set_root_per_replica_; }; @@ -182,6 +194,13 @@ class HloReplicationAnalysis { HloReplication MergeReplications(const HloReplication& replication_a, const HloReplication& replication_b) { + // Merging with a value that is replicated or unique on all devices is a + // copy of one side; only merges of two partially replicated values are + // worth memoizing. + if (!replication_a.IsPartiallyReplicated() || + !replication_b.IsPartiallyReplicated()) { + return replication_a.Merge(replication_b); + } std::pair key = {replication_a, replication_b}; @@ -214,6 +233,47 @@ class HloReplicationAnalysis { bool ComputeHloReplicationOnComputation(const HloComputation* computation, bool mark_everything_not_replicated); + // Records that the replication of `inst` changed: its users, and the callers + // of its computation if it is the root, have to be evaluated again. + void OnReplicationChanged(const HloInstruction* inst); + + // Marks the users of `inst`, and the callers of its computation if it is the + // root, dirty, unless the visit in progress evaluates them anyway. + void MarkDependentsDirty(const HloInstruction* inst); + + // Adds `inst` to dirty_ if it has been evaluated before. The first dirty + // instruction of a computation also makes the callers of that computation + // dirty, since the pending evaluation runs when a caller visits the + // computation again. + void MarkDirty(const HloInstruction* inst); + + // Called after `inst` was evaluated: `inst` stays or becomes dirty if a + // computation it calls has dirty instructions, and leaves dirty_ otherwise. + void MarkEvaluated(const HloInstruction* inst); + + // Returns the replication of `inst`, which must have been computed already. + const TupleTree& GetReplication( + const HloInstruction* inst) const; + + // Merges `source` into `dest` element by element. Returns whether anything + // changed. + bool CombineReplication(const TupleTree& source, + TupleTree* dest); + + // Assigns `replication` to `dest` if it has none yet, or combines it with + // the existing one. Returns whether anything changed. + bool AssignOrCombineReplication(TupleTree replication, + const HloInstruction* dest); + + // Assigns or combines the replication of `source` to `dest`. Returns whether + // anything changed; nothing changes if `source` has no replication yet. + bool PropagateReplication(const HloInstruction* source, + const HloInstruction* dest); + + // Marks `inst` unique on all devices at every index. Returns whether + // anything changed. + bool MarkNotReplicated(const HloInstruction* inst); + // Builds the replica group dedup map that allows caching replication // calculations for all-reduce/all-gather that share the same replica groups. // This can significantly help in compile times when replica groups are very @@ -243,12 +303,48 @@ class HloReplicationAnalysis { // Capture the number of partitions / replicas for the module. const int64_t num_partitions_, replica_count_; - // A map from each analyzed HLO instruction to a shape tree that represents - // whether the instruction outputs the same value across replicas or - // partitions at each shape index. - absl::flat_hash_map> + // A map from each analyzed HLO instruction to a tree that represents whether + // the instruction outputs the same value across replicas or partitions at + // each shape index. The trees hold no Shape, so they stay valid when callers + // change shapes or delete instructions while they hold the analysis. + absl::flat_hash_map> hlo_replication_; + struct ComputationState { + // The module does not change during the analysis, so the post order is + // computed once. + std::vector post_order; + // Whether a visit of the computation has completed. + bool visited = false; + // Number of instructions of the computation in dirty_. + int64_t num_dirty = 0; + }; + // The values must not move when the map grows, since a visit of a nested + // computation adds an entry while the caller iterates over its post order. + absl::node_hash_map + computation_states_; + + // Instructions whose operands, or the roots of the computations they call, + // changed after their last evaluation, and instructions that call a + // computation with dirty instructions. Without partial replication, a repeat + // visit of a computation that is not marked evaluates only these: any other + // instruction would compute its old value again. + absl::flat_hash_set dirty_; + + // The computation of the innermost visit in progress, and whether that visit + // evaluates every instruction. + const HloComputation* visiting_computation_ = nullptr; + bool visiting_all_instructions_ = false; + + // Computations already visited with mark_everything_not_replicated set. + // That visit marks the parameters and every ordinary instruction unique, + // infeeds included; tuples, get-tuple-elements, optimization barriers and + // nested calls derive from those values, and replication only moves towards + // unique, so a repeat visit with the flag set would recompute the same values + // and is skipped. + absl::flat_hash_set + computations_marked_not_replicated_; + // Replications for all-reduce/all-gather that have the same replica groups is // usually identical. We use the following data structures to memoize the // replications for instructions with identical replica groups. diff --git a/third_party/xla/xla/hlo/analysis/hlo_replication_analysis_test.cc b/third_party/xla/xla/hlo/analysis/hlo_replication_analysis_test.cc index 6c601626141e72..e9d6575e3545b6 100644 --- a/third_party/xla/xla/hlo/analysis/hlo_replication_analysis_test.cc +++ b/third_party/xla/xla/hlo/analysis/hlo_replication_analysis_test.cc @@ -15,6 +15,7 @@ limitations under the License. #include "xla/hlo/analysis/hlo_replication_analysis.h" +#include #include #include #include @@ -1078,5 +1079,448 @@ ENTRY entry { FindInstruction(module.get(), "gte.1"), {})); } +// Loop carried values that are partially replicated with two different +// groupings, unique, and replicated, at their own tuple indices. The loop +// needs a second iteration to settle: the partially replicated all-reduce +// results only reach the parameter after the first pass over the body. +TEST_F(HloReplicationAnalysisTest, WhileLoopPartialReplicationPerIndex) { + const std::string module_str = R"hlo( +HloModule WhileLoopPartialReplicationPerIndex + +sum { + a = f32[] parameter(0) + b = f32[] parameter(1) + ROOT add = f32[] add(a, b) +} + +cond { + cond_param = (f32[8], f32[8], f32[8], u32[]) parameter(0) + i = u32[] get-tuple-element(cond_param), index=3 + limit = u32[] constant(4) + ROOT lt = pred[] compare(i, limit), direction=LT +} + +body { + body_param = (f32[8], f32[8], f32[8], u32[]) parameter(0) + x = f32[8] get-tuple-element(body_param), index=0 + y = f32[8] get-tuple-element(body_param), index=1 + u = f32[8] get-tuple-element(body_param), index=2 + i = u32[] get-tuple-element(body_param), index=3 + xu = f32[8] add(x, u) + yu = f32[8] add(y, u) + ar0 = f32[8] all-reduce(xu), to_apply=sum, replica_groups={{0,1},{2,3}} + ar1 = f32[8] all-reduce(yu), to_apply=sum, replica_groups={{0,2},{1,3}} + mixed = f32[8] add(ar0, ar1) + one = u32[] constant(1) + next_i = u32[] add(i, one) + ROOT tuple = (f32[8], f32[8], f32[8], u32[]) tuple(ar0, ar1, mixed, next_i) +} + +ENTRY entry { + p0 = f32[8] parameter(0), parameter_replication={true} + p1 = f32[8] parameter(1), parameter_replication={true} + p2 = f32[8] parameter(2), parameter_replication={false} + zero = u32[] constant(0) + init = (f32[8], f32[8], f32[8], u32[]) tuple(p0, p1, p2, zero) + ROOT while = (f32[8], f32[8], f32[8], u32[]) while(init), condition=cond, body=body +} +)hlo"; + const std::vector pairs01_23 = + CreateReplicaGroups({{0, 1}, {2, 3}}); + const std::vector pairs02_13 = + CreateReplicaGroups({{0, 2}, {1, 3}}); + + ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule( + module_str, /*replica_count=*/4)); + ASSERT_OK_AND_ASSIGN(std::unique_ptr analysis, + HloReplicationAnalysis::RunWithPartialReplication( + module.get(), /*cross_partition_spmd=*/false)); + + for (const char* name : {"while", "body_param", "tuple"}) { + const HloInstruction* inst = FindInstruction(module.get(), name); + EXPECT_FALSE(analysis->HloInstructionIsReplicatedAt(inst, {0})) << name; + EXPECT_TRUE(analysis->HloInstructionIsReplicatedAt(inst, {0}, pairs01_23)) + << name; + EXPECT_FALSE(analysis->HloInstructionIsReplicatedAt(inst, {0}, pairs02_13)) + << name; + EXPECT_FALSE(analysis->HloInstructionIsReplicatedAt(inst, {1})) << name; + EXPECT_FALSE(analysis->HloInstructionIsReplicatedAt(inst, {1}, pairs01_23)) + << name; + EXPECT_TRUE(analysis->HloInstructionIsReplicatedAt(inst, {1}, pairs02_13)) + << name; + EXPECT_FALSE(analysis->HloInstructionIsReplicatedAt(inst, {2})) << name; + EXPECT_FALSE(analysis->HloInstructionIsReplicatedAt(inst, {2}, pairs01_23)) + << name; + EXPECT_FALSE(analysis->HloInstructionIsReplicatedAt(inst, {2}, pairs02_13)) + << name; + EXPECT_TRUE(analysis->HloInstructionIsReplicatedAt(inst, {3})) << name; + EXPECT_TRUE(analysis->HloInstructionIsReplicatedAt(inst, {3}, pairs01_23)) + << name; + } + // u is unique, so the all-reduce operands are unique in every pass and each + // all-reduce result is replicated within its own groups only. + EXPECT_FALSE(analysis->HloInstructionIsReplicatedAt( + FindInstruction(module.get(), "xu"), {}, pairs01_23)); + EXPECT_FALSE(analysis->HloInstructionIsReplicatedAt( + FindInstruction(module.get(), "yu"), {}, pairs02_13)); + EXPECT_TRUE(analysis->HloInstructionIsReplicatedAt( + FindInstruction(module.get(), "ar0"), {}, pairs01_23)); + EXPECT_TRUE(analysis->HloInstructionIsReplicatedAt( + FindInstruction(module.get(), "ar1"), {}, pairs02_13)); + EXPECT_FALSE(analysis->HloInstructionIsReplicatedAt( + FindInstruction(module.get(), "mixed"), {}, pairs01_23)); + EXPECT_FALSE(analysis->HloInstructionIsReplicatedAt( + FindInstruction(module.get(), "mixed"), {}, pairs02_13)); + EXPECT_TRUE(analysis->HloInstructionIsReplicatedAt( + FindInstruction(module.get(), "lt"), {})); +} + +// The loop state enters through a tuple instruction, whose tree is replicated +// at the tuple root as well as at the leaves. The condition turns +// non replicated after the first pass, and the body parameter must then be +// unique at every index, the tuple root included. +TEST_F(HloReplicationAnalysisTest, + NonReplicatedConditionMarksBodyAtEveryIndex) { + const std::string module_str = R"hlo( +HloModule NonReplicatedConditionMarksBodyAtEveryIndex + +cond { + cond_param = (f32[8], u32[]) parameter(0) + i = u32[] get-tuple-element(cond_param), index=1 + limit = u32[] constant(5) + ROOT lt = pred[] compare(i, limit), direction=LT +} + +body { + body_param = (f32[8], u32[]) parameter(0) + x = f32[8] get-tuple-element(body_param), index=0 + i = u32[] get-tuple-element(body_param), index=1 + replica-id = u32[] replica-id() + next_i = u32[] add(i, replica-id) + ROOT tuple = (f32[8], u32[]) tuple(x, next_i) +} + +ENTRY entry { + p0 = f32[8] parameter(0), parameter_replication={true} + zero = u32[] constant(0) + init = (f32[8], u32[]) tuple(p0, zero) + ROOT while = (f32[8], u32[]) while(init), condition=cond, body=body +} +)hlo"; + ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule( + module_str, /*replica_count=*/2)); + ASSERT_OK_AND_ASSIGN(std::unique_ptr analysis, + HloReplicationAnalysis::Run( + module.get(), /*cross_partition_spmd=*/false)); + const HloInstruction* body_param = + FindInstruction(module.get(), "body_param"); + EXPECT_FALSE(analysis->HloInstructionIsReplicatedAt(body_param, {})); + EXPECT_FALSE(analysis->HloInstructionIsReplicatedAt(body_param, {0})); + EXPECT_FALSE(analysis->HloInstructionIsReplicatedAt(body_param, {1})); + EXPECT_FALSE(analysis->HloInstructionIsReplicatedAt( + FindInstruction(module.get(), "x"), {})); + // The body root tuple is unique at its own index too, and so are the loop + // result and the condition parameter that merge it. + EXPECT_FALSE(analysis->HloInstructionIsReplicatedAt( + FindInstruction(module.get(), "tuple"), {})); + EXPECT_FALSE(analysis->HloInstructionIsReplicatedAt( + FindInstruction(module.get(), "while"), {})); + EXPECT_FALSE(analysis->HloInstructionIsReplicatedAt( + FindInstruction(module.get(), "cond_param"), {})); + EXPECT_FALSE(analysis->HloInstructionIsReplicatedAt( + FindInstruction(module.get(), "cond_param"), {0})); + EXPECT_FALSE(analysis->HloInstructionIsReplicatedAt( + FindInstruction(module.get(), "while"), {0})); + EXPECT_FALSE(analysis->HloInstructionIsReplicatedAt( + FindInstruction(module.get(), "while"), {1})); + // The loop input is outside the body and keeps its replicated tuple root. + EXPECT_TRUE(analysis->HloInstructionIsReplicatedAt( + FindInstruction(module.get(), "init"), {})); +} + +// The trip count depends on the partition, so the partitions run the infeed in +// the body a different number of times. Its data is unique although its +// sharding is replicated, and so is the loop result that it feeds. +TEST_F(HloReplicationAnalysisTest, InfeedInDivergentLoopIsNotReplicatedSPMD) { + const std::string module_str = R"hlo( +HloModule InfeedInDivergentLoopIsNotReplicatedSPMD + +cond { + cond_param = (f32[8], u32[]) parameter(0) + i = u32[] get-tuple-element(cond_param), index=1 + pid = u32[] partition-id() + cmp_val = u32[] add(i, pid) + limit = u32[] constant(4) + ROOT lt = pred[] compare(cmp_val, limit), direction=LT +} + +body { + body_param = (f32[8], u32[]) parameter(0) + tok = token[] after-all() + infeed = (f32[8], token[]) infeed(tok), + sharding={{replicated}, {maximal device=0}} + infeed_data = f32[8] get-tuple-element(infeed), index=0 + i = u32[] get-tuple-element(body_param), index=1 + one = u32[] constant(1) + next_i = u32[] add(i, one) + ROOT tuple = (f32[8], u32[]) tuple(infeed_data, next_i) +} + +ENTRY entry { + p0 = f32[8] parameter(0), sharding={replicated} + zero = u32[] constant(0) + init = (f32[8], u32[]) tuple(p0, zero) + ROOT while = (f32[8], u32[]) while(init), condition=cond, body=body +} +)hlo"; + ASSERT_OK_AND_ASSIGN( + auto module, ParseAndReturnVerifiedModule(module_str, /*replica_count=*/1, + /*num_partitions=*/2)); + ASSERT_OK_AND_ASSIGN( + std::unique_ptr analysis, + HloReplicationAnalysis::Run(module.get(), /*cross_partition_spmd=*/true)); + EXPECT_FALSE(analysis->HloInstructionIsReplicatedAt( + FindInstruction(module.get(), "lt"), {})); + EXPECT_FALSE(analysis->HloInstructionIsReplicatedAt( + FindInstruction(module.get(), "infeed"), {0})); + EXPECT_FALSE(analysis->HloInstructionIsReplicatedAt( + FindInstruction(module.get(), "infeed_data"), {})); + for (const char* name : {"tuple", "while"}) { + EXPECT_FALSE(analysis->HloInstructionIsReplicatedAt( + FindInstruction(module.get(), name), {0})) + << name; + } + // The replicated parameter outside the loop stays replicated. + EXPECT_TRUE(analysis->HloInstructionIsReplicatedAt( + FindInstruction(module.get(), "init"), {0})); +} + +// Two loops in a called computation share their body. The second loop makes +// the shared body unique at index 0 after the first loop took its result, and +// nothing that feeds the call changes. The next pass of the outer loop still +// has to visit the call and the first loop again for their results to become +// unique. +TEST_F(HloReplicationAnalysisTest, SharedLoopBodyChangesAfterItsFirstCaller) { + const std::string module_str = R"hlo( +HloModule SharedLoopBodyChangesAfterItsFirstCaller + +cond { + cond_param = (f32[], u32[]) parameter(0) + i = u32[] get-tuple-element(cond_param), index=1 + limit = u32[] constant(4) + ROOT lt = pred[] compare(i, limit), direction=LT +} + +body { + body_param = (f32[], u32[]) parameter(0) + x = f32[] get-tuple-element(body_param), index=0 + i = u32[] get-tuple-element(body_param), index=1 + one = u32[] constant(1) + next_i = u32[] add(i, one) + ROOT tuple = (f32[], u32[]) tuple(x, next_i) +} + +callee { + a = f32[] parameter(0) + b = f32[] parameter(1) + zero = u32[] constant(0) + init1 = (f32[], u32[]) tuple(a, zero) + while1 = (f32[], u32[]) while(init1), condition=cond, body=body + init2 = (f32[], u32[]) tuple(b, zero) + while2 = (f32[], u32[]) while(init2), condition=cond, body=body + r1 = f32[] get-tuple-element(while1), index=0 + r2 = f32[] get-tuple-element(while2), index=0 + ROOT out = (f32[], f32[]) tuple(r1, r2) +} + +outer_cond { + outer_cond_param = (f32[], f32[], u32[]) parameter(0) + j = u32[] get-tuple-element(outer_cond_param), index=2 + limit = u32[] constant(3) + ROOT lt = pred[] compare(j, limit), direction=LT +} + +outer_body { + outer_param = (f32[], f32[], u32[]) parameter(0) + outer_a = f32[] get-tuple-element(outer_param), index=0 + outer_b = f32[] get-tuple-element(outer_param), index=1 + outer_j = u32[] get-tuple-element(outer_param), index=2 + call = (f32[], f32[]) call(outer_a, outer_b), to_apply=callee + call_r1 = f32[] get-tuple-element(call), index=0 + one = u32[] constant(1) + next_j = u32[] add(outer_j, one) + ROOT outer_tuple = (f32[], f32[], u32[]) tuple(call_r1, outer_b, next_j) +} + +ENTRY entry { + p0 = f32[] parameter(0), parameter_replication={true} + p1 = f32[] parameter(1), parameter_replication={false} + zero = u32[] constant(0) + init = (f32[], f32[], u32[]) tuple(p0, p1, zero) + ROOT outer = (f32[], f32[], u32[]) while(init), condition=outer_cond, + body=outer_body +} +)hlo"; + ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule( + module_str, /*replica_count=*/2)); + ASSERT_OK_AND_ASSIGN(std::unique_ptr analysis, + HloReplicationAnalysis::Run( + module.get(), /*cross_partition_spmd=*/false)); + for (const char* name : {"while1", "call", "outer"}) { + EXPECT_FALSE(analysis->HloInstructionIsReplicatedAt( + FindInstruction(module.get(), name), {0})) + << name; + } + EXPECT_FALSE(analysis->HloInstructionIsReplicatedAt( + FindInstruction(module.get(), "outer_a"), {})); + EXPECT_TRUE(analysis->HloInstructionIsReplicatedAt( + FindInstruction(module.get(), "while1"), {1})); + EXPECT_TRUE(analysis->HloInstructionIsReplicatedAt( + FindInstruction(module.get(), "outer_j"), {})); +} + +// Two loops in different called computations share a body that returns its +// parameter, and nothing else. The second loop makes the body unique at index +// 0 after the first loop took its result. Only the body root changes, so the +// first loop and its call have to be visited again when the outer loop +// repeats. +TEST_F(HloReplicationAnalysisTest, IdentityLoopBodySharedAcrossCalls) { + const std::string module_str = R"hlo( +HloModule IdentityLoopBodySharedAcrossCalls + +cond1 { + cond1_param = (f32[], u32[]) parameter(0) + i = u32[] get-tuple-element(cond1_param), index=1 + limit = u32[] constant(4) + ROOT lt = pred[] compare(i, limit), direction=LT +} + +cond2 { + cond2_param = (f32[], u32[]) parameter(0) + i = u32[] get-tuple-element(cond2_param), index=1 + limit = u32[] constant(4) + ROOT lt = pred[] compare(i, limit), direction=LT +} + +body { + ROOT body_param = (f32[], u32[]) parameter(0) +} + +callee1 { + a = f32[] parameter(0) + zero = u32[] constant(0) + init1 = (f32[], u32[]) tuple(a, zero) + while1 = (f32[], u32[]) while(init1), condition=cond1, body=body + ROOT r1 = f32[] get-tuple-element(while1), index=0 +} + +callee2 { + b = f32[] parameter(0) + zero = u32[] constant(0) + init2 = (f32[], u32[]) tuple(b, zero) + while2 = (f32[], u32[]) while(init2), condition=cond2, body=body + ROOT r2 = f32[] get-tuple-element(while2), index=0 +} + +outer_cond { + outer_cond_param = (f32[], f32[], u32[], f32[]) parameter(0) + j = u32[] get-tuple-element(outer_cond_param), index=2 + limit = u32[] constant(3) + ROOT lt = pred[] compare(j, limit), direction=LT +} + +outer_body { + outer_param = (f32[], f32[], u32[], f32[]) parameter(0) + outer_a = f32[] get-tuple-element(outer_param), index=0 + outer_b = f32[] get-tuple-element(outer_param), index=1 + outer_j = u32[] get-tuple-element(outer_param), index=2 + call1 = f32[] call(outer_a), to_apply=callee1 + one = u32[] constant(1) + next_j = u32[] add(outer_j, one) + call2 = f32[] call(outer_b), to_apply=callee2 + ROOT outer_tuple = (f32[], f32[], u32[], f32[]) tuple(call1, outer_b, next_j, + call2) +} + +ENTRY entry { + p0 = f32[] parameter(0), parameter_replication={true} + p1 = f32[] parameter(1), parameter_replication={false} + zero = u32[] constant(0) + init = (f32[], f32[], u32[], f32[]) tuple(p0, p1, zero, p1) + ROOT outer = (f32[], f32[], u32[], f32[]) while(init), condition=outer_cond, + body=outer_body +} +)hlo"; + ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule( + module_str, /*replica_count=*/2)); + ASSERT_OK_AND_ASSIGN(std::unique_ptr analysis, + HloReplicationAnalysis::Run( + module.get(), /*cross_partition_spmd=*/false)); + for (const char* name : {"while1", "outer"}) { + EXPECT_FALSE(analysis->HloInstructionIsReplicatedAt( + FindInstruction(module.get(), name), {0})) + << name; + } + for (const char* name : {"call1", "outer_a"}) { + EXPECT_FALSE(analysis->HloInstructionIsReplicatedAt( + FindInstruction(module.get(), name), {})) + << name; + } + EXPECT_TRUE(analysis->HloInstructionIsReplicatedAt( + FindInstruction(module.get(), "while1"), {1})); +} + +// An all-reduce over groups of one device yields a partially replicated value +// whose device sets are all singletons. Merging it with itself gives unique, so +// in a loop body it turns unique when the body is visited again: with partial +// replication every visit evaluates every instruction. +TEST_F(HloReplicationAnalysisTest, SingletonGroupsInLoopBodyBecomeUnique) { + const std::string module_str = R"hlo( +HloModule SingletonGroupsInLoopBodyBecomeUnique + +sum { + a = f32[] parameter(0) + b = f32[] parameter(1) + ROOT add = f32[] add(a, b) +} + +cond { + cond_param = (f32[], u32[]) parameter(0) + i = u32[] get-tuple-element(cond_param), index=1 + limit = u32[] constant(4) + ROOT lt = pred[] compare(i, limit), direction=LT +} + +body { + body_param = (f32[], u32[]) parameter(0) + x = f32[] get-tuple-element(body_param), index=0 + i = u32[] get-tuple-element(body_param), index=1 + all-reduce = f32[] all-reduce(x), replica_groups={{0},{1}}, to_apply=sum + one = u32[] constant(1) + next_i = u32[] add(i, one) + ROOT tuple = (f32[], u32[]) tuple(all-reduce, next_i) +} + +ENTRY entry { + p0 = f32[] parameter(0), parameter_replication={false} + zero = u32[] constant(0) + init = (f32[], u32[]) tuple(p0, zero) + ROOT while = (f32[], u32[]) while(init), condition=cond, body=body +} +)hlo"; + ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule( + module_str, /*replica_count=*/2)); + ASSERT_OK_AND_ASSIGN(std::unique_ptr analysis, + HloReplicationAnalysis::RunWithPartialReplication( + module.get(), /*cross_partition_spmd=*/false)); + const std::vector singletons = CreateReplicaGroups({{0}, {1}}); + EXPECT_FALSE(analysis->HloInstructionIsReplicatedAt( + FindInstruction(module.get(), "all-reduce"), {}, singletons)); + EXPECT_FALSE(analysis->HloInstructionIsReplicatedAt( + FindInstruction(module.get(), "tuple"), {0}, singletons)); + EXPECT_TRUE(analysis->HloInstructionIsReplicatedAt( + FindInstruction(module.get(), "tuple"), {1})); +} + } // namespace } // namespace xla diff --git a/third_party/xla/xla/hlo/ir/hlo_instructions.h b/third_party/xla/xla/hlo/ir/hlo_instructions.h index d1e46162ccd829..9b316fb58ed935 100644 --- a/third_party/xla/xla/hlo/ir/hlo_instructions.h +++ b/third_party/xla/xla/hlo/ir/hlo_instructions.h @@ -46,7 +46,9 @@ limitations under the License. #include "xla/layout.h" #include "xla/literal.h" #include "xla/literal_pool.h" +#include "xla/primitive_util.h" #include "xla/printer.h" +#include "xla/protobuf_util.h" #include "xla/service/hlo.pb.h" #include "xla/shape.h" #include "xla/shape_util.h" @@ -81,6 +83,10 @@ class HloDimensionsInstruction : public HloInstruction { } } + void HashAdditionalAttributes(absl::HashState h) const override { + absl::HashState::combine(std::move(h), dimensions()); + } + protected: HloDimensionsInstruction(HloOpcode opcode, const Shape& shape, absl::Span dimensions) @@ -462,6 +468,10 @@ class HloCompareInstruction : public HloInstruction { return hlo->opcode() == HloOpcode::kCompare; } + void HashAdditionalAttributes(absl::HashState h) const override { + absl::HashState::combine(std::move(h), direction(), order()); + } + private: void PrintExtraAttributesImpl(AttributePrinter& printer, const HloPrintOptions& options) const override; @@ -1081,6 +1091,10 @@ class HloCollectivePermuteInstruction : public HloChannelInstruction { hlo->opcode() == HloOpcode::kCollectivePermuteStart; } + void HashAdditionalAttributes(absl::HashState h) const override { + absl::HashState::combine(std::move(h), source_target_pairs()); + } + // Whether this is an in-place collective permute (with dynamic slice // operands). Derived from the presence of slice_sizes. bool inplace() const { return !slice_sizes_.empty(); } @@ -1517,6 +1531,11 @@ class HloSliceInstruction : public HloInstruction { return hlo->opcode() == HloOpcode::kSlice; } + void HashAdditionalAttributes(absl::HashState h) const override { + absl::HashState::combine(std::move(h), slice_starts(), slice_limits(), + slice_strides()); + } + private: void PrintExtraAttributesImpl(AttributePrinter& printer, const HloPrintOptions& options) const override; @@ -1579,11 +1598,19 @@ class HloConstantInstruction : public HloInstruction { return false; } - // Add literal to the hash state. + // Hashes the literal the way Identical compares it: by value, ignoring the + // literal's layout, over at most 64 bytes. Literal equality compares sub + // byte types under a bit mask and dynamic arrays only up to their dynamic + // sizes, so their bytes stay out. void HashAdditionalAttributes(absl::HashState h) const override { - if (HasLiteral()) { - absl::HashState::combine(std::move(h), - Literal::AbslHashable(literal())); + if (!HasLiteral()) { + return; + } + const Shape& shape = literal().shape(); + if (shape.IsArray() && shape.is_static() && + !primitive_util::IsSubByteNonPredType(shape.element_type())) { + LiteralBase::Hash(std::move(h), literal()); } } @@ -1976,6 +2003,10 @@ class HloGetTupleElementInstruction : public HloInstruction { return hlo->opcode() == HloOpcode::kGetTupleElement; } + void HashAdditionalAttributes(absl::HashState h) const override { + absl::HashState::combine(std::move(h), tuple_index()); + } + private: void PrintExtraAttributesImpl(AttributePrinter& printer, const HloPrintOptions& options) const override; @@ -2167,6 +2198,11 @@ class HloConvolutionInstruction : public HloInstruction { return hlo->opcode() == HloOpcode::kConvolution; } + void HashAdditionalAttributes(absl::HashState h) const override { + absl::HashState::combine( + std::move(h), protobuf_util::ProtobufHashBySerialization(window())); + } + private: void PrintExtraAttributesImpl(AttributePrinter& printer, const HloPrintOptions& options) const override; @@ -2270,6 +2306,11 @@ class HloReduceWindowInstruction : public HloInstruction { return hlo->opcode() == HloOpcode::kReduceWindow; } + void HashAdditionalAttributes(absl::HashState h) const override { + absl::HashState::combine( + std::move(h), protobuf_util::ProtobufHashBySerialization(window())); + } + private: void PrintExtraAttributesImpl(AttributePrinter& printer, const HloPrintOptions& options) const override; @@ -2461,6 +2502,10 @@ class HloCustomCallInstruction : public HloCallableInstruction { return hlo->opcode() == HloOpcode::kCustomCall; } + void HashAdditionalAttributes(absl::HashState h) const override { + absl::HashState::combine(std::move(h), custom_call_target()); + } + class PerInstructionStorage { // Abstract class for per-instruction storage. public: @@ -2550,6 +2595,12 @@ class HloPadInstruction : public HloInstruction { return hlo->opcode() == HloOpcode::kPad; } + void HashAdditionalAttributes(absl::HashState h) const override { + absl::HashState::combine( + std::move(h), + protobuf_util::ProtobufHashBySerialization(padding_config())); + } + private: void PrintExtraAttributesImpl(AttributePrinter& printer, const HloPrintOptions& options) const override; @@ -2792,6 +2843,10 @@ class HloIotaInstruction : public HloInstruction { return hlo->opcode() == HloOpcode::kIota; } + void HashAdditionalAttributes(absl::HashState h) const override { + absl::HashState::combine(std::move(h), iota_dimension()); + } + private: void PrintExtraAttributesImpl(AttributePrinter& printer, const HloPrintOptions& options) const override; @@ -2861,6 +2916,13 @@ class HloDotInstruction : public HloInstruction { return hlo->opcode() == HloOpcode::kDot; } + void HashAdditionalAttributes(absl::HashState h) const override { + absl::HashState::combine( + std::move(h), + protobuf_util::ProtobufHashBySerialization(dot_dimension_numbers()), + protobuf_util::ProtobufHashBySerialization(precision_config())); + } + private: void PrintExtraAttributesImpl(AttributePrinter& printer, const HloPrintOptions& options) const override; diff --git a/third_party/xla/xla/hlo/transforms/simplifiers/BUILD b/third_party/xla/xla/hlo/transforms/simplifiers/BUILD index e111be91e2c82f..d0c98dc07d5ba0 100644 --- a/third_party/xla/xla/hlo/transforms/simplifiers/BUILD +++ b/third_party/xla/xla/hlo/transforms/simplifiers/BUILD @@ -1414,10 +1414,12 @@ xla_cc_test( deps = [ ":sort_simplifier", "//xla/hlo/parser:hlo_parser", + "//xla/hlo/testlib:filecheck", "//xla/hlo/testlib:hlo_hardware_independent_test_base", "//xla/hlo/testlib:pattern_matcher_gmock", "//xla/hlo/testlib:test", "//xla/service:pattern_matcher", + "@com_google_absl//absl/status:status_matchers", "@com_google_googletest//:gtest_main", "@tsl//tsl/platform:statusor", ], diff --git a/third_party/xla/xla/hlo/transforms/simplifiers/sort_simplifier_test.cc b/third_party/xla/xla/hlo/transforms/simplifiers/sort_simplifier_test.cc index 436c6643f4d9b1..72d8948f43ba83 100644 --- a/third_party/xla/xla/hlo/transforms/simplifiers/sort_simplifier_test.cc +++ b/third_party/xla/xla/hlo/transforms/simplifiers/sort_simplifier_test.cc @@ -17,7 +17,9 @@ limitations under the License. #include +#include "absl/status/status_matchers.h" #include "xla/hlo/parser/hlo_parser.h" +#include "xla/hlo/testlib/filecheck.h" #include "xla/hlo/testlib/hlo_hardware_independent_test_base.h" #include "xla/hlo/testlib/pattern_matcher_gmock.h" #include "xla/hlo/testlib/test.h" @@ -27,6 +29,7 @@ limitations under the License. namespace xla { namespace { +using ::absl_testing::IsOkAndHolds; namespace m = match; using SortSimplifierTest = HloHardwareIndependentTestBase; @@ -50,14 +53,15 @@ TEST_F(SortSimplifierTest, RemoveUnusedSortOperandArrayResult) { dimensions={1}, to_apply=compare ROOT gte = f32[64,8732]{1,0} get-tuple-element(sort), index=0 })"; - TF_ASSERT_OK_AND_ASSIGN(auto module, - ParseAndReturnVerifiedModule(hlo_string)); + ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); SortSimplifier simplifier; uint64_t num_executions = 0; + bool changed = false; do { num_executions++; - } while (simplifier.Run(module.get()).value()); + ASSERT_OK_AND_ASSIGN(changed, simplifier.Run(module.get())); + } while (changed); EXPECT_EQ(num_executions, 2); auto root = module->entry_computation()->root_instruction(); EXPECT_THAT(root, GmockMatch(m::Sort(m::Parameter(0)))); @@ -88,11 +92,10 @@ TEST_F(SortSimplifierTest, RemoveUnusedSortOperandTuple) { gte.1 = u32[64,87] get-tuple-element(sort), index=2 ROOT tuple = (f32[64,87], u32[64,87]) tuple(gte.0, gte.1) })"; - TF_ASSERT_OK_AND_ASSIGN(auto module, - ParseAndReturnVerifiedModule(hlo_string)); + ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); SortSimplifier simplifier; - EXPECT_TRUE(simplifier.Run(module.get()).value()); + EXPECT_THAT(simplifier.Run(module.get()), IsOkAndHolds(true)); auto root = module->entry_computation()->root_instruction(); EXPECT_THAT( root, @@ -119,11 +122,10 @@ TEST_F(SortSimplifierTest, DontRemoveUnusedSortKey) { sort = (f32[64,8732]{1,0}, s32[64,8732]{1,0}) sort(keys, values), dimensions={1}, to_apply=compare ROOT gte = s32[64,8732]{1,0} get-tuple-element(sort), index=1 })"; - TF_ASSERT_OK_AND_ASSIGN(auto module, - ParseAndReturnVerifiedModule(hlo_string)); + ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); SortSimplifier simplifier; - EXPECT_FALSE(simplifier.Run(module.get()).value()); + EXPECT_THAT(simplifier.Run(module.get()), IsOkAndHolds(false)); } TEST_F(SortSimplifierTest, RemoveUnusedFirstOperand) { @@ -145,17 +147,83 @@ TEST_F(SortSimplifierTest, RemoveUnusedFirstOperand) { dimensions={1}, to_apply=compare ROOT gte = s32[64,8732]{1,0} get-tuple-element(sort), index=1 })"; - TF_ASSERT_OK_AND_ASSIGN(auto module, - ParseAndReturnVerifiedModule(hlo_string)); + ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); SortSimplifier simplifier; uint64_t num_executions = 0; + bool changed = false; do { num_executions++; - } while (simplifier.Run(module.get()).value()); + ASSERT_OK_AND_ASSIGN(changed, simplifier.Run(module.get())); + } while (changed); EXPECT_EQ(num_executions, 2); auto root = module->entry_computation()->root_instruction(); EXPECT_THAT(root, GmockMatch(m::Sort(m::Parameter(1)))); } + +TEST_F(SortSimplifierTest, DontRemoveUnusedSortOperandWhenSortIsRoot) { + const char* hlo_string = R"( + HloModule sort_root + + compare { + p.0.lhs = f32[] parameter(0) + p.0.rhs = f32[] parameter(1) + p.1.lhs = s32[] parameter(2) + p.1.rhs = s32[] parameter(3) + ROOT lt = pred[] compare(p.0.lhs, p.0.rhs), direction=LT + } + + ENTRY sort_computation { + keys = f32[64] parameter(0) + values = s32[64] parameter(1) + ROOT sort = (f32[64], s32[64]) sort(keys, values), + dimensions={0}, to_apply=compare + })"; + ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); + + SortSimplifier simplifier; + EXPECT_THAT(simplifier.Run(module.get()), IsOkAndHolds(false)); +} + +TEST_F(SortSimplifierTest, DoesNotRemoveOldComparatorWhenShared) { + const char* hlo_string = R"( + HloModule shared_comparator + + // CHECK-DAG: %compare ({{.*}}p.0.lhs: f32[], {{.*}}p.0.rhs: f32[], {{.*}}p.1.lhs: s32[], {{.*}}p.1.rhs: s32[]) -> pred[] + // CHECK-DAG: %compare.clone ({{.*}}) -> pred[] + compare { + p.0.lhs = f32[] parameter(0) + p.0.rhs = f32[] parameter(1) + p.1.lhs = s32[] parameter(2) + p.1.rhs = s32[] parameter(3) + ROOT lt = pred[] compare(p.0.lhs, p.0.rhs), direction=LT + } + + // CHECK: ENTRY %sort_computation + ENTRY sort_computation { + keys1 = f32[64] parameter(0) + values1 = s32[64] parameter(1) + // CHECK: %[[NEW_SORT:.*]] = f32[64]{{.*}}sort(%keys1), dimensions={0}, to_apply=%compare.clone + sort1 = (f32[64], s32[64]) sort(keys1, values1), + dimensions={0}, to_apply=compare + gte1 = f32[64] get-tuple-element(sort1), index=0 + + keys2 = f32[64] parameter(2) + values2 = s32[64] parameter(3) + // CHECK: %sort2 = (f32[64]{{.*}}, s32[64]{{.*}}) sort(%keys2, %values2), dimensions={0}, to_apply=%compare + sort2 = (f32[64], s32[64]) sort(keys2, values2), + dimensions={0}, to_apply=compare + gte2.0 = f32[64] get-tuple-element(sort2), index=0 + gte2.1 = s32[64] get-tuple-element(sort2), index=1 + + // CHECK: ROOT %tuple = (f32[64]{{.*}}, f32[64]{{.*}}, s32[64]{{.*}}) tuple(%[[NEW_SORT]], %sort2#0, %sort2#1) + ROOT tuple = (f32[64], f32[64], s32[64]) tuple(gte1, gte2.0, gte2.1) + })"; + ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); + + SortSimplifier simplifier; + EXPECT_THAT(simplifier.Run(module.get()), IsOkAndHolds(true)); + EXPECT_THAT(RunFileCheck(module->ToString(), hlo_string), IsOkAndHolds(true)); +} } // namespace } // namespace xla diff --git a/third_party/xla/xla/index_util.cc b/third_party/xla/xla/index_util.cc index 7377db56f9d814..d2a018ad003a1f 100644 --- a/third_party/xla/xla/index_util.cc +++ b/third_party/xla/xla/index_util.cc @@ -17,6 +17,7 @@ limitations under the License. #include +#include "absl/log/check.h" #include "absl/types/span.h" #include "xla/layout_util.h" #include "xla/shape.h" @@ -26,6 +27,12 @@ limitations under the License. namespace xla { +/* static */ int64_t IndexUtil::MultidimensionalIndexToLinearIndex( + const Shape& shape, absl::Span multi_index) { + return MultidimensionalIndexToLinearIndex( + shape, LayoutUtil::MinorToMajor(shape), multi_index); +} + /* static */ DimensionVector IndexUtil::LinearIndexToMultidimensionalIndex( const Shape& shape, int64_t linear_index) { DCHECK_GE(linear_index, 0); diff --git a/third_party/xla/xla/index_util.h b/third_party/xla/xla/index_util.h index ca8ce5e924d1da..9cf51389414087 100644 --- a/third_party/xla/xla/index_util.h +++ b/third_party/xla/xla/index_util.h @@ -42,10 +42,7 @@ class IndexUtil { // on the shape and its layout. The first index in the multi_index is // dimension 0. static int64_t MultidimensionalIndexToLinearIndex( - const Shape& shape, absl::Span multi_index) { - return MultidimensionalIndexToLinearIndex( - shape, LayoutUtil::MinorToMajor(shape), multi_index); - } + const Shape& shape, absl::Span multi_index); // Converts a multidimensional index (eg {x, y, z}) into a linear index based // on the shape and its layout. The first index in the multi_index is @@ -94,9 +91,10 @@ class IndexUtil { // D{L(1)} * ... // Scale factor holding the growing product of D{L(i)} terms. + absl::Span dims = shape.dimensions(); for (size_t i = 0; i < multi_index.size(); ++i) { DCHECK_GE(multi_index[i], 0); - DCHECK_LT(multi_index[i], shape.dimensions(i)) + DCHECK_LT(multi_index[i], dims[i]) << "indexing beyond extent in dimension " << i << ":" << "\n\tindex: " << absl::StrJoin(multi_index, ",") << "\n\tshape: " << ShapeUtil::HumanString(shape); @@ -107,7 +105,7 @@ class IndexUtil { int64_t linear_index = multi_index[minor_to_major[0]]; int64_t scale = 1; for (int i = 1; i < minor_to_major.size(); ++i) { - scale *= shape.dimensions(minor_to_major[i - 1]); + scale *= dims[minor_to_major[i - 1]]; linear_index += scale * multi_index[minor_to_major[i]]; } return linear_index; diff --git a/third_party/xla/xla/inlined_bit_set.h b/third_party/xla/xla/inlined_bit_set.h deleted file mode 100644 index 468c7c692708c5..00000000000000 --- a/third_party/xla/xla/inlined_bit_set.h +++ /dev/null @@ -1,95 +0,0 @@ -/* Copyright 2026 The OpenXLA Authors. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -==============================================================================*/ - -#ifndef XLA_INLINED_BIT_SET_H_ -#define XLA_INLINED_BIT_SET_H_ - -#include -#include -#include - -#include "absl/container/inlined_vector.h" -#include "absl/log/check.h" - -namespace xla { - -// A fixed-size bit set (or bitmap) that stores up to `InlinedWords * 64` bits -// inline without heap allocation, spilling to the heap only when `num_bits` -// exceeds this capacity. Useful for tracking sets of dense integer indices -// over a typically small range `[0, num_bits)`. -template -class InlinedBitSet { - public: - using Word = uint64_t; - static constexpr int kWordSize = sizeof(Word) * 8; - - InlinedBitSet() = default; - explicit InlinedBitSet(int num_bits) : words_(NumWords(num_bits), 0) {} - - void Set(int index) { - DCHECK_GE(index, 0); - const size_t word_idx = IndexToWord(index); - words_[word_idx] |= (Word{1} << (index % kWordSize)); - } - - // Sets all bits in [0, num_bits) to 1 - void SetAll(int num_bits) { - DCHECK_GE(num_bits, 0); - DCHECK_LE(static_cast(num_bits), words_.size() * kWordSize); - const size_t full_words = num_bits / kWordSize; - std::fill_n(words_.begin(), full_words, ~Word{0}); - const size_t remainder = num_bits % kWordSize; - if (remainder != 0) { - words_[full_words] |= (Word{1} << remainder) - 1; - } - } - - void Clear(int index) { - const size_t word_idx = IndexToWord(index); - words_[word_idx] &= ~(Word{1} << (index % kWordSize)); - } - - bool Test(int index) const { - const size_t word_idx = IndexToWord(index); - return (words_[word_idx] & (Word{1} << (index % kWordSize))) != 0; - } - - InlinedBitSet& operator|=(const InlinedBitSet& other) { - DCHECK_EQ(words_.size(), other.words_.size()); - for (size_t i = 0; i < words_.size(); ++i) { - words_[i] |= other.words_[i]; - } - return *this; - } - - private: - static size_t NumWords(int num_bits) { - DCHECK_GE(num_bits, 0); - return (num_bits + kWordSize - 1) / kWordSize; - } - - size_t IndexToWord(int index) const { - DCHECK_GE(index, 0); - const size_t word_idx = index / kWordSize; - DCHECK_LT(word_idx, words_.size()); - return word_idx; - } - - absl::InlinedVector words_; -}; - -} // namespace xla - -#endif // XLA_INLINED_BIT_SET_H_ diff --git a/third_party/xla/xla/inlined_bit_set_test.cc b/third_party/xla/xla/inlined_bit_set_test.cc deleted file mode 100644 index d80413b5320886..00000000000000 --- a/third_party/xla/xla/inlined_bit_set_test.cc +++ /dev/null @@ -1,196 +0,0 @@ -/* Copyright 2026 The OpenXLA Authors. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -==============================================================================*/ - -#include "xla/inlined_bit_set.h" - -#include - -#include "absl/container/flat_hash_set.h" -#include "xla/tsl/platform/test.h" - -namespace xla { -namespace { - -class BruteForceBitSet { - public: - void Set(int index) { bits_.insert(index); } - - void SetAll(int num_bits) { - for (int i = 0; i < num_bits; ++i) { - bits_.insert(i); - } - } - - void Clear(int index) { bits_.erase(index); } - - bool Test(int index) const { return bits_.contains(index); } - - BruteForceBitSet& operator|=(const BruteForceBitSet& other) { - bits_.insert(other.bits_.begin(), other.bits_.end()); - return *this; - } - - private: - absl::flat_hash_set bits_; -}; - -TEST(InlinedBitSetTest, DefaultInitializedAllFalse) { - InlinedBitSet<> bits(130); - - EXPECT_FALSE(bits.Test(0)); - EXPECT_FALSE(bits.Test(63)); - EXPECT_FALSE(bits.Test(64)); - EXPECT_FALSE(bits.Test(127)); - EXPECT_FALSE(bits.Test(128)); -} - -TEST(InlinedBitSetTest, Set) { - InlinedBitSet<> bits(130); - - bits.Set(0); - bits.Set(63); - bits.Set(64); - bits.Set(128); - - EXPECT_TRUE(bits.Test(0)); - EXPECT_TRUE(bits.Test(63)); - EXPECT_TRUE(bits.Test(64)); - EXPECT_FALSE(bits.Test(65)); - EXPECT_FALSE(bits.Test(127)); - EXPECT_TRUE(bits.Test(128)); -} - -TEST(InlinedBitSetTest, Clear) { - InlinedBitSet<> bits(130); - bits.Set(63); - bits.Set(64); - - bits.Clear(63); - - EXPECT_FALSE(bits.Test(63)); - EXPECT_TRUE(bits.Test(64)); -} - -TEST(InlinedBitSetTest, SetAllZeroBits) { - InlinedBitSet<> bits(0); - - bits.SetAll(0); -} - -TEST(InlinedBitSetTest, SetAllExactWordBoundary) { - InlinedBitSet<> bits(64); - - bits.SetAll(64); - - for (int i = 0; i < 64; ++i) { - EXPECT_TRUE(bits.Test(i)); - } -} - -TEST(InlinedBitSetTest, SetAllNonExactBoundary) { - InlinedBitSet<> bits(70); - - bits.SetAll(70); - - for (int i = 0; i < 70; ++i) { - EXPECT_TRUE(bits.Test(i)); - } -} - -TEST(InlinedBitSetTest, SetAllLessThanCapacity) { - InlinedBitSet<> bits(128); - - bits.SetAll(50); - - for (int i = 0; i < 50; ++i) { - EXPECT_TRUE(bits.Test(i)); - } - for (int i = 50; i < 128; ++i) { - EXPECT_FALSE(bits.Test(i)); - } -} - -TEST(InlinedBitSetTest, SetAllZeroOnNonEmpty) { - InlinedBitSet<> bits(64); - - bits.SetAll(0); - - for (int i = 0; i < 64; ++i) { - EXPECT_FALSE(bits.Test(i)); - } -} - -TEST(InlinedBitSetTest, BitwiseOr) { - InlinedBitSet<> a(130); - InlinedBitSet<> b(130); - a.Set(10); - a.Set(100); - b.Set(70); - b.Set(100); - - a |= b; - - EXPECT_TRUE(a.Test(10)); - EXPECT_TRUE(a.Test(70)); - EXPECT_TRUE(a.Test(100)); - EXPECT_FALSE(a.Test(0)); - EXPECT_FALSE(a.Test(129)); -} - -TEST(InlinedBitSetTest, MatchesBruteForceOnRandomOperations) { - constexpr int kNumBits = 200; - constexpr int kNumOps = 500; - std::mt19937 rng(42); - std::uniform_int_distribution op_dist(0, 3); - std::uniform_int_distribution bit_dist(0, kNumBits - 1); - std::uniform_int_distribution count_dist(0, kNumBits); - InlinedBitSet<> bits(kNumBits); - InlinedBitSet<> other_bits(kNumBits); - BruteForceBitSet bf_bits; - BruteForceBitSet bf_other_bits; - - for (int step = 0; step < kNumOps; ++step) { - const int idx = bit_dist(rng); - other_bits.Set(idx); - bf_other_bits.Set(idx); - switch (op_dist(rng)) { - case 0: - bits.Set(idx); - bf_bits.Set(idx); - break; - case 1: - bits.Clear(idx); - bf_bits.Clear(idx); - break; - case 2: { - const int count = count_dist(rng); - bits.SetAll(count); - bf_bits.SetAll(count); - break; - } - case 3: - bits |= other_bits; - bf_bits |= bf_other_bits; - break; - } - } - - for (int i = 0; i < kNumBits; ++i) { - EXPECT_EQ(bits.Test(i), bf_bits.Test(i)) << "Mismatch at index " << i; - } -} - -} // namespace -} // namespace xla diff --git a/third_party/xla/xla/literal.cc b/third_party/xla/xla/literal.cc index 160ad73cd42b76..bc09effad37ee4 100644 --- a/third_party/xla/xla/literal.cc +++ b/third_party/xla/xla/literal.cc @@ -2957,22 +2957,46 @@ LiteralProto LiteralBase::ToProto() const { return proto; } +const void* LiteralBase::untyped_data() const { + return root_piece().untyped_data(); +} + const void* LiteralBase::untyped_data(const ShapeIndex& shape_index) const { return piece(shape_index).untyped_data(); } +void* MutableLiteralBase::untyped_data() { + return mutable_root_piece().untyped_data(); +} + void* MutableLiteralBase::untyped_data(const ShapeIndex& shape_index) { return piece(shape_index).untyped_data(); } +int64_t LiteralBase::size_bytes() const { + return root_piece().size_bytes_dense(); +} + int64_t LiteralBase::size_bytes(const ShapeIndex& shape_index) const { return piece(shape_index).size_bytes_dense(); } +int64_t LiteralBase::total_size_bytes() const { + return root_piece().total_bytes_dense(); +} + int64_t LiteralBase::total_size_bytes(const ShapeIndex& shape_index) const { return piece(shape_index).total_bytes_dense(); } +int64_t LiteralBase::element_count() const { + return root_piece().element_count(); +} + +int64_t LiteralBase::element_count(const ShapeIndex& index) const { + return piece(index).element_count(); +} + std::string LiteralBase::GetR1U8AsString() const { CHECK(shape().IsArray()); CHECK_EQ(shape().dimensions().size(), 1); diff --git a/third_party/xla/xla/literal.h b/third_party/xla/xla/literal.h index 9a90f698d96d76..5c9d91bd2c608c 100644 --- a/third_party/xla/xla/literal.h +++ b/third_party/xla/xla/literal.h @@ -104,16 +104,21 @@ class LiteralBase { // ShapeIndex is not array. See primitive_util.h for the mapping from XLA type // to native type. template - absl::Span data(const ShapeIndex& shape_index = {}) const; + absl::Span data() const; + template + absl::Span data(const ShapeIndex& shape_index) const; // Returns a const pointer to (or size of) the underlying buffer holding the // array at the given shape index. CHECKs if the subshape of the literal at // the given ShapeIndex is not array. - const void* untyped_data(const ShapeIndex& shape_index = {}) const; - int64_t size_bytes(const ShapeIndex& shape_index = {}) const; + const void* untyped_data() const; + const void* untyped_data(const ShapeIndex& shape_index) const; + int64_t size_bytes() const; + int64_t size_bytes(const ShapeIndex& shape_index) const; // total size in bytes of the literal (including pre-allocated dynamic // metadata) - int64_t total_size_bytes(const ShapeIndex& shape_index = {}) const; + int64_t total_size_bytes() const; + int64_t total_size_bytes(const ShapeIndex& shape_index) const; // Computes the size in bytes of the output of the Serialize method. // If `pack_pred` is true, PRED elements are bit-packed (8 elements per byte); @@ -380,13 +385,8 @@ class LiteralBase { // Returns the count of the elements in the array at the given shape index in // this literal. - int64_t element_count(const ShapeIndex& index = {}) const { - if (index.empty()) { - // Common case, avoid GetSubshape(). - return ShapeUtil::ElementsIn(shape()); - } - return ShapeUtil::ElementsIn(ShapeUtil::GetSubshape(shape(), index)); - } + int64_t element_count() const; + int64_t element_count(const ShapeIndex& index) const; // This definition is here to ensure that nobody accidentally implements this // function which would lead to inconsistencies. Use Hash instead. @@ -1309,7 +1309,9 @@ class MutableLiteralBase : public LiteralBase { // given ShapeIndex is not array. See primitive_util.h for the mapping from // XLA type to native type. template - absl::Span data(const ShapeIndex& shape_index = {}); + absl::Span data(); + template + absl::Span data(const ShapeIndex& shape_index); // Unhide const method from parent class. using LiteralBase::data; @@ -1325,7 +1327,8 @@ class MutableLiteralBase : public LiteralBase { // Returns a pointer to the underlying buffer holding the array at the given // shape index. CHECKs if the subshape of the literal at the given ShapeIndex // is not array. - void* untyped_data(const ShapeIndex& shape_index = {}); + void* untyped_data(); + void* untyped_data(const ShapeIndex& shape_index); // Unhide const method from parent class. using LiteralBase::untyped_data; @@ -1944,12 +1947,22 @@ void LiteralBase::Piece::SetLinear(int64_t linear_index, NativeT value) { data().data()[linear_index] = value; } +template +absl::Span LiteralBase::data() const { + return root_piece().data(); +} + template absl::Span LiteralBase::data( const ShapeIndex& shape_index) const { return piece(shape_index).data(); } +template +absl::Span MutableLiteralBase::data() { + return mutable_root_piece().data(); +} + template absl::Span MutableLiteralBase::data(const ShapeIndex& shape_index) { return piece(shape_index).data(); @@ -2208,7 +2221,7 @@ absl::Status MutableLiteralBase::PopulateParallel(Generator&& generator) { TF_RET_CHECK(shape().IsArray()) << __func__ << " is only supported for dense arrays: " << shape(); return PopulateInternal(generator, - /*parallel=*/data().size() > 32); + /*parallel=*/element_count() > 32); } template @@ -2247,9 +2260,8 @@ template ( - std::forward(generator), - /*parallel=*/data().size() > 32); + return PopulateLinearInternal(std::forward(generator), + /*parallel=*/element_count() > 32); } template *> diff --git a/third_party/xla/xla/mlir_hlo/mhlo/IR/hlo_ops_common.td b/third_party/xla/xla/mlir_hlo/mhlo/IR/hlo_ops_common.td index d7ce2fdd56487e..28bba3d301c3d1 100644 --- a/third_party/xla/xla/mlir_hlo/mhlo/IR/hlo_ops_common.td +++ b/third_party/xla/xla/mlir_hlo/mhlo/IR/hlo_ops_common.td @@ -22,6 +22,7 @@ def MHLO_Dialect : Dialect { let useDefaultAttributePrinterParser = 0; let useDefaultTypePrinterParser = 0; + let useStrictPropertiesInAssemblyFormat = 0; } include "mhlo/IR/hlo_base.td" diff --git a/third_party/xla/xla/mosaic/dialect/tpu/tpu.td b/third_party/xla/xla/mosaic/dialect/tpu/tpu.td index 730141fdda7f5d..3833f81c0dfc70 100644 --- a/third_party/xla/xla/mosaic/dialect/tpu/tpu.td +++ b/third_party/xla/xla/mosaic/dialect/tpu/tpu.td @@ -24,6 +24,7 @@ def TPU_Dialect : Dialect { let cppNamespace = "::mlir::tpu"; let useDefaultAttributePrinterParser = 1; let useDefaultTypePrinterParser = 1; + let useStrictPropertiesInAssemblyFormat = 0; let extraClassDeclaration = [{ static constexpr StringRef GetMainKey() { return "tpu.main"; } static constexpr StringRef GetCoreTypeKey() { return "tpu.core_type"; } diff --git a/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_client_test.cc b/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_client_test.cc index efcd52db06227a..971c42c5673635 100644 --- a/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_client_test.cc +++ b/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_client_test.cc @@ -18,6 +18,7 @@ limitations under the License. #include #include #include +#include #include #include #include @@ -1638,6 +1639,42 @@ TEST(StreamExecutorGpuClientTest, CopyFromPinnedHostMemorySpace) { *literal)); } +namespace { + +class FailingHostMemoryAllocator : public HostMemoryAllocator { + public: + OwnedPtr Allocate(size_t size, const AllocateOptions& options) override { + return nullptr; + } +}; + +} // namespace + +TEST(StreamExecutorGpuClientTest, + ToLiteralReturnsResourceExhaustedWhenHostStagingPoolIsFull) { + GpuClientOptions options = GetTestGpuClientOptions(); + options.host_memory_allocator_factory = + [](HostMemoryAllocator::Options options) + -> absl::StatusOr> { + return std::make_unique(); + }; + + ASSERT_OK_AND_ASSIGN(auto client, GetStreamExecutorGpuClient(options)); + ASSERT_OK_AND_ASSIGN( + auto* device_memory_space, + client->addressable_devices()[0]->default_memory_space()); + + constexpr int64_t kTransferBytes = int64_t{1} << 20; + ASSERT_OK_AND_ASSIGN( + auto device_buffer, + client->CreateUninitializedBuffer( + ShapeUtil::MakeShape(U8, {kTransferBytes}), device_memory_space)); + + EXPECT_THAT(device_buffer->ToLiteral().Await(), + StatusIs(absl::StatusCode::kResourceExhausted, + HasSubstr("host staging buffer"))); +} + TEST(StreamExecutorGpuClientTest, CopyToPinnedHostMemorySpaceInt4) { TF_ASSERT_OK_AND_ASSIGN( auto client, GetStreamExecutorGpuClient(GetTestGpuClientOptions())); 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 aa7ae99cc0a108..036dbe82c319f0 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 @@ -424,10 +424,6 @@ absl::StatusOr PjRtStreamExecutorRawClient::AllocateRawBuffer( return absl::InternalError( "se::DeviceAddressAllocator is null in PjRtStreamExecutorRawClient."); } - ABSL_ASSIGN_OR_RETURN( - auto buffer, - allocator_->Allocate(local_device->local_device_id().value(), - on_device_bytes_count, true, layout_memory_space)); tsl::AsyncValueRef mem; if (has_custom_host_memory_allocator_ && layout_memory_space == Layout::kHostMemorySpace) { @@ -437,10 +433,22 @@ absl::StatusOr PjRtStreamExecutorRawClient::AllocateRawBuffer( }; auto buffer = GetHostMemoryAllocator()->Allocate(on_device_bytes_count, alloc_opts); + if (on_device_bytes_count > 0 && buffer == nullptr) { + return ResourceExhausted( + "Failed to allocate a %d-byte buffer from the custom host memory " + "allocator. The allocator may be exhausted or fragmented, or the " + "underlying pinned allocation failed; check earlier allocator " + "warnings for the root cause.", + on_device_bytes_count); + } se::DeviceAddressBase address(buffer.get(), on_device_bytes_count); mem = RawSEDeviceMemory::CreateForeign(address, [buffer = std::move(buffer)]() {}); } else { + ABSL_ASSIGN_OR_RETURN( + auto buffer, + allocator_->Allocate(local_device->local_device_id().value(), + on_device_bytes_count, true, layout_memory_space)); mem = RawSEDeviceMemory::Create(buffer.Release(), local_device, allocator_); } if (client_ != nullptr && local_device->allocation_model() != @@ -582,6 +590,15 @@ PjRtStreamExecutorClient::AllocateLinearizeDest( alloc_opts.local_device_id = local_device->local_device_id(); HostMemoryAllocator::OwnedPtr staging_buffer = GetHostMemoryAllocator()->Allocate(size, alloc_opts); + if (size > 0 && staging_buffer == nullptr) { + return ResourceExhausted( + "Failed to allocate a %d-byte pinned host staging buffer for a " + "host-to-device transfer. The pinned host pool may be exhausted or " + "fragmented (see XLA_PJRT_GPU_HOST_MEMORY_LIMIT_GB), or the " + "underlying pinned allocation failed; check earlier allocator " + "warnings for the root cause.", + size); + } absl::Span span(staging_buffer.get(), size); return PjRtStagingBuffer::Create( 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 3710029a95e619..eba194a6309659 100644 --- a/third_party/xla/xla/pjrt/se/se_raw_buffer.cc +++ b/third_party/xla/xla/pjrt/se/se_raw_buffer.cc @@ -147,6 +147,16 @@ PjRtStreamExecutorRawBuffer::CopyRawHostToDeviceAndReturnEvent( alloc_opts.local_device_id = local_device->local_device_id(); staging_buffer = client->GetHostMemoryAllocator()->Allocate( transfer_size, alloc_opts); + if (staging_buffer == nullptr) { + return ResourceExhausted( + "Failed to allocate a %d-byte pinned host staging buffer for " + "a host-to-device transfer. The pinned host pool may be " + "exhausted or fragmented (see " + "XLA_PJRT_GPU_HOST_MEMORY_LIMIT_GB), or the underlying pinned " + "allocation failed; check earlier allocator warnings for the " + "root cause.", + transfer_size); + } auto copy_to_staging_buffer = [src, transfer_size, staging_buffer]() mutable { tsl::profiler::TraceMe trace("H2D Copy To Staging Buffer"); @@ -218,6 +228,16 @@ PjRtStreamExecutorRawBuffer::CopyRawDeviceToHostAndReturnEvent( alloc_opts.local_device_id = local_device->local_device_id(); staging_buffer = client->GetHostMemoryAllocator()->Allocate( transfer_size, alloc_opts); + if (staging_buffer == nullptr) { + return ResourceExhausted( + "Failed to allocate a %d-byte pinned host staging buffer for " + "a device-to-host transfer. The pinned host pool may be " + "exhausted or fragmented (see " + "XLA_PJRT_GPU_HOST_MEMORY_LIMIT_GB), or the underlying pinned " + "allocation failed; check earlier allocator warnings for the " + "root cause.", + transfer_size); + } ABSL_RETURN_IF_ERROR( stream->Memcpy(staging_buffer.get(), sub_buffer, transfer_size)); auto copy_from_staging_buffer = [dst, transfer_size, @@ -367,6 +387,18 @@ void PjRtStreamExecutorRawBuffer::CopyTo( std::shared_ptr staging_buffer = client_->GetHostMemoryAllocator()->Allocate(GetOnDeviceSizeInBytes(), alloc_opts); + if (GetOnDeviceSizeInBytes() > 0 && staging_buffer == nullptr) { + absl::Status s = ResourceExhausted( + "Failed to allocate a %d-byte pinned host staging buffer for a " + "device-to-device transfer. The pinned host pool may be exhausted or " + "fragmented (see XLA_PJRT_GPU_HOST_MEMORY_LIMIT_GB), or the " + "underlying pinned allocation failed; check earlier allocator " + "warnings for the root cause.", + GetOnDeviceSizeInBytes()); + definition_event_promise.SetError(s); + src_usage_event_promise.SetError(s); + return; + } auto d2h_event = CopyRawDeviceToHostAndReturnEvent( staging_buffer.get(), 0, GetOnDeviceSizeInBytes(), {}); if (!d2h_event.ok()) { diff --git a/third_party/xla/xla/python/_profiler.pyi b/third_party/xla/xla/python/_profiler.pyi index 8b2e09aec6e49d..8f40c019618fec 100644 --- a/third_party/xla/xla/python/_profiler.pyi +++ b/third_party/xla/xla/python/_profiler.pyi @@ -33,7 +33,7 @@ class ProfilerSession: def stop_and_export(self, arg: str, /) -> None: ... def stop(self) -> bytes: ... - def stop_and_get_profile_data() -> ProfileData: ... + def stop_and_get_profile_data(self) -> ProfileData: ... def export(self, arg0: bytes, arg1: str, /) -> None: ... class ProfileOptions: @@ -97,6 +97,11 @@ class TraceMe: @staticmethod def is_enabled() -> bool: ... +def export_to_xprof( + profile_data: "tensorflow::profiler::python::ProfileData", + xprof_dir: str, + session_id: str = ..., +) -> None: ... def get_profiled_instructions_proto(tensorboard_dir: str) -> bytes: ... def get_instructions_profile( tensorboard_dir: str, diff --git a/third_party/xla/xla/python/ifrt/ir/ifrt_dialect.td b/third_party/xla/xla/python/ifrt/ir/ifrt_dialect.td index f5222425d8166d..d8b5b35ab39715 100644 --- a/third_party/xla/xla/python/ifrt/ir/ifrt_dialect.td +++ b/third_party/xla/xla/python/ifrt/ir/ifrt_dialect.td @@ -31,6 +31,7 @@ def Ifrt_Dialect : Dialect { let cppNamespace = "::xla::ifrt"; let useDefaultTypePrinterParser = 1; let useDefaultAttributePrinterParser = 1; + let useStrictPropertiesInAssemblyFormat = 0; let hasOperationAttrVerify = 1; let hasRegionArgAttrVerify = 1; } diff --git a/third_party/xla/xla/python/ifrt/ir/vifrt_dialect.td b/third_party/xla/xla/python/ifrt/ir/vifrt_dialect.td index 299a9f7fe1b8f3..b4a968745a1f75 100644 --- a/third_party/xla/xla/python/ifrt/ir/vifrt_dialect.td +++ b/third_party/xla/xla/python/ifrt/ir/vifrt_dialect.td @@ -42,6 +42,7 @@ def Vifrt_Dialect : Dialect { let useDefaultAttributePrinterParser = 0; let useDefaultTypePrinterParser = 0; + let useStrictPropertiesInAssemblyFormat = 0; } //===--------------------------------------------------------------------------- diff --git a/third_party/xla/xla/python/profiler.cc b/third_party/xla/xla/python/profiler.cc index 3bd55afee08a40..d4ba81a74e9943 100644 --- a/third_party/xla/xla/python/profiler.cc +++ b/third_party/xla/xla/python/profiler.cc @@ -256,13 +256,24 @@ NB_MODULE(_profiler, m) { "stop_and_get_profile_data", [](ProfilerSessionWrapper* sess) -> tensorflow::profiler::python::ProfileData { + // Serialize chunked XSpaces when continuous profiling is enabled to + // avoid hitting the 2GB single-protobuf limit, and pass session_id. + if (sess->session->IsContinuousProfilingEnabled()) { + xla::ThrowIfError(sess->session->Stop()); + auto xspaces = + std::make_shared>( + sess->session->SerializeChunks()); + return tensorflow::profiler::python::ProfileData( + std::move(xspaces), sess->session_id); + } auto xspace = std::make_shared(); // Disables the ProfilerSession xla::ThrowIfError(sess->session->CollectData(xspace.get())); - return tensorflow::profiler::python::ProfileData(xspace); + return tensorflow::profiler::python::ProfileData(std::move(xspace), + sess->session_id); }, nb::call_guard(), - nb::sig("def stop_and_get_profile_data() -> ProfileData")) + nb::sig("def stop_and_get_profile_data(self) -> ProfileData")) .def("export", [](ProfilerSessionWrapper* sess, nb::bytes xspace, const std::string& tensorboard_dir) { tensorflow::profiler::XSpace xspace_proto; @@ -366,6 +377,41 @@ NB_MODULE(_profiler, m) { .def("set_metadata", &TraceMeWrapper::SetMetadata) .def_static("is_enabled", &TraceMeWrapper::IsEnabled); + m.def( + "export_to_xprof", + [](const tensorflow::profiler::python::ProfileData& profile_data, + const std::string& xprof_dir, const std::string& session_id = "") { + // Fall back to profile_data.session_id() when session_id is empty, + // and export chunked XSpaces if continuous profiling was used. + const std::string& effective_session_id = + session_id.empty() ? profile_data.session_id() : session_id; + nb::gil_scoped_release release; + if (profile_data.raw_xspaces() != nullptr) { + if (effective_session_id.empty()) { + xla::ThrowIfError(tsl::profiler::ExportToTensorBoard( + xprof_dir, *profile_data.raw_xspaces())); + } else { + xla::ThrowIfError(tsl::profiler::ExportToTensorBoard( + xprof_dir, effective_session_id, *profile_data.raw_xspaces())); + } + return; + } + if (profile_data.raw_xspace() == nullptr) { + throw xla::XlaRuntimeError("ProfileData does not contain XSpace"); + } + if (effective_session_id.empty()) { + xla::ThrowIfError(tsl::profiler::ExportToTensorBoard( + *profile_data.raw_xspace(), xprof_dir, + /* also_export_trace_json= */ true)); + } else { + xla::ThrowIfError(tsl::profiler::ExportToTensorBoard( + *profile_data.raw_xspace(), xprof_dir, effective_session_id, + /* also_export_trace_json= */ true)); + } + }, + nb::arg("profile_data"), nb::arg("xprof_dir"), + nb::arg("session_id") = ""); + m.def( "get_profiled_instructions_proto", [](std::string tensorboard_dir) { diff --git a/third_party/xla/xla/python/profiler/profile_data_lib.cc b/third_party/xla/xla/python/profiler/profile_data_lib.cc index bbb0a12cef3b0f..08aca8b8eebe72 100644 --- a/third_party/xla/xla/python/profiler/profile_data_lib.cc +++ b/third_party/xla/xla/python/profiler/profile_data_lib.cc @@ -211,9 +211,19 @@ ProfileData::ProfileData(const void* serialized_xspace_ptr, } } -/*explicit*/ ProfileData::ProfileData(std::shared_ptr xspace_ptr) { - xspace_ = xspace_ptr; -} +// Store session_id on ProfileData and support chunked XSpaces while +// aliasing the first chunk so plane visitors remain valid without null derefs. +/*explicit*/ ProfileData::ProfileData(std::shared_ptr xspace_ptr, + std::string session_id) + : xspace_(std::move(xspace_ptr)), session_id_(std::move(session_id)) {} + +/*explicit*/ ProfileData::ProfileData( + std::shared_ptr> xspaces_ptr, std::string session_id) + : xspace_((xspaces_ptr && !xspaces_ptr->empty()) + ? std::shared_ptr(xspaces_ptr, &(*xspaces_ptr)[0]) + : std::make_shared()), + xspaces_(std::move(xspaces_ptr)), + session_id_(std::move(session_id)) {} /*explicit*/ ProfileData::ProfileData(const nb::bytes& serialized_xspace) { if (!xspace_) { diff --git a/third_party/xla/xla/python/profiler/profile_data_lib.h b/third_party/xla/xla/python/profiler/profile_data_lib.h index cdfeee940e30e2..9051d4b143c06c 100644 --- a/third_party/xla/xla/python/profiler/profile_data_lib.h +++ b/third_party/xla/xla/python/profiler/profile_data_lib.h @@ -167,7 +167,13 @@ class ProfileData { ProfileData(const void* serialized_xspace_ptr, size_t serialized_xspace_size); - explicit ProfileData(std::shared_ptr xspace_ptr); + // Supports chunked XSpace vectors (continuous profiling >2GB) and + // preserves session_id on ProfileData for standalone export_profile_data(). + explicit ProfileData(std::shared_ptr xspace_ptr, + std::string session_id = ""); + + explicit ProfileData(std::shared_ptr> xspaces_ptr, + std::string session_id = ""); explicit ProfileData(const nb::bytes& serialized_xspace); @@ -177,11 +183,21 @@ class ProfileData { ProfilePlane* find_plane_with_name(const std::string& name) const; + std::shared_ptr raw_xspace() const { return xspace_; } + + // Exposes chunked XSpaces and session_id so export_to_xprof can save >2GB + // continuous profiles and retain session_id when exporting ProfileData. + std::shared_ptr> raw_xspaces() const { return xspaces_; } + + const std::string& session_id() const { return session_id_; } + private: // The actual XSpace protobuf we are wrapping around. A shared ptr is used so // the different levels of visitors (ProfileData, ProfilePlane, // ProfileLine, etc.) don't depend on the lifetime of others. std::shared_ptr xspace_; + std::shared_ptr> xspaces_; + std::string session_id_; }; ProfileData from_serialized_xspace(const std::string& serialized_xspace); diff --git a/third_party/xla/xla/service/BUILD b/third_party/xla/xla/service/BUILD index 281accbe3b6b0e..81299291f0dcc5 100644 --- a/third_party/xla/xla/service/BUILD +++ b/third_party/xla/xla/service/BUILD @@ -2816,9 +2816,9 @@ cc_library( ":hlo_creation_utils", ":pattern_matcher", ":while_util", + "//xla:bit_set", "//xla:comparison_util", "//xla:frontend_attributes", - "//xla:inlined_bit_set", "//xla:literal_util", "//xla:shape_util", "//xla:status_macros", @@ -2863,6 +2863,7 @@ xla_cc_test( "@com_google_absl//absl/algorithm:container", "@com_google_absl//absl/status:status_matchers", "@com_google_absl//absl/strings", + "@com_google_absl//absl/strings:string_view", "@com_google_googletest//:gtest", ], ) @@ -2903,6 +2904,7 @@ xla_cc_test( "@com_google_absl//absl/algorithm:container", "@com_google_absl//absl/log", "@com_google_absl//absl/status", + "@com_google_absl//absl/status:status_matchers", "@com_google_absl//absl/strings:string_view", "@com_google_googletest//:gtest_main", ], @@ -3541,7 +3543,7 @@ cc_library( deps = [ ":computation_layout", ":hlo_graph_dumper", - ":logical_buffer", + ":hlo_value", ":while_util", "//xla:permutation_util", "//xla:shape_layout", @@ -3549,7 +3551,7 @@ cc_library( "//xla:status_macros", "//xla:util", "//xla:xla_data_proto_cc", - "//xla/hlo/analysis:tuple_points_to_analysis", + "//xla/hlo/analysis:hlo_dataflow_analysis", "//xla/hlo/ir:hlo", "//xla/hlo/pass:hlo_pass", "//xla/hlo/transforms/simplifiers:hlo_dce", @@ -3682,6 +3684,7 @@ xla_cc_test( "@com_google_absl//absl/algorithm:container", "@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:status_matchers", "@com_google_absl//absl/status:statusor", @@ -3871,9 +3874,9 @@ xla_cc_test( srcs = ["layout_assignment_test.cc"], deps = [ ":computation_layout", + ":hlo_value", ":hlo_verifier", ":layout_assignment", - ":logical_buffer", ":pattern_matcher", "//xla:literal", "//xla:literal_util", 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 c2b19168678478..7950039e9d08ad 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 @@ -22,6 +22,7 @@ limitations under the License. #include "absl/algorithm/container.h" #include "absl/log/log.h" #include "absl/status/status.h" +#include "absl/status/status_matchers.h" #include "absl/strings/string_view.h" #include "xla/hlo/ir/hlo_casting_utils.h" #include "xla/hlo/ir/hlo_computation.h" @@ -37,6 +38,8 @@ limitations under the License. namespace xla { namespace { +using ::absl_testing::StatusIs; + class AsyncCollectiveCustomCallRewriterTest : public HloHardwareIndependentTestBase {}; @@ -414,10 +417,10 @@ TEST_F(AsyncCollectiveCustomCallRewriterTest, 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"))); + EXPECT_THAT( + rewriter.Run(module.get()), + StatusIs(absl::StatusCode::kUnimplemented, + ::testing::HasSubstr("contains a negate instruction"))); } // Both modes require the all-gather to have exactly one operand. @@ -428,10 +431,9 @@ TEST_F(AsyncCollectiveCustomCallRewriterTest, 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"))); + EXPECT_THAT(rewriter.Run(module.get()), + StatusIs(absl::StatusCode::kUnimplemented, + ::testing::HasSubstr("has 2 operands"))); } // Both modes require the compute-on computation to contain exactly one @@ -443,9 +445,8 @@ TEST_F(AsyncCollectiveCustomCallRewriterTest, 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"))); + StatusIs(absl::StatusCode::kUnimplemented, + ::testing::HasSubstr("contains at least two"))); } } diff --git a/third_party/xla/xla/service/copy_insertion.cc b/third_party/xla/xla/service/copy_insertion.cc index bb85d255e64bc5..0332767a7e9c9d 100644 --- a/third_party/xla/xla/service/copy_insertion.cc +++ b/third_party/xla/xla/service/copy_insertion.cc @@ -632,19 +632,39 @@ bool IsDisjointInPlaceWhileTupleIndex(const HloInstruction* while_instr, return curr == single_reader || curr == param_gte; } -absl::Status AddCopiesForWhile(const HloAliasAnalysis& alias_analysis, +// Only such a while makes AddCopiesForWhile read HloBuffers. +// LINT.IfChange(disabled_copies_prescan) +bool HasWhileWithDisabledCopies(const HloModule& module) { + for (const HloComputation* computation : module.computations()) { + for (const HloInstruction* instruction : computation->instructions()) { + if (instruction->opcode() == HloOpcode::kWhile && + HasDisableWhileLoopCopiesAttr(instruction)) { + return true; + } + } + } + return false; +} +// LINT.ThenChange(:disabled_copies_buffer_reads) + +// `alias_analysis` is only consulted for a while annotated with +// xla_disable_while_loop_copies and may be null otherwise. +absl::Status AddCopiesForWhile(const HloDataflowAnalysis& dataflow, + const HloAliasAnalysis* alias_analysis, HloInstruction* xla_while) { VLOG(2) << "Adding copies for kWhile instruction " << xla_while->name(); TF_RET_CHECK(xla_while->opcode() == HloOpcode::kWhile); ShapeTree indices_to_copy(xla_while->shape()); - IndicesToCopyForWhile(alias_analysis.dataflow_analysis(), xla_while, - &indices_to_copy); + IndicesToCopyForWhile(dataflow, xla_while, &indices_to_copy); + // LINT.IfChange(disabled_copies_buffer_reads) if (HasDisableWhileLoopCopiesAttr(xla_while)) { VLOG(2) << "While loop copy insertion partially disabled via frontend " "attribute for " << xla_while->name(); + TF_RET_CHECK(alias_analysis != nullptr); + // LINT.ThenChange(:disabled_copies_prescan) // Clear copy requirement for scalars, async pipelined, and disjoint // in-place tuple indices. for (auto& pair : indices_to_copy) { @@ -670,7 +690,7 @@ absl::Status AddCopiesForWhile(const HloAliasAnalysis& alias_analysis, return; } std::vector buffers = - alias_analysis.ComputeBuffersAt(while_init, index); + alias_analysis->ComputeBuffersAt(while_init, index); for (const HloBuffer* buffer : buffers) { if (!seen_buffers.insert(buffer).second) { *init_indices_to_copy.mutable_element(index) = true; @@ -746,9 +766,8 @@ absl::Status AddCopiesForWhile(const HloAliasAnalysis& alias_analysis, // Add copies for the operands of in-place operations. RemoveUnnecessaryCopies // will remove the unnecessary copies. -absl::Status AddCopiesForInPlaceOperation( - const HloAliasAnalysis& alias_analysis, HloInstruction* in_place_op, - int64_t operand_number) { +absl::Status AddCopiesForInPlaceOperation(HloInstruction* in_place_op, + int64_t operand_number) { VLOG(2) << "Adding copies for in-place operation " << in_place_op->name(); HloInstruction* operand = in_place_op->mutable_operand(operand_number); ABSL_ASSIGN_OR_RETURN(HloInstruction * deep_copy, @@ -874,13 +893,12 @@ absl::Status StripControlDependenciesFrom(HloInstruction* instruction) { // roots, in order to resolve interference. We later rely on // RemoveUnnecessaryCopies to drop the unnecessary ones. absl::Status CopyInsertion::AddCopiesForConditional( - const HloAliasAnalysis& alias_analysis, HloInstruction* conditional) { + const HloDataflowAnalysis& dataflow, HloInstruction* conditional) { VLOG(2) << "Adding copies for kConditional instruction " << conditional->name(); ShapeTree indices_to_copy(conditional->shape()); TF_RET_CHECK(conditional->opcode() == HloOpcode::kConditional); - if (!IndicesToCopyForConditional(alias_analysis.dataflow_analysis(), - conditional, &indices_to_copy)) { + if (!IndicesToCopyForConditional(dataflow, conditional, &indices_to_copy)) { VLOG(2) << "No copies necessary for kConditional instruction " << conditional->name(); return absl::OkStatus(); @@ -1236,7 +1254,7 @@ absl::Status AddCopiesForNonCopyableTransitionsRotatedCase( // `chain_start` is the beginning of a non-copyable chain, we add copies to // transition in and out of the chain. absl::Status CopyInsertion::AddCopiesForExplicitNonCopyableTransitions( - const HloAliasAnalysis& alias_analysis, HloInstruction* chain_start) { + HloInstruction* chain_start) { VLOG(2) << "AddCopiesForExplicitNonCopyableTransitions: " << chain_start->ToString(); HloComputation* parent = chain_start->parent(); @@ -1399,12 +1417,11 @@ absl::Status AddCopiesForNonCopyableTransitionsRotatedCase( // (body root) (body root) // absl::Status CopyInsertion::AddCopiesForNonCopyableTransitions( - const HloAliasAnalysis& alias_analysis, HloInstruction* chain_start) { + HloInstruction* chain_start) { if (!IsImplicitNonCopyable(chain_start)) { if (chain_start->IsCustomCall(kPinCustomCallTarget) || chain_start->IsCustomCall(kCreateBufferCustomCallTarget)) { - ABSL_RETURN_IF_ERROR(AddCopiesForExplicitNonCopyableTransitions(alias_analysis, - chain_start)); + ABSL_RETURN_IF_ERROR(AddCopiesForExplicitNonCopyableTransitions(chain_start)); } return absl::OkStatus(); } @@ -1467,8 +1484,27 @@ absl::Status CopyInsertion::AddCopiesForNonCopyableTransitions( absl::Status CopyInsertion::AddCopiesToResolveInterference( HloModule* module, const absl::flat_hash_set& execution_threads) { - ABSL_ASSIGN_OR_RETURN(std::unique_ptr alias_analysis, - HloAliasAnalysis::Run(module, alias_info_)); + // Only the xla_disable_while_loop_copies branch of AddCopiesForWhile reads + // HloBuffers; everything else reads the dataflow analysis. + // HloAliasAnalysis::Run calls HloDataflowAnalysis::Run(*module, + // /*ssa_form=*/true, /*bitcast_defines_value=*/false) and passes alias_info_ + // only to CreateBuffers, so the same two booleans below yield the identical + // dataflow analysis. The analysis is built before the walk on purpose: copies + // inserted by earlier iterations would change the buffers a later build + // would see. + std::unique_ptr alias_analysis; + std::unique_ptr owned_dataflow; + const HloDataflowAnalysis* dataflow = nullptr; + if (HasWhileWithDisabledCopies(*module)) { + ABSL_ASSIGN_OR_RETURN(alias_analysis, + HloAliasAnalysis::Run(module, alias_info_)); + dataflow = &alias_analysis->dataflow_analysis(); + } else { + ABSL_ASSIGN_OR_RETURN(owned_dataflow, + HloDataflowAnalysis::Run(*module, /*ssa_form=*/true, + /*bitcast_defines_value=*/false)); + dataflow = owned_dataflow.get(); + } for (HloComputation* computation : module->MakeNonfusionComputations(execution_threads)) { if (computation->IsAsyncComputation()) { @@ -1477,14 +1513,14 @@ absl::Status CopyInsertion::AddCopiesToResolveInterference( for (HloInstruction* instruction : computation->MakeInstructionPostOrder()) { if (instruction->opcode() == HloOpcode::kWhile) { - ABSL_RETURN_IF_ERROR(AddCopiesForWhile(*alias_analysis, instruction)); + ABSL_RETURN_IF_ERROR( + AddCopiesForWhile(*dataflow, alias_analysis.get(), instruction)); } else if (instruction->opcode() == HloOpcode::kConditional) { - ABSL_RETURN_IF_ERROR(AddCopiesForConditional(*alias_analysis, instruction)); + ABSL_RETURN_IF_ERROR(AddCopiesForConditional(*dataflow, instruction)); } else if (IsNonCopyable(instruction)) { // We currently assume that we don't have a custom-call with // output-to-operand aliases for both buffers and non-buffers. - ABSL_RETURN_IF_ERROR( - AddCopiesForNonCopyableTransitions(*alias_analysis, instruction)); + ABSL_RETURN_IF_ERROR(AddCopiesForNonCopyableTransitions(instruction)); } else { // When an operand is a tuple, we avoid copying the operand multiple // times by recording and checking the operand number of operands that @@ -1583,7 +1619,7 @@ absl::Status CopyInsertion::AddCopiesToResolveInterference( } copied_operands.insert(operand_index_in_this_intr); ABSL_RETURN_IF_ERROR(AddCopiesForInPlaceOperation( - *alias_analysis, instruction, operand_index_in_this_intr)); + instruction, operand_index_in_this_intr)); } } } diff --git a/third_party/xla/xla/service/copy_insertion.h b/third_party/xla/xla/service/copy_insertion.h index 31f928a97b84db..16942891bfe432 100644 --- a/third_party/xla/xla/service/copy_insertion.h +++ b/third_party/xla/xla/service/copy_insertion.h @@ -26,6 +26,7 @@ limitations under the License. #include "absl/strings/string_view.h" #include "xla/hlo/analysis/alias_info.h" #include "xla/hlo/analysis/hlo_alias_analysis.h" +#include "xla/hlo/analysis/hlo_dataflow_analysis.h" #include "xla/hlo/ir/hlo_computation.h" #include "xla/hlo/ir/hlo_instruction.h" #include "xla/hlo/ir/hlo_module.h" @@ -136,15 +137,14 @@ class CopyInsertion : public HloModulePass { // Add copies for conditional instructions. virtual absl::Status AddCopiesForConditional( - const HloAliasAnalysis& alias_analysis, HloInstruction* conditional); + const HloDataflowAnalysis& dataflow, HloInstruction* conditional); // Adds copies for transitioning into and out of non-copyable values. - absl::Status AddCopiesForNonCopyableTransitions( - const HloAliasAnalysis& alias_analysis, HloInstruction* chain_start); + absl::Status AddCopiesForNonCopyableTransitions(HloInstruction* chain_start); // Adds copies for transitioning into and out of non-copyable values for a // explicit non-copyable chain. absl::Status AddCopiesForExplicitNonCopyableTransitions( - const HloAliasAnalysis& alias_analysis, HloInstruction* chain_start); + HloInstruction* chain_start); // Backend specific information about whether an instruction can share buffer // with its operand. @@ -160,6 +160,8 @@ class CopyInsertion : public HloModulePass { const absl::flat_hash_set& execution_threads) override; private: + friend class CopyInsertionTestPeer; + absl::Status AddCopiesToResolveInterference( HloModule* module, const absl::flat_hash_set& execution_threads); diff --git a/third_party/xla/xla/service/copy_insertion_test.cc b/third_party/xla/xla/service/copy_insertion_test.cc index ce2da447ae0a58..0c20d1bfce6910 100644 --- a/third_party/xla/xla/service/copy_insertion_test.cc +++ b/third_party/xla/xla/service/copy_insertion_test.cc @@ -28,6 +28,7 @@ limitations under the License. #include "absl/algorithm/container.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/status_matchers.h" #include "absl/status/statusor.h" @@ -67,6 +68,17 @@ limitations under the License. namespace op = xla::testing::opcode_matchers; namespace xla { + +// Reaches the private interference step so a test can run it on its own. +class CopyInsertionTestPeer { + public: + static absl::Status AddCopiesToResolveInterference( + CopyInsertion& copy_insertion, HloModule* module) { + return copy_insertion.AddCopiesToResolveInterference( + module, /*execution_threads=*/{}); + } +}; + namespace { using ::testing::NotNull; @@ -1167,6 +1179,125 @@ ENTRY %WhileEntry () -> (f32[16], f32[16]) { EXPECT_EQ(CountCopies(*module), 1); } +// %init_buf reaches the annotated loop twice: directly at index 2 and through +// the output of %first at index 1, which aliases the init of %first. The +// duplicate is only visible on the unmodified module: once %first has its init +// copied, %init_buf no longer shares a buffer with the output of %first. +constexpr absl::string_view kDuplicateInitBuffersThroughEarlierWhileHlo = R"( +HloModule DuplicateInitBuffersThroughEarlierWhile + +%FirstBody (loop_state: (s32[], f32[16])) -> (s32[], f32[16]) { + %loop_state = (s32[], f32[16]) parameter(0) + %indvar = s32[] get-tuple-element(%loop_state), index=0 + %c1 = s32[] constant(1) + %next_indvar = s32[] add(%indvar, %c1) + %data = f32[16] get-tuple-element(%loop_state), index=1 + %f1 = f32[] constant(1.0) + %c_add = f32[16] broadcast(%f1) + %next_data = f32[16] add(%data, %c_add) + ROOT %tuple = (s32[], f32[16]) tuple(%next_indvar, %next_data) +} + +%FirstCondition (loop_state: (s32[], f32[16])) -> pred[] { + %loop_state = (s32[], f32[16]) parameter(0) + %indvar = s32[] get-tuple-element(%loop_state), index=0 + %limit = s32[] constant(10) + ROOT %cmp = pred[] compare(%indvar, %limit), direction=LT +} + +%Body (loop_state: (s32[], f32[16], f32[16])) -> (s32[], f32[16], f32[16]) { + %loop_state = (s32[], f32[16], f32[16]) parameter(0) + %indvar = s32[] get-tuple-element(%loop_state), index=0 + %c1 = s32[] constant(1) + %next_indvar = s32[] add(%indvar, %c1) + %acc1 = f32[16] get-tuple-element(%loop_state), index=1 + %acc2 = f32[16] get-tuple-element(%loop_state), index=2 + %f1 = f32[] constant(1.0) + %update = f32[4] broadcast(%f1) + %c0 = s32[] constant(0) + %dus1 = f32[16] dynamic-update-slice(%acc1, %update, %c0), + frontend_attributes={xla_disjoint_read_write_regions="true"} + %dus2 = f32[16] dynamic-update-slice(%acc2, %update, %c0), + frontend_attributes={xla_disjoint_read_write_regions="true"} + ROOT %tuple = (s32[], f32[16], f32[16]) tuple(%next_indvar, %dus1, %dus2) +} + +%Condition (loop_state: (s32[], f32[16], f32[16])) -> pred[] { + %loop_state = (s32[], f32[16], f32[16]) parameter(0) + %indvar = s32[] get-tuple-element(%loop_state), index=0 + %limit = s32[] constant(10) + ROOT %cmp = pred[] compare(%indvar, %limit), direction=LT +} + +ENTRY %WhileEntry () -> (f32[16], f32[16], f32[16]) { + %c0 = s32[] constant(0) + %indvar_init = s32[] negate(%c0) + %zero = f32[] constant(0.0) + %init_buf = f32[16] broadcast(%zero) + %first_init = (s32[], f32[16]) tuple(%indvar_init, %init_buf) + %first = (s32[], f32[16]) while(%first_init), + condition=%FirstCondition, body=%FirstBody + %first_data = f32[16] get-tuple-element(%first), index=1 + %init_tuple = (s32[], f32[16], f32[16]) tuple(%indvar_init, %first_data, %init_buf) + %while = (s32[], f32[16], f32[16]) while(%init_tuple), + condition=%Condition, body=%Body, + frontend_attributes={xla_disable_while_loop_copies="true"} + %out1 = f32[16] get-tuple-element(%while), index=1 + %out2 = f32[16] get-tuple-element(%while), index=2 + ROOT %root = (f32[16], f32[16], f32[16]) tuple(%out1, %out2, %init_buf) +} +)"; + +TEST_F(WhileCopyInsertionTest, + DisableWhileLoopCopiesDuplicateInitBuffersThroughEarlierWhile) { + ASSERT_OK_AND_ASSIGN(auto module, + ParseAndReturnVerifiedModule( + kDuplicateInitBuffersThroughEarlierWhileHlo)); + InsertCopies(module.get()); + + EXPECT_EQ(CountCopies(*module->GetComputationWithName("Body")), 0); + // The duplicate at index 2 is copied; %init_buf is live out, so the in-place + // update of index 2 could not run on its buffer. + const HloInstruction* while_hlo = + module->entry_computation()->root_instruction()->operand(0)->operand(0); + EXPECT_THAT(while_hlo->operand(0), + op::Tuple(op::Negate(), op::GetTupleElement(op::While()), + op::Copy(op::Broadcast()))); +} + +// Same module through the interference step alone, so the copy cannot come +// from the alias analysis that later steps build. +TEST_F(WhileCopyInsertionTest, + DisableWhileLoopCopiesDuplicateInitBuffersInterferenceStepOnly) { + ASSERT_OK_AND_ASSIGN(auto module, + ParseAndReturnVerifiedModule( + kDuplicateInitBuffersThroughEarlierWhileHlo)); + CopyInsertion copy_insertion(&alias_info_); + ASSERT_OK(CopyInsertionTestPeer::AddCopiesToResolveInterference( + copy_insertion, module.get())); + + EXPECT_EQ(CountCopies(*module->GetComputationWithName("Body")), 0); + // The step deep copies the init, which leaves get-tuple-element/tuple + // rewraps for the tuple simplifier of the full pass; look through them. + auto through_rewraps = [](const HloInstruction* instr) { + while (instr->opcode() == HloOpcode::kGetTupleElement && + instr->operand(0)->opcode() == HloOpcode::kTuple) { + instr = instr->operand(0)->operand(instr->tuple_index()); + } + return instr; + }; + const HloInstruction* while_hlo = + module->entry_computation()->root_instruction()->operand(0)->operand(0); + const HloInstruction* init = while_hlo->operand(0); + ASSERT_EQ(init->opcode(), HloOpcode::kTuple); + EXPECT_THAT(through_rewraps(init->operand(0)), op::Negate()); + EXPECT_THAT(through_rewraps(init->operand(1)), + op::GetTupleElement(op::While())); + const HloInstruction* copy = through_rewraps(init->operand(2)); + ASSERT_THAT(copy, op::Copy()); + EXPECT_THAT(through_rewraps(copy->operand(0)), op::Broadcast()); +} + // Tests Copy Insertion when a while feeds another while // PARAMETER // | | diff --git a/third_party/xla/xla/service/gpu/BUILD b/third_party/xla/xla/service/gpu/BUILD index 73d9ad7dc320c8..f93d35c2da47d0 100644 --- a/third_party/xla/xla/service/gpu/BUILD +++ b/third_party/xla/xla/service/gpu/BUILD @@ -712,6 +712,7 @@ cc_library( deps = [ ":buffer_allocations", ":gpu_constants", + ":gpu_memory_space_assignment", "//xla:shape_util", "//xla:util", "//xla:xla_proto_cc", @@ -746,6 +747,35 @@ cc_library( ], ) +xla_cc_test( + name = "gpu_executable_buffer_allocator_test", + srcs = ["gpu_executable_buffer_allocator_test.cc"], + deps = [ + ":buffer_allocations", + ":gpu_constants", + ":gpu_executable_buffer_allocator", + ":gpu_memory_space_assignment", + "//xla:executable_run_options", + "//xla:shape_util", + "//xla:xla_data_proto_cc", + "//xla:xla_proto_cc", + "//xla/service:buffer_assignment", + "//xla/service:service_executable_run_options", + "//xla/stream_executor:device_address", + "//xla/stream_executor:device_address_allocator", + "//xla/stream_executor:mock_platform", + "//xla/stream_executor:mock_stream", + "//xla/stream_executor:platform", + "//xla/stream_executor:stream", + "@com_google_absl//absl/status", + "@com_google_absl//absl/status:status_macros", + "@com_google_absl//absl/status:status_matchers", + "@com_google_absl//absl/status:statusor", + "@com_google_absl//absl/types:span", + "@com_google_googletest//:gtest_main", + ], +) + xla_cc_test( name = "gpu_executable_va_remap_allocator_test", srcs = ["gpu_executable_va_remap_allocator_test.cc"], diff --git a/third_party/xla/xla/service/gpu/gpu_compiler.cc b/third_party/xla/xla/service/gpu/gpu_compiler.cc index b2c99e37d950ae..6a828612bf9bb7 100644 --- a/third_party/xla/xla/service/gpu/gpu_compiler.cc +++ b/third_party/xla/xla/service/gpu/gpu_compiler.cc @@ -2193,7 +2193,8 @@ absl::Status GpuCompiler::OptimizeHloPostLayoutAssignment( alias_info, mlir_context, /*only_fuse_if_profitable=*/true, /*use_experimental_tiling=*/ - debug_options.xla_gpu_experimental_enable_tiling_propagation()); + debug_options.xla_gpu_experimental_enable_tiling_propagation(), + thread_pool, &mlir_context_pool_); } pipeline.AddPass(); diff --git a/third_party/xla/xla/service/gpu/gpu_executable_buffer_allocator.cc b/third_party/xla/xla/service/gpu/gpu_executable_buffer_allocator.cc index bd90c461c433c2..a86b4d8cb936e5 100644 --- a/third_party/xla/xla/service/gpu/gpu_executable_buffer_allocator.cc +++ b/third_party/xla/xla/service/gpu/gpu_executable_buffer_allocator.cc @@ -38,6 +38,7 @@ limitations under the License. #include "xla/service/gpu/buffer_allocations.h" #include "xla/service/gpu/gpu_constants.h" #include "xla/service/gpu/gpu_executable_va_remap_allocator.h" +#include "xla/service/gpu/gpu_memory_space_assignment.h" #include "xla/service/service_executable_run_options.h" #include "xla/shape.h" #include "xla/shape_util.h" @@ -52,6 +53,21 @@ limitations under the License. namespace xla::gpu { namespace { +// True for allocations that GenerateBufferAllocations obtains from the device +// memory allocator: everything that is not a parameter, a constant, or a +// thread-local buffer. +bool IsTransientAllocation(const BufferAllocation& allocation) { + return !allocation.is_thread_local() && + !allocation.is_entry_computation_parameter() && + !allocation.is_constant(); +} + +// True for allocations in the collective memory space S(1). +bool IsCollectiveMemoryAllocation(const BufferAllocation& allocation) { + return allocation.color() == + static_cast(MemorySpaceColor::kCollective); +} + absl::Status CheckAlignment(const BufferAllocation& allocation, se::DeviceAddressBase buffer, int arg_idx) { const int64_t expected_alignment = [&] { @@ -230,15 +246,36 @@ GpuExecutableBufferAllocator::ExecutionScope::GenerateBufferAllocations( const int64_t num_buffers = owner_->allocations_.size(); ABSL_RETURN_IF_ERROR(Prepare(run_options, device_ordinal)); - std::vector buffers; - buffers.reserve(num_buffers); + // Resolve allocations in two passes so that every collective memory (S(1)) + // buffer of this execution is allocated before any default memory (S(0)) + // buffer. S(1) buffers back NCCL symmetric windows and can only live in the + // fixed, preallocated part of the shared BFC arena, while S(0) may later be + // allowed to grow past it. Allocating S(1) first guarantees that S(0) + // pressure within the same execution can never take the space S(1) needs. + // Both passes keep allocation-index order, so S(1) placement stays a + // deterministic function of the program, and `buffers` remains indexed by + // allocation index. + std::vector buffers(num_buffers); + auto resolve = [&](int64_t i) -> absl::Status { + const BufferAllocation& allocation = *owner_->allocations_[i]; + ABSL_ASSIGN_OR_RETURN(buffers[i], BufferForAllocation( + get_parameter_buffer, globals, allocation, + memory_allocator, device_ordinal, i)); + return CheckAlignment(allocation, buffers[i], i); + }; + + std::vector deferred_indices; for (int64_t i = 0; i < num_buffers; ++i) { const BufferAllocation& allocation = *owner_->allocations_[i]; - ABSL_ASSIGN_OR_RETURN( - buffers.emplace_back(), - BufferForAllocation(get_parameter_buffer, globals, allocation, - memory_allocator, device_ordinal, i)); - ABSL_RETURN_IF_ERROR(CheckAlignment(allocation, buffers.back(), i)); + if (IsTransientAllocation(allocation) && + !IsCollectiveMemoryAllocation(allocation)) { + deferred_indices.push_back(i); + continue; + } + ABSL_RETURN_IF_ERROR(resolve(i)); + } + for (int64_t i : deferred_indices) { + ABSL_RETURN_IF_ERROR(resolve(i)); } return BufferAllocations(buffers, device_ordinal, memory_allocator); } diff --git a/third_party/xla/xla/service/gpu/gpu_executable_buffer_allocator.h b/third_party/xla/xla/service/gpu/gpu_executable_buffer_allocator.h index e71e1533cc2ea7..a8a8b821e97172 100644 --- a/third_party/xla/xla/service/gpu/gpu_executable_buffer_allocator.h +++ b/third_party/xla/xla/service/gpu/gpu_executable_buffer_allocator.h @@ -94,7 +94,9 @@ class GpuExecutableBufferAllocator { // Builds the BufferAllocations for an execution. Entry-computation // parameter buffers are obtained from `get_parameter_buffer`; all other // allocations are resolved internally, including alignment checking and - // any subclass allocation-address policy. + // any subclass allocation-address policy. Collective memory (S(1)) + // buffers are allocated before all other transient buffers, in allocation + // index order within each group. absl::StatusOr GenerateBufferAllocations( const ServiceExecutableRunOptions* run_options, ParameterBufferResolver get_parameter_buffer, diff --git a/third_party/xla/xla/service/gpu/gpu_executable_buffer_allocator_test.cc b/third_party/xla/xla/service/gpu/gpu_executable_buffer_allocator_test.cc new file mode 100644 index 00000000000000..38f9f46121ad86 --- /dev/null +++ b/third_party/xla/xla/service/gpu/gpu_executable_buffer_allocator_test.cc @@ -0,0 +1,241 @@ +/* Copyright 2026 The OpenXLA Authors. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#include "xla/service/gpu/gpu_executable_buffer_allocator.h" + +#include +#include +#include +#include +#include + +#include +#include +#include "absl/status/status.h" +#include "absl/status/status_macros.h" +#include "absl/status/status_matchers.h" +#include "absl/status/statusor.h" +#include "absl/types/span.h" +#include "xla/executable_run_options.h" +#include "xla/service/buffer_assignment.h" +#include "xla/service/gpu/buffer_allocations.h" +#include "xla/service/gpu/gpu_constants.h" +#include "xla/service/gpu/gpu_memory_space_assignment.h" +#include "xla/service/service_executable_run_options.h" +#include "xla/shape_util.h" +#include "xla/stream_executor/device_address.h" +#include "xla/stream_executor/device_address_allocator.h" +#include "xla/stream_executor/mock_platform.h" +#include "xla/stream_executor/mock_stream.h" +#include "xla/stream_executor/platform.h" +#include "xla/stream_executor/stream.h" +#include "xla/xla.pb.h" +#include "xla/xla_data.pb.h" + +namespace xla::gpu { +namespace { + +using ::absl_testing::StatusIs; +using ::testing::ElementsAre; +using ::testing::NiceMock; +using ::testing::Pair; + +constexpr int64_t kDefault = static_cast(MemorySpaceColor::kDefault); +constexpr int64_t kCollective = + static_cast(MemorySpaceColor::kCollective); +constexpr int64_t kTempBuffer = + static_cast(MemorySpaceColor::kTempBuffer); + +// Host-backed buffers must satisfy the alignment check on every allocation +// kind. +constexpr uint64_t kAlignment = kXlaAllocatedBufferAlignBytes; + +// Serves every request from aligned host memory and records the memory space +// and size of each request in the order it was made. +class RecordingAllocator final : public se::DeviceAddressAllocator { + public: + RecordingAllocator(const se::Platform* platform, se::Stream* stream) + : se::DeviceAddressAllocator(platform), stream_(stream) {} + + absl::StatusOr> Allocate( + int device_ordinal, uint64_t size, bool retry_on_failure, + int64_t memory_space) override { + if (failing_memory_space_.has_value() && + *failing_memory_space_ == memory_space) { + return absl::ResourceExhaustedError("injected allocation failure"); + } + requests_.push_back({memory_space, size}); + auto storage = std::make_unique(size + kAlignment); + uint8_t* aligned = reinterpret_cast( + (reinterpret_cast(storage.get()) + kAlignment - 1) & + ~uintptr_t{kAlignment - 1}); + storage_.push_back(std::move(storage)); + return se::ScopedDeviceAddress( + se::DeviceAddressBase(aligned, size), device_ordinal, this); + } + + absl::Status Deallocate(int device_ordinal, + se::DeviceAddressBase mem) override { + if (!mem.is_null()) { + ++num_deallocations_; + } + return absl::OkStatus(); + } + + absl::StatusOr GetStream(int device_ordinal) override { + return stream_; + } + + // Makes every subsequent request for `memory_space` fail. + void FailRequestsFor(int64_t memory_space) { + failing_memory_space_ = memory_space; + } + + // (memory_space, size) of every successful request, in request order. + const std::vector>& requests() const { + return requests_; + } + + int num_deallocations() const { return num_deallocations_; } + + private: + se::Stream* stream_; + std::optional failing_memory_space_; + std::vector> requests_; + std::vector> storage_; + int num_deallocations_ = 0; +}; + +class GpuExecutableBufferAllocatorTest : public ::testing::Test { + protected: + GpuExecutableBufferAllocatorTest() : allocator_(&platform_, &stream_) { + run_options_.set_stream(&stream_); + service_run_options_ = ServiceExecutableRunOptions(run_options_); + } + + // Builds the BufferAllocations for one execution of `allocations` with the + // base (ALWAYS_UPDATE) buffer allocator and the recording device allocator. + absl::StatusOr Generate( + absl::Span allocations) { + GpuExecutableBufferAllocator executable_allocator( + "test", allocations, ShapeUtil::MakeShape(F32, {64}), &debug_options_, + /*thunk_executor=*/nullptr); + auto get_parameter_buffer = [&](const BufferAllocation& allocation) + -> absl::StatusOr { + return GpuExecutableBufferAllocator::ParameterBuffer{ + se::DeviceAddressBase(parameter_storage_, sizeof(parameter_storage_)), + allocation.parameter_number()}; + }; + ABSL_ASSIGN_OR_RETURN( + std::unique_ptr scope, + executable_allocator.CreateExecutionScope(&service_run_options_, + &allocator_, + /*device_ordinal=*/0)); + return scope->GenerateBufferAllocations( + &service_run_options_, get_parameter_buffer, &globals_, &allocator_, + /*device_ordinal=*/0); + } + + NiceMock platform_; + NiceMock stream_; + RecordingAllocator allocator_; + ExecutableRunOptions run_options_; + ServiceExecutableRunOptions service_run_options_; + DebugOptions debug_options_; + GpuExecutableBufferAllocator::BufferAllocToDeviceMemoryMap globals_; + alignas(kAlignment) uint8_t parameter_storage_[1024]; + alignas(kAlignment) uint8_t constant_storage_[256]; +}; + +TEST_F(GpuExecutableBufferAllocatorTest, + AllocatesCollectiveMemoryBeforeOtherTransientBuffers) { + BufferAllocation parameter(/*index=*/0, /*size=*/1024, kDefault); + parameter.set_entry_computation_parameter( + /*parameter_number=*/0, /*param_shape_index=*/{}, + /*parameter_aliased_with_output=*/false); + BufferAllocation default_temp(/*index=*/1, /*size=*/1024, kDefault); + BufferAllocation collective_temp(/*index=*/2, /*size=*/2048, kCollective); + BufferAllocation constant(/*index=*/3, /*size=*/256, kDefault); + constant.set_constant(true); + BufferAllocation collective_output(/*index=*/4, /*size=*/512, kCollective); + collective_output.set_maybe_live_out(true); + BufferAllocation default_output(/*index=*/5, /*size=*/768, kDefault); + default_output.set_maybe_live_out(true); + BufferAllocation thread_local_buffer(/*index=*/6, /*size=*/128, kDefault); + thread_local_buffer.set_is_thread_local(true); + BufferAllocation separate_temp(/*index=*/7, /*size=*/64, kTempBuffer); + std::vector allocations = { + ¶meter, &default_temp, + &collective_temp, &constant, + &collective_output, &default_output, + &thread_local_buffer, &separate_temp}; + globals_[3] = se::DeviceAddressBase(constant_storage_, 256); + + ASSERT_OK_AND_ASSIGN(BufferAllocations buffers, Generate(allocations)); + + // Every collective request precedes every other request, and each group + // keeps allocation index order. + EXPECT_THAT(allocator_.requests(), + ElementsAre(Pair(kCollective, 2048), Pair(kCollective, 512), + Pair(kDefault, 1024), Pair(kDefault, 768), + Pair(kTempBuffer, 64))); + + // The result is still indexed by allocation index. + EXPECT_EQ(buffers.GetDeviceAddress(0).opaque(), + static_cast(parameter_storage_)); + EXPECT_EQ(buffers.GetDeviceAddress(1).size(), 1024); + EXPECT_EQ(buffers.GetDeviceAddress(2).size(), 2048); + EXPECT_EQ(buffers.GetDeviceAddress(3).opaque(), + static_cast(constant_storage_)); + EXPECT_EQ(buffers.GetDeviceAddress(4).size(), 512); + EXPECT_EQ(buffers.GetDeviceAddress(5).size(), 768); + EXPECT_TRUE(buffers.GetDeviceAddress(6).is_null()); + EXPECT_EQ(buffers.GetDeviceAddress(7).size(), 64); + + ASSERT_OK(buffers.TearDown(/*live_addresses=*/{}, allocations)); + EXPECT_EQ(allocator_.num_deallocations(), 5); +} + +TEST_F(GpuExecutableBufferAllocatorTest, + KeepsAllocationIndexOrderWithoutCollectiveMemory) { + BufferAllocation first(/*index=*/0, /*size=*/256, kDefault); + BufferAllocation second(/*index=*/1, /*size=*/512, kDefault); + BufferAllocation third(/*index=*/2, /*size=*/1024, kDefault); + std::vector allocations = {&first, &second, &third}; + + ASSERT_OK_AND_ASSIGN(BufferAllocations buffers, Generate(allocations)); + + EXPECT_THAT(allocator_.requests(), + ElementsAre(Pair(kDefault, 256), Pair(kDefault, 512), + Pair(kDefault, 1024))); + ASSERT_OK(buffers.TearDown(/*live_addresses=*/{}, allocations)); +} + +TEST_F(GpuExecutableBufferAllocatorTest, + ReturnsDefaultMemoryFailureAfterCollectiveMemoryIsAllocated) { + allocator_.FailRequestsFor(kDefault); + BufferAllocation default_temp(/*index=*/0, /*size=*/1024, kDefault); + BufferAllocation collective_temp(/*index=*/1, /*size=*/2048, kCollective); + std::vector allocations = {&default_temp, + &collective_temp}; + + EXPECT_THAT(Generate(allocations).status(), + StatusIs(absl::StatusCode::kResourceExhausted)); + // The collective buffer was requested before the default one failed. + EXPECT_THAT(allocator_.requests(), ElementsAre(Pair(kCollective, 2048))); +} + +} // namespace +} // namespace xla::gpu diff --git a/third_party/xla/xla/service/gpu/model/BUILD b/third_party/xla/xla/service/gpu/model/BUILD index d221700bdd30fa..b7800f761d509a 100644 --- a/third_party/xla/xla/service/gpu/model/BUILD +++ b/third_party/xla/xla/service/gpu/model/BUILD @@ -447,6 +447,8 @@ cc_library( "//xla/service/gpu:gpu_fusible", "//xla/service/gpu:launch_dimensions", "//xla/stream_executor:device_description", + "//xla/tsl/concurrency:executor", + "//xla/tsl/concurrency:future", "@com_google_absl//absl/base:core_headers", "@com_google_absl//absl/log", "@com_google_absl//absl/log:check", @@ -580,11 +582,13 @@ cc_library( "//xla/codegen/tiling:tiling_specification", "//xla/codegen/tiling/experimental:tile", "//xla/codegen/tiling/experimental:tiled_hlo", + "//xla/codegen/tiling/experimental:tiling_space_utils", "//xla/codegen/xtile:block_level_parameters", "//xla/codegen/xtile:tiling_from_block_parameters", "//xla/codegen/xtile/codegen:emitter_helpers", "//xla/hlo/ir:hlo", "//xla/hlo/utils:hlo_traversal", + "//xla/runtime:object_pool", "//xla/service:decision", "//xla/service:hlo_cost_analysis", "//xla/service:instruction_fusion", @@ -592,6 +596,8 @@ cc_library( "//xla/service/gpu:hlo_fusion_analysis", "//xla/service/gpu:ir_emission_utils", "//xla/stream_executor:device_description", + "//xla/tsl/concurrency:executor", + "//xla/tsl/concurrency:future", "//xla/tsl/platform:env", "@com_google_absl//absl/algorithm:container", "@com_google_absl//absl/container:flat_hash_map", @@ -640,6 +646,7 @@ xla_cc_test( "//xla/service/gpu:ir_emission_utils", "//xla/stream_executor:device_description", "//xla/tests:xla_internal_test_main", + "//xla/tsl/concurrency:future", "//xla/tsl/platform:env", "//xla/tsl/testing:temporary_directory", "@com_google_absl//absl/container:flat_hash_map", diff --git a/third_party/xla/xla/service/gpu/model/combined_gpu_performance_model.cc b/third_party/xla/xla/service/gpu/model/combined_gpu_performance_model.cc index 8dc83ba48a520c..1eabdb08f877ee 100644 --- a/third_party/xla/xla/service/gpu/model/combined_gpu_performance_model.cc +++ b/third_party/xla/xla/service/gpu/model/combined_gpu_performance_model.cc @@ -39,6 +39,8 @@ limitations under the License. #include "xla/service/gpu/model/gpu_performance_model_base.h" #include "xla/service/hlo_cost_analysis.h" #include "xla/stream_executor/device_description.h" +#include "xla/tsl/concurrency/executor.h" +#include "xla/tsl/concurrency/future.h" #include "xla/xla_data.pb.h" namespace xla { @@ -49,12 +51,13 @@ CombinedGpuPerformanceModel::CombinedGpuPerformanceModel( HloFusionAnalysisCache& fusion_analysis_cache, mlir::MLIRContext& mlir_context, HloCostAnalysis::ShapeSizeFunction shape_size, bool use_experimental_tiling, - bool enable_same_shape_multi_output_fusion) + bool enable_same_shape_multi_output_fusion, + MlirContextPool* mlir_context_pool) : device_info_(device_info), fusion_analysis_cache_(fusion_analysis_cache), indexing_model_(&device_info, &fusion_analysis_cache, shape_size, &mlir_context, use_experimental_tiling, - enable_same_shape_multi_output_fusion), + enable_same_shape_multi_output_fusion, mlir_context_pool), model_(device_info, fusion_analysis_cache, cache_) {} absl::StatusOr @@ -262,10 +265,19 @@ void CombinedGpuPerformanceModel::Invalidate( cache_.Invalidate(instruction); } -absl::StatusOr -CombinedGpuPerformanceModel::TryFindBestTilingForFusion( - const HloFusionAdaptor& fusion_adaptor) { - return indexing_model_.TryFindBestTilingForFusion(fusion_adaptor); +tsl::Future +CombinedGpuPerformanceModel::TryFindBestTilingForFusionAsync( + const HloFusionAdaptor& fusion_adaptor, tsl::Executor* executor) { + return indexing_model_.TryFindBestTilingForFusionAsync(fusion_adaptor, + executor); +} + +tsl::Future +CombinedGpuPerformanceModel::TryFindTopKBestTilingsForFusionAsync( + const HloFusionAdaptor& fusion_adaptor, int top_k, + tsl::Executor* executor) { + return indexing_model_.TryFindTopKBestTilingsForFusionAsync(fusion_adaptor, + top_k, executor); } } // namespace gpu diff --git a/third_party/xla/xla/service/gpu/model/combined_gpu_performance_model.h b/third_party/xla/xla/service/gpu/model/combined_gpu_performance_model.h index 94970dfed3a524..80831d7a1b9df2 100644 --- a/third_party/xla/xla/service/gpu/model/combined_gpu_performance_model.h +++ b/third_party/xla/xla/service/gpu/model/combined_gpu_performance_model.h @@ -32,6 +32,8 @@ limitations under the License. #include "xla/service/gpu/model/gpu_performance_model_base.h" #include "xla/service/hlo_cost_analysis.h" #include "xla/stream_executor/device_description.h" +#include "xla/tsl/concurrency/executor.h" +#include "xla/tsl/concurrency/future.h" namespace xla { namespace gpu { @@ -51,7 +53,8 @@ class CombinedGpuPerformanceModel : public GpuPerformanceModelBase { ABSL_ATTRIBUTE_LIFETIME_BOUND, mlir::MLIRContext& mlir_context ABSL_ATTRIBUTE_LIFETIME_BOUND, HloCostAnalysis::ShapeSizeFunction shape_size, - bool use_experimental_tiling, bool enable_same_shape_multi_output_fusion); + bool use_experimental_tiling, bool enable_same_shape_multi_output_fusion, + MlirContextPool* mlir_context_pool = nullptr); // Returns runtime data analysis results of a single instruction. // @@ -81,8 +84,25 @@ class CombinedGpuPerformanceModel : public GpuPerformanceModelBase { // Returns the best tiling for a fusion. // // DOES NOT cache the result. - absl::StatusOr TryFindBestTilingForFusion( - const HloFusionAdaptor& fusion_adaptor); + // NOTE: `fusion_adaptor` and `this` are captured; callers must ensure both + // `fusion_adaptor` and this model instance outlive the returned + // `tsl::Future`. If executor is nullptr, defaults to + // tsl::InlineExecutor::Instance(). + [[nodiscard]] tsl::Future + TryFindBestTilingForFusionAsync(const HloFusionAdaptor& fusion_adaptor, + tsl::Executor* executor = nullptr); + + // Returns top_k best tilings for a fusion. + // + // DOES NOT cache the result. + // NOTE: `fusion_adaptor` and `this` are captured; callers must ensure both + // `fusion_adaptor` and this model instance outlive the returned + // `tsl::Future`. If executor is nullptr, defaults to + // tsl::InlineExecutor::Instance(). + [[nodiscard]] tsl::Future + TryFindTopKBestTilingsForFusionAsync(const HloFusionAdaptor& fusion_adaptor, + int top_k, + tsl::Executor* executor = nullptr); // Invalidates all cache entries related to given instruction. // diff --git a/third_party/xla/xla/service/gpu/model/gpu_indexing_performance_model.cc b/third_party/xla/xla/service/gpu/model/gpu_indexing_performance_model.cc index 636570437775f0..3438e5ca5045e3 100644 --- a/third_party/xla/xla/service/gpu/model/gpu_indexing_performance_model.cc +++ b/third_party/xla/xla/service/gpu/model/gpu_indexing_performance_model.cc @@ -44,9 +44,11 @@ limitations under the License. #include "llvm/ADT/STLExtras.h" #include "llvm/ADT/SmallVector.h" #include "llvm/Support/MathExtras.h" +#include "mlir/IR/MLIRContext.h" #include "xla/backends/gpu/codegen/triton/fusion.h" #include "xla/codegen/tiling/experimental/tiled_hlo.h" #include "xla/codegen/tiling/experimental/tiling_space.h" +#include "xla/codegen/tiling/experimental/tiling_space_utils.h" #include "xla/codegen/tiling/symbolic_tile_analysis.h" #include "xla/codegen/tiling/tiled_hlo_computation.h" #include "xla/codegen/tiling/tiling_specification.h" @@ -73,6 +75,8 @@ limitations under the License. #include "xla/shape.h" #include "xla/shape_util.h" #include "xla/stream_executor/device_description.h" +#include "xla/tsl/concurrency/executor.h" +#include "xla/tsl/concurrency/future.h" #include "xla/tsl/platform/env.h" #include "xla/util.h" #include "tsl/platform/path.h" @@ -642,6 +646,231 @@ absl::Status DumpAndLogTopKCandidates( return absl::OkStatus(); } +absl::StatusOr> +FilterAndSortCandidates( + std::vector> raw_candidates, int top_k, + const HloFusionAdaptor& fusion_adaptor) { + absl::InlinedVector candidates; + for (auto& res : raw_candidates) { + if (res.has_value()) { + candidates.push_back(std::move(*res)); + } + } + VLOG(1) << absl::StrCat("Found ", candidates.size(), + " valid tiling candidates."); + absl::c_stable_sort( + candidates, [](const TiledRunTimeData& a, const TiledRunTimeData& b) { + return a.runtime_data.exec_time < b.runtime_data.exec_time; + }); + if (candidates.size() > top_k) { + candidates.resize(top_k); + } + ABSL_RETURN_IF_ERROR(DumpAndLogTopKCandidates(fusion_adaptor, candidates)); + return candidates; +} + +// Evaluates a single tiling candidate: clones `base_tiling_space`, assigns +// `padded_tile_sizes` to the clone, tiles the fusion and estimates its run +// time. Returns std::nullopt if the candidate cannot be tiled or if it +// violates Triton constraints. +// +// If `mlir_context_pool` is not null, the clone is rebound to a context +// borrowed from the pool, so that concurrent evaluations do not contend on a +// single mlir::MLIRContext. +absl::StatusOr> EvaluateCandidate( + const HloFusionAdaptor& fusion_adaptor, + const experimental::TilingSpace& base_tiling_space, + absl::Span padded_tile_sizes, + MlirContextPool* mlir_context_pool, + const se::DeviceDescription& device_info, + HloCostAnalysis::ShapeSizeFunction shape_size, + absl::FunctionRef flops_fn) { + // Cloning is faster than calling TilingSpace::Create() for each tiling + // candidate. + std::optional borrowed_context; + mlir::MLIRContext* target_context = nullptr; + if (mlir_context_pool != nullptr) { + ABSL_ASSIGN_OR_RETURN(borrowed_context, mlir_context_pool->GetOrCreate()); + target_context = (*borrowed_context)->get(); + } + std::unique_ptr tiling_space = + base_tiling_space.Clone(target_context); + ABSL_RETURN_IF_ERROR(tiling_space->AssignTileSizes(padded_tile_sizes)); + VLOG(3) << "Trying tile sizes " << absl::StrJoin(padded_tile_sizes, ","); + + const absl::StatusOr tiled_computation = + experimental::TiledHloComputation::Tile(fusion_adaptor, + std::move(tiling_space)); + if (!tiled_computation.ok()) { + // TODO: b/511080616 - GetValidTilings() must return only tilings that can + // be tiled and we should treat all errors here as a failure. + VLOG(3) << "Tiling failed for " << absl::StrJoin(padded_tile_sizes, ",") + << " with error: " << tiled_computation.status().message(); + return std::nullopt; + } + if (const Decision valid = experimental::VerifyTritonConstraints( + *tiled_computation, device_info); + !valid) { + VLOG(3) << "Triton constraints violated for tiling " << valid.Explain(); + return std::nullopt; + } + + ABSL_ASSIGN_OR_RETURN( + std::optional tiled_run_time_data, + EstimateTiledRunTimeDataImpl(fusion_adaptor, *tiled_computation, + device_info, shape_size, flops_fn)); + + if (tiled_run_time_data.has_value()) { + VLOG(2) << "Accepted tile sizes [" + << absl::StrJoin(tiled_run_time_data->block_level_parameters + .output_tile_sizes.front(), + ",") + << "], exec_time " << tiled_run_time_data->runtime_data.exec_time; + } + return tiled_run_time_data; +} + +tsl::Future TryFindTopKBestTilingsWithTilingSpace( + const HloFusionAdaptor& fusion_adaptor, int top_k, tsl::Executor* executor, + std::function flops_fn, + mlir::MLIRContext* mlir_context, MlirContextPool* mlir_context_pool, + const se::DeviceDescription& device_info, + HloCostAnalysis::ShapeSizeFunction shape_size) { + CHECK(executor == &tsl::InlineExecutor::Instance() || + mlir_context_pool != nullptr || mlir_context->isMultithreadingEnabled()) + << "Concurrent tiling evaluation requires either an MlirContextPool or " + "an MLIRContext with multithreading enabled."; + + ABSL_ASSIGN_OR_RETURN( + std::unique_ptr base_tiling_space_unique, + experimental::TilingSpace::Create(fusion_adaptor, mlir_context)); + + ABSL_ASSIGN_OR_RETURN(auto tilings, base_tiling_space_unique->GetValidTilings()); + VLOG(1) << absl::StrCat( + "TryFindTopKBestTilingsForFusionAsync tiling_space evaluating ", + tilings.size(), " tilings."); + + std::shared_ptr base_tiling_space = + std::move(base_tiling_space_unique); + + std::vector>> candidate_futures; + candidate_futures.reserve(tilings.size()); + for (const llvm::SmallVector& tiling : tilings) { + // Assign padded tile size as Triton emitter will require that. + // Symbolic analysis route hides this assumption. + llvm::SmallVector padded_tile_sizes = + xla::xtile::GetPaddedTileSizes(tiling); + + // Early pruning invalid candidates as an optimization. + if (Decision decision = experimental::VerifySubsetOfTritonConstraints( + padded_tile_sizes, *base_tiling_space, device_info); + !decision) { + VLOG(3) << "Pre-filtering candidate violating Triton constraints: " + << absl::StrJoin(padded_tile_sizes, ",") << ": " + << decision.Explain(); + continue; + } + + if (!DoesTilingFitInRegisters(fusion_adaptor, *base_tiling_space, + padded_tile_sizes, device_info)) { + VLOG(3) << "Pre-filtering candidate exceeding register limit: " + << absl::StrJoin(padded_tile_sizes, ","); + continue; + } + + candidate_futures.push_back(tsl::MakeFutureOn( + *executor, + [base_space = base_tiling_space, &fusion_adaptor, + device_info = &device_info, shape_size, flops_fn, mlir_context_pool, + padded_tile_sizes = std::move(padded_tile_sizes)]() { + return EvaluateCandidate(fusion_adaptor, *base_space, + padded_tile_sizes, mlir_context_pool, + *device_info, shape_size, flops_fn); + })); + } + + if (candidate_futures.empty()) { + return TopKTiledRunTimeDataOrError( + absl::InlinedVector{}); + } + + return tsl::JoinFutures(absl::MakeSpan(candidate_futures)) + .Map(*executor, + [&fusion_adaptor, + top_k](std::vector> results) + -> absl::StatusOr { + ABSL_ASSIGN_OR_RETURN(auto candidates, + FilterAndSortCandidates(std::move(results), top_k, + fusion_adaptor)); + return TopKTiledRunTimeDataOrError(std::move(candidates)); + }); +} + +tsl::Future +TryFindTopKBestTilingsWithSymbolicAnalysis( + const HloFusionAdaptor& fusion_adaptor, int top_k, + std::function flops_fn, + mlir::MLIRContext* mlir_context, const se::DeviceDescription& device_info, + HloCostAnalysis::ShapeSizeFunction shape_size) { + SymbolicTileAnalysisOrError analysis_or_error = + SymbolicTileAnalysis::AnalyzeFusion( + fusion_adaptor, mlir_context, + TritonEmitterConstraints::GetBuilder(device_info)); + + if (const auto* fusion_decision = + std::get_if(&analysis_or_error)) { + VLOG(2) << "TryFindTopKBestTilingsForFusionAsync SymbolicTileAnalysis " + "rejected: " + << fusion_decision->Explain(); + return *fusion_decision; + } + + const auto& analysis = std::get(analysis_or_error); + + ABSL_ASSIGN_OR_RETURN(auto tilings, analysis.GetValidTilings()); + VLOG(1) << absl::StrCat( + "TryFindTopKBestTilingsForFusionAsync symbolic analysis evaluating ", + tilings.size(), " tilings."); + + std::vector> raw_candidates; + raw_candidates.reserve(tilings.size()); + for (const auto& tiling : tilings) { + // TODO(b/372454662): This needs to be adjusted if we want to support more + // than one "real root" (i.e. a root without users). + // Currently ComputeTiledComputation() may fail and return an + // Unimplemented error for cases of multi-output fusion that we do not + // support yet. + auto maybe_tiled_hlo_computation = analysis.ComputeTiledComputation(tiling); + if (!maybe_tiled_hlo_computation.ok()) { + if (maybe_tiled_hlo_computation.status().code() == + absl::StatusCode::kUnimplemented && + absl::StrContains(maybe_tiled_hlo_computation.status().message(), + "multi-output fusion")) { + continue; + } + return maybe_tiled_hlo_computation.status(); + } + + ABSL_ASSIGN_OR_RETURN(std::optional tiled_run_time_data, + EstimateTiledRunTimeDataImpl( + fusion_adaptor, *maybe_tiled_hlo_computation, + device_info, shape_size, flops_fn)); + + if (tiled_run_time_data.has_value()) { + ABSL_ASSIGN_OR_RETURN(FlatTiling flat_tiling, + tiling.Flatten(analysis.GetTilingSpecification())); + VLOG(2) << "Accepted tile sizes [" << absl::StrJoin(flat_tiling, ",") + << "], exec_time " << tiled_run_time_data->runtime_data.exec_time; + raw_candidates.push_back(std::move(tiled_run_time_data)); + } + } + + ABSL_ASSIGN_OR_RETURN(auto candidates, + FilterAndSortCandidates(std::move(raw_candidates), top_k, + fusion_adaptor)); + return TopKTiledRunTimeDataOrError(std::move(candidates)); +} + } // namespace int64_t GpuPerformanceModelWithIndexingAnalysis::FlopsPerElement( @@ -827,207 +1056,90 @@ absl::flat_hash_map PrecomputeFlopsMap( } // namespace internal -absl::StatusOr -GpuPerformanceModelWithIndexingAnalysis::TryFindTopKBestTilingsForFusion( - const HloFusionAdaptor& fusion_adaptor, int top_k) { +tsl::Future +GpuPerformanceModelWithIndexingAnalysis::TryFindTopKBestTilingsForFusionAsync( + const HloFusionAdaptor& fusion_adaptor, int top_k, + tsl::Executor* executor) { XLA_SCOPED_LOGGING_TIMER( "GpuPerformanceModelWithIndexingAnalysis::" - "TryFindTopKBestTilingsForFusion"); + "TryFindTopKBestTilingsForFusionAsync"); if (!fusion_adaptor.GetRoots().empty() && fusion_adaptor.GetRoots()[0].instruction().parent() != nullptr) { const HloInstruction& root = fusion_adaptor.GetRoots()[0].instruction(); - VLOG(1) << "TryFindTopKBestTilingsForFusion adaptor root " + VLOG(1) << "TryFindTopKBestTilingsForFusionAsync adaptor root " << root.ToString(HloPrintOptions::ShortParsable()) << " parent " << root.parent()->name(); } - absl::InlinedVector candidates; - - if (use_experimental_tiling_) { - using experimental::TiledHloComputation; - using experimental::TilingSpace; - - ABSL_ASSIGN_OR_RETURN(std::unique_ptr base_tiling_space, - TilingSpace::Create(fusion_adaptor, mlir_context_)); - ABSL_ASSIGN_OR_RETURN(auto tilings, base_tiling_space->GetValidTilings()); - VLOG(1) << absl::StrCat( - "TryFindTopKBestTilingsForFusion tiling_space evaluating ", - tilings.size(), " tilings."); - - const absl::flat_hash_map flops_map = - internal::PrecomputeFlopsMap( - fusion_adaptor, - [this](const HloInstruction* hlo) { return FlopsPerElement(hlo); }); + if (executor == nullptr) { + executor = &tsl::InlineExecutor::Instance(); + } - auto flops_fn = [&flops_map](const HloInstruction* hlo) -> int64_t { + if (use_experimental_tiling_) { + auto flops_map = std::make_shared< + const absl::flat_hash_map>( + internal::PrecomputeFlopsMap(fusion_adaptor, + [this](const HloInstruction* hlo) { + return FlopsPerElement(hlo); + })); + + auto flops_fn = [flops_map](const HloInstruction* hlo) -> int64_t { if (hlo->opcode() == HloOpcode::kParameter || hlo->opcode() == HloOpcode::kTuple || hlo->opcode() == HloOpcode::kGetTupleElement) { return 0; } - auto it = flops_map.find(hlo); - CHECK(it != flops_map.end()) + auto it = flops_map->find(hlo); + CHECK(it != flops_map->end()) << "Missing precomputed FLOPs for " << hlo->name(); return it->second; }; - for (const llvm::SmallVector& tiling : tilings) { - // Assign padded tile size as Triton emitter will require that. - // Symbolic analysis route hides this assumption. - llvm::SmallVector padded_tile_sizes = - xla::xtile::GetPaddedTileSizes(tiling); - - // Early pruning invalid candidates as an optimization. - if (Decision decision = experimental::VerifySubsetOfTritonConstraints( - padded_tile_sizes, *base_tiling_space, *device_info_); - !decision) { - VLOG(3) << "Pre-filtering candidate violating Triton constraints: " - << absl::StrJoin(padded_tile_sizes, ",") << ": " - << decision.Explain(); - continue; - } - - if (!DoesTilingFitInRegisters(fusion_adaptor, *base_tiling_space, - padded_tile_sizes, *device_info_)) { - VLOG(3) << "Pre-filtering candidate exceeding register limit: " - << absl::StrJoin(padded_tile_sizes, ","); - continue; - } - - // Cloning is faster than calling TilingSpace::Create() for each - // tiling candidate. - std::unique_ptr tiling_space = base_tiling_space->Clone(); - ABSL_RETURN_IF_ERROR(tiling_space->AssignTileSizes(padded_tile_sizes)); - VLOG(3) << "Trying tile sizes " << absl::StrJoin(padded_tile_sizes, ","); - - const absl::StatusOr tiled_computation = - TiledHloComputation::Tile(fusion_adaptor, std::move(tiling_space)); - if (!tiled_computation.ok()) { - // TODO: b/511080616 - GetValidTilings() must return only tilings that - // can be tiled and we should treat all errors here as a failure. - VLOG(3) << "Tiling failed for " << absl::StrJoin(tiling, ",") - << " with error: " << tiled_computation.status().message(); - continue; - } - if (const Decision valid = experimental::VerifyTritonConstraints( - *tiled_computation, *device_info_); - !valid) { - VLOG(3) << "Triton constraints violated for tiling " << valid.Explain(); - continue; - } - - ABSL_ASSIGN_OR_RETURN( - std::optional tiled_run_time_data, - EstimateTiledRunTimeDataImpl(fusion_adaptor, *tiled_computation, - *device_info_, shape_size_, flops_fn)); - - if (tiled_run_time_data.has_value()) { - VLOG(2) << "Accepted tile sizes [" - << absl::StrJoin(tiled_run_time_data->block_level_parameters - .output_tile_sizes.front(), - ",") - << "], exec_time " - << tiled_run_time_data->runtime_data.exec_time; - candidates.push_back(*tiled_run_time_data); - } - } - } else { - SymbolicTileAnalysisOrError analysis_or_error = - SymbolicTileAnalysis::AnalyzeFusion( - fusion_adaptor, mlir_context_, - TritonEmitterConstraints::GetBuilder(*device_info_)); - - if (const auto* fusion_decision = - std::get_if(&analysis_or_error)) { - VLOG(2) - << "TryFindTopKBestTilingsForFusion SymbolicTileAnalysis rejected: " - << fusion_decision->Explain(); - return *fusion_decision; - } - - SymbolicTileAnalysis analysis = - std::get(std::move(analysis_or_error)); - - ABSL_ASSIGN_OR_RETURN(auto tilings, analysis.GetValidTilings()); - VLOG(1) << absl::StrCat( - "TryFindTopKBestTilingsForFusion symbolic analysis evaluating ", - tilings.size(), " tilings."); - for (const auto& tiling : tilings) { - // TODO(b/372454662): This needs to be adjusted if we want to support more - // than one "real root" (i.e. a root without users). - // Currently ComputeTiledComputation() may fail and return an - // Unimplemented error for cases of multi-output fusion that we do not - // support yet. - auto maybe_tiled_hlo_computation = - analysis.ComputeTiledComputation(tiling); - if (!maybe_tiled_hlo_computation.ok()) { - if (maybe_tiled_hlo_computation.status().code() == - absl::StatusCode::kUnimplemented && - absl::StrContains(maybe_tiled_hlo_computation.status().message(), - "multi-output fusion")) { - continue; - } - return maybe_tiled_hlo_computation.status(); - } - - ABSL_ASSIGN_OR_RETURN( - std::optional tiled_run_time_data, - EstimateTiledRunTimeDataImpl( - fusion_adaptor, *maybe_tiled_hlo_computation, *device_info_, - shape_size_, [this](const HloInstruction* hlo) { - return FlopsPerElement(hlo); - })); - - if (tiled_run_time_data.has_value()) { - ABSL_ASSIGN_OR_RETURN(FlatTiling flat_tiling, - tiling.Flatten(analysis.GetTilingSpecification())); - VLOG(2) << "Accepted tile sizes [" << absl::StrJoin(flat_tiling, ",") - << "], exec_time " - << tiled_run_time_data->runtime_data.exec_time; - candidates.push_back(*tiled_run_time_data); - } - } + return TryFindTopKBestTilingsWithTilingSpace( + fusion_adaptor, top_k, executor, flops_fn, mlir_context_, + mlir_context_pool_, *device_info_, shape_size_); } - VLOG(1) << absl::StrCat("Found ", candidates.size(), - " valid tiling candidates."); - absl::c_stable_sort( - candidates, [](const TiledRunTimeData& a, const TiledRunTimeData& b) { - return a.runtime_data.exec_time < b.runtime_data.exec_time; - }); - - if (candidates.size() > top_k) { - candidates.resize(top_k); - } - - ABSL_RETURN_IF_ERROR(DumpAndLogTopKCandidates(fusion_adaptor, candidates)); - - return candidates; + auto flops_fn = [this](const HloInstruction* hlo) -> int64_t { + return FlopsPerElement(hlo); + }; + return TryFindTopKBestTilingsWithSymbolicAnalysis(fusion_adaptor, top_k, + flops_fn, mlir_context_, + *device_info_, shape_size_); } -absl::StatusOr -GpuPerformanceModelWithIndexingAnalysis::TryFindBestTilingForFusion( - const HloFusionAdaptor& fusion_adaptor) { - ABSL_ASSIGN_OR_RETURN(auto top_k_result, TryFindTopKBestTilingsForFusion( - fusion_adaptor, /*top_k=*/1)); - if (std::holds_alternative(top_k_result)) { - return std::get(top_k_result); - } - auto& tilings = - std::get>(top_k_result); - if (tilings.empty()) { - return FusionDecision::Forbid("No valid tilings found."); +tsl::Future +GpuPerformanceModelWithIndexingAnalysis::TryFindBestTilingForFusionAsync( + const HloFusionAdaptor& fusion_adaptor, tsl::Executor* executor) { + if (executor == nullptr) { + executor = &tsl::InlineExecutor::Instance(); } - VLOG(1) - << "TryFindBestTilingForFusion " - << fusion_adaptor.GetRoots().front().instruction().ToString( - HloPrintOptions::ShortParsable()) - << " best output_tile_sizes [" - << absl::StrJoin( - tilings.front().block_level_parameters.output_tile_sizes.front(), - ",") - << "], exec_time " << tilings.front().runtime_data.exec_time; - return tilings.front(); + return TryFindTopKBestTilingsForFusionAsync(fusion_adaptor, /*top_k=*/1, + executor) + .Map(*executor, + [&fusion_adaptor](const TopKTiledRunTimeDataOrError& top_k_result) + -> TiledRunTimeDataOrError { + if (std::holds_alternative(top_k_result)) { + return std::get(top_k_result); + } + const auto& tilings = + std::get>( + top_k_result); + if (tilings.empty()) { + return FusionDecision::Forbid("No valid tilings found."); + } + VLOG(1) + << "TryFindBestTilingForFusionAsync " + << fusion_adaptor.GetRoots().front().instruction().ToString( + HloPrintOptions::ShortParsable()) + << " best output_tile_sizes [" + << absl::StrJoin( + tilings.front() + .block_level_parameters.output_tile_sizes.front(), + ",") + << "], exec_time " << tilings.front().runtime_data.exec_time; + return tilings.front(); + }); } } // namespace gpu diff --git a/third_party/xla/xla/service/gpu/model/gpu_indexing_performance_model.h b/third_party/xla/xla/service/gpu/model/gpu_indexing_performance_model.h index 50cc2989f3cf83..be9d1e10d2e26e 100644 --- a/third_party/xla/xla/service/gpu/model/gpu_indexing_performance_model.h +++ b/third_party/xla/xla/service/gpu/model/gpu_indexing_performance_model.h @@ -17,6 +17,8 @@ limitations under the License. #define XLA_SERVICE_GPU_MODEL_GPU_INDEXING_PERFORMANCE_MODEL_H_ #include +#include +#include #include #include @@ -31,6 +33,7 @@ limitations under the License. #include "xla/codegen/xtile/block_level_parameters.h" #include "xla/hlo/ir/hlo_instruction.h" #include "xla/hlo/utils/hlo_traversal.h" +#include "xla/runtime/object_pool.h" #include "xla/service/gpu/model/fusion_analysis_cache.h" #include "xla/service/gpu/model/gpu_hlo_cost_analysis.h" #include "xla/service/gpu/model/gpu_performance_model_base.h" @@ -38,10 +41,14 @@ limitations under the License. #include "xla/service/hlo_cost_analysis.h" #include "xla/service/instruction_fusion.h" #include "xla/stream_executor/device_description.h" +#include "xla/tsl/concurrency/executor.h" +#include "xla/tsl/concurrency/future.h" namespace xla { namespace gpu { +using MlirContextPool = ObjectPool>; + // Contains informations about block level parameters and run time of a fusion. struct TiledRunTimeData { EstimateRunTimeData runtime_data; @@ -72,7 +79,8 @@ class GpuPerformanceModelWithIndexingAnalysis : public GpuPerformanceModelBase { HloFusionAnalysisCache* fusion_analysis_cache, HloCostAnalysis::ShapeSizeFunction shape_size, mlir::MLIRContext* mlir_context, bool use_experimental_tiling, - bool enable_same_shape_multi_output_fusion) + bool enable_same_shape_multi_output_fusion, + MlirContextPool* mlir_context_pool = nullptr) : hlo_op_profile_(&HloOpProfiles::Singleton().GetProfile(*device_info)), device_info_(device_info), fusion_analysis_cache_(fusion_analysis_cache), @@ -84,6 +92,7 @@ class GpuPerformanceModelWithIndexingAnalysis : public GpuPerformanceModelBase { /*count_multiple_input_accesses=*/true}, *device_info_), mlir_context_(mlir_context), + mlir_context_pool_(mlir_context_pool), use_experimental_tiling_(use_experimental_tiling), enable_same_shape_multi_output_fusion_( enable_same_shape_multi_output_fusion) {} @@ -125,13 +134,26 @@ class GpuPerformanceModelWithIndexingAnalysis : public GpuPerformanceModelBase { // Returns FusionDecision if the fusion can't be tiled or there are no valid // block level parameters. // Otherwise returns block level parameters that give the best execution time. - absl::StatusOr TryFindBestTilingForFusion( - const HloFusionAdaptor& fusion_adaptor); + // + // NOTE: `fusion_adaptor` and `this` are captured; callers must ensure both + // `fusion_adaptor` and this model instance outlive the returned + // `tsl::Future`. If executor is nullptr, defaults to + // tsl::InlineExecutor::Instance(). + [[nodiscard]] tsl::Future + TryFindBestTilingForFusionAsync(const HloFusionAdaptor& fusion_adaptor, + tsl::Executor* executor = nullptr); // Returns top_k (possibly fewer if not enough valid tilings are found) block // level parameters for the given fusion. - absl::StatusOr TryFindTopKBestTilingsForFusion( - const HloFusionAdaptor& fusion_adaptor, int top_k); + // + // NOTE: `fusion_adaptor` and `this` are captured; callers must ensure both + // `fusion_adaptor` and this model instance outlive the returned + // `tsl::Future`. If executor is nullptr, defaults to + // tsl::InlineExecutor::Instance(). + [[nodiscard]] tsl::Future + TryFindTopKBestTilingsForFusionAsync(const HloFusionAdaptor& fusion_adaptor, + int top_k, + tsl::Executor* executor = nullptr); // Returns an estimate how many FLOPs will be used to produce one element of // the output. @@ -144,6 +166,7 @@ class GpuPerformanceModelWithIndexingAnalysis : public GpuPerformanceModelBase { HloCostAnalysis::ShapeSizeFunction shape_size_; GpuHloCostAnalysis cost_analysis_; mlir::MLIRContext* mlir_context_; + MlirContextPool* mlir_context_pool_; bool use_experimental_tiling_; bool enable_same_shape_multi_output_fusion_; }; diff --git a/third_party/xla/xla/service/gpu/model/gpu_indexing_performance_model_test.cc b/third_party/xla/xla/service/gpu/model/gpu_indexing_performance_model_test.cc index 711a6ad045107c..03554a5406806c 100644 --- a/third_party/xla/xla/service/gpu/model/gpu_indexing_performance_model_test.cc +++ b/third_party/xla/xla/service/gpu/model/gpu_indexing_performance_model_test.cc @@ -53,7 +53,9 @@ limitations under the License. #include "xla/shape.h" #include "xla/shape_util.h" #include "xla/stream_executor/device_description.h" +#include "xla/tsl/concurrency/future.h" #include "xla/tsl/platform/env.h" +#include "xla/tsl/platform/threadpool.h" #include "xla/tsl/testing/temporary_directory.h" #include "xla/util.h" #include "tsl/platform/path.h" @@ -75,7 +77,14 @@ class GpuIndexingPerformanceModelTest bool use_experimental_tiling() const { return GetParam(); } + static std::unique_ptr CreateMlirContext() { + auto ctx = std::make_unique(); + ctx->disableMultithreading(); + return ctx; + } + mlir::MLIRContext mlir_context_; + MlirContextPool mlir_context_pool_{CreateMlirContext, 4}; // The reference times in the test cases below are measured // on A6000 by profiling the execution of the HLOs. se::DeviceDescription device_info_{TestGpuDeviceInfo::RTXA6000DeviceInfo()}; @@ -86,7 +95,8 @@ class GpuIndexingPerformanceModelTest HloCostAnalysis::DefaultShapeSize, &mlir_context_, use_experimental_tiling(), - /*enable_same_shape_multi_output_fusion=*/false}; + /*enable_same_shape_multi_output_fusion=*/false, + &mlir_context_pool_}; size_t WarpSize() const { return ::xla::gpu::WarpSize(device_info_); } @@ -442,7 +452,8 @@ ENTRY entry_computation { ASSERT_OK_AND_ASSIGN( TiledRunTimeDataOrError tiling_result, - indexing_cost_model_.TryFindBestTilingForFusion(*fusion_adaptor)); + indexing_cost_model_.TryFindBestTilingForFusionAsync(*fusion_adaptor) + .Await()); ASSERT_TRUE(std::holds_alternative(tiling_result)); @@ -495,7 +506,8 @@ ENTRY entry_computation { ASSERT_OK_AND_ASSIGN( TiledRunTimeDataOrError tiling_result, - indexing_cost_model_.TryFindBestTilingForFusion(*fusion_adaptor)); + indexing_cost_model_.TryFindBestTilingForFusionAsync(*fusion_adaptor) + .Await()); ASSERT_TRUE(std::holds_alternative(tiling_result)); @@ -538,8 +550,10 @@ ENTRY main { ASSERT_OK_AND_ASSIGN( auto top_k_result, - indexing_cost_model_.TryFindTopKBestTilingsForFusion(*fusion_adaptor, - /*top_k=*/3)); + indexing_cost_model_ + .TryFindTopKBestTilingsForFusionAsync(*fusion_adaptor, + /*top_k=*/3) + .Await()); ASSERT_TRUE((std::holds_alternative>( top_k_result))); @@ -583,7 +597,8 @@ ENTRY main { ASSERT_OK_AND_ASSIGN( TiledRunTimeDataOrError tiling_result, - indexing_cost_model_.TryFindBestTilingForFusion(*fusion_adaptor)); + indexing_cost_model_.TryFindBestTilingForFusionAsync(*fusion_adaptor) + .Await()); ASSERT_TRUE(std::holds_alternative(tiling_result)); @@ -1330,6 +1345,61 @@ ENTRY entry { EXPECT_EQ(flops_map.at(exp), indexing_cost_model_.FlopsPerElement(exp)); } +TEST_P(GpuIndexingPerformanceModelTest, + TryFindTopKBestTilingsForFusionAsync_ParallelMatchesSequential) { + ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(R"( +HloModule m + +fused_computation { + param_0.1 = f32[1024,1024] parameter(0) + ROOT negate.1 = f32[1024,1024] negate(param_0.1) +} + +ENTRY main { + param_0.3 = f32[1024,1024] parameter(0) + ROOT fusion = f32[1024,1024] fusion(param_0.3), kind=kCustom, calls=fused_computation +} +)")); + + auto fusion_adaptor = HloFusionAdaptor::ForInstruction( + module->entry_computation()->root_instruction()); + + constexpr int top_k = 5; + ASSERT_OK_AND_ASSIGN( + auto seq_result, + indexing_cost_model_ + .TryFindTopKBestTilingsForFusionAsync(*fusion_adaptor, top_k) + .Await()); + + tsl::thread::ThreadPool thread_pool(tsl::Env::Default(), "test", 8); + ASSERT_OK_AND_ASSIGN(auto parallel_result, + indexing_cost_model_ + .TryFindTopKBestTilingsForFusionAsync( + *fusion_adaptor, top_k, thread_pool.AsExecutor()) + .Await()); + + ASSERT_TRUE((std::holds_alternative>( + seq_result))); + const auto& seq_candidates = + std::get>(seq_result); + ASSERT_FALSE(seq_candidates.empty()); + + ASSERT_TRUE((std::holds_alternative>( + parallel_result))); + const auto& parallel_candidates = + std::get>(parallel_result); + + ASSERT_EQ(parallel_candidates.size(), seq_candidates.size()); + for (int i = 0; i < seq_candidates.size(); ++i) { + EXPECT_EQ(parallel_candidates[i].runtime_data.exec_time, + seq_candidates[i].runtime_data.exec_time); + EXPECT_EQ(parallel_candidates[i].block_level_parameters.output_tile_sizes, + seq_candidates[i].block_level_parameters.output_tile_sizes); + EXPECT_EQ(parallel_candidates[i].block_level_parameters.num_warps, + seq_candidates[i].block_level_parameters.num_warps); + } +} + } // namespace } // namespace gpu } // namespace xla diff --git a/third_party/xla/xla/service/layout_assignment.cc b/third_party/xla/xla/service/layout_assignment.cc index 88fa67ba987317..3238f50e4a377e 100644 --- a/third_party/xla/xla/service/layout_assignment.cc +++ b/third_party/xla/xla/service/layout_assignment.cc @@ -19,12 +19,14 @@ limitations under the License. #include #include #include +#include #include #include #include #include #include "absl/algorithm/container.h" +#include "absl/container/flat_hash_map.h" #include "absl/container/flat_hash_set.h" #include "absl/log/check.h" #include "absl/log/log.h" @@ -36,7 +38,7 @@ limitations under the License. #include "absl/strings/str_join.h" #include "absl/strings/string_view.h" #include "absl/types/span.h" -#include "xla/hlo/analysis/tuple_points_to_analysis.h" +#include "xla/hlo/analysis/hlo_dataflow_analysis.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" @@ -51,7 +53,7 @@ limitations under the License. #include "xla/map_util.h" #include "xla/permutation_util.h" #include "xla/service/computation_layout.h" -#include "xla/service/logical_buffer.h" +#include "xla/service/hlo_value.h" #include "xla/service/while_util.h" #include "xla/shape.h" #include "xla/shape_layout.h" @@ -78,7 +80,7 @@ std::ostream& operator<<(std::ostream& out, } BufferLayoutConstraint::BufferLayoutConstraint(const Layout& layout, - const LogicalBuffer& buffer, + const HloValue& buffer, bool mandatory, bool dfs, int64_t priority) : LayoutConstraint(mandatory, dfs, priority), buffer_(&buffer) { @@ -220,75 +222,72 @@ std::string ComputationLayoutConstraint::ToString() const { layout_state_, computation_layout_.ToString()); } -void LayoutAssignment::SetPointsToAnalysis( - std::unique_ptr analysis) { - points_to_analysis_ = std::move(analysis); +void LayoutAssignment::SetDataflowAnalysis( + std::unique_ptr analysis) { + dataflow_analysis_ = std::move(analysis); buffer_sets_cache_.clear(); - buffer_sets_cache_analysis_ = points_to_analysis_.get(); + buffer_sets_cache_analysis_ = dataflow_analysis_.get(); } -PointsToSet::BufferSet* LayoutAssignment::GetBufferSet( +LayoutAssignment::BufferSet* LayoutAssignment::GetBufferSet( const HloInstruction* instruction) const { - DCHECK_EQ(buffer_sets_cache_analysis_, points_to_analysis_.get()) - << "points_to_analysis_ changed behind buffer_sets_cache_; use " - "SetPointsToAnalysis"; + DCHECK_EQ(buffer_sets_cache_analysis_, dataflow_analysis_.get()) + << "dataflow_analysis_ changed behind buffer_sets_cache_; use " + "SetDataflowAnalysis"; auto it = buffer_sets_cache_.find(instruction); if (it != buffer_sets_cache_.end()) { return it->second.get(); } auto& buffer_set = - buffer_sets_cache_ - .emplace(instruction, std::make_unique()) + buffer_sets_cache_.emplace(instruction, std::make_unique()) .first->second; - const auto& points_to_set = points_to_analysis_->GetPointsToSet(instruction); - points_to_set.ForEachElement( - [&buffer_set](const ShapeIndex& /*index*/, - const PointsToSet::BufferList& buffers) { - buffer_set->insert(buffers.begin(), buffers.end()); + const auto& instruction_value_set = + dataflow_analysis_->GetInstructionValueSet(instruction); + instruction_value_set.ForEachElement( + [&buffer_set](const ShapeIndex& /*index*/, const HloValueSet& value_set) { + buffer_set->insert(value_set.values().begin(), + value_set.values().end()); }); return buffer_set.get(); } bool LayoutAssignment::AnyOperandBufferForwarded( const HloInstruction* instruction, int64_t operand_no) const { - // The operand is potentially forwarded if the intersection of points-to sets + // The operand is potentially forwarded if the intersection of value sets // of the operand and the instruction is non-empty. - PointsToSet::BufferSet* output_buffers = GetBufferSet(instruction); - PointsToSet::BufferSet* operand_buffers = - GetBufferSet(instruction->operand(operand_no)); + BufferSet* output_buffers = GetBufferSet(instruction); + BufferSet* operand_buffers = GetBufferSet(instruction->operand(operand_no)); // Probe the larger set with elements of the smaller one; intersection // emptiness is symmetric. - PointsToSet::BufferSet* small = output_buffers; - PointsToSet::BufferSet* large = operand_buffers; + BufferSet* small = output_buffers; + BufferSet* large = operand_buffers; if (small->size() > large->size()) { std::swap(small, large); } - return absl::c_any_of( - *small, [&](const LogicalBuffer* b) { return large->count(b) > 0; }); + return absl::c_any_of(*small, + [&](const HloValue* b) { return large->count(b) > 0; }); } bool LayoutAssignment::AllOperandBuffersForwarded( const HloInstruction* instruction, int64_t operand_no) const { - // The operand is potentially forwarded if the intersection of points-to sets + // The operand is potentially forwarded if the intersection of value sets // of the operand and the instruction is non-empty. - PointsToSet::BufferSet* output_buffers = GetBufferSet(instruction); - PointsToSet::BufferSet* operand_buffers = - GetBufferSet(instruction->operand(operand_no)); + BufferSet* output_buffers = GetBufferSet(instruction); + BufferSet* operand_buffers = GetBufferSet(instruction->operand(operand_no)); // Each buffer in operand_buffers should also occur in output_buffers. - return absl::c_all_of(*operand_buffers, [&](const LogicalBuffer* b) { + return absl::c_all_of(*operand_buffers, [&](const HloValue* b) { return output_buffers->count(b) > 0; }); } absl::Status LayoutAssignment::SetBufferLayout(const Layout& layout, - const LogicalBuffer& buffer, + const HloValue& buffer, bool mandatory, bool dfs, int64_t priority, const HloInstruction* user) { VLOG(3) << "SetBufferLayout : " << buffer << " : " << LayoutUtil::HumanString(layout) << " with priority " << priority << "; mandatory = " << mandatory << "; dfs = " << dfs << "\n"; - ABSL_RETURN_IF_ERROR(points_to_analysis_->VerifyBuffer(buffer)); if (unconstrained_buffer_ids_.erase(buffer.id()) > 0) { VLOG(3) << "Erase buffer from unconstrained ids\n"; } @@ -466,8 +465,8 @@ absl::Status LayoutAssignment::SetInstructionLayout( instruction->shape(), [this, layout, instruction, mandatory, allow_alias, priority]( const Shape& subshape, const ShapeIndex& index) -> absl::Status { - auto buffers = - points_to_analysis_->GetPointsToSet(instruction).element(index); + const auto& buffers = + dataflow_analysis_->GetValueSet(instruction, index).values(); CHECK_EQ(1, buffers.size()); if (!allow_alias) { CHECK_EQ(buffers[0]->instruction(), instruction); @@ -506,8 +505,8 @@ absl::Status LayoutAssignment::SetInstructionLayout( if (!subshape_index.empty() && index != subshape_index) { return absl::OkStatus(); } - auto buffers = - points_to_analysis_->GetPointsToSet(instruction).element(index); + const auto& buffers = + dataflow_analysis_->GetValueSet(instruction, index).values(); CHECK_EQ(1, buffers.size()); if (!allow_alias) { CHECK_EQ(buffers[0]->instruction(), instruction); @@ -538,7 +537,7 @@ absl::Status LayoutAssignment::SetInstructionLayout( } const BufferLayoutConstraint* LayoutAssignment::GetBufferLayoutConstraint( - const LogicalBuffer& buffer) const { + const HloValue& buffer) const { return buffer_constraints_.Find(buffer.id()); } @@ -586,13 +585,18 @@ std::string LayoutAssignment::ToString( "): ", constraints.OperandLayout(instruction, i)->ToString(), "\n"); } } - for (const LogicalBuffer* buffer : - points_to_analysis_->GetBuffersDefinedByInstruction(instruction)) { - auto* buffer_constraint = GetBufferLayoutConstraint(*buffer); - if (buffer_constraint != nullptr) { - absl::StrAppend(&output, " ", buffer->ToString(), " : ", - LayoutUtil::HumanString(buffer_constraint->layout()), - "\n"); + for (const auto& [index, value_set] : + dataflow_analysis_->GetInstructionValueSet(instruction)) { + for (const HloValue* buffer : value_set.values()) { + if (buffer->instruction() != instruction) { + continue; + } + auto* buffer_constraint = GetBufferLayoutConstraint(*buffer); + if (buffer_constraint != nullptr) { + absl::StrAppend(&output, " ", buffer->ToString(), " : ", + LayoutUtil::HumanString(buffer_constraint->layout()), + "\n"); + } } } } @@ -740,7 +744,13 @@ absl::Status LayoutAssignment::AddParameterConstraints( absl::Status LayoutAssignment::AddCollectiveConstraints( HloInstruction* instruction) { - ABSL_RETURN_IF_ERROR(SetInstructionLayout(instruction->shape(), instruction)); + // HloDataflowAnalysis forwards operand HloValues into tuple element {0} of + // async collective start instructions and forwards start HloValues into + // collective done instructions, so allow_alias is required when constraining + // the instruction output shape. + ABSL_RETURN_IF_ERROR(SetInstructionLayout(instruction->shape(), instruction, + /*mandatory=*/true, /*dfs=*/true, + /*allow_alias=*/true)); for (int64_t i = 0; i < instruction->operand_count(); ++i) { CHECK(instruction->shape().IsArray() || (instruction->shape().IsTuple() && @@ -2017,7 +2027,13 @@ void LayoutAssignment::SetupCopiedInstruction(const HloInstruction& instruction, absl::Status LayoutAssignment::CheckLayouts( HloModule* module, const absl::flat_hash_set& execution_threads) { - ABSL_ASSIGN_OR_RETURN(auto points_to_analysis, TuplePointsToAnalysis::Run(module)); + ABSL_ASSIGN_OR_RETURN(auto dataflow_analysis, + HloDataflowAnalysis::Run( + *module, /*ssa_form=*/false, + /*bitcast_defines_value=*/false, execution_threads, + /*propagate_through_calls=*/false, + /*precompute_uses=*/std::nullopt, + /*propagate_through_control_flow=*/false)); for (auto* computation : module->MakeNonfusionComputations(execution_threads)) { for (auto* instruction : computation->instructions()) { @@ -2026,19 +2042,18 @@ absl::Status LayoutAssignment::CheckLayouts( TF_RET_CHECK(LayoutUtil::HasLayout(instruction->shape())); ABSL_RETURN_IF_ERROR(ShapeUtil::ValidateShape(instruction->shape())); - // Use points-to analysis to verify that every subshape element in the - // output of the instruction matches the layout of the logical buffer + // Use dataflow analysis to verify that every subshape element in the + // output of the instruction matches the layout of the HloValue // which could be the source of the subshape value. - const PointsToSet& points_to_set = - points_to_analysis->GetPointsToSet(instruction); - ABSL_RETURN_IF_ERROR(points_to_set.ForEachElementWithStatus( - [&instruction]( - ShapeIndex index, - const PointsToSet::BufferList& buffers) -> absl::Status { + const InstructionValueSet& value_set_tree = + dataflow_analysis->GetInstructionValueSet(instruction); + ABSL_RETURN_IF_ERROR(value_set_tree.ForEachElementWithStatus( + [&instruction](ShapeIndex index, + const HloValueSet& value_set) -> absl::Status { if (ShapeUtil::IsLeafIndex(instruction->shape(), index)) { const Shape& instruction_subshape = ShapeUtil::GetSubshape(instruction->shape(), index); - for (const LogicalBuffer* buffer : buffers) { + for (const HloValue* buffer : value_set.values()) { if (!Shape::Equal() .IgnoreDynamicDimension() .IgnoreBuffer() @@ -2046,7 +2061,7 @@ absl::Status LayoutAssignment::CheckLayouts( buffer->shape())) { return Internal( "Layout of instruction %s at index {%s} does not match " - "source LogicalBuffer %s: %s vs %s", + "source HloValue %s: %s vs %s", instruction->name(), absl::StrJoin(index, ","), buffer->ToString(), ShapeUtil::HumanStringWithLayout(instruction_subshape), @@ -2406,22 +2421,21 @@ absl::Status LayoutAssignment::PropagateConstraints( namespace { // Returns a vector containing all array-shaped uses (instruction and operand -// number) of the given logical buffer or its aliases. +// number) of the given HloValue or its aliases. std::vector> GetArrayUsesOfBuffer( - const TuplePointsToAnalysis::BufferAliasVector& aliases) { + const HloValue& buffer) { std::vector> uses; - for (const auto& buffer_alias : aliases) { - if (!buffer_alias.instruction()->shape().IsArray()) { + for (const HloPosition& position : buffer.positions()) { + if (!position.instruction->shape().IsArray()) { continue; } // This alias must be the top-level (index == {}) of the instruction's // result because the instruction produces an array. - CHECK(buffer_alias.index().empty()); + CHECK(position.index.empty()); // Add all uses of the instruction's output. - for (const HloInstruction* user : buffer_alias.instruction()->users()) { - for (int64_t operand_no : - user->OperandIndices(buffer_alias.instruction())) { + for (const HloInstruction* user : position.instruction->users()) { + for (int64_t operand_no : user->OperandIndices(position.instruction)) { uses.emplace_back(user, operand_no); } } @@ -2437,17 +2451,17 @@ absl::Status LayoutAssignment::PropagateUseConstraintToDefs( const HloInstruction* user) { // Try to set all logical buffers which may be sources of the given operand to // match the given layout. - const PointsToSet& points_to_set = - points_to_analysis_->GetPointsToSet(instruction); - return points_to_set.ForEachElementWithStatus( + const InstructionValueSet& value_set_tree = + dataflow_analysis_->GetInstructionValueSet(instruction); + return value_set_tree.ForEachElementWithStatus( [&shape_layout, this, priority, user]( const ShapeIndex& index, - const PointsToSet::BufferList& buffers) -> absl::Status { + const HloValueSet& value_set) -> absl::Status { const auto& subshape = ShapeUtil::GetSubshape(shape_layout.shape(), index); if (ShapeUtil::IsLeafIndex(shape_layout.shape(), index) && subshape.has_layout()) { - for (const LogicalBuffer* buffer : buffers) { + for (const HloValue* buffer : value_set.values()) { if (buffer->shape().IsArray() && (buffer->instruction()->opcode() != HloOpcode::kReduce || !buffer->instruction()->shape().IsTuple())) { @@ -2493,16 +2507,14 @@ absl::Status LayoutAssignment::PropagateOperandConstraintToResultForCustomCall( continue; } ShapeIndex shape_index = output_operand_pair.first; - if (!points_to_analysis_->InstructionDefinesBufferAtIndex(user, - shape_index)) { + if (!dataflow_analysis_->ValueIsDefinedAt(user, shape_index)) { return absl::OkStatus(); } - ABSL_ASSIGN_OR_RETURN( - const LogicalBuffer* buffer, - points_to_analysis_->GetBufferDefinedAt(user, shape_index)); + const HloValue& buffer = + dataflow_analysis_->GetValueDefinedAt(user, shape_index); return SetBufferLayout( - operand_constraint.shape_layout().layout(), *buffer, + operand_constraint.shape_layout().layout(), buffer, /*mandatory=*/OperandLayoutAlwaysPropagateForward(user), /*dfs=*/InstructionShouldPropagateDepthFirst(*user), operand_constraint.priority()); @@ -2552,11 +2564,10 @@ absl::Status LayoutAssignment::PropagateOperandConstraint( user->operand_count() == 1 ? ShapeIndex() : ShapeIndex({operand_constraint.operand_no()}); - ABSL_ASSIGN_OR_RETURN( - const LogicalBuffer* buffer, - points_to_analysis_->GetBufferDefinedAt(user, shape_index)); + const HloValue& buffer = + dataflow_analysis_->GetValueDefinedAt(user, shape_index); ABSL_RETURN_IF_ERROR(SetBufferLayout(operand_constraint.shape_layout().layout(), - *buffer, + buffer, /*mandatory=*/true, /*dfs=*/true)); } @@ -2631,22 +2642,20 @@ absl::Status LayoutAssignment::PropagateOperandConstraint( operand->shape().dimensions().size()) { return absl::OkStatus(); } - if (!points_to_analysis_->InstructionDefinesBufferAtIndex( - user, shape_index)) { + if (!dataflow_analysis_->ValueIsDefinedAt(user, shape_index)) { return absl::OkStatus(); } // TODO(b/67641796): Are there cases except fusion that use this code // path? - ABSL_ASSIGN_OR_RETURN( - const LogicalBuffer* buffer, - points_to_analysis_->GetBufferDefinedAt(user, shape_index)); + const HloValue& buffer = + dataflow_analysis_->GetValueDefinedAt(user, shape_index); // If we already have a constraint for the buffer it was assigned but // hasn't propagated yet. This can happen with diamond-shaped graphs // where one path is first evaluated in depth-first order (we're here) // and the other path is propagated later. We don't set the layout // here as it will always be overwritten later. ABSL_RETURN_IF_ERROR(SetBufferLayout( - operand_constraint.shape_layout().layout(), *buffer, + operand_constraint.shape_layout().layout(), buffer, /*mandatory=*/true, /*dfs=*/true, operand_constraint.priority())); return absl::OkStatus(); })); @@ -2662,19 +2671,17 @@ absl::Status LayoutAssignment::PropagateOperandConstraint( if (subshape.dimensions().size() <= 1) { return absl::OkStatus(); } - if (!points_to_analysis_->InstructionDefinesBufferAtIndex( - user, shape_index)) { + if (!dataflow_analysis_->ValueIsDefinedAt(user, shape_index)) { return absl::OkStatus(); } - ABSL_ASSIGN_OR_RETURN( - const LogicalBuffer* buffer, - points_to_analysis_->GetBufferDefinedAt(user, shape_index)); + const HloValue& buffer = + dataflow_analysis_->GetValueDefinedAt(user, shape_index); std::unique_ptr layout = ChooseOutputLayoutFromOperandLayout( operand_constraint.shape_layout().layout(), user, operand_constraint.operand_no()); if (layout != nullptr) { ABSL_RETURN_IF_ERROR(SetBufferLayout( - *layout, *buffer, + *layout, buffer, /*mandatory=*/OperandLayoutAlwaysPropagateForward(user), /*dfs=*/InstructionShouldPropagateDepthFirst(*user), operand_constraint.priority())); @@ -2687,7 +2694,7 @@ absl::Status LayoutAssignment::PropagateOperandConstraint( absl::Status LayoutAssignment::PropagateBufferConstraintToOperands( const BufferLayoutConstraint& buffer_constraint, LayoutConstraints* constraints) { - const LogicalBuffer& buffer = buffer_constraint.buffer(); + const HloValue& buffer = buffer_constraint.buffer(); const HloInstruction* instruction = buffer.instruction(); if (IsAtMostRank1(instruction->shape())) { return absl::OkStatus(); @@ -2768,7 +2775,7 @@ absl::Status LayoutAssignment::PropagateBufferConstraint( const BufferLayoutConstraint& buffer_constraint, LayoutConstraints* constraints) { // Only propagate array layouts. - const LogicalBuffer& buffer = buffer_constraint.buffer(); + const HloValue& buffer = buffer_constraint.buffer(); if (!buffer.IsArray()) { return absl::OkStatus(); } @@ -2782,13 +2789,12 @@ absl::Status LayoutAssignment::PropagateBufferConstraintToUses( LayoutConstraints* constraints) { VLOG(5) << "PropagateBufferConstraintToUses: " << buffer_constraint.ToString(); - const LogicalBuffer& buffer = buffer_constraint.buffer(); + const HloValue& buffer = buffer_constraint.buffer(); TF_RET_CHECK(buffer.IsArray()); // Propagate the layout to all array uses of the logical buffer. This skips // uses of the buffer where the buffer is the element of a tuple. - for (const auto& user_operand_no : - GetArrayUsesOfBuffer(points_to_analysis_->GetBufferAliases(buffer))) { + for (const auto& user_operand_no : GetArrayUsesOfBuffer(buffer)) { const HloInstruction* user = user_operand_no.first; int64_t operand_no = user_operand_no.second; // Only add an operand constraint if the user does not forward the buffer @@ -2825,11 +2831,12 @@ absl::Status LayoutAssignment::PropagateBufferConstraintToUses( ShapeIndex used_index = buffer.index(); used_index.push_front(index); - ABSL_ASSIGN_OR_RETURN(auto buffer, points_to_analysis_->GetBufferDefinedAt( - inputs, used_index)); - - ABSL_RETURN_IF_ERROR(SetBufferLayout(buffer_constraint.layout(), *buffer, - /*mandatory=*/false)); + for (const HloValue* input_buffer : + dataflow_analysis_->GetValueSet(inputs, used_index).values()) { + ABSL_RETURN_IF_ERROR(SetBufferLayout(buffer_constraint.layout(), + *input_buffer, + /*mandatory=*/false)); + } } } @@ -2861,30 +2868,30 @@ absl::Status LayoutAssignment::PropagateResultConstraint( absl::StatusOr LayoutAssignment::InferArrayLayout( const HloInstruction* instruction, const ShapeIndex& index) { const auto& source_buffers = - points_to_analysis_->GetPointsToSet(instruction).element(index); + dataflow_analysis_->GetValueSet(instruction, index).values(); TF_RET_CHECK(!source_buffers.empty()); - // Verify the layout is the same for every LogicalBuffer which this location + // Verify the layout is the same for every HloValue which this location // ('instruction' and 'index') points to. const Layout* first_buffer_layout = nullptr; - for (const LogicalBuffer* source_buffer : source_buffers) { + for (const HloValue* source_buffer : source_buffers) { VLOG(5) << "Logical buffer: " << source_buffer->ToString() << "\n"; auto* source_buffer_constraint = GetBufferLayoutConstraint(*source_buffer); if (source_buffer_constraint == nullptr) { // This should not happen because we've assigned layouts to all // instructions preceding this one. - return Internal("LogicalBuffer %s does not have a layout", + return Internal("HloValue %s does not have a layout", source_buffer->ToString()); } if (first_buffer_layout == nullptr) { first_buffer_layout = &source_buffer_constraint->layout(); } else if (!Layout::Equal().MinorToMajorOnly()( - source_buffer->shape().layout(), *first_buffer_layout)) { - // The points-to set is ambiguous for this index and the different source + source_buffer_constraint->layout(), *first_buffer_layout)) { + // The value set is ambiguous for this index and the different source // buffers have different layouts. This case is possible in valid XLA // computations because we do not propagate BufferLayoutConstraints to all - // LogicalBuffers which may alias the constrained LogicalBuffer at some + // HloValues which may alias the constrained HloValue at some // point in the computation. return FailedPrecondition( "Array at index {%s} in instruction %s aliases buffers %s " @@ -2967,23 +2974,26 @@ absl::Status LayoutAssignment::AssignLayouts(LayoutConstraints& constraints) { // by the respective BufferLayoutConstraints. Any array shapes in the output // of the instruction which are not defined by the instruction (eg, array // elements in a Tuple instruction) will be assigned below via inference. - for (const LogicalBuffer* buffer : - points_to_analysis_->GetBuffersDefinedByInstruction(instruction)) { - if (!buffer->shape().IsArray()) { - continue; - } - TF_RET_CHECK(buffer->instruction() == instruction); - auto* buffer_layout_constraint = GetBufferLayoutConstraint(*buffer); - TF_RET_CHECK(buffer_layout_constraint != nullptr); - if (instruction->opcode() == HloOpcode::kConstant) { - // For constants, we also need to change the layout of the internal - // literal. - instruction->RelayoutConstant(buffer_layout_constraint->layout(), - buffer->index()); - } else { - Shape* buffer_subshape = ShapeUtil::GetMutableSubshape( - instruction->mutable_shape(), buffer->index()); - *buffer_subshape->mutable_layout() = buffer_layout_constraint->layout(); + for (const auto& [index, value_set] : + dataflow_analysis_->GetInstructionValueSet(instruction)) { + for (const HloValue* buffer : value_set.values()) { + if (!buffer->shape().IsArray() || + buffer->instruction() != instruction) { + continue; + } + auto* buffer_layout_constraint = GetBufferLayoutConstraint(*buffer); + TF_RET_CHECK(buffer_layout_constraint != nullptr); + if (instruction->opcode() == HloOpcode::kConstant) { + // For constants, we also need to change the layout of the internal + // literal. + instruction->RelayoutConstant(buffer_layout_constraint->layout(), + buffer->index()); + } else { + Shape* buffer_subshape = ShapeUtil::GetMutableSubshape( + instruction->mutable_shape(), buffer->index()); + *buffer_subshape->mutable_layout() = + buffer_layout_constraint->layout(); + } } } @@ -3261,10 +3271,10 @@ absl::Status LayoutAssignment::ClearComputationLayouts( void LayoutAssignment::InitUnconstrainedBuffers(HloComputation* computation) { for (HloInstruction* inst : computation->instructions()) { - points_to_analysis_->GetPointsToSet(inst).ForEachElement( - [&](const ShapeIndex&, const PointsToSet::BufferList& buffers) { - for (const LogicalBuffer* buffer : buffers) { - // The points to analysis is computed per module, restrict + dataflow_analysis_->GetInstructionValueSet(inst).ForEachElement( + [&](const ShapeIndex&, const HloValueSet& value_set) { + for (const HloValue* buffer : value_set.values()) { + // The dataflow analysis is computed per module, restrict // constraints to array buffers in this computation. if (buffer->IsArray() && buffer->instruction()->parent() == computation) { @@ -3311,13 +3321,9 @@ absl::Status LayoutAssignment::AddCustomCallConstraints( } void LayoutAssignment::RecordUnconstrainedLayoutInstructions() { - for (LogicalBuffer::Id buffer_id : unconstrained_buffer_ids_) { - VLOG(5) - << "unconstrained instruction:" - << points_to_analysis_->GetBuffer(buffer_id).instruction()->ToString() - << "\n"; + for (HloValue::Id buffer_id : unconstrained_buffer_ids_) { unconstrained_layout_instructions_.insert( - points_to_analysis_->GetBuffer(buffer_id).instruction()); + dataflow_analysis_->GetValue(buffer_id).instruction()); } } @@ -3328,9 +3334,9 @@ absl::Status LayoutAssignment::AssignLayoutsToUnconstrainedBuffers( // Arbitrarily pick the first unconstrained buffer and assign a layout. // Constants use their literal's layout; other buffers get a default layout. - // unconstrained_buffer_ids_ has a stable sort based on LogicalBuffer::Id. - const LogicalBuffer& buffer = - points_to_analysis_->GetBuffer(*unconstrained_buffer_ids_.begin()); + // unconstrained_buffer_ids_ has a stable sort based on HloValue::Id. + const HloValue& buffer = + dataflow_analysis_->GetValue(*unconstrained_buffer_ids_.begin()); const HloInstruction* instruction = buffer.instruction(); Layout new_layout = instruction->opcode() == HloOpcode::kConstant @@ -3598,8 +3604,14 @@ absl::Status LayoutAssignment::VerifyEntryComputationLayout( absl::StatusOr> LayoutAssignment::SetupPropagation( HloModule* module, const absl::flat_hash_set& execution_threads) { - ABSL_ASSIGN_OR_RETURN(auto points_to_analysis, TuplePointsToAnalysis::Run(module)); - SetPointsToAnalysis(std::move(points_to_analysis)); + ABSL_ASSIGN_OR_RETURN(auto dataflow_analysis, + HloDataflowAnalysis::Run( + *module, /*ssa_form=*/false, + /*bitcast_defines_value=*/false, execution_threads, + /*propagate_through_calls=*/false, + /*precompute_uses=*/std::nullopt, + /*propagate_through_control_flow=*/false)); + SetDataflowAnalysis(std::move(dataflow_analysis)); auto computations_to_work = module->MakeNonfusionComputations(execution_threads); // If the reverse_computation_order_ flag is set, reverse the ordering of @@ -3962,8 +3974,8 @@ absl::Status LayoutAssignment::ClearPreviousPassSideEffects( unconstrained_layout_instructions_.clear(); unconstrained_buffer_ids_.clear(); buffer_constraints_.Clear(); - // buffer_sets_cache_ stays: it depends only on points_to_analysis_, which - // SetPointsToAnalysis owns. + // buffer_sets_cache_ stays: it depends only on dataflow_analysis_, which + // SetDataflowAnalysis owns. return absl::OkStatus(); } absl::Status LayoutAssignment::AddCopyForOperand(HloInstruction* instruction, diff --git a/third_party/xla/xla/service/layout_assignment.h b/third_party/xla/xla/service/layout_assignment.h index abcfe3eb5a23a6..934b35780a8590 100644 --- a/third_party/xla/xla/service/layout_assignment.h +++ b/third_party/xla/xla/service/layout_assignment.h @@ -33,7 +33,7 @@ limitations under the License. #include "absl/status/status.h" #include "absl/status/statusor.h" #include "absl/strings/string_view.h" -#include "xla/hlo/analysis/tuple_points_to_analysis.h" +#include "xla/hlo/analysis/hlo_dataflow_analysis.h" #include "xla/hlo/ir/hlo_computation.h" #include "xla/hlo/ir/hlo_instruction.h" #include "xla/hlo/ir/hlo_module.h" @@ -43,7 +43,7 @@ limitations under the License. #include "xla/layout_util.h" #include "xla/map_util.h" #include "xla/service/computation_layout.h" -#include "xla/service/logical_buffer.h" +#include "xla/service/hlo_value.h" #include "xla/shape.h" #include "xla/shape_layout.h" #include "xla/shape_util.h" @@ -137,12 +137,12 @@ std::ostream& operator<<(std::ostream& out, const LayoutConstraint& constraint); // or induced this buffer layout constraint. class BufferLayoutConstraint : public LayoutConstraint { public: - // Constructs a BufferLayoutConstraint for the given LogicalBuffer. - BufferLayoutConstraint(const Layout& layout, const LogicalBuffer& buffer, + // Constructs a BufferLayoutConstraint for the given HloValue. + BufferLayoutConstraint(const Layout& layout, const HloValue& buffer, bool mandatory, bool dfs, int64_t priority); - // Returns the constrained LogicalBuffer. - const LogicalBuffer& buffer() const { return *buffer_; } + // Returns the constrained HloValue. + const HloValue& buffer() const { return *buffer_; } // Returns the target Layout for the buffer. const Layout& layout() const { return layout_[0]; } @@ -159,8 +159,8 @@ class BufferLayoutConstraint : public LayoutConstraint { // The constrained layout(s) for the buffer. absl::InlinedVector layout_; - // The LogicalBuffer being constrained. - const LogicalBuffer* buffer_; + // The HloValue being constrained. + const HloValue* buffer_; // The consumer instruction that induced this constraint, if any. const HloInstruction* from_user_ = nullptr; @@ -400,8 +400,8 @@ class LayoutAssignment : public HloModulePass { ChannelLayoutConstraints* channel_constraints = nullptr, bool reverse_computation_order = false); ~LayoutAssignment() override {} - const TuplePointsToAnalysis& points_to_analysis() const { - return *points_to_analysis_; + const HloDataflowAnalysis& dataflow_analysis() const { + return *dataflow_analysis_; } absl::string_view name() const override { return "layout-assignment"; } @@ -574,14 +574,12 @@ class LayoutAssignment : public HloModulePass { // Add a constraint on the layout of a LogicalBuffer, the layout of the // operand of the instruction, or the layout of the result of the computation, // respectively. - absl::Status SetBufferLayout(const Layout& layout, - const LogicalBuffer& buffer, + absl::Status SetBufferLayout(const Layout& layout, const HloValue& buffer, bool mandatory = true, bool dfs = true) { return SetBufferLayout(layout, buffer, mandatory, dfs, current_priority_); } - absl::Status SetBufferLayout(const Layout& layout, - const LogicalBuffer& buffer, bool mandatory, - bool dfs, int64_t priority, + absl::Status SetBufferLayout(const Layout& layout, const HloValue& buffer, + bool mandatory, bool dfs, int64_t priority, const HloInstruction* from_user = nullptr); absl::Status SetOperandLayout(const Shape& shape_with_layout, const HloInstruction* instruction, @@ -647,7 +645,7 @@ class LayoutAssignment : public HloModulePass { const ComputationLayoutConstraint& layout_constraint, LayoutConstraints* constraints); - virtual Layout GetUnconstrainedLayout(const LogicalBuffer& buffer) { + virtual Layout GetUnconstrainedLayout(const HloValue& buffer) { Layout layout = LayoutUtil::GetDefaultLayoutForShape(buffer.shape()); if (buffer.shape().has_layout()) { layout.set_memory_space(buffer.shape().layout().memory_space()); @@ -662,13 +660,14 @@ class LayoutAssignment : public HloModulePass { absl::Status PropagateUnconstraintedBuffers(LayoutConstraints* constraints); const BufferLayoutConstraint* GetBufferLayoutConstraint( - const LogicalBuffer& buffer) const; + const HloValue& buffer) const; absl::StatusOr GetInstructionBufferLayoutConstraint(const HloInstruction* instruction) const; + using BufferSet = absl::flat_hash_set; // Find a bufferset in the bufferset cache. This is useful since we can // currently create the flattened buffer set for the same instruction many // times, which is often slow. - PointsToSet::BufferSet* GetBufferSet(const HloInstruction* instruction) const; + BufferSet* GetBufferSet(const HloInstruction* instruction) const; // Similar to above, but returns true only if all buffers associated with that // operand are forwarded. bool AllOperandBuffersForwarded(const HloInstruction* instruction, @@ -871,8 +870,8 @@ class LayoutAssignment : public HloModulePass { // Constrains layouts for custom calls that have specific layout requirements. absl::Status AddCustomCallConstraints(LayoutConstraints* constraints); - // Initializes unconstrained_buffer_ids_ with all array-shaped logical buffers - // in the given computation. + // Initializes unconstrained_buffer_ids_ with all array HloValues in the + // given computation. void InitUnconstrainedBuffers(HloComputation* computation); // Records instructions that lack layout constraints before applying default @@ -1023,12 +1022,12 @@ class LayoutAssignment : public HloModulePass { // Adds constraints related to host Send/Recv instructions. absl::Status BuildHostChannelConstraints(HloComputation* computation); - // Replaces points_to_analysis_ and drops the buffer sets memoized from the + // Replaces dataflow_analysis_ and drops the buffer sets memoized from the // previous analysis. - void SetPointsToAnalysis(std::unique_ptr analysis); + void SetDataflowAnalysis(std::unique_ptr analysis); - // Module points to analysis that can be updated for cloned computations. - std::unique_ptr points_to_analysis_; + // Module dataflow analysis that can be updated for cloned computations. + std::unique_ptr dataflow_analysis_; // The set of HLO instructions which lacked any layout constraint, thus // receiving propagated default layouts. @@ -1072,25 +1071,24 @@ class LayoutAssignment : public HloModulePass { ChannelLayoutConstraints host_channel_constraints_; // Array-shaped buffers which have not yet been constrained, in id order. - absl::btree_set unconstrained_buffer_ids_; + absl::btree_set unconstrained_buffer_ids_; - // Buffer sets of points_to_analysis_ memoized by GetBufferSet. The entries + // Buffer sets of dataflow_analysis_ memoized by GetBufferSet. The entries // point at buffers owned by the analysis, so the cache is only valid for the - // analysis it was filled from: SetPointsToAnalysis, the only place the + // analysis it was filled from: SetDataflowAnalysis, the only place the // analysis changes, clears it, and GetBufferSet checks that pairing. - mutable absl::flat_hash_map> + mutable absl::flat_hash_map> buffer_sets_cache_; - const TuplePointsToAnalysis* buffer_sets_cache_analysis_ = nullptr; + const HloDataflowAnalysis* buffer_sets_cache_analysis_ = nullptr; - // Buffer layout constraints stored densely by LogicalBuffer::Id, which the - // points-to analysis assigns sequentially. Clear() invalidates every entry + // Buffer layout constraints stored densely by HloValue::Id, which the + // dataflow analysis assigns sequentially. Clear() invalidates every entry // in O(1) by bumping a generation, so no table is destroyed and rebuilt per // propagation round. class BufferConstraintTable { public: // The live constraint of `id`, or nullptr. - BufferLayoutConstraint* Find(LogicalBuffer::Id id) const { + BufferLayoutConstraint* Find(HloValue::Id id) const { const size_t slot = static_cast(id); return slot < constraints_.size() && generations_[slot] == generation_ ? constraints_[slot].get() @@ -1099,8 +1097,7 @@ class LayoutAssignment : public HloModulePass { // Makes `constraint` the live constraint of `id`. BufferLayoutConstraint* Insert( - LogicalBuffer::Id id, - std::unique_ptr constraint) { + HloValue::Id id, std::unique_ptr constraint) { const size_t slot = static_cast(id); if (slot >= constraints_.size()) { constraints_.resize(slot + 1); diff --git a/third_party/xla/xla/service/layout_assignment_test.cc b/third_party/xla/xla/service/layout_assignment_test.cc index b7268e3234e460..c49e7d4bf1fe2d 100644 --- a/third_party/xla/xla/service/layout_assignment_test.cc +++ b/third_party/xla/xla/service/layout_assignment_test.cc @@ -43,8 +43,8 @@ limitations under the License. #include "xla/literal.h" #include "xla/literal_util.h" #include "xla/service/computation_layout.h" +#include "xla/service/hlo_value.h" #include "xla/service/hlo_verifier.h" -#include "xla/service/logical_buffer.h" #include "xla/service/pattern_matcher.h" #include "xla/shape.h" #include "xla/shape_layout.h" @@ -525,7 +525,7 @@ class OperandsMustBeTheSameLayoutAssignment : public LayoutAssignment { absl::Status PropagateBufferConstraint( const BufferLayoutConstraint& buffer_constraint, LayoutConstraints* constraints) override { - const LogicalBuffer& buffer = buffer_constraint.buffer(); + const HloValue& buffer = buffer_constraint.buffer(); const HloInstruction* instruction = buffer.instruction(); // Force the operands' layout to the output layout. diff --git a/third_party/xla/xla/service/spmd/shardy/sdy_round_trip/BUILD b/third_party/xla/xla/service/spmd/shardy/sdy_round_trip/BUILD index 8b05f417940548..2ca1c537bfc8d0 100644 --- a/third_party/xla/xla/service/spmd/shardy/sdy_round_trip/BUILD +++ b/third_party/xla/xla/service/spmd/shardy/sdy_round_trip/BUILD @@ -74,8 +74,6 @@ cc_library( deps = [ "//xla:side_effect_util", "//xla/hlo/ir:hlo_sharding", - "//xla/hlo/ir:mesh_and_axis", - "//xla/hlo/ir:named_sharding", "//xla/hlo/parser:hlo_parser", "//xla/hlo/translate/hlo_to_mhlo:hlo_utils", "//xla/service/spmd/shardy:constants", diff --git a/third_party/xla/xla/service/spmd/shardy/sdy_round_trip/import_shardy_attrs.cc b/third_party/xla/xla/service/spmd/shardy/sdy_round_trip/import_shardy_attrs.cc index 4b3f5b085e2f92..30228198eaa7e2 100644 --- a/third_party/xla/xla/service/spmd/shardy/sdy_round_trip/import_shardy_attrs.cc +++ b/third_party/xla/xla/service/spmd/shardy/sdy_round_trip/import_shardy_attrs.cc @@ -55,8 +55,6 @@ limitations under the License. #include "shardy/dialect/sdy/ir/utils.h" #include "stablehlo/dialect/StablehloOps.h" #include "xla/hlo/ir/hlo_sharding.h" -#include "xla/hlo/ir/mesh_and_axis.h" -#include "xla/hlo/ir/named_sharding.h" #include "xla/hlo/parser/hlo_parser.h" #include "xla/hlo/translate/hlo_to_mhlo/hlo_utils.h" #include "xla/service/spmd/shardy/constants.h" @@ -84,7 +82,6 @@ using ::mlir::func::FuncOp; using ::mlir::sdy::getSharding; using ::mlir::sdy::kShardingAttr; using ::mlir::sdy::kShardingRuleAttr; -using ::mlir::sdy::ManualAxesAttr; using ::mlir::sdy::MeshAttr; using ::mlir::sdy::OpShardingRuleAttr; using ::mlir::sdy::PropagationBarrierOp; @@ -209,37 +206,6 @@ bool handleFuncResultSharding( return anyChanged; } -ManualAxesAttr extractManualAxes(const xla::HloSharding& hloSharding, - mlir::MLIRContext* context) { - if (hloSharding.IsTuple()) { - // Manual axes are only attached to ranked tensors, so a leading token - // element can legitimately have none. Take the first element that does. - for (const xla::HloSharding& element : hloSharding.tuple_elements()) { - if (ManualAxesAttr manualAxes = extractManualAxes(element, context)) { - return manualAxes; - } - } - return nullptr; - } - if (hloSharding.UseNamedShardingLeaf()) { - const xla::NamedSharding& namedSharding = hloSharding.named_sharding(); - llvm::SmallVector manualAxes; - const xla::Mesh& mesh = namedSharding.mesh(); - for (const auto& axisRef : namedSharding.manual_axes()) { - // Silently dropping an axis here would yield a manual computation with - // fewer manual axes than it actually has, which miscompiles rather than - // fails, so treat an out-of-range index as a bug. - CHECK_LT(axisRef.mesh_axis_index(), mesh.axis_names().size()); - manualAxes.push_back(mlir::StringAttr::get( - context, mesh.axis_names()[axisRef.mesh_axis_index()])); - } - if (!manualAxes.empty()) { - return ManualAxesAttr::get(context, manualAxes); - } - } - return nullptr; -} - // The sharding information is in the `kXlaShardingAttr` attribute. void convertShardyAttrsWithHloShardingV3(FuncOp funcOp) { for (auto [argNum, argType] : llvm::enumerate(funcOp.getArgumentTypes())) { @@ -306,38 +272,12 @@ void convertShardyAttrsWithHloShardingV3(FuncOp funcOp) { StringRef targetName = customCallOp.getCallTargetName(); if (targetName == kShardingCustomCallTargetName || targetName == "X64Combine" || - targetName == kGlobalToLocalShapeCallTargetName || - targetName == kLocalToGlobalShapeCallTargetName || isPythonCallbackCustomCall(customCallOp)) { - HloSharding hloSharding = parseShardingFromString(shardingAttr); if (auto sdySharding = convertToSdySharding( - hloSharding, customCallOp->getResultTypes(), - customCallOp->getContext())) { + parseShardingFromString(shardingAttr), + customCallOp->getResultTypes(), customCallOp->getContext())) { customCallOp->setAttr(kShardingAttr, sdySharding); } - // Only `xla.sdy.GlobalToLocalShape` (when there are operands) and the - // manual computation CallOp (when there are no operands) may carry - // manual axes. - if (targetName == kGlobalToLocalShapeCallTargetName) { - if (auto manualAxesAttr = - extractManualAxes(hloSharding, customCallOp->getContext())) { - customCallOp->setAttr(kManualAxes, manualAxesAttr); - } - } - } - } else if (auto callOp = mlir::dyn_cast(op); - callOp && isManualComputation(callOp)) { - // The body call op carries the local out shardings and the manual axes. - // It is the only carrier of the manual axes when the manual computation - // has no operands, since there is no `xla.sdy.GlobalToLocalShape` then. - HloSharding hloSharding = parseShardingFromString(shardingAttr); - if (auto sdySharding = convertToSdySharding( - hloSharding, callOp->getResultTypes(), callOp->getContext())) { - callOp->setAttr(kShardingAttr, sdySharding); - } - if (auto manualAxesAttr = - extractManualAxes(hloSharding, callOp->getContext())) { - callOp->setAttr(kManualAxes, manualAxesAttr); } } diff --git a/third_party/xla/xla/service/spmd/shardy/sdy_round_trip/shard_map_export.cc b/third_party/xla/xla/service/spmd/shardy/sdy_round_trip/shard_map_export.cc index cb87f6e39b80e0..6b15a224661c57 100644 --- a/third_party/xla/xla/service/spmd/shardy/sdy_round_trip/shard_map_export.cc +++ b/third_party/xla/xla/service/spmd/shardy/sdy_round_trip/shard_map_export.cc @@ -17,10 +17,7 @@ limitations under the License. #include -#include "llvm/ADT/STLExtras.h" -#include "llvm/ADT/SmallVector.h" #include "llvm/ADT/StringRef.h" -#include "llvm/Support/CommandLine.h" #include "llvm/Support/ErrorHandling.h" #include "mlir/Dialect/Func/IR/FuncOps.h" #include "mlir/IR/AffineMap.h" @@ -43,7 +40,6 @@ limitations under the License. #include "mlir/Support/LLVM.h" #include "mlir/Support/TypeID.h" #include "mlir/Transforms/DialectConversion.h" -#include "shardy/dialect/sdy/ir/constants.h" #include "shardy/dialect/sdy/ir/dialect.h" #include "shardy/dialect/sdy/ir/utils.h" #include "stablehlo/dialect/StablehloOps.h" @@ -63,44 +59,14 @@ using ::mlir::func::FuncOp; namespace stablehlo = ::mlir::stablehlo; namespace sdy = ::mlir::sdy; -// Returns `shardings` with `manualAxes` stripped from each element, i.e. the -// local view of tensors that are sharded over those axes. -// -// A `NamedSharding` may not mention the same axis twice, so an op whose result -// is a local tensor carries the manual axes in its `xla.sdy.manual_axes` -// attribute and the remaining axes in its dimension shardings. -sdy::TensorShardingPerValueAttr localizeShardings( - sdy::TensorShardingPerValueAttr shardings, - mlir::ArrayRef manualAxes, mlir::MLIRContext* context) { - llvm::SmallVector localShardings; - localShardings.reserve(shardings.size()); - for (sdy::TensorShardingAttr sharding : shardings.getShardings()) { - localShardings.push_back(sdy::eraseManualAxes(sharding, manualAxes)); - } - return sdy::TensorShardingPerValueAttr::get(context, localShardings); -} - class SdyRoundTripShardMapExportPass : public mlir::PassWrapper> { public: MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(SdyRoundTripShardMapExportPass) - SdyRoundTripShardMapExportPass() = default; - explicit SdyRoundTripShardMapExportPass(bool enableHloShardingV3) { - this->enableHloShardingV3 = enableHloShardingV3; - } - SdyRoundTripShardMapExportPass(const SdyRoundTripShardMapExportPass& other) - : mlir::PassWrapper>(other) { - this->enableHloShardingV3 = other.enableHloShardingV3.getValue(); - } - - Option enableHloShardingV3{ - *this, "enable-hlo-sharding-v3", - llvm::cl::desc("Whether to enable HloShardingV3 which is the mesh and " - "axis based sharding representation."), - llvm::cl::init(false)}; + SdyRoundTripShardMapExportPass(bool enableHloShardingV3) + : enableHloShardingV3(enableHloShardingV3) {} void runOnOperation() final { ModuleOp moduleOp = getOperation(); @@ -124,81 +90,28 @@ class SdyRoundTripShardMapExportPass rewriter.setInsertionPoint(manualComputation); mlir::ValueRange operands = manualComputation->getOperands(); if (!operands.empty()) { - llvm::SmallVector globalOperands(operands.begin(), - operands.end()); - sdy::TensorShardingPerValueAttr inShardings = - manualComputation.getInShardings(); - if (enableHloShardingV3) { - inShardings = sdy::inlineMesh(symbolTable, inShardings); - // `GlobalToLocalShape`'s result is the *local* tensor, so it can only - // carry the local in shardings. Materialize the reshard that - // `ManualComputationOp` performs implicitly on its operands as an - // explicit `@Sharding` custom call. This keeps the exported module - // sharding consistent - every op's sharding describes its own result - // - and lets the import side read the global in shardings straight - // off the operand instead of inferring them. - for (auto [globalOperand, inSharding] : - llvm::zip_equal(globalOperands, inShardings.getShardings())) { - auto shardingOp = stablehlo::CustomCallOp::create( - rewriter, loc, globalOperand.getType(), globalOperand); - shardingOp.setCallTargetName(kShardingCustomCallTargetName); - // Import reads the global in shardings back off these wrappers, so - // they have to survive the HLO round trip intact. Without a side - // effect a wrapper of a zero-element tensor is replaced by a - // constant (see `ZeroSizedHloElimination`, which likewise spares - // the side-effecting `xla.sdy.GlobalToLocalShape`), which would - // lose that operand's sharding. The side effect also stops CSE from - // merging wrappers that carry different shardings. - shardingOp.setHasSideEffect(true); - shardingOp->setAttr(mlir::sdy::kShardingAttr, - sdy::TensorShardingPerValueAttr::get( - moduleOp.getContext(), inSharding)); - globalOperand = shardingOp.getResult(0); - } - } auto globalToLocalShape = stablehlo::CustomCallOp::create( - rewriter, loc, manualCompBodyArgTypes, globalOperands); + rewriter, loc, manualCompBodyArgTypes, operands); globalToLocalShape.setCallTargetName(kGlobalToLocalShapeCallTargetName); // We mark `xla.sdy.GlobalToLocalShape` as side-effecting to avoid // CSE deduping it with another taking the same operands, as it would // ignore the frontend attributes that could be different. globalToLocalShape.setHasSideEffect(true); + sdy::TensorShardingPerValueAttr inShardings = + manualComputation.getInShardings(); if (enableHloShardingV3) { - globalToLocalShape->setAttr( - mlir::sdy::kShardingAttr, - localizeShardings(inShardings, manualComputation.getManualAxes(), - moduleOp.getContext())); - globalToLocalShape->setAttr(kManualAxes, - manualComputation.getManualAxesAttr()); - } else { - setFrontendAttribute(globalToLocalShape, kInShardings, inShardings); - setFrontendAttribute(globalToLocalShape, kManualAxes, - manualComputation.getManualAxesAttr()); + inShardings = sdy::inlineMesh(symbolTable, inShardings); } + setFrontendAttribute(globalToLocalShape, kInShardings, inShardings); + setFrontendAttribute(globalToLocalShape, kManualAxes, + manualComputation.getManualAxesAttr()); operands = globalToLocalShape->getResults(); } - sdy::TensorShardingPerValueAttr outShardings = - manualComputation.getOutShardings(); - if (enableHloShardingV3) { - outShardings = sdy::inlineMesh(symbolTable, outShardings); - } - auto callOp = CallOp::create(rewriter, loc, localResultTypes, funcName, operands); setFrontendAttribute(callOp, kXlaInlineableAttr, rewriter.getStringAttr("xla_late")); - if (enableHloShardingV3 && !outShardings.getShardings().empty()) { - // The call op produces the local tensors, so it carries the local view - // of the out shardings together with the manual axes. This is also the - // only carrier of the manual axes when the manual computation has no - // operands, i.e. when there is no `xla.sdy.GlobalToLocalShape`. - callOp->setAttr( - mlir::sdy::kShardingAttr, - localizeShardings(outShardings, manualComputation.getManualAxes(), - moduleOp.getContext())); - callOp->setAttr(kManualAxes, manualComputation.getManualAxesAttr()); - } mlir::ResultRange results = manualComputation->getResults(); if (!results.empty()) { @@ -209,13 +122,14 @@ class SdyRoundTripShardMapExportPass // We mark `xla.sdy.LocalToGlobalShape` as side-effecting to avoid // CSE removing it if it has no users. localToGlobalShape.setHasSideEffect(true); + sdy::TensorShardingPerValueAttr outShardings = + manualComputation.getOutShardings(); if (enableHloShardingV3) { - localToGlobalShape->setAttr(mlir::sdy::kShardingAttr, outShardings); - } else { - setFrontendAttribute(localToGlobalShape, kOutShardings, outShardings); - setFrontendAttribute(localToGlobalShape, kManualAxes, - manualComputation.getManualAxesAttr()); + outShardings = sdy::inlineMesh(symbolTable, outShardings); } + setFrontendAttribute(localToGlobalShape, kOutShardings, outShardings); + setFrontendAttribute(localToGlobalShape, kManualAxes, + manualComputation.getManualAxesAttr()); results = localToGlobalShape->getResults(); } sdy::inlineRegionAndConvertTerminatorOp( @@ -233,16 +147,15 @@ class SdyRoundTripShardMapExportPass "1. A separate function for the body of the `ManualComputationOp`." "2. A `CallOp` calling the function in #1, marked as not inlinable." "3. A pair of `CustomCallOp`s that change the shape of the " - " arguments/results." - " Under HloShardingV3 the shardings and manual axes are attached as " - "native attributes, and each operand is additionally wrapped in a " - "`@Sharding` custom call holding the global in sharding so that " - "every op's sharding describes its own result. Otherwise they are " - "saved as frontend attrs on the pair of `CustomCallOp`s."; + " arguments/results. They save the in/out shardings and manual " + " axes as frontend attrs."; } void getDependentDialects(mlir::DialectRegistry& registry) const final { registry.insert(); } + + private: + bool enableHloShardingV3 = false; }; } // namespace diff --git a/third_party/xla/xla/service/spmd/shardy/sdy_round_trip/shard_map_import.cc b/third_party/xla/xla/service/spmd/shardy/sdy_round_trip/shard_map_import.cc index 29ea30dd68e693..2e2151e8d11f62 100644 --- a/third_party/xla/xla/service/spmd/shardy/sdy_round_trip/shard_map_import.cc +++ b/third_party/xla/xla/service/spmd/shardy/sdy_round_trip/shard_map_import.cc @@ -75,50 +75,6 @@ using ::mlir::stablehlo::CustomCallOp; namespace sdy = ::mlir::sdy; -// Recovers the global in shardings of a manual computation from the -// `@Sharding` custom calls that export wraps around each operand of -// `globalToLocalShape`, and removes those wrappers. -// -// Export materializes the reshard that `sdy.manual_computation` performs -// implicitly on entry as an explicit `@Sharding` op, so that every op's -// sharding describes its own result. Undoing it here makes the round trip an -// identity. -// -// Returns an empty vector, and leaves the IR untouched, if any operand is not -// wrapped - the shardings are only meaningful as a complete set. -SmallVector takeGlobalInShardings( - CustomCallOp globalToLocalShape, mlir::IRRewriter& rewriter) { - SmallVector globalInShardings; - SmallVector shardingOps; - for (mlir::Value globalOperand : globalToLocalShape.getOperands()) { - auto shardingOp = globalOperand.getDefiningOp(); - if (!shardingOp || - shardingOp.getCallTargetName() != kShardingCustomCallTargetName) { - return {}; - } - sdy::TensorShardingAttr sharding = - sdy::getSharding(shardingOp.getResult(0)); - if (!sharding) { - return {}; - } - globalInShardings.push_back(sharding); - shardingOps.push_back(shardingOp); - } - - for (auto [i, shardingOp] : llvm::enumerate(shardingOps)) { - globalToLocalShape->setOperand(i, shardingOp.getOperand(0)); - } - // The same wrapper can feed several operands, so only visit it once - after - // the first erase the handle is dangling. - llvm::SmallDenseSet seen; - for (CustomCallOp shardingOp : shardingOps) { - if (seen.insert(shardingOp).second && shardingOp.use_empty()) { - rewriter.eraseOp(shardingOp); - } - } - return globalInShardings; -} - mlir::LogicalResult rewriteManualComputation( CallOp callOp, mlir::IRRewriter& rewriter, const mlir::SymbolTable& symbolTable) { @@ -140,7 +96,7 @@ mlir::LogicalResult rewriteManualComputation( // `ManualComputationOp` differently depending on whether the original had // operands/results. CustomCallOp globalToLocalShape; - SmallVector operands(callOp.getOperands()); + mlir::ValueRange operands = callOp.getOperands(); if (!operands.empty()) { // An input to `sdy.manual_computation` can have a dimension of size 0 // (i.e. 0 num-elements), in which case, the corresponding result of @@ -158,8 +114,7 @@ mlir::LogicalResult rewriteManualComputation( globalToLocalShape = (*customCallResIt).getDefiningOp(); CHECK_EQ(globalToLocalShape.getCallTargetName(), kGlobalToLocalShapeCallTargetName); - operands.assign(globalToLocalShape->operand_begin(), - globalToLocalShape->operand_end()); + operands = globalToLocalShape->getOperands(); } mlir::TypeRange resultTypes = callOp->getResultTypes(); @@ -195,13 +150,10 @@ mlir::LogicalResult rewriteManualComputation( if (!customCallOp) { return; } - if (mlir::DictionaryAttr frontendAttrs = getFrontendAttrs(customCallOp)) { - if (hasKey(frontendAttrs, shardingAttrName)) { - shardings = parseStringAttr( - frontendAttrs, shardingAttrName); - } - if (manualAxes.empty() && hasKey(frontendAttrs, kManualAxes)) { + shardings = parseStringAttr( + frontendAttrs, shardingAttrName); + if (manualAxes.empty()) { manualAxes = parseStringAttr(frontendAttrs, kManualAxes); } @@ -210,45 +162,6 @@ mlir::LogicalResult rewriteManualComputation( setShardingAttrs(globalToLocalShape, inShardings, kInShardings); setShardingAttrs(localToGlobalShape, outShardings, kOutShardings); - - // Under HloShardingV3 the shardings and manual axes round trip as native - // attributes rather than frontend attributes, so read them directly. - if (outShardings.empty() && localToGlobalShape) { - // `LocalToGlobalShape` produces the global tensors, so its own sharding is - // the global out shardings. - if (auto sharding = - localToGlobalShape->getAttrOfType( - sdy::kShardingAttr)) { - outShardings = sharding; - } - } - - if (manualAxes.empty()) { - // A manual computation without operands has no `GlobalToLocalShape`, in - // which case the body call op is the only carrier of the manual axes. - for (Operation* op : - {globalToLocalShape.getOperation(), callOp.getOperation()}) { - auto axes = - op ? op->getAttrOfType(kManualAxes) : nullptr; - if (axes && !axes.empty()) { - manualAxes = axes; - break; - } - } - } - - if (inShardings.empty() && globalToLocalShape) { - SmallVector globalInShardings = - takeGlobalInShardings(globalToLocalShape, rewriter); - if (!globalInShardings.empty()) { - inShardings = - sdy::TensorShardingPerValueAttr::get(context, globalInShardings); - // `takeGlobalInShardings` rewired `globalToLocalShape` past the wrappers. - operands.assign(globalToLocalShape->operand_begin(), - globalToLocalShape->operand_end()); - } - } - auto manualComputationOp = rewriter.replaceOpWithNewOp( callOp, resultTypes, operands, inShardings, outShardings, manualAxes); @@ -363,12 +276,10 @@ class SdyRoundTripShardMapImportPass } StringRef getDescription() const override { - return "converts a CallOp calling a @xla.sdy.manual_computation_body func, " - "wrapped with a pair of `CustomCallOps` that change the shape of " - "the arguments/results, to a ManualComputationOp. The in/out " - "shardings and manual axes are read from frontend attrs, or, under " - "HloShardingV3, from native attributes on the boundary ops and on " - "the `@Sharding` custom calls wrapping the operands"; + return "converts a CallOp calling a @xla.sdy.manual_computation_body func " + "with in/out shardings and manual axes as frontend attrs, wrapped " + "with a pair of `CustomCallOps` that change the shape of the " + "arguments/results, to a ManualComputationOp"; } void getDependentDialects(mlir::DialectRegistry& registry) const final { registry.insert(); diff --git a/third_party/xla/xla/service/spmd/shardy/stablehlo_round_trip/export_shardings.cc b/third_party/xla/xla/service/spmd/shardy/stablehlo_round_trip/export_shardings.cc index e1f91c750f6882..0c76e727f68379 100644 --- a/third_party/xla/xla/service/spmd/shardy/stablehlo_round_trip/export_shardings.cc +++ b/third_party/xla/xla/service/spmd/shardy/stablehlo_round_trip/export_shardings.cc @@ -136,11 +136,6 @@ void exportFunc(FuncOp funcOp, const SymbolTable& symbolTable, if (ManualAxesAttr manualAxesAttr = funcOp.getArgAttrOfType(argNum, kManualAxes)) { manualAxes = manualAxesAttr.getValue(); - if (mlir::isa( - funcOp.getArgument(argNum).getType()) && - getMeshAttr(sdySharding).getAxes().size() != manualAxes.size()) { - manualAxes = {}; - } attrs.erase(kManualAxes); } attrs.set(kXlaShardingAttr, @@ -166,10 +161,6 @@ void exportFunc(FuncOp funcOp, const SymbolTable& symbolTable, if (auto manualAxesAttr = mlir::dyn_cast_or_null(attrs.get(kManualAxes))) { manualAxes = manualAxesAttr.getValue(); - if (mlir::isa(funcOp.getResultTypes()[resNum]) && - getMeshAttr(sdySharding).getAxes().size() != manualAxes.size()) { - manualAxes = {}; - } attrs.erase(kManualAxes); } attrs.set(kXlaShardingAttr, @@ -328,38 +319,24 @@ HloSharding getHloShardingForOp( std::function getMeshAttr, ArrayRef manualAxes, bool enableHloShardingV3, bool simplifyReplicatedShardings) { - auto getManualAxesForType = [&](mlir::Type type, - TensorShardingAttr sdySharding) { - if (mlir::isa(type) && - getMeshAttr(sdySharding).getAxes().size() != manualAxes.size()) { - return ArrayRef(); - } - return manualAxes; - }; - bool isNoResultMaximal = op->getNumResults() == 0 && shardings.size() == 1 && (getMeshAttr(shardings.front()).isMaximal() || shardings.front().isFullyReplicated()); CHECK(shardings.size() == op->getNumResults() || isNoResultMaximal); if (op->getNumResults() == 1 || isNoResultMaximal) { - ArrayRef resultManualAxes = - op->getNumResults() == 1 - ? getManualAxesForType(op->getResultTypes().front(), - shardings.front()) - : manualAxes; - return convertToHloSharding(shardings.front(), getMeshAttr, - resultManualAxes, enableHloShardingV3, + return convertToHloSharding(shardings.front(), getMeshAttr, manualAxes, + enableHloShardingV3, simplifyReplicatedShardings); } std::vector newShardings; newShardings.reserve(shardings.size()); - for (auto [type, sdySharding] : - llvm::zip_equal(op->getResultTypes(), shardings)) { - newShardings.push_back(convertToHloSharding( - sdySharding, getMeshAttr, getManualAxesForType(type, sdySharding), - enableHloShardingV3, simplifyReplicatedShardings)); - } + llvm::transform(shardings, std::back_inserter(newShardings), + [&](TensorShardingAttr sdySharding) { + return convertToHloSharding(sdySharding, getMeshAttr, + manualAxes, enableHloShardingV3, + simplifyReplicatedShardings); + }); std::vector shapes; shapes.reserve(op->getNumResults()); diff --git a/third_party/xla/xla/service/spmd/shardy/test/round_trip_pipeline_manual_computation.mlir b/third_party/xla/xla/service/spmd/shardy/test/round_trip_pipeline_manual_computation.mlir index 597d749388fdf6..9a05f216ae8a8e 100644 --- a/third_party/xla/xla/service/spmd/shardy/test/round_trip_pipeline_manual_computation.mlir +++ b/third_party/xla/xla/service/spmd/shardy/test/round_trip_pipeline_manual_computation.mlir @@ -154,42 +154,3 @@ func.func @main( } : (!stablehlo.token, tensor<2xi64>) -> (!stablehlo.token, tensor<2xi64>) return %0#0, %0#1 : !stablehlo.token, tensor<2xi64> } - -// ----- - -// ***************** Zero-element operand test ***************** - -// A zero-element operand is replaced by a constant during the HLO round trip, -// which drops any non side-effecting op producing it. Make sure the remaining -// operands keep their in shardings rather than the manual computation coming -// back with none at all. - -// Make sure this temp attr doesn't exist anymore. -// CHECK-NOT: sharding_hlo_string - -// CHECK-V2: sdy.mesh @mesh_1 = <["a"=4]> -// CHECK-V3: sdy.mesh @mesh = <["a"=4]> -sdy.mesh @mesh_1 = <["a"=4]> - -// CHECK-LABEL: func.func @main -func.func @main(%arg0: tensor<4x8xi32>) -> tensor<4x8xi32> { - // CHECK-NEXT: %[[MANUAL_COMP:.*]] = sdy.manual_computation(%arg0) - // CHECK-V2-SAME{LITERAL}: in_shardings=[<@mesh_1, [{"a"}, {}]>] out_shardings=[<@mesh_1, [{"a"}, {}]>] manual_axes={"a"} (%arg1: tensor<1x8xi32>) { - // CHECK-V3-SAME{LITERAL}: in_shardings=[<@mesh, [{"a"}, {}]>] out_shardings=[<@mesh, [{"a"}, {}]>] manual_axes={"a"} (%arg1: tensor<1x8xi32>) { - // CHECK-NEXT: sdy.return %arg1 : tensor<1x8xi32> - // CHECK-NEXT: } : (tensor<4x8xi32>) -> tensor<4x8xi32> - // CHECK-NEXT: return %[[MANUAL_COMP]] : tensor<4x8xi32> - %c = stablehlo.constant dense<> : tensor<0xi32> - %0 = sdy.manual_computation(%c, %arg0) - in_shardings=[<@mesh_1, [{}], replicated={"a"}>, <@mesh_1, [{"a"}, {}]>] - out_shardings=[<@mesh_1, [{"a"}, {}]>] - manual_axes={"a"} (%arg1: tensor<0xi32>, %arg2: tensor<1x8xi32>) { - %init = stablehlo.constant dense<0> : tensor - %red = stablehlo.reduce(%arg1 init: %init) applies stablehlo.add across dimensions = [0] : (tensor<0xi32>, tensor) -> tensor - %bcast = stablehlo.broadcast_in_dim %red, dims = [] : (tensor) -> tensor<1x8xi32> - %add = stablehlo.add %arg2, %bcast : tensor<1x8xi32> - sdy.return %add : tensor<1x8xi32> - } : (tensor<0xi32>, tensor<4x8xi32>) -> tensor<4x8xi32> - return %0 : tensor<4x8xi32> -} - diff --git a/third_party/xla/xla/service/spmd/shardy/test/sdy_round_trip_import_pipeline_hlo_sharding_v3.mlir b/third_party/xla/xla/service/spmd/shardy/test/sdy_round_trip_import_pipeline_hlo_sharding_v3.mlir index 1b0f94f6667750..58b7012ab5b707 100644 --- a/third_party/xla/xla/service/spmd/shardy/test/sdy_round_trip_import_pipeline_hlo_sharding_v3.mlir +++ b/third_party/xla/xla/service/spmd/shardy/test/sdy_round_trip_import_pipeline_hlo_sharding_v3.mlir @@ -204,30 +204,6 @@ module @module_1 { return %3 : tensor<1xi32> } - // Tests importing shard map custom calls annotated with HloShardingV3 mhlo.sharding. - // CHECK-LABEL: func @manual_computation_hlo_sharding_v3 - // CHECK-SAME: (%arg0: tensor<8x16xf32>) -> tensor<8x16xf32> { - // CHECK-NEXT: %[[MAN_COMP:.*]] = sdy.manual_computation(%arg0) - // CHECK-SAME{LITERAL}: in_shardings=[<@mesh_2, [{"a"}, {}]>] out_shardings=[<@mesh_2, [{"a"}, {}]>] manual_axes={"a"} - // CHECK-SAME: (%arg1: tensor<2x16xf32>) { - // CHECK-NEXT: %[[ADD:.*]] = stablehlo.add %arg1, %arg1 : tensor<2x16xf32> - // CHECK-NEXT: sdy.return %[[ADD]] : tensor<2x16xf32> - // CHECK-NEXT: } : (tensor<8x16xf32>) -> tensor<8x16xf32> - // CHECK-NEXT: return %[[MAN_COMP]] : tensor<8x16xf32> - // CHECK-NEXT: } - func.func @manual_computation_hlo_sharding_v3(%arg0: tensor<8x16xf32>) -> tensor<8x16xf32> { - %s0 = stablehlo.custom_call @Sharding(%arg0) {mhlo.sharding = "{mesh['a'=4,'b'=2], [{'a'}, {}]}"} : (tensor<8x16xf32>) -> tensor<8x16xf32> - %0 = stablehlo.custom_call @xla.sdy.GlobalToLocalShape(%s0) {has_side_effect = true, mhlo.sharding = "{mesh['a'=4,'b'=2], [{}, {}], manual={'a'}}"} : (tensor<8x16xf32>) -> tensor<2x16xf32> - %1 = call @xla.sdy.manual_computation_body_hlo_v3(%0) : (tensor<2x16xf32>) -> tensor<2x16xf32> - %2 = stablehlo.custom_call @xla.sdy.LocalToGlobalShape(%1) {has_side_effect = true, mhlo.sharding = "{mesh['a'=4,'b'=2], [{'a'}, {}]}"} : (tensor<2x16xf32>) -> tensor<8x16xf32> - return %2 : tensor<8x16xf32> - } - - func.func private @xla.sdy.manual_computation_body_hlo_v3(%arg0: tensor<2x16xf32>) -> tensor<2x16xf32> { - %0 = stablehlo.add %arg0, %arg0 : tensor<2x16xf32> - return %0 : tensor<2x16xf32> - } - // CHECK-LABEL: func @frontend_attr_not_sharding // CHECK-SAME: %arg0: tensor<16x8xf32> {sdy.sharding = #sdy.sharding<@mesh_0, [{"a", "b"}, {?}]>}, // CHECK-SAME: %arg1: tensor<16x8xf32> {mhlo.frontend_attributes = {baz = 1 : i32, foo = "bar"}}, diff --git a/third_party/xla/xla/service/spmd/shardy/test/sdy_round_trip_shard_map_export_hlo_sharding_v3.mlir b/third_party/xla/xla/service/spmd/shardy/test/sdy_round_trip_shard_map_export_hlo_sharding_v3.mlir deleted file mode 100644 index fc83e872e68335..00000000000000 --- a/third_party/xla/xla/service/spmd/shardy/test/sdy_round_trip_shard_map_export_hlo_sharding_v3.mlir +++ /dev/null @@ -1,285 +0,0 @@ -// 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. -// ============================================================================== -// RUN: sdy_opt %s -xla-sdy-round-trip-shard-map-export='enable-hlo-sharding-v3=true' 2>&1 | FileCheck %s - -// Tests exporting `sdy.manual_computation` under HloShardingV3 (`enable-hlo-sharding-v3=true`). -// Verifies that: -// 1. Every op carries a native `sdy.sharding` that describes its own result: -// - a `@Sharding` custom call per operand holds the global in-shardings, -// making the reshard `sdy.manual_computation` performs implicitly on entry -// explicit. It is side-effecting so that the HLO round trip cannot drop it -// (see `ZeroSizedHloElimination` for a zero-element operand) or CSE two -// wrappers carrying different shardings; -// - `GlobalToLocalShape` and the body call op hold the local shardings plus -// typed `xla.sdy.manual_axes`; -// - `LocalToGlobalShape` holds the global out-shardings. -// 2. Legacy string `mhlo.frontend_attributes` (xla.sdy.in_shardings, etc.) are NOT emitted. - -sdy.mesh @mesh_0 = <["a"=4, "b"=2]> -sdy.mesh @mesh_1 = <["a"=2, "b"=2, "c"=2, "d"=2]> - -// Ensure string frontend attributes are completely absent in V3 export: -// CHECK-NOT: xla.sdy.in_shardings -// CHECK-NOT: xla.sdy.out_shardings - -// Tests basic export with multiple inputs, outputs, and all-reduce. -// CHECK-LABEL: func @single_manual_comp -func.func @single_manual_comp(%arg0: tensor<8x16xf32> {sdy.sharding = #sdy.sharding<@mesh_0, [{"a", ?}, {"b", ?}]>}, %arg1: tensor<16x32xf32> {sdy.sharding = #sdy.sharding<@mesh_0, [{"b", ?}, {?}]>}) -> (tensor<8x32xf32> {sdy.sharding = #sdy.sharding<@mesh_0, [{"a"}, {}]>}) { - // CHECK-NEXT: %[[SHARDING_0:.*]] = stablehlo.custom_call @Sharding(%arg0) - // CHECK-SAME: has_side_effect = true - // CHECK-SAME{LITERAL}: sdy.sharding = #sdy.sharding_per_value<[, [{"a"}, {"b"}]>]> - // CHECK-SAME: (tensor<8x16xf32>) -> tensor<8x16xf32> - // CHECK-NEXT: %[[SHARDING_1:.*]] = stablehlo.custom_call @Sharding(%arg1) - // CHECK-SAME: has_side_effect = true - // CHECK-SAME{LITERAL}: sdy.sharding = #sdy.sharding_per_value<[, [{"b"}, {}], replicated={"a"}>]> - // CHECK-SAME: (tensor<16x32xf32>) -> tensor<16x32xf32> - // CHECK-NEXT: %[[GLOBAL_TO_LOCAL:.*]]:2 = stablehlo.custom_call @xla.sdy.GlobalToLocalShape(%[[SHARDING_0]], %[[SHARDING_1]]) - // CHECK-SAME{LITERAL}: sdy.sharding = #sdy.sharding_per_value<[, [{}, {}]>, , [{}, {}]>]> - // CHECK-SAME{LITERAL}: xla.sdy.manual_axes = #sdy - // CHECK-SAME: (tensor<8x16xf32>, tensor<16x32xf32>) -> (tensor<2x8xf32>, tensor<8x32xf32>) - // CHECK-NEXT: %[[SHMAP:.*]] = call @xla.sdy.manual_computation_body(%[[GLOBAL_TO_LOCAL]]#0, %[[GLOBAL_TO_LOCAL]]#1) - // CHECK-SAME: {mhlo.frontend_attributes = {inlineable = "xla_late"} - // CHECK-SAME{LITERAL}: sdy.sharding = #sdy.sharding_per_value<[, [{}, {}]>]> - // CHECK-SAME{LITERAL}: xla.sdy.manual_axes = #sdy - // CHECK-SAME: (tensor<2x8xf32>, tensor<8x32xf32>) -> tensor<2x32xf32> - // CHECK-NEXT: %[[LOCAL_TO_GLOBAL:.*]] = stablehlo.custom_call @xla.sdy.LocalToGlobalShape(%[[SHMAP]]) - // CHECK-SAME{LITERAL}: sdy.sharding = #sdy.sharding_per_value<[, [{"a"}, {}], replicated={"b"}>]> - // CHECK-SAME: (tensor<2x32xf32>) -> tensor<8x32xf32> - // CHECK-NEXT: return %[[LOCAL_TO_GLOBAL]] : tensor<8x32xf32> - %0 = sdy.manual_computation(%arg0, %arg1) in_shardings=[<@mesh_0, [{"a"}, {"b"}]>, <@mesh_0, [{"b"}, {}], replicated={"a"}>] out_shardings=[<@mesh_0, [{"a"}, {}], replicated={"b"}>] manual_axes={"a", "b"} (%arg2: tensor<2x8xf32>, %arg3: tensor<8x32xf32>) { - %1 = stablehlo.add %arg2, %arg2 : tensor<2x8xf32> - %2 = stablehlo.dot %1, %arg3 : (tensor<2x8xf32>, tensor<8x32xf32>) -> tensor<2x32xf32> - %3 = "stablehlo.all_reduce"(%2) <{channel_handle = #stablehlo.channel_handle, replica_groups = dense<[[0, 1], [2, 3], [4, 5], [6, 7]]> : tensor<4x2xi64>, use_global_device_ids}> ({ - ^bb0(%arg4: tensor, %arg5: tensor): - %4 = stablehlo.add %arg4, %arg5 : tensor - stablehlo.return %4 : tensor - }) : (tensor<2x32xf32>) -> tensor<2x32xf32> - sdy.return %3 : tensor<2x32xf32> - } : (tensor<8x16xf32>, tensor<16x32xf32>) -> tensor<8x32xf32> - return %0 : tensor<8x32xf32> -} - -// Tests chained manual computations where the output of one feeds the input of -// another with a *different* sharding. The `@Sharding` custom call on the -// second boundary is what preserves the reshard between them; without it the -// `in_shardings` of the second manual computation would be lost. -// CHECK-LABEL: func @manual_comp_using_another -func.func @manual_comp_using_another(%arg0: tensor<8x8xf32> {sdy.sharding = #sdy.sharding<@mesh_0, [{"a"}, {}]>}) - -> (tensor<8x8xf32> {sdy.sharding = #sdy.sharding<@mesh_0, [{}, {"b"}]>}) { - // CHECK-NEXT: %[[SHARDING_0:.*]] = stablehlo.custom_call @Sharding(%arg0) - // CHECK-SAME{LITERAL}: sdy.sharding = #sdy.sharding_per_value<[, [{"a"}, {}]>]> - // CHECK-NEXT: %[[GLOBAL_TO_LOCAL_0:.*]] = stablehlo.custom_call @xla.sdy.GlobalToLocalShape(%[[SHARDING_0]]) - // CHECK-SAME{LITERAL}: sdy.sharding = #sdy.sharding_per_value<[, [{}, {}]>]> - // CHECK-SAME{LITERAL}: xla.sdy.manual_axes = #sdy - // CHECK-SAME: (tensor<8x8xf32>) -> tensor<2x8xf32> - // CHECK-NEXT: %[[SHMAP_0:.*]] = call @xla.sdy.manual_computation_body_0(%[[GLOBAL_TO_LOCAL_0]]) - // CHECK-SAME: {mhlo.frontend_attributes = {inlineable = "xla_late"} - // CHECK-SAME{LITERAL}: xla.sdy.manual_axes = #sdy - // CHECK-SAME: (tensor<2x8xf32>) -> tensor<2x8xf32> - // CHECK-NEXT: %[[LOCAL_TO_GLOBAL_0:.*]] = stablehlo.custom_call @xla.sdy.LocalToGlobalShape(%[[SHMAP_0]]) - // CHECK-SAME{LITERAL}: sdy.sharding = #sdy.sharding_per_value<[, [{"a"}, {}]>]> - // CHECK-SAME: (tensor<2x8xf32>) -> tensor<8x8xf32> - // CHECK-NEXT: %[[SHARDING_1:.*]] = stablehlo.custom_call @Sharding(%[[LOCAL_TO_GLOBAL_0]]) - // CHECK-SAME{LITERAL}: sdy.sharding = #sdy.sharding_per_value<[, [{}, {"b"}]>]> - // CHECK-NEXT: %[[GLOBAL_TO_LOCAL_1:.*]] = stablehlo.custom_call @xla.sdy.GlobalToLocalShape(%[[SHARDING_1]]) - // CHECK-SAME{LITERAL}: sdy.sharding = #sdy.sharding_per_value<[, [{}, {}]>]> - // CHECK-SAME{LITERAL}: xla.sdy.manual_axes = #sdy - // CHECK-SAME: (tensor<8x8xf32>) -> tensor<8x4xf32> - // CHECK-NEXT: %[[SHMAP_1:.*]] = call @xla.sdy.manual_computation_body_1(%[[GLOBAL_TO_LOCAL_1]]) - // CHECK-SAME: {mhlo.frontend_attributes = {inlineable = "xla_late"} - // CHECK-SAME{LITERAL}: xla.sdy.manual_axes = #sdy - // CHECK-SAME: (tensor<8x4xf32>) -> tensor<8x4xf32> - // CHECK-NEXT: %[[LOCAL_TO_GLOBAL_1:.*]] = stablehlo.custom_call @xla.sdy.LocalToGlobalShape(%[[SHMAP_1]]) - // CHECK-SAME{LITERAL}: sdy.sharding = #sdy.sharding_per_value<[, [{}, {"b"}]>]> - // CHECK-SAME: (tensor<8x4xf32>) -> tensor<8x8xf32> - // CHECK-NEXT: return %[[LOCAL_TO_GLOBAL_1]] : tensor<8x8xf32> - %0 = sdy.manual_computation(%arg0) in_shardings=[<@mesh_0, [{"a"}, {}]>] out_shardings=[<@mesh_0, [{"a"}, {}]>] manual_axes={"a"} (%arg1: tensor<2x8xf32>) { - sdy.return %arg1 : tensor<2x8xf32> - } : (tensor<8x8xf32>) -> tensor<8x8xf32> - - %1 = sdy.manual_computation(%0) in_shardings=[<@mesh_0, [{}, {"b"}]>] out_shardings=[<@mesh_0, [{}, {"b"}]>] manual_axes={"b"} (%arg1: tensor<8x4xf32>) { - sdy.return %arg1 : tensor<8x4xf32> - } : (tensor<8x8xf32>) -> tensor<8x8xf32> - return %1 : tensor<8x8xf32> -} - -// Tests nested manual computations: an inner shard_map inside an outer shard_map body. -// CHECK-LABEL: func @nested_shmaps -func.func @nested_shmaps(%arg0: tensor<4x8xf32> {sdy.sharding = #sdy.sharding<@mesh_1, [{"a"}, {"b"}]>}) -> (tensor<4x8xf32> {sdy.sharding = #sdy.sharding<@mesh_1, [{"a", ?}, {?}]>}) { - // CHECK-NEXT: %[[SHARDING:.*]] = stablehlo.custom_call @Sharding(%arg0) - // CHECK-SAME{LITERAL}: sdy.sharding = #sdy.sharding_per_value<[, [{"a"}, {}]>]> - // CHECK-NEXT: %[[GLOBAL_TO_LOCAL:.*]] = stablehlo.custom_call @xla.sdy.GlobalToLocalShape(%[[SHARDING]]) - // CHECK-SAME{LITERAL}: sdy.sharding = #sdy.sharding_per_value<[, [{}, {}]>]> - // CHECK-SAME{LITERAL}: xla.sdy.manual_axes = #sdy - // CHECK-SAME: (tensor<4x8xf32>) -> tensor<2x8xf32> - // CHECK-NEXT: %[[SHMAP:.*]] = call @xla.sdy.manual_computation_body_3(%[[GLOBAL_TO_LOCAL]]) - // CHECK-SAME: {mhlo.frontend_attributes = {inlineable = "xla_late"} - // CHECK-SAME{LITERAL}: xla.sdy.manual_axes = #sdy - // CHECK-SAME: (tensor<2x8xf32>) -> tensor<2x8xf32> - // CHECK-NEXT: %[[LOCAL_TO_GLOBAL:.*]] = stablehlo.custom_call @xla.sdy.LocalToGlobalShape(%[[SHMAP]]) - // CHECK-SAME{LITERAL}: sdy.sharding = #sdy.sharding_per_value<[, [{"a"}, {}]>]> - // CHECK-SAME: (tensor<2x8xf32>) -> tensor<4x8xf32> - // CHECK-NEXT: return %[[LOCAL_TO_GLOBAL]] : tensor<4x8xf32> - %0 = sdy.manual_computation(%arg0) in_shardings=[<@mesh_1, [{"a"}, {}]>] out_shardings=[<@mesh_1, [{"a"}, {}]>] manual_axes={"a"} (%arg1: tensor<2x8xf32>) { - %1 = sdy.manual_computation(%arg1) in_shardings=[<@mesh_1, [{}, {"b"}]>] out_shardings=[<@mesh_1, [{}, {"b"}]>] manual_axes={"b"} (%arg2: tensor<2x4xf32>) { - %2 = stablehlo.multiply %arg2, %arg2 : tensor<2x4xf32> - sdy.return %2 : tensor<2x4xf32> - } : (tensor<2x8xf32>) -> tensor<2x8xf32> - sdy.return %1 : tensor<2x8xf32> - } : (tensor<4x8xf32>) -> tensor<4x8xf32> - return %0 : tensor<4x8xf32> -} - -// Tests nested manual computation where the outer body performs additional math on the inner result. -// CHECK-LABEL: func @nested_shmaps_extra_op -func.func @nested_shmaps_extra_op(%arg0: tensor<4x8xf32> {sdy.sharding = #sdy.sharding<@mesh_1, [{"a"}, {"b"}]>}) -> (tensor<4x8xf32> {sdy.sharding = #sdy.sharding<@mesh_1, [{"a", ?}, {?}]>}) { - // CHECK-NEXT: %[[SHARDING:.*]] = stablehlo.custom_call @Sharding(%arg0) - // CHECK-SAME{LITERAL}: sdy.sharding = #sdy.sharding_per_value<[, [{"a"}, {}]>]> - // CHECK-NEXT: %[[GLOBAL_TO_LOCAL:.*]] = stablehlo.custom_call @xla.sdy.GlobalToLocalShape(%[[SHARDING]]) - // CHECK-SAME{LITERAL}: sdy.sharding = #sdy.sharding_per_value<[, [{}, {}]>]> - // CHECK-SAME{LITERAL}: xla.sdy.manual_axes = #sdy - // CHECK-SAME: (tensor<4x8xf32>) -> tensor<2x8xf32> - // CHECK-NEXT: %[[SHMAP:.*]] = call @xla.sdy.manual_computation_body_5(%[[GLOBAL_TO_LOCAL]]) - // CHECK-SAME: {mhlo.frontend_attributes = {inlineable = "xla_late"} - // CHECK-SAME{LITERAL}: xla.sdy.manual_axes = #sdy - // CHECK-SAME: (tensor<2x8xf32>) -> tensor<2x8xf32> - // CHECK-NEXT: %[[LOCAL_TO_GLOBAL:.*]] = stablehlo.custom_call @xla.sdy.LocalToGlobalShape(%[[SHMAP]]) - // CHECK-SAME{LITERAL}: sdy.sharding = #sdy.sharding_per_value<[, [{"a"}, {}]>]> - // CHECK-SAME: (tensor<2x8xf32>) -> tensor<4x8xf32> - // CHECK-NEXT: return %[[LOCAL_TO_GLOBAL]] : tensor<4x8xf32> - %0 = sdy.manual_computation(%arg0) in_shardings=[<@mesh_1, [{"a"}, {}]>] out_shardings=[<@mesh_1, [{"a"}, {}]>] manual_axes={"a"} (%arg1: tensor<2x8xf32>) { - %1 = sdy.manual_computation(%arg1) in_shardings=[<@mesh_1, [{}, {"b"}]>] out_shardings=[<@mesh_1, [{}, {"b"}]>] manual_axes={"b"} (%arg2: tensor<2x4xf32>) { - %2 = stablehlo.multiply %arg2, %arg2 : tensor<2x4xf32> - sdy.return %2 : tensor<2x4xf32> - } : (tensor<2x8xf32>) -> tensor<2x8xf32> - %3 = stablehlo.add %1, %1 : tensor<2x8xf32> - sdy.return %3 : tensor<2x8xf32> - } : (tensor<4x8xf32>) -> tensor<4x8xf32> - return %0 : tensor<4x8xf32> -} - -// Edge case: Manual computation producing an output but taking no inputs. There -// is no `GlobalToLocalShape`, so the body call op is the only carrier of the -// manual axes. -// CHECK-LABEL: func @manual_computation_no_inputs -func.func @manual_computation_no_inputs() -> tensor<4xi64> { - // CHECK-NEXT: %[[SHMAP:.*]] = call @xla.sdy.manual_computation_body_6() - // CHECK-SAME: {mhlo.frontend_attributes = {inlineable = "xla_late"} - // CHECK-SAME{LITERAL}: sdy.sharding = #sdy.sharding_per_value<[, [{}]>]> - // CHECK-SAME{LITERAL}: xla.sdy.manual_axes = #sdy - // CHECK-SAME: () -> tensor<2xi64> - // CHECK-NEXT: %[[LOCAL_TO_GLOBAL:.*]] = stablehlo.custom_call @xla.sdy.LocalToGlobalShape(%[[SHMAP]]) - // CHECK-SAME{LITERAL}: sdy.sharding = #sdy.sharding_per_value<[, [{"b"}]>]> - // CHECK-SAME: (tensor<2xi64>) -> tensor<4xi64> - // CHECK-NEXT: return %[[LOCAL_TO_GLOBAL]] : tensor<4xi64> - %0 = sdy.manual_computation() in_shardings=[] out_shardings=[<@mesh_0, [{"b"}]>] manual_axes={"b"} () { - %1 = stablehlo.constant dense<[2, 3]> : tensor<2xi64> - sdy.return %1 : tensor<2xi64> - } : () -> tensor<4xi64> - func.return %0 : tensor<4xi64> -} - -// Edge case: Manual computation consuming an input but producing no outputs. -// The body call op has no results, so `GlobalToLocalShape` is the only carrier -// of the manual axes. -// CHECK-LABEL: func @manual_computation_no_outputs -func.func @manual_computation_no_outputs(%arg0: tensor<4xi64>) { - // CHECK-NEXT: %[[SHARDING:.*]] = stablehlo.custom_call @Sharding(%arg0) - // CHECK-SAME{LITERAL}: sdy.sharding = #sdy.sharding_per_value<[, [{"b"}]>]> - // CHECK-NEXT: %[[GLOBAL_TO_LOCAL:.*]] = stablehlo.custom_call @xla.sdy.GlobalToLocalShape(%[[SHARDING]]) - // CHECK-SAME{LITERAL}: sdy.sharding = #sdy.sharding_per_value<[, [{}]>]> - // CHECK-SAME{LITERAL}: xla.sdy.manual_axes = #sdy - // CHECK-SAME: (tensor<4xi64>) -> tensor<2xi64> - // CHECK-NEXT: call @xla.sdy.manual_computation_body_7(%[[GLOBAL_TO_LOCAL]]) - // CHECK-SAME: {mhlo.frontend_attributes = {inlineable = "xla_late"}} - // CHECK-SAME: (tensor<2xi64>) -> () - // CHECK-NEXT: return - sdy.manual_computation(%arg0) in_shardings=[<@mesh_0, [{"b"}]>] out_shardings=[] manual_axes={"b"} (%arg1: tensor<2xi64>) { - stablehlo.custom_call @sdy_testonly(%arg1) : (tensor<2xi64>) -> () - sdy.return - } : (tensor<4xi64>) -> () - func.return -} - -// Edge case: Degenerate manual computation with no inputs and no outputs. -// CHECK-LABEL: func @manual_computation_no_inputs_no_outputs -func.func @manual_computation_no_inputs_no_outputs() { - // CHECK-NEXT: call @xla.sdy.manual_computation_body_8() {mhlo.frontend_attributes = {inlineable = "xla_late"}} : () -> () - sdy.manual_computation() in_shardings=[] out_shardings=[] manual_axes={} () { - sdy.return - } : () -> () - func.return -} - -// Type preservation: Token arguments should not have manual axes attached in V3. -// CHECK-LABEL: func @manual_computation_tokens -func.func @manual_computation_tokens(%arg0: !stablehlo.token, %arg1: tensor<2xi64>) -> (!stablehlo.token, tensor<2xi64>) { - // CHECK-NEXT: %[[SHARDING_0:.*]] = stablehlo.custom_call @Sharding(%arg0) - // CHECK-SAME{LITERAL}: sdy.sharding = #sdy.sharding_per_value<[, []>]> - // CHECK-SAME: (!stablehlo.token) -> !stablehlo.token - // CHECK-NEXT: %[[SHARDING_1:.*]] = stablehlo.custom_call @Sharding(%arg1) - // CHECK-SAME{LITERAL}: sdy.sharding = #sdy.sharding_per_value<[, [{"b"}]>]> - // CHECK-SAME: (tensor<2xi64>) -> tensor<2xi64> - // CHECK-NEXT: %[[GLOBAL_TO_LOCAL:.*]]:2 = stablehlo.custom_call @xla.sdy.GlobalToLocalShape(%[[SHARDING_0]], %[[SHARDING_1]]) - // CHECK-SAME{LITERAL}: sdy.sharding = #sdy.sharding_per_value<[, []>, , [{}]>]> - // CHECK-SAME{LITERAL}: xla.sdy.manual_axes = #sdy - // CHECK-SAME: (!stablehlo.token, tensor<2xi64>) -> (!stablehlo.token, tensor<1xi64>) - // CHECK-NEXT: %[[SHMAP:.*]]:2 = call @xla.sdy.manual_computation_body_9(%[[GLOBAL_TO_LOCAL]]#0, %[[GLOBAL_TO_LOCAL]]#1) - // CHECK-SAME: {mhlo.frontend_attributes = {inlineable = "xla_late"} - // CHECK-SAME{LITERAL}: xla.sdy.manual_axes = #sdy - // CHECK-SAME: (!stablehlo.token, tensor<1xi64>) -> (!stablehlo.token, tensor<1xi64>) - // CHECK-NEXT: %[[LOCAL_TO_GLOBAL:.*]]:2 = stablehlo.custom_call @xla.sdy.LocalToGlobalShape(%[[SHMAP]]#0, %[[SHMAP]]#1) - // CHECK-SAME{LITERAL}: sdy.sharding = #sdy.sharding_per_value<[, []>, , [{"b"}]>]> - // CHECK-SAME: (!stablehlo.token, tensor<1xi64>) -> (!stablehlo.token, tensor<2xi64>) - // CHECK-NEXT: return %[[LOCAL_TO_GLOBAL]]#0, %[[LOCAL_TO_GLOBAL]]#1 : !stablehlo.token, tensor<2xi64> - %0:2 = sdy.manual_computation(%arg0, %arg1) in_shardings=[<@mesh_0, []>, <@mesh_0, [{"b"}]>] out_shardings=[<@mesh_0, []>, <@mesh_0, [{"b"}]>] manual_axes={"b"} (%arg2: !stablehlo.token, %arg3: tensor<1xi64>) { - sdy.return %arg2, %arg3 : !stablehlo.token, tensor<1xi64> - } : (!stablehlo.token, tensor<2xi64>) -> (!stablehlo.token, tensor<2xi64>) - func.return %0#0, %0#1 : !stablehlo.token, tensor<2xi64> -} - - -// CHECK-LABEL: func @xla.sdy.manual_computation_body_3(%arg0: tensor<2x8xf32>) -> tensor<2x8xf32> -// CHECK-NEXT: %[[SHARDING:.*]] = stablehlo.custom_call @Sharding(%arg0) -// CHECK-SAME{LITERAL}: sdy.sharding = #sdy.sharding_per_value<[, [{}, {"b"}]>]> -// CHECK-NEXT: %[[GLOBAL_TO_LOCAL:.*]] = stablehlo.custom_call @xla.sdy.GlobalToLocalShape(%[[SHARDING]]) -// CHECK-SAME{LITERAL}: sdy.sharding = #sdy.sharding_per_value<[, [{}, {}]>]> -// CHECK-SAME{LITERAL}: xla.sdy.manual_axes = #sdy -// CHECK-SAME: (tensor<2x8xf32>) -> tensor<2x4xf32> -// CHECK-NEXT: %[[SHMAP:.*]] = call @xla.sdy.manual_computation_body_2(%[[GLOBAL_TO_LOCAL]]) -// CHECK-SAME: {mhlo.frontend_attributes = {inlineable = "xla_late"} -// CHECK-SAME{LITERAL}: xla.sdy.manual_axes = #sdy -// CHECK-SAME: (tensor<2x4xf32>) -> tensor<2x4xf32> -// CHECK-NEXT: %[[LOCAL_TO_GLOBAL:.*]] = stablehlo.custom_call @xla.sdy.LocalToGlobalShape(%[[SHMAP]]) -// CHECK-SAME{LITERAL}: sdy.sharding = #sdy.sharding_per_value<[, [{}, {"b"}]>]> -// CHECK-SAME: (tensor<2x4xf32>) -> tensor<2x8xf32> -// CHECK-NEXT: return %[[LOCAL_TO_GLOBAL]] : tensor<2x8xf32> - -// CHECK-LABEL: func @xla.sdy.manual_computation_body_5(%arg0: tensor<2x8xf32>) -> tensor<2x8xf32> -// CHECK-NEXT: %[[SHARDING:.*]] = stablehlo.custom_call @Sharding(%arg0) -// CHECK-SAME{LITERAL}: sdy.sharding = #sdy.sharding_per_value<[, [{}, {"b"}]>]> -// CHECK-NEXT: %[[GLOBAL_TO_LOCAL:.*]] = stablehlo.custom_call @xla.sdy.GlobalToLocalShape(%[[SHARDING]]) -// CHECK-SAME{LITERAL}: sdy.sharding = #sdy.sharding_per_value<[, [{}, {}]>]> -// CHECK-SAME{LITERAL}: xla.sdy.manual_axes = #sdy -// CHECK-SAME: (tensor<2x8xf32>) -> tensor<2x4xf32> -// CHECK-NEXT: %[[SHMAP:.*]] = call @xla.sdy.manual_computation_body_4(%[[GLOBAL_TO_LOCAL]]) -// CHECK-SAME: {mhlo.frontend_attributes = {inlineable = "xla_late"} -// CHECK-SAME{LITERAL}: xla.sdy.manual_axes = #sdy -// CHECK-SAME: (tensor<2x4xf32>) -> tensor<2x4xf32> -// CHECK-NEXT: %[[LOCAL_TO_GLOBAL:.*]] = stablehlo.custom_call @xla.sdy.LocalToGlobalShape(%[[SHMAP]]) -// CHECK-SAME{LITERAL}: sdy.sharding = #sdy.sharding_per_value<[, [{}, {"b"}]>]> -// CHECK-SAME: (tensor<2x4xf32>) -> tensor<2x8xf32> -// CHECK-NEXT: %[[ADD:.*]] = stablehlo.add %[[LOCAL_TO_GLOBAL]], %[[LOCAL_TO_GLOBAL]] : tensor<2x8xf32> -// CHECK-NEXT: return %[[ADD]] : tensor<2x8xf32> diff --git a/third_party/xla/xla/service/spmd/shardy/test/sdy_round_trip_shard_map_import_failure.mlir b/third_party/xla/xla/service/spmd/shardy/test/sdy_round_trip_shard_map_import_failure.mlir index 15b7eec24bdca0..47d1ba107a8bf5 100644 --- a/third_party/xla/xla/service/spmd/shardy/test/sdy_round_trip_shard_map_import_failure.mlir +++ b/third_party/xla/xla/service/spmd/shardy/test/sdy_round_trip_shard_map_import_failure.mlir @@ -45,38 +45,3 @@ func.func @xla.sdy.manual_computation_body(%arg0: tensor<0x8xf32>) -> tensor<0x8 return %arg0 : tensor<0x8xf32> } -// ----- - -sdy.mesh @mesh = <["a"=2, "b"=2]> - -// Tests import error when manual computation body call operand is not produced by GlobalToLocalShape. -func.func @manual_computation_missing_global_to_local_shape_direct_attrs(%arg0: tensor<0x16xf32>) -> (tensor<0x16xf32>) { - %c = stablehlo.constant dense<0.000000e+00> : tensor<0x8xf32> - // expected-error @+2 {{'func.call' op expected at least one operand of @xla.sdy.manual_computation_body to be produced by a xla.sdy.GlobalToLocalShape CustomCallOp}} - // expected-error @+1 {{failed to rewrite func.call to manual computation}} - %0 = call @xla.sdy.manual_computation_body(%c) : (tensor<0x8xf32>) -> tensor<0x8xf32> - %1 = stablehlo.custom_call @xla.sdy.LocalToGlobalShape(%0) {has_side_effect = true, sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{}, {"b"}]>]>, xla.sdy.manual_axes = #sdy} : (tensor<0x8xf32>) -> (tensor<0x16xf32>) - return %1 : tensor<0x16xf32> -} - -func.func @xla.sdy.manual_computation_body(%arg0: tensor<0x8xf32>) -> tensor<0x8xf32> { - return %arg0 : tensor<0x8xf32> -} - -// ----- - -sdy.mesh @mesh = <["a"=2, "b"=2]> - -// Tests import error when manual computation body call result is not consumed by LocalToGlobalShape. -func.func @manual_computation_missing_local_to_global_shape_direct_attrs(%arg0: tensor<0x16xf32>) -> (tensor<0x8xf32>) { - %0 = stablehlo.custom_call @xla.sdy.GlobalToLocalShape(%arg0) {has_side_effect = true, sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{}, {}]>]>, xla.sdy.manual_axes = #sdy} : (tensor<0x16xf32>) -> tensor<0x8xf32> - // expected-error @+2 {{'func.call' op expected the first use of @xla.sdy.manual_computation_body to be a xla.sdy.LocalToGlobalShape CustomCallOp}} - // expected-error @+1 {{failed to rewrite func.call to manual computation}} - %1 = call @xla.sdy.manual_computation_body(%0) : (tensor<0x8xf32>) -> tensor<0x8xf32> - return %1 : tensor<0x8xf32> -} - -func.func @xla.sdy.manual_computation_body(%arg0: tensor<0x8xf32>) -> tensor<0x8xf32> { - return %arg0 : tensor<0x8xf32> -} - diff --git a/third_party/xla/xla/service/spmd/shardy/test/sdy_round_trip_shard_map_import_hlo_sharding_v3.mlir b/third_party/xla/xla/service/spmd/shardy/test/sdy_round_trip_shard_map_import_hlo_sharding_v3.mlir deleted file mode 100644 index 6e7060cb48ff30..00000000000000 --- a/third_party/xla/xla/service/spmd/shardy/test/sdy_round_trip_shard_map_import_hlo_sharding_v3.mlir +++ /dev/null @@ -1,362 +0,0 @@ -// 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. -// ============================================================================== -// RUN: sdy_opt %s --split-input-file -xla-sdy-test-flatten-call-graph -xla-sdy-round-trip-shard-map-import 2>&1 | FileCheck %s - -// Tests importing shard map custom calls back into `sdy.manual_computation` under -// HloShardingV3 (where frontend_attributes are omitted). -// Verifies that: -// 1. Shardings and manual axes are read directly off the native attributes of the -// boundary ops, without any inference: -// - global `in_shardings` come from the `@Sharding` custom call wrapping each -// operand of `GlobalToLocalShape`, which is then unwired and erased; -// - global `out_shardings` come from `LocalToGlobalShape`; -// - `manual_axes` come from `GlobalToLocalShape`, or from the body call op -// when the manual computation has no operands. -// 2. Outlined manual computation call bodies are inlined back into `sdy.manual_computation`. - -sdy.mesh @mesh_0 = <["a"=4, "b"=2]> -sdy.mesh @mesh_1 = <["a"=2, "b"=2, "c"=2, "d"=2]> - -// Tests single manual computation with multiple inputs/outputs, all-reduce, and -// in/out shardings read off the `@Sharding` wrappers and `LocalToGlobalShape`. -// CHECK-LABEL: func @single_manual_comp -func.func @single_manual_comp(%arg0: tensor<8x16xf32>, %arg1: tensor<16x32xf32>) -> (tensor<8x32xf32>) { - // CHECK-NOT: call @xla.sdy.manual_computation_body - // CHECK: %[[MAN_COMP:.*]] = sdy.manual_computation(%arg0, %arg1) - // CHECK-SAME{LITERAL}: in_shardings=[<@mesh_0, [{"a"}, {"b"}]>, <@mesh_0, [{"b"}, {}], replicated={"a"}>] - // CHECK-SAME{LITERAL}: out_shardings=[<@mesh_0, [{"a"}, {}], replicated={"b"}>] - - // CHECK-SAME{LITERAL}: manual_axes={"a", "b"} - // CHECK-SAME: (%arg2: tensor<2x8xf32>, %arg3: tensor<8x32xf32>) { - // CHECK-NEXT: %[[ADD_0:.*]] = stablehlo.add %arg2, %arg2 : tensor<2x8xf32> - // CHECK-NEXT: %[[DOT:.*]] = stablehlo.dot %[[ADD_0]], %arg3 : (tensor<2x8xf32>, tensor<8x32xf32>) -> tensor<2x32xf32> - // CHECK-NEXT: %[[REDUCE:.*]] = "stablehlo.all_reduce"(%[[DOT]]) - // CHECK-NEXT: ^bb0(%arg4: tensor, %arg5: tensor): - // CHECK-NEXT: %[[ADD_1:.*]] = stablehlo.add %arg4, %arg5 : tensor - // CHECK-NEXT: stablehlo.return %[[ADD_1]] : tensor - // CHECK-NEXT: }) : (tensor<2x32xf32>) -> tensor<2x32xf32> - // CHECK-NEXT: sdy.return %[[REDUCE]] : tensor<2x32xf32> - // CHECK-NEXT: } : (tensor<8x16xf32>, tensor<16x32xf32>) -> tensor<8x32xf32> - // CHECK-NEXT: return %[[MAN_COMP]] : tensor<8x32xf32> - %s0 = stablehlo.custom_call @Sharding(%arg0) {sdy.sharding = #sdy.sharding_per_value<[<@mesh_0, [{"a"}, {"b"}]>]>} : (tensor<8x16xf32>) -> tensor<8x16xf32> - %s1 = stablehlo.custom_call @Sharding(%arg1) {sdy.sharding = #sdy.sharding_per_value<[<@mesh_0, [{"b"}, {}], replicated={"a"}>]>} : (tensor<16x32xf32>) -> tensor<16x32xf32> - %0:2 = stablehlo.custom_call @xla.sdy.GlobalToLocalShape(%s0, %s1) {has_side_effect = true, sdy.sharding = #sdy.sharding_per_value<[<@mesh_0, [{}, {}]>, <@mesh_0, [{}, {}]>]>, xla.sdy.manual_axes = #sdy} : (tensor<8x16xf32>, tensor<16x32xf32>) -> (tensor<2x8xf32>, tensor<8x32xf32>) - %1 = call @xla.sdy.manual_computation_body(%0#0, %0#1) : (tensor<2x8xf32>, tensor<8x32xf32>) -> tensor<2x32xf32> - %2 = stablehlo.custom_call @xla.sdy.LocalToGlobalShape(%1) {has_side_effect = true, sdy.sharding = #sdy.sharding_per_value<[<@mesh_0, [{"a"}, {}], replicated={"b"}>]>} : (tensor<2x32xf32>) -> tensor<8x32xf32> - return %2 : tensor<8x32xf32> -} - -// CHECK-LABEL: func @single_manual_comp_name_is_not_prefix_nor_suffix -func.func @single_manual_comp_name_is_not_prefix_nor_suffix(%arg0: tensor<8x8xf32>) -> (tensor<8x8xf32>) { - // CHECK-NOT: call @my_model.___call__.fwd.xla.sdy.manual_computation_body_14.1234 - // CHECK: %[[MAN_COMP:.*]] = sdy.manual_computation(%arg0) - // CHECK-SAME{LITERAL}: in_shardings=[<@mesh_0, [{"a"}, {}]>] - // CHECK-SAME{LITERAL}: out_shardings=[<@mesh_0, [{"a"}, {}]>] - // CHECK-SAME{LITERAL}: manual_axes={"a"} - // CHECK-SAME: (%arg1: tensor<2x8xf32>) { - // CHECK-NEXT: sdy.return %arg1 : tensor<2x8xf32> - // CHECK-NEXT: } : (tensor<8x8xf32>) -> tensor<8x8xf32> - // CHECK-NEXT: return %[[MAN_COMP]] : tensor<8x8xf32> - %s0 = stablehlo.custom_call @Sharding(%arg0) {sdy.sharding = #sdy.sharding_per_value<[<@mesh_0, [{"a"}, {}]>]>} : (tensor<8x8xf32>) -> tensor<8x8xf32> - %0 = stablehlo.custom_call @xla.sdy.GlobalToLocalShape(%s0) {has_side_effect = true, sdy.sharding = #sdy.sharding_per_value<[<@mesh_0, [{}, {}]>]>, xla.sdy.manual_axes = #sdy} : (tensor<8x8xf32>) -> tensor<2x8xf32> - %1 = call @my_model.___call__.fwd.xla.sdy.manual_computation_body_14.1234(%0) : (tensor<2x8xf32>) -> tensor<2x8xf32> - %2 = stablehlo.custom_call @xla.sdy.LocalToGlobalShape(%1) {has_side_effect = true, sdy.sharding = #sdy.sharding_per_value<[<@mesh_0, [{"a"}, {}]>]>} : (tensor<2x8xf32>) -> tensor<8x8xf32> - return %2 : tensor<8x8xf32> -} - -// Tests chained manual computations: outputs of first shard map feed second shard map. -// CHECK-LABEL: func @manual_comp_using_another -func.func @manual_comp_using_another(%arg0: tensor<8x8xf32>) -> tensor<8x8xf32> { - // CHECK-NOT: call @xla.sdy.manual_computation_body_0 - // CHECK: %[[MAN_COMP_0:.*]] = sdy.manual_computation(%arg0) - // CHECK-SAME{LITERAL}: in_shardings=[<@mesh_0, [{"a"}, {}]>] - // CHECK-SAME{LITERAL}: out_shardings=[<@mesh_0, [{"a"}, {}]>] - // CHECK-SAME{LITERAL}: manual_axes={"a"} - // CHECK-SAME: (%arg1: tensor<2x8xf32>) { - // CHECK-NEXT: sdy.return %arg1 : tensor<2x8xf32> - // CHECK-NEXT: } : (tensor<8x8xf32>) -> tensor<8x8xf32> - // CHECK-NOT: call @xla.sdy.manual_computation_body_1 - // CHECK-NEXT: %[[MAN_COMP_1:.*]] = sdy.manual_computation(%[[MAN_COMP_0]]) - // CHECK-SAME{LITERAL}: in_shardings=[<@mesh_0, [{}, {"b"}]>] - // CHECK-SAME{LITERAL}: out_shardings=[<@mesh_0, [{}, {"b"}]>] - // CHECK-SAME{LITERAL}: manual_axes={"b"} - // CHECK-SAME: (%arg1: tensor<8x4xf32>) { - // CHECK-NEXT: sdy.return %arg1 : tensor<8x4xf32> - // CHECK-NEXT: } : (tensor<8x8xf32>) -> tensor<8x8xf32> - // CHECK-NEXT: return %[[MAN_COMP_1]] : tensor<8x8xf32> - %s0 = stablehlo.custom_call @Sharding(%arg0) {sdy.sharding = #sdy.sharding_per_value<[<@mesh_0, [{"a"}, {}]>]>} : (tensor<8x8xf32>) -> tensor<8x8xf32> - %0 = stablehlo.custom_call @xla.sdy.GlobalToLocalShape(%s0) {has_side_effect = true, sdy.sharding = #sdy.sharding_per_value<[<@mesh_0, [{}, {}]>]>, xla.sdy.manual_axes = #sdy} : (tensor<8x8xf32>) -> tensor<2x8xf32> - %1 = call @xla.sdy.manual_computation_body_0(%0) : (tensor<2x8xf32>) -> tensor<2x8xf32> - %2 = stablehlo.custom_call @xla.sdy.LocalToGlobalShape(%1) {has_side_effect = true, sdy.sharding = #sdy.sharding_per_value<[<@mesh_0, [{"a"}, {}]>]>} : (tensor<2x8xf32>) -> tensor<8x8xf32> - %s1 = stablehlo.custom_call @Sharding(%2) {sdy.sharding = #sdy.sharding_per_value<[<@mesh_0, [{}, {"b"}]>]>} : (tensor<8x8xf32>) -> tensor<8x8xf32> - %3 = stablehlo.custom_call @xla.sdy.GlobalToLocalShape(%s1) {has_side_effect = true, sdy.sharding = #sdy.sharding_per_value<[<@mesh_0, [{}, {}]>]>, xla.sdy.manual_axes = #sdy} : (tensor<8x8xf32>) -> tensor<8x4xf32> - %4 = call @xla.sdy.manual_computation_body_1(%3) : (tensor<8x4xf32>) -> tensor<8x4xf32> - %5 = stablehlo.custom_call @xla.sdy.LocalToGlobalShape(%4) {has_side_effect = true, sdy.sharding = #sdy.sharding_per_value<[<@mesh_0, [{}, {"b"}]>]>} : (tensor<8x4xf32>) -> tensor<8x8xf32> - return %5 : tensor<8x8xf32> -} - -// Tests nested manual computations: reconstructing inner and outer sdy.manual_computation regions. -// CHECK-LABEL: func @nested_shmaps -func.func @nested_shmaps(%arg0: tensor<4x8xf32>) -> tensor<4x8xf32> { - // CHECK-NOT: call @xla.sdy.manual_computation_body_2 - // CHECK-NOT: call @xla.sdy.manual_computation_body_3 - // CHECK: %[[MAN_COMP_0:.*]] = sdy.manual_computation(%arg0) - // CHECK-SAME{LITERAL}: in_shardings=[<@mesh_1, [{"a"}, {}]>] - // CHECK-SAME{LITERAL}: out_shardings=[<@mesh_1, [{"a"}, {}]>] - // CHECK-SAME{LITERAL}: manual_axes={"a"} - // CHECK-SAME: (%arg1: tensor<2x8xf32>) { - // CHECK-NEXT: %[[MAN_COMP_1:.*]] = sdy.manual_computation(%arg1) - // CHECK-SAME{LITERAL}: in_shardings=[<@mesh_1, [{}, {"b"}]>] - // CHECK-SAME{LITERAL}: out_shardings=[<@mesh_1, [{}, {"b"}]>] - // CHECK-SAME{LITERAL}: manual_axes={"b"} - // CHECK-SAME: (%arg2: tensor<2x4xf32>) { - // CHECK-NEXT: %[[MUL:.*]] = stablehlo.multiply %arg2, %arg2 : tensor<2x4xf32> - // CHECK-NEXT: sdy.return %[[MUL]] : tensor<2x4xf32> - // CHECK-NEXT: } : (tensor<2x8xf32>) -> tensor<2x8xf32> - // CHECK-NEXT: sdy.return %[[MAN_COMP_1]] : tensor<2x8xf32> - // CHECK-NEXT: } : (tensor<4x8xf32>) -> tensor<4x8xf32> - // CHECK-NEXT: return %[[MAN_COMP_0]] : tensor<4x8xf32> - %s0 = stablehlo.custom_call @Sharding(%arg0) {sdy.sharding = #sdy.sharding_per_value<[<@mesh_1, [{"a"}, {}]>]>} : (tensor<4x8xf32>) -> tensor<4x8xf32> - %0 = stablehlo.custom_call @xla.sdy.GlobalToLocalShape(%s0) {has_side_effect = true, sdy.sharding = #sdy.sharding_per_value<[<@mesh_1, [{}, {}]>]>, xla.sdy.manual_axes = #sdy} : (tensor<4x8xf32>) -> tensor<2x8xf32> - %1 = call @xla.sdy.manual_computation_body_3(%0) : (tensor<2x8xf32>) -> tensor<2x8xf32> - %2 = stablehlo.custom_call @xla.sdy.LocalToGlobalShape(%1) {has_side_effect = true, sdy.sharding = #sdy.sharding_per_value<[<@mesh_1, [{"a"}, {}]>]>} : (tensor<2x8xf32>) -> tensor<4x8xf32> - return %2 : tensor<4x8xf32> -} - -// Tests nested manual computation with operations in the outer region around the inner region. -// CHECK-LABEL: func @nested_shmaps_extra_op -func.func @nested_shmaps_extra_op(%arg0: tensor<4x8xf32>) -> tensor<4x8xf32> { - // CHECK-NOT: call @xla.sdy.manual_computation_body_4 - // CHECK-NOT: call @xla.sdy.manual_computation_body_5 - // CHECK: %[[MAN_COMP_0:.*]] = sdy.manual_computation(%arg0) - // CHECK-SAME{LITERAL}: in_shardings=[<@mesh_1, [{"a"}, {}]>] - // CHECK-SAME{LITERAL}: out_shardings=[<@mesh_1, [{"a"}, {}]>] - // CHECK-SAME{LITERAL}: manual_axes={"a"} - // CHECK-SAME: (%arg1: tensor<2x8xf32>) { - // CHECK-NEXT: %[[MAN_COMP_1:.*]] = sdy.manual_computation(%arg1) - // CHECK-SAME{LITERAL}: in_shardings=[<@mesh_1, [{}, {"b"}]>] - // CHECK-SAME{LITERAL}: out_shardings=[<@mesh_1, [{}, {"b"}]>] - // CHECK-SAME{LITERAL}: manual_axes={"b"} - // CHECK-SAME: (%arg2: tensor<2x4xf32>) { - // CHECK-NEXT: %[[MUL:.*]] = stablehlo.multiply %arg2, %arg2 : tensor<2x4xf32> - // CHECK-NEXT: sdy.return %[[MUL]] : tensor<2x4xf32> - // CHECK-NEXT: } : (tensor<2x8xf32>) -> tensor<2x8xf32> - // CHECK-NEXT: %[[ADD:.*]] = stablehlo.add %[[MAN_COMP_1]], %[[MAN_COMP_1]] : tensor<2x8xf32> - // CHECK-NEXT: sdy.return %[[ADD]] : tensor<2x8xf32> - // CHECK-NEXT: } : (tensor<4x8xf32>) -> tensor<4x8xf32> - // CHECK-NEXT: return %[[MAN_COMP_0]] : tensor<4x8xf32> - %s0 = stablehlo.custom_call @Sharding(%arg0) {sdy.sharding = #sdy.sharding_per_value<[<@mesh_1, [{"a"}, {}]>]>} : (tensor<4x8xf32>) -> tensor<4x8xf32> - %0 = stablehlo.custom_call @xla.sdy.GlobalToLocalShape(%s0) {has_side_effect = true, sdy.sharding = #sdy.sharding_per_value<[<@mesh_1, [{}, {}]>]>, xla.sdy.manual_axes = #sdy} : (tensor<4x8xf32>) -> tensor<2x8xf32> - %1 = call @xla.sdy.manual_computation_body_5(%0) : (tensor<2x8xf32>) -> tensor<2x8xf32> - %2 = stablehlo.custom_call @xla.sdy.LocalToGlobalShape(%1) {has_side_effect = true, sdy.sharding = #sdy.sharding_per_value<[<@mesh_1, [{"a"}, {}]>]>} : (tensor<2x8xf32>) -> tensor<4x8xf32> - return %2 : tensor<4x8xf32> -} - -// Tests manual computation with no inputs: there is no `GlobalToLocalShape`, so -// the body call op is the only carrier of the manual axes. -// CHECK-LABEL: func @manual_computation_no_inputs -func.func @manual_computation_no_inputs() -> tensor<4xi64> { - - // CHECK-NOT: call @xla.sdy.manual_computation_body_6 - // CHECK: %[[SHMAP:.*]] = sdy.manual_computation() - // CHECK-SAME{LITERAL}: in_shardings=[] - // CHECK-SAME{LITERAL}: out_shardings=[<@mesh_0, [{"b"}]>] - // CHECK-SAME{LITERAL}: manual_axes={"b"} - // CHECK-SAME{LITERAL}: () { - // CHECK-NEXT: %[[C:.*]] = stablehlo.constant dense<[2, 3]> : tensor<2xi64> - // CHECK-NEXT: sdy.return %[[C]] : tensor<2xi64> - // CHECK-NEXT: } : () -> tensor<4xi64> - // CHECK-NEXT: return %[[SHMAP]] : tensor<4xi64> - %0 = call @xla.sdy.manual_computation_body_6() {sdy.sharding = #sdy.sharding_per_value<[<@mesh_0, [{}]>]>, xla.sdy.manual_axes = #sdy} : () -> tensor<2xi64> - %1 = stablehlo.custom_call @xla.sdy.LocalToGlobalShape(%0) {has_side_effect = true, sdy.sharding = #sdy.sharding_per_value<[<@mesh_0, [{"b"}]>]>} : (tensor<2xi64>) -> tensor<4xi64> - return %1 : tensor<4xi64> -} - -// Edge case: Manual computation consuming an input but producing no outputs. -// CHECK-LABEL: func @manual_computation_no_outputs -func.func @manual_computation_no_outputs(%arg0: tensor<4xi64>) { - // CHECK-NOT: call @xla.sdy.manual_computation_body_7 - // CHECK: sdy.manual_computation(%arg0) - // CHECK-SAME{LITERAL}: in_shardings=[<@mesh_0, [{"b"}]>] - // CHECK-SAME{LITERAL}: out_shardings=[] - // CHECK-SAME{LITERAL}: manual_axes={"b"} - // CHECK-SAME{LITERAL}: (%arg1: tensor<2xi64>) { - // CHECK-NEXT: stablehlo.custom_call @sdy_testonly(%arg1) : (tensor<2xi64>) -> () - // CHECK-NEXT: sdy.return - // CHECK-NEXT: } : (tensor<4xi64>) -> () - // CHECK-NEXT: return - %s0 = stablehlo.custom_call @Sharding(%arg0) {sdy.sharding = #sdy.sharding_per_value<[<@mesh_0, [{"b"}]>]>} : (tensor<4xi64>) -> tensor<4xi64> - %0 = stablehlo.custom_call @xla.sdy.GlobalToLocalShape(%s0) {has_side_effect = true, sdy.sharding = #sdy.sharding_per_value<[<@mesh_0, [{}]>]>, xla.sdy.manual_axes = #sdy} : (tensor<4xi64>) -> tensor<2xi64> - call @xla.sdy.manual_computation_body_7(%0) : (tensor<2xi64>) -> () - return -} - -// Edge case: Degenerate manual computation with no inputs and no outputs. -// CHECK-LABEL: func @manual_computation_no_inputs_no_outputs -func.func @manual_computation_no_inputs_no_outputs() { - // CHECK-NEXT: sdy.manual_computation() in_shardings=[] out_shardings=[] manual_axes={} () { - // CHECK-NEXT: sdy.return - // CHECK-NEXT: } : () -> () - // CHECK-NEXT: return - call @xla.sdy.manual_computation_body_8() {mhlo.frontend_attributes = {inlineable = "false"}} : () -> () - return -} - -// Edge case: Manual computation with 0-sized dimension tensors. -// CHECK-LABEL: func @manual_computation_some_zero_dim_inputs_outputs -func.func @manual_computation_some_zero_dim_inputs_outputs(%arg0: tensor<0x16xf32>, %arg1: tensor<16x32xf32>) -> (tensor<0x32xf32>, tensor<16x32xf32>) { - - // CHECK-NOT: call @xla.sdy.manual_computation_body - // CHECK: %[[CONST_0_32:.*]] = stablehlo.constant dense<0.000000e+00> : tensor<0x32xf32> - // CHECK: %[[MAN_COMP:.*]]:2 = sdy.manual_computation(%arg0, %arg1) - // CHECK-SAME{LITERAL}: in_shardings=[<@mesh_0, [{}, {"b"}]>, <@mesh_0, [{"b"}, {}]>] - // CHECK-SAME{LITERAL}: out_shardings=[<@mesh_0, [{}, {}], replicated={"b"}>, <@mesh_0, [{"b"}, {}]>] - // CHECK-SAME{LITERAL}: manual_axes={"b"} - // CHECK-SAME: (%arg2: tensor<0x8xf32>, %arg3: tensor<8x32xf32>) { - // CHECK-NEXT: %[[DOT:.*]] = stablehlo.dot %arg2, %arg3 - // CHECK-NEXT: sdy.return %[[DOT]], %arg3 - // CHECK-NEXT: } : (tensor<0x16xf32>, tensor<16x32xf32>) -> (tensor<0x32xf32>, tensor<16x32xf32>) - // CHECK-NEXT: return %[[CONST_0_32]], %[[MAN_COMP]]#1 - %c1 = stablehlo.constant dense<0.000000e+00> : tensor<0x8xf32> - %c2 = stablehlo.constant dense<0.000000e+00> : tensor<0x32xf32> - %s0 = stablehlo.custom_call @Sharding(%arg0) {sdy.sharding = #sdy.sharding_per_value<[<@mesh_0, [{}, {"b"}]>]>} : (tensor<0x16xf32>) -> tensor<0x16xf32> - %s1 = stablehlo.custom_call @Sharding(%arg1) {sdy.sharding = #sdy.sharding_per_value<[<@mesh_0, [{"b"}, {}]>]>} : (tensor<16x32xf32>) -> tensor<16x32xf32> - %0:2 = stablehlo.custom_call @xla.sdy.GlobalToLocalShape(%s0, %s1) {has_side_effect = true, sdy.sharding = #sdy.sharding_per_value<[<@mesh_0, [{}, {}]>, <@mesh_0, [{}, {}]>]>, xla.sdy.manual_axes = #sdy} : (tensor<0x16xf32>, tensor<16x32xf32>) -> (tensor<0x8xf32>, tensor<8x32xf32>) - %1:2 = call @xla.sdy.manual_computation_body_9(%c1, %0#1) : (tensor<0x8xf32>, tensor<8x32xf32>) -> (tensor<0x32xf32>, tensor<8x32xf32>) - %2:2 = stablehlo.custom_call @xla.sdy.LocalToGlobalShape(%c2, %1#1) {has_side_effect = true, sdy.sharding = #sdy.sharding_per_value<[<@mesh_0, [{}, {}], replicated={"b"}>, <@mesh_0, [{"b"}, {}]>]>} : (tensor<0x32xf32>, tensor<8x32xf32>) -> (tensor<0x32xf32>, tensor<16x32xf32>) - return %c2, %2#1 : tensor<0x32xf32>, tensor<16x32xf32> -} - -// Type preservation: Token arguments are preserved with empty sharding and not assigned manual axes. -// CHECK-LABEL: func @manual_computation_tokens -func.func @manual_computation_tokens(%arg0: !stablehlo.token, %arg1: tensor<2xi64>) -> (!stablehlo.token, tensor<2xi64>) { - // CHECK-NOT: call @xla.sdy.manual_computation_body_11 - // CHECK: %[[MAN_COMP:.*]]:2 = sdy.manual_computation(%arg0, %arg1) - // CHECK-SAME{LITERAL}: in_shardings=[<@mesh_0, []>, <@mesh_0, [{"b"}]>] - // CHECK-SAME{LITERAL}: out_shardings=[<@mesh_0, []>, <@mesh_0, [{"b"}]>] - // CHECK-SAME{LITERAL}: manual_axes={"b"} - // CHECK-SAME: (%arg2: !stablehlo.token, %arg3: tensor<1xi64>) { - // CHECK-NEXT: sdy.return %arg2, %arg3 : !stablehlo.token, tensor<1xi64> - // CHECK-NEXT: } : (!stablehlo.token, tensor<2xi64>) -> (!stablehlo.token, tensor<2xi64>) - // CHECK-NEXT: return %[[MAN_COMP]]#0, %[[MAN_COMP]]#1 : !stablehlo.token, tensor<2xi64> - %s0 = stablehlo.custom_call @Sharding(%arg0) {sdy.sharding = #sdy.sharding_per_value<[<@mesh_0, []>]>} : (!stablehlo.token) -> !stablehlo.token - %s1 = stablehlo.custom_call @Sharding(%arg1) {sdy.sharding = #sdy.sharding_per_value<[<@mesh_0, [{"b"}]>]>} : (tensor<2xi64>) -> tensor<2xi64> - %0:2 = stablehlo.custom_call @xla.sdy.GlobalToLocalShape(%s0, %s1) {has_side_effect = true, sdy.sharding = #sdy.sharding_per_value<[<@mesh_0, []>, <@mesh_0, [{}]>]>, xla.sdy.manual_axes = #sdy} : (!stablehlo.token, tensor<2xi64>) -> (!stablehlo.token, tensor<1xi64>) - %1:2 = call @xla.sdy.manual_computation_body_11(%0#0, %0#1) : (!stablehlo.token, tensor<1xi64>) -> (!stablehlo.token, tensor<1xi64>) - %2:2 = stablehlo.custom_call @xla.sdy.LocalToGlobalShape(%1#0, %1#1) {has_side_effect = true, sdy.sharding = #sdy.sharding_per_value<[<@mesh_0, []>, <@mesh_0, [{"b"}]>]>} : (!stablehlo.token, tensor<1xi64>) -> (!stablehlo.token, tensor<2xi64>) - return %2#0, %2#1 : !stablehlo.token, tensor<2xi64> -} - -// Tests stray/dangling boundary custom calls where results are unused: safely stripped. -// CHECK-LABEL: func @stray_unused_manual_computation_custom_calls -func.func @stray_unused_manual_computation_custom_calls(%arg0: tensor<0x16xf32>) -> tensor<0x16xf32> { - // CHECK-NEXT: %[[CONST_0_8:.*]] = stablehlo.constant dense<0.000000e+00> : tensor<0x8xf32> - // CHECK-NEXT: %[[CONST_0_16:.*]] = stablehlo.constant dense<0.000000e+00> : tensor<0x16xf32> - // CHECK-NOT: call @xla.sdy.manual_computation_body - // CHECK-NOT: call @xla.sdy.GlobalToLocalShape - // CHECK-NOT: call @xla.sdy.LocalToGlobalShape - // CHECK-NEXT: return %[[CONST_0_16]] - %c1 = stablehlo.constant dense<0.000000e+00> : tensor<0x8xf32> - %c2 = stablehlo.constant dense<0.000000e+00> : tensor<0x16xf32> - %0 = stablehlo.custom_call @xla.sdy.GlobalToLocalShape(%arg0) {has_side_effect = true, sdy.sharding = #sdy.sharding_per_value<[<@mesh_0, [{}, {}]>]>, xla.sdy.manual_axes = #sdy} : (tensor<0x16xf32>) -> tensor<0x8xf32> - %1 = stablehlo.custom_call @xla.sdy.LocalToGlobalShape(%c1) {has_side_effect = true, sdy.sharding = #sdy.sharding_per_value<[<@mesh_0, [{"b"}, {}]>]>} : (tensor<0x8xf32>) -> tensor<0x16xf32> - return %c2 : tensor<0x16xf32> -} - -// CHECK-NOT: func private @xla.sdy.manual_computation_body( -func.func private @xla.sdy.manual_computation_body(%arg0: tensor<2x8xf32>, %arg1: tensor<8x32xf32>) -> tensor<2x32xf32> { - %0 = stablehlo.add %arg0, %arg0 : tensor<2x8xf32> - %1 = stablehlo.dot %0, %arg1 : (tensor<2x8xf32>, tensor<8x32xf32>) -> tensor<2x32xf32> - %2 = "stablehlo.all_reduce"(%1) <{replica_groups = dense<[[0], [1]]> : tensor<2x1xi64>}> ({ - ^bb0(%arg2: tensor, %arg3: tensor): - %3 = stablehlo.add %arg2, %arg3 : tensor - stablehlo.return %3 : tensor - }) : (tensor<2x32xf32>) -> tensor<2x32xf32> - return %2 : tensor<2x32xf32> -} - -// CHECK-NOT: func private @my_model.___call__.fwd.xla.sdy.manual_computation_body_14.1234( -func.func private @my_model.___call__.fwd.xla.sdy.manual_computation_body_14.1234(%arg0: tensor<2x8xf32>) -> tensor<2x8xf32> { - return %arg0 : tensor<2x8xf32> -} - -// CHECK-NOT: func private @xla.sdy.manual_computation_body_0( -func.func private @xla.sdy.manual_computation_body_0(%arg0: tensor<2x8xf32>) -> tensor<2x8xf32> { - return %arg0 : tensor<2x8xf32> -} - -// CHECK-NOT: func private @xla.sdy.manual_computation_body_1( -func.func private @xla.sdy.manual_computation_body_1(%arg0: tensor<8x4xf32>) -> tensor<8x4xf32> { - return %arg0 : tensor<8x4xf32> -} - -// CHECK-NOT: func private @xla.sdy.manual_computation_body_2( -func.func private @xla.sdy.manual_computation_body_2(%arg0: tensor<2x4xf32>) -> tensor<2x4xf32> { - %0 = stablehlo.multiply %arg0, %arg0 : tensor<2x4xf32> - return %0 : tensor<2x4xf32> -} - -// CHECK-NOT: func private @xla.sdy.manual_computation_body_3( -func.func private @xla.sdy.manual_computation_body_3(%arg0: tensor<2x8xf32>) -> tensor<2x8xf32> { - %s0 = stablehlo.custom_call @Sharding(%arg0) {sdy.sharding = #sdy.sharding_per_value<[<@mesh_1, [{}, {"b"}]>]>} : (tensor<2x8xf32>) -> tensor<2x8xf32> - %0 = stablehlo.custom_call @xla.sdy.GlobalToLocalShape(%s0) {has_side_effect = true, sdy.sharding = #sdy.sharding_per_value<[<@mesh_1, [{}, {}]>]>, xla.sdy.manual_axes = #sdy} : (tensor<2x8xf32>) -> tensor<2x4xf32> - %1 = call @xla.sdy.manual_computation_body_2(%0) : (tensor<2x4xf32>) -> tensor<2x4xf32> - %2 = stablehlo.custom_call @xla.sdy.LocalToGlobalShape(%1) {has_side_effect = true, sdy.sharding = #sdy.sharding_per_value<[<@mesh_1, [{}, {"b"}]>]>} : (tensor<2x4xf32>) -> tensor<2x8xf32> - return %2 : tensor<2x8xf32> -} - -// CHECK-NOT: func private @xla.sdy.manual_computation_body_4( -func.func private @xla.sdy.manual_computation_body_4(%arg0: tensor<2x4xf32>) -> tensor<2x4xf32> { - %0 = stablehlo.multiply %arg0, %arg0 : tensor<2x4xf32> - return %0 : tensor<2x4xf32> -} - -// CHECK-NOT: func private @xla.sdy.manual_computation_body_5( -func.func private @xla.sdy.manual_computation_body_5(%arg0: tensor<2x8xf32>) -> tensor<2x8xf32> { - %s0 = stablehlo.custom_call @Sharding(%arg0) {sdy.sharding = #sdy.sharding_per_value<[<@mesh_1, [{}, {"b"}]>]>} : (tensor<2x8xf32>) -> tensor<2x8xf32> - %0 = stablehlo.custom_call @xla.sdy.GlobalToLocalShape(%s0) {has_side_effect = true, sdy.sharding = #sdy.sharding_per_value<[<@mesh_1, [{}, {}]>]>, xla.sdy.manual_axes = #sdy} : (tensor<2x8xf32>) -> tensor<2x4xf32> - %1 = call @xla.sdy.manual_computation_body_4(%0) : (tensor<2x4xf32>) -> tensor<2x4xf32> - %2 = stablehlo.custom_call @xla.sdy.LocalToGlobalShape(%1) {has_side_effect = true, sdy.sharding = #sdy.sharding_per_value<[<@mesh_1, [{}, {"b"}]>]>} : (tensor<2x4xf32>) -> tensor<2x8xf32> - %3 = stablehlo.add %2, %2 : tensor<2x8xf32> - return %3 : tensor<2x8xf32> -} - -// CHECK-NOT: func private @xla.sdy.manual_computation_body_6( -func.func private @xla.sdy.manual_computation_body_6() -> tensor<2xi64> { - %0 = stablehlo.constant dense<[2, 3]> : tensor<2xi64> - return %0 : tensor<2xi64> -} - -// CHECK-NOT: func private @xla.sdy.manual_computation_body_7( -func.func private @xla.sdy.manual_computation_body_7(%arg0: tensor<2xi64>) { - stablehlo.custom_call @sdy_testonly(%arg0) : (tensor<2xi64>) -> () - return -} - -// CHECK-NOT: func private @xla.sdy.manual_computation_body_8( -func.func private @xla.sdy.manual_computation_body_8() { - return -} - -// CHECK-NOT: func private @xla.sdy.manual_computation_body_9( -func.func private @xla.sdy.manual_computation_body_9(%arg0: tensor<0x8xf32>, %arg1: tensor<8x32xf32>) -> (tensor<0x32xf32>, tensor<8x32xf32>) { - %0 = stablehlo.dot %arg0, %arg1 : (tensor<0x8xf32>, tensor<8x32xf32>) -> tensor<0x32xf32> - return %0, %arg1 : tensor<0x32xf32>, tensor<8x32xf32> -} - -// CHECK-NOT: func private @xla.sdy.manual_computation_body_11( -func.func private @xla.sdy.manual_computation_body_11(%arg0: !stablehlo.token, %arg1: tensor<1xi64>) -> (!stablehlo.token, tensor<1xi64>) { - return %arg0, %arg1 : !stablehlo.token, tensor<1xi64> -} diff --git a/third_party/xla/xla/service/while_loop_simplifier.cc b/third_party/xla/xla/service/while_loop_simplifier.cc index aa1cb7b60455cb..3e99ac8dcbc57d 100644 --- a/third_party/xla/xla/service/while_loop_simplifier.cc +++ b/third_party/xla/xla/service/while_loop_simplifier.cc @@ -30,6 +30,7 @@ limitations under the License. #include "absl/status/status_macros.h" #include "absl/strings/string_view.h" #include "absl/types/span.h" +#include "xla/bit_set.h" #include "xla/comparison_util.h" #include "xla/frontend_attributes.h" #include "xla/hlo/analysis/while_loop_analysis.h" @@ -43,7 +44,6 @@ limitations under the License. #include "xla/hlo/ir/hlo_print_options.h" #include "xla/hlo/transforms/simplifiers/hlo_dce.h" #include "xla/hlo/utils/hlo_query.h" -#include "xla/inlined_bit_set.h" #include "xla/literal_util.h" #include "xla/primitive_util.h" #include "xla/service/call_inliner.h" @@ -143,19 +143,15 @@ void CopyMetadata(HloInstruction* old_while_op, HloInstruction* new_while_op) { new_while_op->set_metadata(old_while_op->metadata()); } -// This is a utility function that removes the given tuple indices from the -// while loop init, body, and condition. The final shape returned is still the -// same as before. If set index_for_replaced will replace any use of the removed -// indices in the final shape with a copy of the removed index. +// This is a utility function that removes the tuple indices whose bit is not +// set in used_tuple_indices from the while loop init, body, and condition. The +// final shape returned is still the same as before. If set index_for_replaced +// will replace any use of the removed indices in the final shape with a copy of +// the removed index. static absl::StatusOr RemoveDeadTupleIndices( - HloInstruction* while_op, absl::flat_hash_set& used_tuple_indices, + HloInstruction* while_op, const InlinedBitSet<>& used_tuple_indices, std::optional> dead_to_surviving_index = std::nullopt) { - // Build up maps from the old/new to the new/old tuple indices. - std::vector new_to_old_tuple_idx(used_tuple_indices.begin(), - used_tuple_indices.end()); - absl::c_sort(new_to_old_tuple_idx); - HloModule* module = while_op->GetModule(); HloComputation* computation = while_op->parent(); HloInstruction* while_init = while_op->mutable_operand(0); @@ -163,6 +159,15 @@ static absl::StatusOr RemoveDeadTupleIndices( HloComputation* while_body = while_op->while_body(); HloInstruction* while_body_root = while_body->root_instruction(); + // Build up maps from the old/new to the new/old tuple indices. + std::vector new_to_old_tuple_idx; + const int64_t old_tuple_size = while_init->shape().tuple_shapes().size(); + for (int64_t old_idx = 0; old_idx < old_tuple_size; ++old_idx) { + if (used_tuple_indices.Test(old_idx)) { + new_to_old_tuple_idx.push_back(old_idx); + } + } + auto print_no_metadata = HloPrintOptions().set_print_metadata(false); absl::flat_hash_map old_to_new_tuple_idx; @@ -355,6 +360,253 @@ bool AllWhileParamConsumersGte(const HloInstruction* while_op) { return true; } +// The tuple indices of a while loop that are read after the loop: all of them +// when the loop is the root of its computation or has a user other than +// get-tuple-element. +InlinedBitSet<> TupleIndicesUsedAfterLoop(const HloInstruction* while_op, + int64_t tuple_size) { + InlinedBitSet<> used(tuple_size); + if (while_op == while_op->parent()->root_instruction()) { + used.SetAll(tuple_size); + return used; + } + for (const HloInstruction* user : while_op->users()) { + if (user->opcode() != HloOpcode::kGetTupleElement) { + used.SetAll(tuple_size); + return used; + } + used.Set(user->tuple_index()); + } + return used; +} + +// Whether the body passes the tuple element at 'index' of its parameter +// through unmodified. +bool IsPassedThroughUnmodified(const HloComputation* while_body, + int64_t index) { + const HloInstruction* output = while_body->root_instruction()->operand(index); + return output->opcode() == HloOpcode::kGetTupleElement && + output->operand(0) == while_body->parameter_instruction(0) && + output->tuple_index() == index; +} + +// How a while loop's body and condition depend on the loop's candidate tuple +// indices, the indices TryRemoveDeadWhileParams may remove: candidate_bit +// numbers them and holds -1 for every other index. Per instruction, the +// candidates it depends on within one iteration; across instructions, the +// connected components of the operand graph. +// +// The instructions of the body and the condition are addressed by slot, the +// body's local ids first, then the condition's. The slots map the instructions +// to a dense index range, so the per instruction state lives in flat arrays +// instead of hash maps keyed by instruction. Local ids stay stable during the +// analysis: nothing mutates the graph before RemoveDeadTupleIndices. +class WhileInputDependencies { + public: + WhileInputDependencies(const HloInstruction* while_op, + absl::Span candidate_bit, + int num_candidates) + : while_body_(while_op->while_body()), + while_cond_(while_op->while_condition()), + while_body_root_(while_body_->root_instruction()), + candidate_bit_(candidate_bit), + body_slots_(while_body_->next_unique_instruction_internal_id()), + num_slots_(body_slots_ + + while_cond_->next_unique_instruction_internal_id()), + side_effecting_row_(num_slots_), + affecting_others_row_(num_slots_ + 1), + input_deps_(num_slots_ + 2, num_candidates), + sets_(num_slots_) { + PropagateThrough(while_body_); + PropagateThrough(while_cond_); + AccumulateOutputDependencies(); + } + + int64_t tuple_size() const { return candidate_bit_.size(); } + + // Whether the index may be removed at all. + bool IsCandidate(int64_t index) const { return candidate_bit_[index] >= 0; } + + // Whether a side effecting instruction or the loop condition depends on the + // candidate index. + bool ReachesSink(int64_t index) const { + DCHECK(IsCandidate(index)); + return input_deps_.Test(side_effecting_row_, candidate_bit_[index]); + } + + // Whether a body output other than its own depends on the candidate index. + bool AffectsOtherOutputs(int64_t index) const { + DCHECK(IsCandidate(index)); + return input_deps_.Test(affecting_others_row_, candidate_bit_[index]); + } + + // The connected component of the body output at the index, numbered below + // num_slots(). The get-tuple-elements of the index belong to the same + // component, so indices that depend on each other share it. + int ComponentOf(int64_t index) { + return sets_.Find(SlotOf(while_body_root_->operand(index))); + } + + int num_slots() const { return num_slots_; } + + private: + int SlotOf(const HloInstruction* inst) const { + DCHECK(inst->parent() == while_body_ || inst->parent() == while_cond_); + return (inst->parent() == while_body_ ? 0 : body_slots_) + inst->local_id(); + } + + // The parameters and the body root take no part in the dependencies. + bool IsTracked(const HloInstruction* inst) const { + return inst != while_body_->parameter_instruction(0) && + inst != while_cond_->parameter_instruction(0) && + inst != while_body_root_; + } + + // One post order pass over a computation. Operands come before their users, + // so the dependencies of an instruction are the union of its operands' by + // the time it is visited. A get-tuple-element of the parameter seeds its own + // bit and joins the component of the output it feeds back into. A side + // effecting instruction and the condition root are sinks: the inputs they + // depend on must stay. + void PropagateThrough(const HloComputation* comp) { + const HloInstruction* while_input = comp->parameter_instruction(0); + comp->ForEachInstructionPostOrder([&](HloInstruction* inst) { + if (!IsTracked(inst)) { + return; + } + const int slot = SlotOf(inst); + if (inst->opcode() == HloOpcode::kGetTupleElement && + inst->operand(0) == while_input) { + if (candidate_bit_[inst->tuple_index()] >= 0) { + input_deps_.Set(slot, candidate_bit_[inst->tuple_index()]); + } + const HloInstruction* output = + while_body_root_->operand(inst->tuple_index()); + if (output != inst) { + sets_.Merge(SlotOf(output), slot); + } + } else { + for (const HloInstruction* operand : inst->operands()) { + if (!IsTracked(operand)) { + continue; + } + sets_.Merge(SlotOf(operand), slot); + if (const std::optional index = + ParameterIndex(operand, while_input)) { + // The row of a parameter get-tuple-element holds only its own + // bit, so one Set replaces the OR of a whole row. + if (candidate_bit_[*index] >= 0) { + input_deps_.Set(slot, candidate_bit_[*index]); + } + } else { + input_deps_.Or(slot, SlotOf(operand)); + } + } + } + if (inst->HasSideEffect() || inst == while_cond_->root_instruction()) { + input_deps_.Or(side_effecting_row_, slot); + } + }); + } + + // The tuple index when inst is a get-tuple-element of the given parameter. + static std::optional ParameterIndex(const HloInstruction* inst, + const HloInstruction* param) { + if (inst->opcode() == HloOpcode::kGetTupleElement && + inst->operand(0) == param) { + return inst->tuple_index(); + } + return std::nullopt; + } + + // Collects the inputs that reach some body output other than their own; a + // self update does not count. + void AccumulateOutputDependencies() { + const HloInstruction* while_input = while_body_->parameter_instruction(0); + for (int64_t i = 0; i < while_body_root_->operand_count(); ++i) { + const HloInstruction* output = while_body_root_->operand(i); + if (const std::optional index = + ParameterIndex(output, while_input)) { + // The output copies input *index; the row holds only that bit, and a + // pass through (*index == i) is a self update. + if (*index != i && candidate_bit_[*index] >= 0) { + input_deps_.Set(affecting_others_row_, candidate_bit_[*index]); + } + continue; + } + const int output_slot = SlotOf(output); + if (candidate_bit_[i] >= 0) { + input_deps_.OrIgnoringBit(affecting_others_row_, output_slot, + candidate_bit_[i]); + } else { + input_deps_.Or(affecting_others_row_, output_slot); + } + } + } + + const HloComputation* const while_body_; + const HloComputation* const while_cond_; + const HloInstruction* const while_body_root_; + const absl::Span candidate_bit_; + const int body_slots_; + const int num_slots_; + // Two extra rows of input_deps_: the inputs that reach a side effect or the + // condition, and the inputs that reach another output. + const int side_effecting_row_; + const int affecting_others_row_; + DenseBitSets input_deps_; + DenseUnionFind sets_; +}; + +// Clears index i of the surviving indices; returns 1 when it was still set. +int64_t EraseIndex(int64_t i, InlinedBitSet<>& used_tuple_indices) { + if (!used_tuple_indices.Test(i)) { + return 0; + } + used_tuple_indices.Clear(i); + return 1; +} + +// Case 1: an index that is unused after the loop can go on its own when no +// other output, side effect or condition depends on its input. Returns the +// number of indices newly removed from used_tuple_indices. +int64_t EraseIndicesDeadOnTheirOwn(const InlinedBitSet<>& used_after_loop, + const WhileInputDependencies& deps, + InlinedBitSet<>& used_tuple_indices) { + int64_t num_erased = 0; + for (int64_t i = 0; i < deps.tuple_size(); ++i) { + if (!used_after_loop.Test(i) && !deps.AffectsOtherOutputs(i) && + !deps.ReachesSink(i)) { + VLOG(2) << "Remove with dependencies " << i; + num_erased += EraseIndex(i, used_tuple_indices); + } + } + return num_erased; +} + +// Case 2: a connected group of indices can go when none of its indices blocks +// it. An index blocks its group when it is not a candidate, that is used after +// the loop and not passed through, or when a side effect or the condition +// depends on it. Returns the number of indices newly removed from +// used_tuple_indices. +int64_t EraseIndicesDeadAsGroup(WhileInputDependencies& deps, + InlinedBitSet<>& used_tuple_indices) { + InlinedBitSet<> blocked_components(deps.num_slots()); + for (int64_t i = 0; i < deps.tuple_size(); ++i) { + if (!deps.IsCandidate(i) || deps.ReachesSink(i)) { + blocked_components.Set(deps.ComponentOf(i)); + } + } + int64_t num_erased = 0; + for (int64_t i = 0; i < deps.tuple_size(); ++i) { + if (!blocked_components.Test(deps.ComponentOf(i))) { + VLOG(2) << "Remove with groups: index " << i; + num_erased += EraseIndex(i, used_tuple_indices); + } + } + return num_erased; +} + } // namespace absl::StatusOr TryRemoveDeadWhileParams(HloInstruction* while_op) { @@ -372,7 +624,6 @@ absl::StatusOr TryRemoveDeadWhileParams(HloInstruction* while_op) { } HloInstruction* while_init = while_op->mutable_operand(0); - HloComputation* while_cond = while_op->while_condition(); HloComputation* while_body = while_op->while_body(); HloInstruction* while_body_root = while_body->root_instruction(); @@ -388,12 +639,6 @@ absl::StatusOr TryRemoveDeadWhileParams(HloInstruction* while_op) { const int64_t tuple_size = ShapeUtil::TupleElementCount(while_init->shape()); auto print_no_metadata = HloPrintOptions().set_print_metadata(false); - // A set that stores all unused indices. Initialize to all indices and then - // remove elements from it. - absl::flat_hash_set used_tuple_indices; - for (int64_t i = 0; i < tuple_size; ++i) { - used_tuple_indices.insert(i); - } // Bail if param0 of while_cond or while_body has users which aren't of type // get-tuple-element. @@ -406,146 +651,49 @@ absl::StatusOr TryRemoveDeadWhileParams(HloInstruction* while_op) { "empty."; return false; } - InlinedBitSet<> used_indices_after_loop(tuple_size); - if (while_op == while_op->parent()->root_instruction()) { - used_indices_after_loop.SetAll(tuple_size); - } - for (auto user : while_op->users()) { - if (user->opcode() != HloOpcode::kGetTupleElement) { - used_indices_after_loop.SetAll(tuple_size); - break; - } - used_indices_after_loop.Set(user->tuple_index()); - } + const InlinedBitSet<> used_after_loop = + TupleIndicesUsedAfterLoop(while_op, tuple_size); // We identify dead loop parameters in two cases (provided they do not cause // side effects or affect the loop condition): // 1) There is no use after the loop, and the input does not affect other // outputs (though it may update itself). // 2) A group of mutually dependent elements whose outputs are either unused - // after the loop or are passed-through unmodified. We use UnionFind to - // approximate connected components (which has false negatives for directed - // dependencies). - - // Tracks the set of inputs that each instruction depends on (in one - // iteration). Used to find inputs affecting other outputs (case 1) and - // inputs affecting the condition or side effects (cases 1 & 2). - absl::flat_hash_map> inst_input_deps; - // Find disjoint sets of connected instruction groups. This helps finding a - // group of inter-dependent indices that can be removed together. For case 2). - absl::flat_hash_map> - disjoint_sets; - // Initialize. - for (HloComputation* comp : {while_body, while_cond}) { - HloInstruction* while_input = comp->parameter_instruction(0); - for (HloInstruction* inst : comp->instructions()) { - if (inst == while_input || inst == while_body_root) { - continue; - } - disjoint_sets[inst].Get() = inst; - } - } - // Track the dependencies and merge the disjoint sets. - InlinedBitSet<> side_effecting_indices(tuple_size); - for (HloComputation* comp : {while_body, while_cond}) { - HloInstruction* while_input = comp->parameter_instruction(0); - for (HloInstruction* inst : comp->MakeInstructionPostOrder()) { - if (inst == while_input || inst == while_body_root) { - continue; - } - auto& deps = inst_input_deps.try_emplace(inst, tuple_size).first->second; - auto& my_set = disjoint_sets[inst]; - if (inst->opcode() == HloOpcode::kGetTupleElement && - inst->operand(0) == while_input) { - deps.Set(inst->tuple_index()); - HloInstruction* output = - while_body_root->mutable_operand(inst->tuple_index()); - if (output != inst) { - disjoint_sets[output].Merge(&my_set); - } - } else { - for (HloInstruction* operand : inst->operands()) { - if (operand == while_body_root) { - // The root is skipped above, so it has no map entries. Looking it - // up here would insert one and could dangle deps and my_set. - continue; - } - disjoint_sets[operand].Merge(&my_set); - auto it = inst_input_deps.find(operand); - if (it != inst_input_deps.end()) { - deps |= it->second; - } - } - } - if (inst->HasSideEffect() || inst == while_cond->root_instruction()) { - side_effecting_indices |= deps; - } - } - } - // Find inputs that can be removed because they don't affect others. - InlinedBitSet<> indices_affecting_others(tuple_size); - for (int64_t i = 0; i < tuple_size; ++i) { - HloInstruction* output = while_body_root->mutable_operand(i); - auto it = inst_input_deps.find(output); - if (it != inst_input_deps.end()) { - InlinedBitSet<> deps = it->second; - // Ignore self-updates. - deps.Clear(i); - indices_affecting_others |= deps; - } - } + // after the loop or are passed through unmodified. We use a union find data + // structure to approximate connected components (which has false negatives + // for directed dependencies). + + // Only an index that is unused after the loop or passed through unmodified + // can be removed, in either case. The dependency analysis tracks just those + // candidate indices, numbered by candidate_bit; every other index keeps -1 + // there. + std::vector candidate_bit(tuple_size, -1); + int num_candidates = 0; for (int64_t i = 0; i < tuple_size; ++i) { - if (!indices_affecting_others.Test(i) && !used_indices_after_loop.Test(i) && - !side_effecting_indices.Test(i)) { - VLOG(2) << "Remove with dependencies " << i; - used_tuple_indices.erase(i); + if (!used_after_loop.Test(i) || IsPassedThroughUnmodified(while_body, i)) { + candidate_bit[i] = num_candidates++; } } - // Find the connected groups of input/output indices. - absl::flat_hash_map> groups; - for (int64_t i = 0; i < tuple_size; ++i) { - HloInstruction* output = while_body_root->mutable_operand(i); - groups[disjoint_sets[output].Get()].insert(i); - } - for (HloComputation* comp : {while_body, while_cond}) { - HloInstruction* while_input = comp->parameter_instruction(0); - for (HloInstruction* gte : while_input->users()) { - groups[disjoint_sets[gte].Get()].insert(gte->tuple_index()); - } - } - for (const auto& group : groups) { - if (absl::c_any_of(group.second, [&](int64_t index) { - // We cannot remove this index causes side effects, or if its output - // is not passed through from input and it is used after the while op. - const HloInstruction* output = while_body_root->operand(index); - return side_effecting_indices.Test(index) || - (used_indices_after_loop.Test(index) && - !(output->opcode() == HloOpcode::kGetTupleElement && - output->operand(0) == - while_body->parameter_instruction(0) && - output->tuple_index() == index)); - })) { - continue; - } - VLOG(2) << "Remove with groups:"; - for (int64_t index : group.second) { - VLOG(2) << " index " << index; - used_tuple_indices.erase(index); - } + if (num_candidates == 0) { + VLOG(2) << "Loop " << while_op->ToString(print_no_metadata) + << " has no removable candidate index."; + return false; } + WhileInputDependencies deps(while_op, candidate_bit, num_candidates); + // The indices that survive: all of them until a case removes some. + InlinedBitSet<> used_tuple_indices(tuple_size); + used_tuple_indices.SetAll(tuple_size); + const int64_t num_removed = + EraseIndicesDeadOnTheirOwn(used_after_loop, deps, used_tuple_indices) + + EraseIndicesDeadAsGroup(deps, used_tuple_indices); - if (used_tuple_indices.size() == tuple_size) { + if (num_removed == 0) { VLOG(2) << "Loop " << while_op->ToString(print_no_metadata) << " uses all of its inputs; no simplification possible."; return false; } - // If we got here, used_tuple_indices.size() < tuple_size, meaning some - // elements of the loop's tuple aren't used by while_body or while_cond. - CHECK_LT(used_tuple_indices.size(), tuple_size); - - VLOG(1) << "Eliminating " << tuple_size - used_tuple_indices.size() - << " elements from tuple of " + VLOG(1) << "Eliminating " << num_removed << " elements from tuple of " << while_op->ToString(print_no_metadata); ABSL_ASSIGN_OR_RETURN(while_op, @@ -627,12 +775,12 @@ absl::StatusOr RemoveRepeatedWhileTupleIndices( } // And finally, remove the dead indices. - absl::flat_hash_set used_tuple_indices; - for (int index = 0; - index < while_body->root_instruction()->shape().tuple_shapes().size(); - ++index) { - if (dead_to_surviving_index.find(index) == dead_to_surviving_index.end()) { - used_tuple_indices.insert(index); + const int64_t tuple_size = + while_body->root_instruction()->shape().tuple_shapes().size(); + InlinedBitSet<> used_tuple_indices(tuple_size); + for (int64_t index = 0; index < tuple_size; ++index) { + if (!dead_to_surviving_index.contains(index)) { + used_tuple_indices.Set(index); } } diff --git a/third_party/xla/xla/service/while_loop_simplifier_test.cc b/third_party/xla/xla/service/while_loop_simplifier_test.cc index 37d7f9a1a3c853..dc348d5027c9c4 100644 --- a/third_party/xla/xla/service/while_loop_simplifier_test.cc +++ b/third_party/xla/xla/service/while_loop_simplifier_test.cc @@ -25,6 +25,7 @@ limitations under the License. #include "absl/status/status_matchers.h" #include "absl/strings/str_cat.h" #include "absl/strings/str_replace.h" +#include "absl/strings/string_view.h" #include "xla/hlo/ir/hlo_computation.h" #include "xla/hlo/ir/hlo_instruction.h" #include "xla/hlo/ir/hlo_opcode.h" @@ -45,6 +46,9 @@ namespace { using ::testing::_; namespace op = xla::testing::opcode_matchers; +using op::Constant; +using op::Tuple; +using op::While; // Returns the first kWhile instruction within m's entry computation. HloInstruction* FindFirstWhile(HloModule* m) { @@ -2342,5 +2346,70 @@ TEST_F(WhileLoopSimplifierTest, SimplifierWithDisabledWhileLoopDceAttr) { EXPECT_FALSE(changed); } +// An async-update or async-done carries no called computation of its own; its +// side effect is the wrapped instruction's, reached through the chain. The +// loop below feeds x to the async-start and y to the async-update only, and +// SIDE_EFFECT selects whether the wrapped custom call has a side effect. +constexpr absl::string_view kAsyncChainHlo = R"hlo( + HloModule AsyncChain + + async_comp { + a = f32[] parameter(0) + b = f32[] parameter(1) + ROOT cc = f32[] custom-call(a, b), custom_call_target="Foo", custom_call_has_side_effect=SIDE_EFFECT + } + cond { + p = (s32[], f32[], f32[]) parameter(0) + i = s32[] get-tuple-element(p), index=0 + ten = s32[] constant(10) + ROOT lt = pred[] compare(i, ten), direction=LT + } + body { + p = (s32[], f32[], f32[]) parameter(0) + i = s32[] get-tuple-element(p), index=0 + x = f32[] get-tuple-element(p), index=1 + y = f32[] get-tuple-element(p), index=2 + one = s32[] constant(1) + i.next = s32[] add(i, one) + start = ((f32[]), f32[], s32[]) async-start(x), calls=async_comp + update = ((f32[], f32[]), f32[], s32[]) async-update(start, y) + done = f32[] async-done(update) + ROOT t = (s32[], f32[], f32[]) tuple(i.next, x, y) + } + ENTRY main { + zero = s32[] constant(0) + one = f32[] constant(1) + two = f32[] constant(2) + init = (s32[], f32[], f32[]) tuple(zero, one, two) + loop = (s32[], f32[], f32[]) while(init), condition=cond, body=body + ROOT r = s32[] get-tuple-element(loop), index=0 + } + )hlo"; + +// A side effect behind the async chain keeps x and y alive, as it would +// behind the start itself. +TEST_F(WhileLoopSimplifierTest, AsyncChainWithSideEffectKeepsItsInputsAlive) { + ASSERT_OK_AND_ASSIGN(auto m, ParseAndReturnVerifiedModule(absl::StrReplaceAll( + kAsyncChainHlo, {{"SIDE_EFFECT", "true"}}))); + HloInstruction* loop = FindFirstWhile(m.get()); + + ASSERT_OK_AND_ASSIGN(bool changed, TryRemoveDeadWhileParams(loop)); + + EXPECT_FALSE(changed); +} + +// Without the side effect both x and y are dead. +TEST_F(WhileLoopSimplifierTest, AsyncChainWithoutSideEffectIsDead) { + ASSERT_OK_AND_ASSIGN(auto m, + ParseAndReturnVerifiedModule(absl::StrReplaceAll( + kAsyncChainHlo, {{"SIDE_EFFECT", "false"}}))); + HloInstruction* loop = FindFirstWhile(m.get()); + + ASSERT_OK_AND_ASSIGN(bool changed, TryRemoveDeadWhileParams(loop)); + + EXPECT_TRUE(changed); + EXPECT_THAT(FindFirstWhile(m.get()), While(Tuple(Constant()))); +} + } // namespace } // namespace xla diff --git a/third_party/xla/xla/stream_executor/BUILD b/third_party/xla/xla/stream_executor/BUILD index 2d863d1a66fc41..ba5e63312be9f5 100644 --- a/third_party/xla/xla/stream_executor/BUILD +++ b/third_party/xla/xla/stream_executor/BUILD @@ -1093,7 +1093,7 @@ cc_library( name = "command_buffer", srcs = ["command_buffer.cc"], hdrs = ["command_buffer.h"], - visibility = [":internal"], + visibility = internal_visibility([":internal"]), deps = [ ":bit_pattern", ":device_address", diff --git a/third_party/xla/xla/stream_executor/cuda/BUILD b/third_party/xla/xla/stream_executor/cuda/BUILD index d23960e276d69b..5d69ba482808a8 100644 --- a/third_party/xla/xla/stream_executor/cuda/BUILD +++ b/third_party/xla/xla/stream_executor/cuda/BUILD @@ -84,21 +84,6 @@ config_setting( }, ) -bool_flag( - name = "enable_libnvjitlink_support", - build_setting_default = if_google( - True, - oss_value = False, - ), -) - -config_setting( - name = "libnvjitlink_support_enabled", - flag_values = { - ":enable_libnvjitlink_support": "True", - }, -) - cc_library( name = "cuda_platform_id", srcs = ["cuda_platform_id.cc"], @@ -1308,51 +1293,32 @@ cc_library( name = "nvjitlink_support", srcs = ["nvjitlink_support.cc"], hdrs = ["nvjitlink_support.h"], - local_defines = select({ - ":libnvjitlink_support_enabled": [ - "LIBNVJITLINK_SUPPORT=true", - ], - "//conditions:default": [ - "LIBNVJITLINK_SUPPORT=false", - ], - }) + if_cuda_newer_than( - "12_0", - ["CUDA_SUPPORTS_NVJITLINK=true"], - ["CUDA_SUPPORTS_NVJITLINK=false"], - ), ) cc_library( - name = "nvjitlink_stub", - srcs = [ + name = "nvjitlink", + srcs = ["nvjitlink.cc"], + hdrs = [ "nvjitlink.h", - "nvjitlink_stub.cc", ], - deps = [ - ":compilation_provider", - ":cuda_compute_capability", - "//xla/stream_executor/gpu:gpu_asm_opts", - "@com_google_absl//absl/status", - "@com_google_absl//absl/status:statusor", - "@com_google_absl//absl/types:span", + local_defines = [ + # By default, nvJitLink.h defines static inline wrappers that forward to + # versioned __nvJitLink*__ symbols. Defining + # NVJITLINK_NO_INLINE makes nvjitlink.cc call the unversioned + # nvJitLink* symbols directly so it is compatible with nvjitlink_stub. + "NVJITLINK_NO_INLINE", ], -) - -cc_library( - name = "nvjitlink_impl", - srcs = [ - "nvjitlink.h", - "nvjitlink_impl.cc", + tags = [ + "cuda-only", + "gpu", ], - tags = ["manual"], deps = [ ":compilation_provider", ":cuda_compute_capability", ":ptx_compiler_helpers", "//xla/stream_executor:kernel_stats", "//xla/stream_executor/gpu:gpu_asm_opts", - "//xla/tsl/platform:errors", - "//xla/tsl/platform:statusor", + "//xla/tsl/cuda:nvjitlink", # buildcleaner: keep "@com_google_absl//absl/algorithm:container", "@com_google_absl//absl/base:core_headers", "@com_google_absl//absl/cleanup", @@ -1366,37 +1332,6 @@ cc_library( "@com_google_absl//absl/strings", "@com_google_absl//absl/types:span", "@local_config_cuda//cuda:cuda_headers", - "@local_config_cuda//cuda:nvjitlink", # buildcleaner: keep - ], -) - -# Since select() can't be nested, we need to wrap the cuda_newer_than check in a separate -# library target. -cc_library( - name = "nvjitlink_cuda_supported", - # Even though the macro is called `*_newer_than`, it does a greater-than-or-equal-to comparison. - deps = if_cuda_newer_than( - "12_0", - [":nvjitlink_impl"], - [":nvjitlink_stub"], - ), -) - -cc_library( - name = "nvjitlink", - hdrs = [ - "nvjitlink.h", - ], - deps = select({ - ":libnvjitlink_support_enabled": [":nvjitlink_cuda_supported"], - "//conditions:default": [":nvjitlink_stub"], - }) + [ - ":compilation_provider", - ":cuda_compute_capability", - "//xla/stream_executor:device_description", - "//xla/stream_executor/gpu:gpu_asm_opts", - "@com_google_absl//absl/status:statusor", - "@com_google_absl//absl/types:span", ], ) @@ -1412,19 +1347,20 @@ xla_cc_test( ]), # The test fails with msan/tsan since its introduction. tags = [ + "cuda-only", + "gpu", "nomsan", "notsan", ], deps = [ + ":compilation_provider", ":cuda_compute_capability", ":nvjitlink", ":nvjitlink_support", "//xla/stream_executor:kernel_stats", "//xla/stream_executor/gpu:gpu_asm_opts", - "//xla/tsl/platform:statusor", "@com_google_absl//absl/status", "@com_google_absl//absl/status:status_matchers", - "@com_google_absl//absl/strings", "@com_google_absl//absl/types:span", "@com_google_googletest//:gtest", "@com_google_googletest//:gtest_main", @@ -1435,6 +1371,10 @@ cc_library( name = "nvjitlink_known_issues", srcs = ["nvjitlink_known_issues.cc"], hdrs = ["nvjitlink_known_issues.h"], + tags = [ + "cuda-only", + "gpu", + ], deps = [ ":nvjitlink", "@com_google_absl//absl/status:statusor", @@ -1446,15 +1386,17 @@ xla_cc_test( srcs = ["nvjitlink_known_issues_test.cc"], # LibNvJitLink is a binary-only library. Therefore is not compatible with msan/tsan. tags = [ + "cuda-only", + "gpu", "nomsan", "notsan", ], deps = [ + ":nvjitlink", ":nvjitlink_known_issues", ":nvjitlink_support", "@com_google_googletest//:gtest", "@com_google_googletest//:gtest_main", - "@tsl//tsl/platform:test", ], ) @@ -2448,6 +2390,10 @@ cc_library( name = "nvjitlink_compilation_provider", srcs = ["nvjitlink_compilation_provider.cc"], hdrs = ["nvjitlink_compilation_provider.h"], + tags = [ + "cuda-only", + "gpu", + ], deps = [ ":compilation_options", ":compilation_provider", @@ -2669,6 +2615,7 @@ cc_library( ":composite_compilation_provider", ":defer_relocatable_compilation_compilation_provider", ":driver_compilation_provider", + ":nvjitlink", ":nvjitlink_compilation_provider", ":nvjitlink_known_issues", ":nvjitlink_support", diff --git a/third_party/xla/xla/stream_executor/cuda/assemble_compilation_provider.cc b/third_party/xla/xla/stream_executor/cuda/assemble_compilation_provider.cc index 618c597b393d10..8b2956fe7e4381 100644 --- a/third_party/xla/xla/stream_executor/cuda/assemble_compilation_provider.cc +++ b/third_party/xla/xla/stream_executor/cuda/assemble_compilation_provider.cc @@ -32,6 +32,7 @@ limitations under the License. #include "xla/stream_executor/cuda/composite_compilation_provider.h" #include "xla/stream_executor/cuda/defer_relocatable_compilation_compilation_provider.h" #include "xla/stream_executor/cuda/driver_compilation_provider.h" +#include "xla/stream_executor/cuda/nvjitlink.h" #include "xla/stream_executor/cuda/nvjitlink_compilation_provider.h" #include "xla/stream_executor/cuda/nvjitlink_known_issues.h" #include "xla/stream_executor/cuda/nvjitlink_support.h" @@ -64,6 +65,17 @@ absl::Status HasNvJitLinkSupport(const CompilationProviderOptions& options) { return absl::OkStatus(); } + // When built with the nvjitlink stub library, libnvJitLink.so is loaded + // dynamically at runtime and may not be present in the environment. Probe it + // via GetNvJitLinkVersion() first so auto mode can fall back gracefully. + if (absl::StatusOr version = GetNvJitLinkVersion(); + !version.ok()) { + return absl::UnavailableError( + absl::StrCat("LibNvJitLink is disabled in auto mode because the " + "library could not be loaded: ", + version.status().message())); + } + if (LoadedNvJitLinkHasKnownIssues()) { return absl::UnavailableError( "LibNvJitLink is disabled since the loaded library version has known " diff --git a/third_party/xla/xla/stream_executor/cuda/nvjitlink_impl.cc b/third_party/xla/xla/stream_executor/cuda/nvjitlink.cc similarity index 91% rename from third_party/xla/xla/stream_executor/cuda/nvjitlink_impl.cc rename to third_party/xla/xla/stream_executor/cuda/nvjitlink.cc index 30056ad6b96853..9b043eab05fa1f 100644 --- a/third_party/xla/xla/stream_executor/cuda/nvjitlink_impl.cc +++ b/third_party/xla/xla/stream_executor/cuda/nvjitlink.cc @@ -13,6 +13,8 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ +#include "xla/stream_executor/cuda/nvjitlink.h" + #include #include #include @@ -39,12 +41,29 @@ limitations under the License. #include "third_party/gpus/cuda/include/nvJitLink.h" #include "xla/stream_executor/cuda/compilation_provider.h" #include "xla/stream_executor/cuda/cuda_compute_capability.h" -#include "xla/stream_executor/cuda/nvjitlink.h" #include "xla/stream_executor/cuda/ptx_compiler_helpers.h" #include "xla/stream_executor/gpu/gpu_asm_opts.h" #include "xla/stream_executor/kernel_stats.h" -#include "xla/tsl/platform/errors.h" -#include "xla/tsl/platform/statusor.h" + +// Workaround for a bug in nvJitLink.h prior to CUDA 13.1: when +// NVJITLINK_NO_INLINE is defined, the header omits the `extern` declarations +// for the unversioned nvJitLink* functions instead of declaring them. This bug +// was fixed in CUDA 13.1. +// FIXME: Remove this workaround once the minimum supported CUDA version is at +// least 13.1. +extern "C" { +nvJitLinkResult nvJitLinkCreate(nvJitLinkHandle*, uint32_t, const char**); +nvJitLinkResult nvJitLinkDestroy(nvJitLinkHandle*); +nvJitLinkResult nvJitLinkAddData(nvJitLinkHandle, nvJitLinkInputType, + const void*, size_t, const char*); +nvJitLinkResult nvJitLinkComplete(nvJitLinkHandle); +nvJitLinkResult nvJitLinkGetLinkedCubinSize(nvJitLinkHandle, size_t*); +nvJitLinkResult nvJitLinkGetLinkedCubin(nvJitLinkHandle, void*); +nvJitLinkResult nvJitLinkGetErrorLogSize(nvJitLinkHandle, size_t*); +nvJitLinkResult nvJitLinkGetErrorLog(nvJitLinkHandle, char*); +nvJitLinkResult nvJitLinkGetInfoLogSize(nvJitLinkHandle, size_t*); +nvJitLinkResult nvJitLinkGetInfoLog(nvJitLinkHandle, char*); +} // extern "C" namespace stream_executor { diff --git a/third_party/xla/xla/stream_executor/cuda/nvjitlink_known_issues_test.cc b/third_party/xla/xla/stream_executor/cuda/nvjitlink_known_issues_test.cc index 73b56b04fbe05e..03fe8a6e12f7ad 100644 --- a/third_party/xla/xla/stream_executor/cuda/nvjitlink_known_issues_test.cc +++ b/third_party/xla/xla/stream_executor/cuda/nvjitlink_known_issues_test.cc @@ -15,19 +15,24 @@ limitations under the License. #include "xla/stream_executor/cuda/nvjitlink_known_issues.h" +#include #include +#include "xla/stream_executor/cuda/nvjitlink.h" #include "xla/stream_executor/cuda/nvjitlink_support.h" -#include "tsl/platform/test.h" namespace { -TEST(NvJitLinkKnownIssuesTest, ReturnsFalseWhenNvJitLinkIsNotAvailable) { - if (stream_executor::IsLibNvJitLinkSupported()) { - GTEST_SKIP(); +TEST(NvJitLinkKnownIssuesTest, ReturnsExpectedResultForLoadedVersion) { + if (!stream_executor::IsLibNvJitLinkSupported()) { + EXPECT_FALSE(stream_executor::LoadedNvJitLinkHasKnownIssues()); + return; } - // This is the only invariance we can test without writing a pointless change - // detector test. - EXPECT_FALSE(stream_executor::LoadedNvJitLinkHasKnownIssues()); + ASSERT_OK_AND_ASSIGN(stream_executor::NvJitLinkVersion version, + stream_executor::GetNvJitLinkVersion()); + constexpr stream_executor::NvJitLinkVersion kMinVersionWithoutKnownIssues{12, + 5}; + EXPECT_EQ(stream_executor::LoadedNvJitLinkHasKnownIssues(), + version < kMinVersionWithoutKnownIssues); } } // namespace diff --git a/third_party/xla/xla/stream_executor/cuda/nvjitlink_stub.cc b/third_party/xla/xla/stream_executor/cuda/nvjitlink_stub.cc deleted file mode 100644 index 4c3ddc4b5efd32..00000000000000 --- a/third_party/xla/xla/stream_executor/cuda/nvjitlink_stub.cc +++ /dev/null @@ -1,40 +0,0 @@ -/* Copyright 2024 The OpenXLA Authors. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -==============================================================================*/ - -#include "absl/status/status.h" -#include "absl/status/statusor.h" -#include "absl/types/span.h" -#include "xla/stream_executor/cuda/compilation_provider.h" -#include "xla/stream_executor/cuda/cuda_compute_capability.h" -#include "xla/stream_executor/cuda/nvjitlink.h" -#include "xla/stream_executor/gpu/gpu_asm_opts.h" - -namespace stream_executor { - -absl::StatusOr GetNvJitLinkVersion() { - return absl::UnimplementedError("libnvjitlink is not supported"); -} - -absl::StatusOr CompileAndLinkUsingLibNvJitLink( - const CudaComputeCapability& cc, absl::Span inputs, - GpuAsmOpts options, bool cancel_if_reg_spill, bool dump_compilation_log) { - return absl::UnimplementedError("libnvjitlink is not supported"); -} - -absl::StatusOr GetLatestPtxIsaVersionForLibNvJitLink() { - return absl::UnimplementedError("libnvjitlink is not supported"); -} - -} // namespace stream_executor diff --git a/third_party/xla/xla/stream_executor/cuda/nvjitlink_support.cc b/third_party/xla/xla/stream_executor/cuda/nvjitlink_support.cc index 0a1a870ff181da..48fbee2cf4cbf4 100644 --- a/third_party/xla/xla/stream_executor/cuda/nvjitlink_support.cc +++ b/third_party/xla/xla/stream_executor/cuda/nvjitlink_support.cc @@ -24,7 +24,7 @@ bool IsLibNvJitLinkSupported() { #ifdef MEMORY_SANITIZER return false; #else - return LIBNVJITLINK_SUPPORT && CUDA_SUPPORTS_NVJITLINK; + return true; #endif } } // namespace stream_executor diff --git a/third_party/xla/xla/stream_executor/cuda/nvjitlink_test.cc b/third_party/xla/xla/stream_executor/cuda/nvjitlink_test.cc index 305b078601322d..4785528a129f78 100644 --- a/third_party/xla/xla/stream_executor/cuda/nvjitlink_test.cc +++ b/third_party/xla/xla/stream_executor/cuda/nvjitlink_test.cc @@ -15,24 +15,20 @@ limitations under the License. #include "xla/stream_executor/cuda/nvjitlink.h" -#include - #include #include -#include #include #include #include #include "absl/status/status.h" #include "absl/status/status_matchers.h" // IWYU pragma: keep -#include "absl/strings/str_replace.h" #include "absl/types/span.h" +#include "xla/stream_executor/cuda/compilation_provider.h" #include "xla/stream_executor/cuda/cuda_compute_capability.h" #include "xla/stream_executor/cuda/nvjitlink_support.h" #include "xla/stream_executor/gpu/gpu_asm_opts.h" #include "xla/stream_executor/kernel_stats.h" -#include "xla/tsl/platform/statusor.h" namespace { @@ -43,7 +39,7 @@ namespace { // constexpr const char kDependeePtx[] = R"( .version 8.0 -.target sm_52 +.target sm_80 .address_size 64 // .globl _Z5magicv @@ -64,7 +60,7 @@ constexpr const char kDependeePtx[] = R"( // constexpr const char kDependentPtx[] = R"( .version 8.0 -.target sm_52 +.target sm_80 .address_size 64 // .globl _Z6kernelPi @@ -75,7 +71,6 @@ constexpr const char kDependentPtx[] = R"( .visible .entry _Z6kernelPi( .param .u64 _Z6kernelPi_param_0 ) -// Insert .maxnreg directive here! { .reg .b32 %r<2>; .reg .b64 %rd<3>; @@ -102,7 +97,7 @@ constexpr const char kDependentPtx[] = R"( // constexpr const char kStandalonePtx[] = R"( .version 8.0 -.target sm_52 +.target sm_80 .address_size 64 // .globl _Z6kernelPi @@ -123,8 +118,112 @@ constexpr const char kStandalonePtx[] = R"( })"; -constexpr stream_executor::CudaComputeCapability kDefaultComputeCapability{5, - 2}; +// Generated by the following command: +// +// echo "__global__ void kernel(int* in) { for (int i=0; i < 25; i++) \ +// { in[i] += i; } for (int i=0; i < 25; i++) { in[24-i] += i; }}" \ +// | nvcc -arch=sm_80 -o - -rdc true --ptx --x cu - +// +// The `.maxnreg` directive was added manually afterwards. +constexpr const char kSpillingPtx[] = R"( +.version 8.0 +.target sm_80 +.address_size 64 + + // .globl _Z6kernelPi + +.visible .entry _Z6kernelPi( + .param .u64 _Z6kernelPi_param_0 +) + .maxnreg 16 +{ + .reg .b32 %r<51>; + .reg .b64 %rd<3>; + + + ld.param.u64 %rd1, [_Z6kernelPi_param_0]; + cvta.to.global.u64 %rd2, %rd1; + ld.global.u32 %r1, [%rd2+4]; + ld.global.u32 %r2, [%rd2+8]; + ld.global.u32 %r3, [%rd2+12]; + ld.global.u32 %r4, [%rd2+16]; + ld.global.u32 %r5, [%rd2+20]; + ld.global.u32 %r6, [%rd2+24]; + ld.global.u32 %r7, [%rd2+28]; + ld.global.u32 %r8, [%rd2+32]; + ld.global.u32 %r9, [%rd2+36]; + ld.global.u32 %r10, [%rd2+40]; + ld.global.u32 %r11, [%rd2+44]; + ld.global.u32 %r12, [%rd2+48]; + ld.global.u32 %r13, [%rd2+52]; + ld.global.u32 %r14, [%rd2+56]; + ld.global.u32 %r15, [%rd2+60]; + ld.global.u32 %r16, [%rd2+64]; + ld.global.u32 %r17, [%rd2+68]; + ld.global.u32 %r18, [%rd2+72]; + ld.global.u32 %r19, [%rd2+76]; + ld.global.u32 %r20, [%rd2+80]; + ld.global.u32 %r21, [%rd2+84]; + ld.global.u32 %r22, [%rd2+88]; + ld.global.u32 %r23, [%rd2+92]; + ld.global.u32 %r24, [%rd2+96]; + add.s32 %r25, %r24, 24; + st.global.u32 [%rd2+96], %r25; + add.s32 %r26, %r23, 24; + st.global.u32 [%rd2+92], %r26; + add.s32 %r27, %r22, 24; + st.global.u32 [%rd2+88], %r27; + add.s32 %r28, %r21, 24; + st.global.u32 [%rd2+84], %r28; + add.s32 %r29, %r20, 24; + st.global.u32 [%rd2+80], %r29; + add.s32 %r30, %r19, 24; + st.global.u32 [%rd2+76], %r30; + add.s32 %r31, %r18, 24; + st.global.u32 [%rd2+72], %r31; + add.s32 %r32, %r17, 24; + st.global.u32 [%rd2+68], %r32; + add.s32 %r33, %r16, 24; + st.global.u32 [%rd2+64], %r33; + add.s32 %r34, %r15, 24; + st.global.u32 [%rd2+60], %r34; + add.s32 %r35, %r14, 24; + st.global.u32 [%rd2+56], %r35; + add.s32 %r36, %r13, 24; + st.global.u32 [%rd2+52], %r36; + add.s32 %r37, %r12, 24; + st.global.u32 [%rd2+48], %r37; + add.s32 %r38, %r11, 24; + st.global.u32 [%rd2+44], %r38; + add.s32 %r39, %r10, 24; + st.global.u32 [%rd2+40], %r39; + add.s32 %r40, %r9, 24; + st.global.u32 [%rd2+36], %r40; + add.s32 %r41, %r8, 24; + st.global.u32 [%rd2+32], %r41; + add.s32 %r42, %r7, 24; + st.global.u32 [%rd2+28], %r42; + add.s32 %r43, %r6, 24; + st.global.u32 [%rd2+24], %r43; + add.s32 %r44, %r5, 24; + st.global.u32 [%rd2+20], %r44; + add.s32 %r45, %r4, 24; + st.global.u32 [%rd2+16], %r45; + add.s32 %r46, %r3, 24; + st.global.u32 [%rd2+12], %r46; + add.s32 %r47, %r2, 24; + st.global.u32 [%rd2+8], %r47; + add.s32 %r48, %r1, 24; + st.global.u32 [%rd2+4], %r48; + ld.global.u32 %r49, [%rd2]; + add.s32 %r50, %r49, 24; + st.global.u32 [%rd2], %r50; + ret; +} +)"; + +constexpr stream_executor::CudaComputeCapability kDefaultComputeCapability{8, + 0}; // Just a helper function that wraps `CompileAndLinkUsingLibNvJitLink`. It makes // it easier to deal with C string PTX literals. @@ -187,13 +286,9 @@ TEST_F(NvJitLinkTest, CanAlsoJustCompileSingleCompilationUnit) { } TEST_F(NvJitLinkTest, CancelsOnRegSpill) { - std::string dependent_ptx = absl::StrReplaceAll( - kDependentPtx, {{"// Insert .maxnreg directive here!", ".maxnreg 16"}}); - // We have to disable optimization here, otherwise PTXAS will optimize our // trivial register usages away and we don't spill as intended. - EXPECT_THAT(CompileAndLinkHelper(kDefaultComputeCapability, - {dependent_ptx.c_str(), kDependeePtx}, + EXPECT_THAT(CompileAndLinkHelper(kDefaultComputeCapability, {kSpillingPtx}, /*disable_gpuasm_optimizations=*/true, /*cancel_if_reg_spill=*/true), absl_testing::StatusIs(absl::StatusCode::kCancelled)); @@ -201,8 +296,7 @@ TEST_F(NvJitLinkTest, CancelsOnRegSpill) { // We also test the converse to ensure our test case isn't broken. ASSERT_OK_AND_ASSIGN( stream_executor::cuda::Assembly assembly, - CompileAndLinkHelper(kDefaultComputeCapability, - {dependent_ptx.c_str(), kDependeePtx}, + CompileAndLinkHelper(kDefaultComputeCapability, {kSpillingPtx}, /*disable_gpuasm_optimizations=*/true, /*cancel_if_reg_spill=*/false)); ASSERT_EQ(assembly.module_stats.size(), 1); diff --git a/third_party/xla/xla/stream_executor/cuda/topk_kernel_cuda_common.cu.h b/third_party/xla/xla/stream_executor/cuda/topk_kernel_cuda_common.cu.h index 046e65812860d8..3bb05e6822e6dc 100644 --- a/third_party/xla/xla/stream_executor/cuda/topk_kernel_cuda_common.cu.h +++ b/third_party/xla/xla/stream_executor/cuda/topk_kernel_cuda_common.cu.h @@ -81,6 +81,20 @@ __device__ __forceinline__ auto ToOrdered(T x) { return x; } } + +template +__device__ __forceinline__ T FromOrdered(OrderedT val) { + if constexpr (sizeof(T) == 4 && !std::is_integral_v) { + uint32_t u = (val & 0x80000000u) ? (val ^ 0x80000000u) : ~val; + return absl::bit_cast(u); + } else if constexpr (sizeof(T) == 2 && !std::is_integral_v) { + uint16_t u = (val & 0x8000u) ? static_cast(val ^ 0x8000u) + : static_cast(~val); + return absl::bit_cast(u); + } else { + return val; + } +} } // namespace details // Default implementation for KV holder. Useful for testing while adding support @@ -88,15 +102,14 @@ __device__ __forceinline__ auto ToOrdered(T x) { // implementations below. template struct Descending { + using OrderedKey = decltype(details::ToOrdered(T{})); struct KVT { - T key; + OrderedKey key; V idx; }; __device__ __forceinline__ static bool cmp(const KVT& lhs, const KVT& rhs) { - auto l = details::ToOrdered(lhs.key); - auto r = details::ToOrdered(rhs.key); - return l == r ? lhs.idx < rhs.idx : l > r; + return lhs.key == rhs.key ? lhs.idx < rhs.idx : lhs.key > rhs.key; } }; @@ -189,7 +202,7 @@ struct TopK { // TODO(doak): Use bitonic sort. #pragma unroll for (int i = 0; i < K; i++) { - tmp[i] = {key[Idx(i)], VT(Idx(i))}; + tmp[i] = {details::ToOrdered(key[Idx(i)]), VT(Idx(i))}; } #pragma unroll for (int i = 0; i < K; i++) { @@ -205,7 +218,7 @@ struct TopK { constexpr uint32_t WarpSize = WAVEFRONT_SIZE; for (int idx = K; idx < n; idx++) { - KVT kv{key[Idx(idx)], VT(Idx(idx))}; + KVT kv{details::ToOrdered(key[Idx(idx)]), VT(Idx(idx))}; Push(tmp, kv); } Reduce(tmp, WarpSize); @@ -239,7 +252,7 @@ struct TopK { return; } for (int i = 0; i < num_outputs_; ++i) { - keys[i] = tmp[i].key; + keys[i] = details::FromOrdered(tmp[i].key); idxs[i] = tmp[i].idx; } } diff --git a/third_party/xla/xla/stream_executor/rocm/rocm_dnn.cc b/third_party/xla/xla/stream_executor/rocm/rocm_dnn.cc index 39f1ff7fa0afe3..c8935d0a0308bb 100644 --- a/third_party/xla/xla/stream_executor/rocm/rocm_dnn.cc +++ b/third_party/xla/xla/stream_executor/rocm/rocm_dnn.cc @@ -2966,10 +2966,11 @@ absl::Status MIOpenSupport::GetConvolveRunners( } std::vector profile_results; - if (!GetMIOpenConvolveAlgorithms( + if (!GetMIOpenConvolveAlgorithmsInternal( kind, input_type, output_type, stream, input_descriptor, input_data, filter_descriptor, filter_data, output_descriptor, output_data, - convolution_descriptor, scratch_allocator, &profile_results)) + convolution_descriptor, scratch_allocator, &profile_results, + use_fallback)) return absl::InternalError("GetMIOpenConvolveAlgorithms failure"); for (const auto& profile_result : profile_results) { @@ -3092,13 +3093,33 @@ bool MIOpenSupport::GetMIOpenConvolveAlgorithms( const dnn::ConvolutionDescriptor& convolution_descriptor, ScratchAllocator* scratch_allocator, std::vector* out_algorithms) { + return GetMIOpenConvolveAlgorithmsInternal( + kind, input_type, output_type, stream, input_descriptor, input_data, + filter_descriptor, filter_data, output_descriptor, output_data, + convolution_descriptor, scratch_allocator, out_algorithms, + /*use_fallback=*/false); +} + +bool MIOpenSupport::GetMIOpenConvolveAlgorithmsInternal( + dnn::ConvolutionKind kind, dnn::DataType input_type, + dnn::DataType output_type, Stream* stream, + const dnn::BatchDescriptor& input_descriptor, DeviceAddressBase input_data, + const dnn::FilterDescriptor& filter_descriptor, + DeviceAddressBase filter_data, + const dnn::BatchDescriptor& output_descriptor, + DeviceAddressBase output_data, + const dnn::ConvolutionDescriptor& convolution_descriptor, + ScratchAllocator* scratch_allocator, + std::vector* out_algorithms, bool use_fallback) { // TODO(rocm): Create handles only once and reuse them between the methods - if (!PopulateMIOpenFindDb(kind, input_type, output_type, stream, - input_descriptor, input_data, filter_descriptor, - filter_data, output_descriptor, output_data, - convolution_descriptor, scratch_allocator) - .ok()) { - return false; + if (!use_fallback) { + if (!PopulateMIOpenFindDb(kind, input_type, output_type, stream, + input_descriptor, input_data, filter_descriptor, + filter_data, output_descriptor, output_data, + convolution_descriptor, scratch_allocator) + .ok()) { + return false; + } } return GetMIOpenConvolveAlgorithmsImmediateMode( kind, input_type, output_type, stream, input_descriptor, diff --git a/third_party/xla/xla/stream_executor/rocm/rocm_dnn.h b/third_party/xla/xla/stream_executor/rocm/rocm_dnn.h index c293839253e1f9..e2f3044166096f 100644 --- a/third_party/xla/xla/stream_executor/rocm/rocm_dnn.h +++ b/third_party/xla/xla/stream_executor/rocm/rocm_dnn.h @@ -256,6 +256,19 @@ class MIOpenSupport : public dnn::DnnSupport { ScratchAllocator* scratch_allocator, std::vector* out_algorithms) override; + bool GetMIOpenConvolveAlgorithmsInternal( + dnn::ConvolutionKind kind, dnn::DataType input_type, + dnn::DataType output_type, Stream* stream, + const dnn::BatchDescriptor& input_descriptor, + DeviceAddressBase input_data, + const dnn::FilterDescriptor& filter_descriptor, + DeviceAddressBase filter_data, + const dnn::BatchDescriptor& output_descriptor, + DeviceAddressBase output_data, + const dnn::ConvolutionDescriptor& convolution_descriptor, + ScratchAllocator* scratch_allocator, + std::vector* out_algorithms, bool use_fallback); + absl::Status GetMIOpenConvolveAlgorithmsImmediateMode( dnn::ConvolutionKind kind, dnn::DataType input_type, dnn::DataType output_type, Stream* stream, diff --git a/third_party/xla/xla/stream_executor/rocm/topk_kernel_rocm_common.cu.h b/third_party/xla/xla/stream_executor/rocm/topk_kernel_rocm_common.cu.h index df2b8db853fa68..6742958e8d85c8 100644 --- a/third_party/xla/xla/stream_executor/rocm/topk_kernel_rocm_common.cu.h +++ b/third_party/xla/xla/stream_executor/rocm/topk_kernel_rocm_common.cu.h @@ -85,6 +85,20 @@ __device__ __forceinline__ auto ToOrdered(T x) { return x; } } + +template +__device__ __forceinline__ T FromOrdered(OrderedT val) { + if constexpr (sizeof(T) == 4 && !std::is_integral_v) { + uint32_t u = (val & 0x80000000u) ? (val ^ 0x80000000u) : ~val; + return absl::bit_cast(u); + } else if constexpr (sizeof(T) == 2 && !std::is_integral_v) { + uint16_t u = (val & 0x8000u) ? static_cast(val ^ 0x8000u) + : static_cast(~val); + return absl::bit_cast(u); + } else { + return val; + } +} } // namespace details // Default implementation for KV holder. Useful for testing while adding support @@ -92,15 +106,14 @@ __device__ __forceinline__ auto ToOrdered(T x) { // implementations below. template struct Descending { + using OrderedKey = decltype(details::ToOrdered(T{})); struct KVT { - T key; + OrderedKey key; V idx; }; __device__ __forceinline__ static bool cmp(const KVT& lhs, const KVT& rhs) { - auto l = details::ToOrdered(lhs.key); - auto r = details::ToOrdered(rhs.key); - return l == r ? lhs.idx < rhs.idx : l > r; + return lhs.key == rhs.key ? lhs.idx < rhs.idx : lhs.key > rhs.key; } }; @@ -193,7 +206,7 @@ struct TopK { // TODO(doak): Use bitonic sort. #pragma unroll for (int i = 0; i < K; i++) { - tmp[i] = {key[Idx(i)], VT(Idx(i))}; + tmp[i] = {details::ToOrdered(key[Idx(i)]), VT(Idx(i))}; } #pragma unroll for (int i = 0; i < K; i++) { @@ -209,7 +222,7 @@ struct TopK { constexpr uint32_t WarpSize = WAVEFRONT_SIZE; for (int idx = K; idx < n; idx++) { - KVT kv{key[Idx(idx)], VT(Idx(idx))}; + KVT kv{details::ToOrdered(key[Idx(idx)]), VT(Idx(idx))}; Push(tmp, kv); } Reduce(tmp, WarpSize); @@ -237,7 +250,7 @@ struct TopK { Reduce(tmp, blockDim.x / WarpSize); if (threadIdx.x != 0) return; for (int i = 0; i < num_outputs_; ++i) { - keys[i] = tmp[i].key; + keys[i] = details::FromOrdered(tmp[i].key); idxs[i] = tmp[i].idx; } } diff --git a/third_party/xla/xla/tests/constraint_propagator.cc b/third_party/xla/xla/tests/constraint_propagator.cc index efc9489b1e2602..ca487d25a098a7 100644 --- a/third_party/xla/xla/tests/constraint_propagator.cc +++ b/third_party/xla/xla/tests/constraint_propagator.cc @@ -329,7 +329,7 @@ std::optional GetSymmetricMagnitudeBound( const ConstraintInterval& interval) { if ((interval.min == ConstraintInterval::kMin && interval.max == ConstraintInterval::kMax) || - interval.min > 0.0 || interval.max < 0.0) { + interval.min >= 0.0 || interval.max <= 0.0) { return std::nullopt; } double m = std::min(std::abs(interval.min), std::abs(interval.max)); diff --git a/third_party/xla/xla/tests/constraint_propagator_test.cc b/third_party/xla/xla/tests/constraint_propagator_test.cc index cd35a5753dd931..e7f9f1e32677cf 100644 --- a/third_party/xla/xla/tests/constraint_propagator_test.cc +++ b/third_party/xla/xla/tests/constraint_propagator_test.cc @@ -157,10 +157,8 @@ ENTRY main { auto y_int = states[module->entry_computation()->parameter_instruction(1)] .GetConstraintInterval(); - EXPECT_TRUE(x_int.IsPositive()); - EXPECT_FALSE(x_int.exclude_zero); - EXPECT_TRUE(y_int.IsNegative()); - EXPECT_FALSE(y_int.exclude_zero); + EXPECT_EQ(x_int, ConstraintInterval::Positive()); + EXPECT_EQ(y_int, ConstraintInterval::Negative()); } TEST_F(ConstraintPropagatorTest, SubtractNegative) { @@ -183,10 +181,8 @@ ENTRY main { auto y_int = states[module->entry_computation()->parameter_instruction(1)] .GetConstraintInterval(); - EXPECT_TRUE(x_int.IsNegative()); - EXPECT_FALSE(x_int.exclude_zero); - EXPECT_TRUE(y_int.IsPositive()); - EXPECT_FALSE(y_int.exclude_zero); + EXPECT_EQ(x_int, ConstraintInterval::Negative()); + EXPECT_EQ(y_int, ConstraintInterval::Positive()); } TEST_F(ConstraintPropagatorTest, RsqrtMulForcedNegative) { @@ -961,9 +957,11 @@ ENTRY main { auto p1_int = states[module->entry_computation()->parameter_instruction(1)] .GetConstraintInterval(); EXPECT_FALSE(p0_int.IsEmpty()); - EXPECT_LE(p0_int.max, 0.0); + EXPECT_DOUBLE_EQ(p0_int.min, -4.0); + EXPECT_DOUBLE_EQ(p0_int.max, 0.0); EXPECT_FALSE(p1_int.IsEmpty()); - EXPECT_GE(p1_int.min, 0.0); + EXPECT_DOUBLE_EQ(p1_int.min, 0.0); + EXPECT_DOUBLE_EQ(p1_int.max, 4.0); } // Tests that for Power(base, exp) with base in (0, 1) (decay schedule, e.g. diff --git a/third_party/xla/xla/tools/multihost_hlo_runner/functional_hlo_runner.cc b/third_party/xla/xla/tools/multihost_hlo_runner/functional_hlo_runner.cc index 89f7202922dd19..a98bd795160457 100644 --- a/third_party/xla/xla/tools/multihost_hlo_runner/functional_hlo_runner.cc +++ b/third_party/xla/xla/tools/multihost_hlo_runner/functional_hlo_runner.cc @@ -373,9 +373,6 @@ absl::StatusOr FetchAndLogOutput( const std::vector>>& output_buffers, ModuleOutputMode module_output_mode, bool log_output) { CHECK(!output_buffers.empty()); - absl::Mutex mu; - absl::Status status; - size_t num_pending_transfers = 0; bool device_0_is_local = false; for (PjRtDevice* device : GetLocalDevices(client)) { if (device->id() == 0) { @@ -383,67 +380,81 @@ absl::StatusOr FetchAndLogOutput( } } - if (module_output_mode == ModuleOutputMode::kReturnDevice0Outputs && - device_0_is_local) { - num_pending_transfers = output_buffers[0].size(); - } else if (module_output_mode == ModuleOutputMode::kReturnOutputs) { - for (const auto& bs : output_buffers) { - num_pending_transfers += bs.size(); - } - } - PerDeviceLiteralVecType outputs; - for (int i = 0; i < output_buffers.size(); ++i) { - if (output_buffers[i].empty()) { - continue; - } - const int device_id = output_buffers[i][0]->device()->id(); - std::vector& output_slice = outputs[device_id]; - if (module_output_mode == ModuleOutputMode::kReturnOutputs || - (module_output_mode == ModuleOutputMode::kReturnDevice0Outputs && - device_id == 0)) { - output_slice.reserve(output_buffers[i].size()); - for (const auto& buffer : output_buffers[i]) { - TF_RET_CHECK(buffer->device() == output_buffers[i][0]->device()) - << "All outputs from a given vector of outputs should be for the " - "same device"; - ABSL_ASSIGN_OR_RETURN(auto logical_shape, buffer->logical_on_device_shape()); - output_slice.emplace_back( - ShapeUtil::DeviceShapeToHostShape(logical_shape)); - buffer->ToLiteral(&output_slice.back()).OnReady([&](absl::Status s) { - absl::MutexLock lock(mu); - --num_pending_transfers; - status.Update(s); - }); + absl::Mutex mu; + absl::Status status; + size_t num_pending_transfers = 0; + + absl::Status issue_status = [&]() -> absl::Status { + for (int i = 0; i < output_buffers.size(); ++i) { + if (output_buffers[i].empty()) { + continue; } - } else { - for (const auto& buffer : output_buffers[i]) { - TF_RET_CHECK(buffer->device() == output_buffers[i][0]->device()) - << "All outputs from a given vector of outputs should be for the " - "same device"; - ABSL_RETURN_IF_ERROR(buffer->GetReadyFuture().Await()); + const int device_id = output_buffers[i][0]->device()->id(); + std::vector& output_slice = outputs[device_id]; + if (module_output_mode == ModuleOutputMode::kReturnOutputs || + (module_output_mode == ModuleOutputMode::kReturnDevice0Outputs && + device_id == 0)) { + output_slice.reserve(output_buffers[i].size()); + for (const auto& buffer : output_buffers[i]) { + if (buffer->device() != output_buffers[i][0]->device()) { + return absl::InternalError( + "All outputs from a given vector of outputs should be for the " + "same device"); + } + ABSL_ASSIGN_OR_RETURN(auto logical_shape, + buffer->logical_on_device_shape()); + output_slice.emplace_back( + ShapeUtil::DeviceShapeToHostShape(logical_shape)); + { + absl::MutexLock lock(mu); + ++num_pending_transfers; + } + buffer->ToLiteral(&output_slice.back()).OnReady([&](absl::Status s) { + absl::MutexLock lock(mu); + --num_pending_transfers; + status.Update(s); + }); + } + } else { + for (const auto& buffer : output_buffers[i]) { + if (buffer->device() != output_buffers[i][0]->device()) { + return absl::InternalError( + "All outputs from a given vector of outputs should be for the " + "same device"); + } + ABSL_RETURN_IF_ERROR(buffer->GetReadyFuture().Await()); + } } } - } - if (module_output_mode == ModuleOutputMode::kReturnOutputs || - (module_output_mode == ModuleOutputMode::kReturnDevice0Outputs && - device_0_is_local)) { - auto cond = [&]() { return !status.ok() || num_pending_transfers == 0; }; + return absl::OkStatus(); + }(); + + // The ToLiteral callbacks above reference `outputs`, `mu`, `status` and + // `num_pending_transfers`, so every issued transfer must complete before this + // function returns, including when a transfer fails or issuing one fails. + { + auto all_transfers_done = [&]() { return num_pending_transfers == 0; }; absl::MutexLock lock(mu); - mu.Await(absl::Condition(&cond)); + mu.Await(absl::Condition(&all_transfers_done)); + ABSL_RETURN_IF_ERROR(issue_status); ABSL_RETURN_IF_ERROR(status); - if (log_output) { - for (const PjRtDevice* device : GetLocalDevices(client)) { - int device_id = device->id(); - if (module_output_mode == ModuleOutputMode::kReturnDevice0Outputs && - device_id != 0) { - continue; - } - LOG(INFO) << "Outputs for device_id: " << device_id; - const std::vector& output_slice = outputs[device_id]; - for (int i = 0; i < output_slice.size(); ++i) { - LOG(INFO) << "output[" << i << "]: " << output_slice[i].ToString(); - } + } + + if (log_output && + (module_output_mode == ModuleOutputMode::kReturnOutputs || + (module_output_mode == ModuleOutputMode::kReturnDevice0Outputs && + device_0_is_local))) { + for (const PjRtDevice* device : GetLocalDevices(client)) { + int device_id = device->id(); + if (module_output_mode == ModuleOutputMode::kReturnDevice0Outputs && + device_id != 0) { + continue; + } + LOG(INFO) << "Outputs for device_id: " << device_id; + const std::vector& output_slice = outputs[device_id]; + for (int i = 0; i < output_slice.size(); ++i) { + LOG(INFO) << "output[" << i << "]: " << output_slice[i].ToString(); } } } diff --git a/third_party/xla/xla/tsl/cuda/BUILD.bazel b/third_party/xla/xla/tsl/cuda/BUILD.bazel index 5b2105c7a56ba6..12b161e31ae14a 100644 --- a/third_party/xla/xla/tsl/cuda/BUILD.bazel +++ b/third_party/xla/xla/tsl/cuda/BUILD.bazel @@ -439,6 +439,55 @@ alias( visibility = ["//visibility:public"], ) +cuda_stub( + name = "nvjitlink", + srcs = ["nvjitlink.symbols"], +) + +cc_library( + name = "nvjitlink_stub", # buildifier: disable=duplicated-name + srcs = select({ + "@xla//xla/tsl:is_cuda_enabled_and_oss": [ + "nvjitlink.tramp.S", + "nvjitlink_stub.cc", + ], + "//conditions:default": [], + }), + defines = [ + # By default, nvJitLink.h defines static inline wrappers that forward to + # versioned __nvJitLink*__ symbols. Defining + # NVJITLINK_NO_INLINE makes consumers call the unversioned nvJitLink* + # symbols exported by nvjitlink.symbols instead. + "NVJITLINK_NO_INLINE", + ], + linkopts = if_cuda_is_configured(if_cuda_newer_than( + "13_0", + if_false = cuda_rpath_flags("nvidia/nvjitlink/lib"), + if_true = cuda_rpath_flags("nvidia/cu13/lib"), + )), + local_defines = [ + "IMPLIB_EXPORT_SHIMS=1", + ], + textual_hdrs = ["nvjitlink.inc"], + visibility = ["//visibility:public"], + deps = select({ + "@xla//xla/tsl:is_cuda_enabled_and_oss": [ + # keep sorted + "//xla/tsl/platform:logging", + "@local_config_cuda//cuda:cuda_headers", + "@tsl//tsl/platform:dso_loader", + "@tsl//tsl/platform:load_library", + ], + "//conditions:default": [], + }), +) + +alias( + name = "nvjitlink", # buildifier: disable=duplicated-name + actual = if_cuda_libs("@cuda_nvjitlink//:nvjitlink", ":nvjitlink_stub"), + visibility = ["//visibility:public"], +) + cuda_stub( name = "nvml", srcs = ["nvml.symbols"], diff --git a/third_party/xla/xla/tsl/cuda/nvjitlink.symbols b/third_party/xla/xla/tsl/cuda/nvjitlink.symbols new file mode 100644 index 00000000000000..805d74fb1f72c4 --- /dev/null +++ b/third_party/xla/xla/tsl/cuda/nvjitlink.symbols @@ -0,0 +1,16 @@ +nvJitLinkAddData +nvJitLinkAddFile +nvJitLinkComplete +nvJitLinkCreate +nvJitLinkDestroy +nvJitLinkGetErrorLog +nvJitLinkGetErrorLogSize +nvJitLinkGetInfoLog +nvJitLinkGetInfoLogSize +nvJitLinkGetLinkedCubin +nvJitLinkGetLinkedCubinSize +nvJitLinkGetLinkedLTOIR +nvJitLinkGetLinkedLTOIRSize +nvJitLinkGetLinkedPtx +nvJitLinkGetLinkedPtxSize +nvJitLinkVersion diff --git a/third_party/xla/xla/tsl/cuda/nvjitlink_stub.cc b/third_party/xla/xla/tsl/cuda/nvjitlink_stub.cc new file mode 100644 index 00000000000000..db753abb93e1c4 --- /dev/null +++ b/third_party/xla/xla/tsl/cuda/nvjitlink_stub.cc @@ -0,0 +1,73 @@ +/* Copyright 2026 The OpenXLA Authors. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#include + +#include "third_party/gpus/cuda/include/nvJitLink.h" +#include "xla/tsl/platform/logging.h" +#include "tsl/platform/dso_loader.h" +#include "tsl/platform/load_library.h" + +// Implements the nvJitLink API by forwarding to nvJitLink loaded from the DSO. + +namespace { +// Returns DSO handle or null if loading the DSO fails. +void* GetDsoHandle() { + static auto handle = []() -> void* { + auto handle_or = tsl::internal::DsoLoader::GetNvJitLinkDsoHandle(); + if (!handle_or.ok()) { + return nullptr; + } + return handle_or.value(); + }(); + return handle; +} + +void* LoadSymbol(const char* symbol_name) { + void* symbol = nullptr; + if (auto handle = GetDsoHandle()) { + tsl::internal::GetSymbolFromLibrary(handle, symbol_name, &symbol) + .IgnoreError(); + } + return symbol; +} + +const char* kSymbols[] = { +#include "xla/tsl/cuda/nvjitlink.inc" +}; + +constexpr size_t kNumSymbols = sizeof(kSymbols) / sizeof(const char*); + +} // namespace + +extern "C" { + +static nvJitLinkResult GetSymbolNotFoundError() { + return NVJITLINK_ERROR_INTERNAL; +} + +extern void* _nvjitlink_tramp_table[]; + +void _nvjitlink_tramp_resolve(int i) { + CHECK_LE(0, i); + CHECK_LT(i, kNumSymbols); + void* p = LoadSymbol(kSymbols[i]); + if (!p) { + p = reinterpret_cast(&GetSymbolNotFoundError); + } + _nvjitlink_tramp_table[i] = p; +} + +} // extern "C" diff --git a/third_party/xla/xla/tsl/platform/cloud/gcs_file_system_test.cc b/third_party/xla/xla/tsl/platform/cloud/gcs_file_system_test.cc index 76f2d24379504f..4e0fc2e020e358 100644 --- a/third_party/xla/xla/tsl/platform/cloud/gcs_file_system_test.cc +++ b/third_party/xla/xla/tsl/platform/cloud/gcs_file_system_test.cc @@ -634,7 +634,8 @@ TEST(GcsFileSystemTest, NewRandomAccessFile_WithBlockCache) { std::unique_ptr( new FakeHttpRequestFactory(&requests)), std::unique_ptr(new FakeZoneProvider), 9 /* block size */, - 18 /* max bytes */, 0 /* max staleness */, 3600 /* stat cache max age */, + 1 << 20 /* max bytes: large enough that nothing is evicted */, + 0 /* max staleness */, 3600 /* stat cache max age */, 0 /* stat cache max entries */, 0 /* matching paths cache max age */, 0 /* matching paths cache max entries */, kTestRetryConfig, kTestTimeoutConfig, *kAllowedLocationsDefault, @@ -723,7 +724,8 @@ TEST(GcsFileSystemTest, NewRandomAccessFile_WithBlockCache_Flush) { std::unique_ptr( new FakeHttpRequestFactory(&requests)), std::unique_ptr(new FakeZoneProvider), 9 /* block size */, - 18 /* max bytes */, 0 /* max staleness */, 3600 /* stat cache max age */, + 1 << 20 /* max bytes: large enough that nothing is evicted */, + 0 /* max staleness */, 3600 /* stat cache max age */, 0 /* stat cache max entries */, 0 /* matching paths cache max age */, 0 /* matching paths cache max entries */, kTestRetryConfig, kTestTimeoutConfig, *kAllowedLocationsDefault, @@ -843,7 +845,8 @@ TEST(GcsFileSystemTest, std::unique_ptr( new FakeHttpRequestFactory(&requests)), std::unique_ptr(new FakeZoneProvider), 9 /* block size */, - 18 /* max bytes */, 0 /* max staleness */, 0 /* stat cache max age */, + 1 << 20 /* max bytes: large enough that nothing is evicted */, + 0 /* max staleness */, 0 /* stat cache max age */, 0 /* stat cache max entries */, 0 /* matching paths cache max age */, 0 /* matching paths cache max entries */, kTestRetryConfig, kTestTimeoutConfig, *kAllowedLocationsDefault, diff --git a/third_party/xla/xla/tsl/platform/default/dso_loader.cc b/third_party/xla/xla/tsl/platform/default/dso_loader.cc index ed9ecb415e5322..88e8c8afb36e5a 100644 --- a/third_party/xla/xla/tsl/platform/default/dso_loader.cc +++ b/third_party/xla/xla/tsl/platform/default/dso_loader.cc @@ -86,7 +86,8 @@ absl::StatusOr GetCudaDsoHandle(const std::string& name, if (name == "cufft" || name == "cusolver") { alt_version = (version == "12") ? "11" : (version == "11" ? "12" : ""); } else if (name == "cublas" || name == "cublasLt" || name == "cudart" || - name == "nvrtc" || name == "cusparse" || name == "cupti") { + name == "nvrtc" || name == "cusparse" || name == "cupti" || + name == "nvJitLink") { alt_version = (version == "13") ? "12" : (version == "12" ? "13" : ""); } @@ -128,6 +129,10 @@ absl::StatusOr GetNvrtcDsoHandle() { return GetCudaDsoHandle("nvrtc", GetCudaRtVersion()); } +absl::StatusOr GetNvJitLinkDsoHandle() { + return GetCudaDsoHandle("nvJitLink", GetCudaRtVersion()); +} + absl::StatusOr GetCudaRuntimeDsoHandle() { return GetCudaDsoHandle("cudart", GetCudaRtVersion()); } @@ -232,6 +237,11 @@ absl::StatusOr GetCudnnDsoHandle() { return *result; } +absl::StatusOr GetNvJitLinkDsoHandle() { + static auto result = new auto(DsoLoader::GetNvJitLinkDsoHandle()); + return *result; +} + } // namespace CachedDsoLoader } // namespace internal } // namespace tsl diff --git a/third_party/xla/xla/tsl/platform/default/dso_loader.h b/third_party/xla/xla/tsl/platform/default/dso_loader.h index 32e80ef473ec78..ffbfcb48d876d8 100644 --- a/third_party/xla/xla/tsl/platform/default/dso_loader.h +++ b/third_party/xla/xla/tsl/platform/default/dso_loader.h @@ -16,6 +16,10 @@ limitations under the License. // Common DSO loading functionality: exposes callables that dlopen DSOs // in either the runfiles directories +// IWYU pragma: private, include "tsl/platform/dso_loader.h" +// IWYU pragma: friend ".*tsl/platform/.*dso_loader\.h" +// IWYU pragma: friend ".*tsl/platform/default/.*" + #ifndef XLA_TSL_PLATFORM_DEFAULT_DSO_LOADER_H_ #define XLA_TSL_PLATFORM_DEFAULT_DSO_LOADER_H_ @@ -41,6 +45,7 @@ absl::StatusOr GetNcclDsoHandle(); absl::StatusOr GetNvshmemDsoHandle(); absl::StatusOr GetNvInferDsoHandle(); absl::StatusOr GetNvInferPluginDsoHandle(); +absl::StatusOr GetNvJitLinkDsoHandle(); absl::StatusOr GetNvmlDsoHandle(); absl::StatusOr GetNvrtcDsoHandle(); @@ -69,6 +74,7 @@ absl::StatusOr GetCusolverDsoHandle(); absl::StatusOr GetCusparseDsoHandle(); absl::StatusOr GetCuptiDsoHandle(); absl::StatusOr GetCudnnDsoHandle(); +absl::StatusOr GetNvJitLinkDsoHandle(); } // namespace CachedDsoLoader } // namespace internal diff --git a/third_party/xla/xla/union_find.h b/third_party/xla/xla/union_find.h index c3001d9fdcb5fb..2f241a3737af13 100644 --- a/third_party/xla/xla/union_find.h +++ b/third_party/xla/xla/union_find.h @@ -16,6 +16,12 @@ limitations under the License. #ifndef XLA_UNION_FIND_H_ #define XLA_UNION_FIND_H_ +#include +#include +#include + +#include "absl/log/check.h" + namespace xla { // Union-Find data structure. @@ -77,6 +83,75 @@ UnionFind* UnionFind::FindRoot() { return parent_; } +// Union-Find optimized for a dense index range [0, num_elements). Use it +// instead of UnionFind when the elements already have dense ids (for +// example the local ids of a computation's instructions). Avoids per element +// allocations and the need for a hash_map from element to cluster: the whole +// structure is two flat arrays. +class DenseUnionFind { + public: + // Every element starts in its own cluster: both arrays zero initialized. + explicit DenseUnionFind(int num_elements) + : parent_plus_one_(CheckNonNegative(num_elements), 0), + rank_(num_elements, 0) {} + + int num_elements() const { return parent_plus_one_.size(); } + + // Returns the representative of the cluster containing 'element'. Compresses + // the path on the way up by path halving, the one pass variant of path + // compression: every other node on the path is pointed at its grandparent. + int Find(int element) { + DCHECK_GE(element, 0); + DCHECK_LT(element, num_elements()); + while (true) { + const int parent = Parent(element); + if (parent == element) { + return element; + } + const int grandparent = Parent(parent); + parent_plus_one_[element] = grandparent + 1; + element = grandparent; + } + } + + // Merges the clusters of 'a' and 'b' by rank: the root of higher rank + // becomes the representative of the merged cluster, and equal ranks grow the + // surviving root's rank by one. + void Merge(int a, int b) { + a = Find(a); + b = Find(b); + if (a == b) { + return; + } + if (rank_[a] < rank_[b]) { + std::swap(a, b); + } + if (rank_[a] == rank_[b]) { + ++rank_[a]; + } + parent_plus_one_[b] = a + 1; + } + + private: + static int CheckNonNegative(int count) { + DCHECK_GE(count, 0); + return count; + } + + // A root is its own parent. + int Parent(int element) const { + const int stored = parent_plus_one_[element]; + return stored == 0 ? element : stored - 1; + } + + // The parent of each element plus one; 0 marks a root, so the zero + // initialized array holds singletons and needs no iota. + std::vector parent_plus_one_; + // The rank of a root bounds the height of its tree and stays below + // log2(num_elements), so a byte holds it. + std::vector rank_; +}; + } // namespace xla #endif // XLA_UNION_FIND_H_ diff --git a/third_party/xla/xla/union_find_test.cc b/third_party/xla/xla/union_find_test.cc new file mode 100644 index 00000000000000..331be60ec12dd8 --- /dev/null +++ b/third_party/xla/xla/union_find_test.cc @@ -0,0 +1,209 @@ +/* Copyright 2026 The OpenXLA Authors. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#include "xla/union_find.h" + +#include +#include + +#include "xla/tsl/platform/test.h" +#include "xla/tsl/platform/test_benchmark.h" + +namespace xla { +namespace { + +using ::benchmark::DoNotOptimize; +using ::testing::benchmark::State; + +TEST(UnionFindTest, MergeKeepsTheCallersValue) { + UnionFind a(1); + UnionFind b(2); + UnionFind c(3); + + a.Merge(&b); + c.Merge(&a); + + EXPECT_EQ(a.Get(), 3); + EXPECT_EQ(b.Get(), 3); + EXPECT_EQ(c.Get(), 3); + EXPECT_EQ(a.Size(), 3); +} + +TEST(DenseUnionFindTest, ZeroElements) { + DenseUnionFind sets(0); + EXPECT_EQ(sets.num_elements(), 0); +} + +TEST(DenseUnionFindTest, StartsAsSingletons) { + DenseUnionFind sets(4); + + EXPECT_EQ(sets.num_elements(), 4); + for (int i = 0; i < 4; ++i) { + EXPECT_EQ(sets.Find(i), i); + } +} + +TEST(DenseUnionFindTest, MergeJoinsClusters) { + DenseUnionFind sets(6); + + sets.Merge(0, 1); + sets.Merge(2, 3); + sets.Merge(1, 3); + sets.Merge(4, 4); + + EXPECT_EQ(sets.Find(0), sets.Find(3)); + EXPECT_EQ(sets.Find(1), sets.Find(2)); + EXPECT_NE(sets.Find(0), sets.Find(4)); + EXPECT_NE(sets.Find(4), sets.Find(5)); + EXPECT_EQ(sets.Find(4), 4); + EXPECT_EQ(sets.Find(5), 5); +} + +// Union by rank: the root of higher rank represents the merged cluster, and a +// tie keeps the first argument's root and raises its rank. +TEST(DenseUnionFindTest, RepresentativeIsTheRootOfHigherRank) { + DenseUnionFind sets(5); + + sets.Merge(0, 1); // Tie: 0 stays the root, now of rank 1. + EXPECT_EQ(sets.Find(1), 0); + sets.Merge(2, 0); // 2 has rank 0, so 0 keeps representing the cluster. + EXPECT_EQ(sets.Find(2), 0); + sets.Merge(3, 4); // Tie: 3 becomes a root of rank 1. + sets.Merge(3, 0); // Tie between ranks 1: 3 stays the root, now of rank 2. + for (int i = 0; i < 5; ++i) { + EXPECT_EQ(sets.Find(i), 3); + } +} + +// Reference implementation: one label per element, relabelled on every merge. +class BruteForceSets { + public: + explicit BruteForceSets(int n) : label_(n) { + for (int i = 0; i < n; ++i) { + label_[i] = i; + } + } + bool SameSet(int a, int b) const { return label_[a] == label_[b]; } + void Merge(int a, int b) { + const int from = label_[b]; + const int to = label_[a]; + for (int& label : label_) { + if (label == from) { + label = to; + } + } + } + + private: + std::vector label_; +}; + +TEST(DenseUnionFindTest, MatchesBruteForceOnRandomMerges) { + constexpr int kNumInstances = 100; + constexpr int kNumMerges = 200; + std::mt19937 rng(11); + std::uniform_int_distribution size_dist(1, 60); + + for (int instance = 0; instance < kNumInstances; ++instance) { + const int n = size_dist(rng); + std::uniform_int_distribution element(0, n - 1); + DenseUnionFind sets(n); + BruteForceSets reference(n); + + for (int merge = 0; merge < kNumMerges; ++merge) { + const int a = element(rng); + const int b = element(rng); + sets.Merge(a, b); + reference.Merge(a, b); + const int x = element(rng); + const int y = element(rng); + ASSERT_EQ(sets.Find(x) == sets.Find(y), reference.SameSet(x, y)) + << "instance " << instance << " merge " << merge; + } + for (int x = 0; x < n; ++x) { + for (int y = 0; y < n; ++y) { + ASSERT_EQ(sets.Find(x) == sets.Find(y), reference.SameSet(x, y)) + << "instance " << instance << " pair " << x << ", " << y; + } + } + } +} + +// The access pattern of TryRemoveDeadWhileParams over the instruction slots of +// a loop, including the construction of the structure: +// - one merge per operand edge, +// - then one find per tuple index. +std::vector RandomEdges(int n) { + std::mt19937 rng(5); + std::uniform_int_distribution element(0, n - 1); + std::vector edges(2 * n); + for (int& e : edges) { + e = element(rng); + } + return edges; +} + +void BM_DenseUnionFindTryRemoveDeadWhileParamsPattern(State& state) { + const int n = state.range(0); + const std::vector edges = RandomEdges(n); + for (auto s : state) { + DenseUnionFind sets(n); + for (int i = 0; i + 1 < edges.size(); i += 2) { + sets.Merge(edges[i], edges[i + 1]); + } + int sink = 0; + for (int i = 0; i < n; ++i) { + sink += sets.Find(i); + } + DoNotOptimize(sink); + } + state.SetItemsProcessed(state.iterations() * 2 * n); +} +BENCHMARK(BM_DenseUnionFindTryRemoveDeadWhileParamsPattern) + ->Arg(1024) + ->Arg(100000) + ->Arg(1 << 20) + ->Arg(1 << 22); + +// The same pattern on UnionFind: one node per element, holding its index as +// the value, so a find is a Get. +void BM_UnionFindTryRemoveDeadWhileParamsPattern(State& state) { + const int n = state.range(0); + const std::vector edges = RandomEdges(n); + for (auto s : state) { + std::vector> sets; + sets.reserve(n); + for (int i = 0; i < n; ++i) { + sets.emplace_back(i); + } + for (int i = 0; i + 1 < edges.size(); i += 2) { + sets[edges[i]].Merge(&sets[edges[i + 1]]); + } + int sink = 0; + for (int i = 0; i < n; ++i) { + sink += sets[i].Get(); + } + DoNotOptimize(sink); + } + state.SetItemsProcessed(state.iterations() * 2 * n); +} +BENCHMARK(BM_UnionFindTryRemoveDeadWhileParamsPattern) + ->Arg(1024) + ->Arg(100000) + ->Arg(1 << 20) + ->Arg(1 << 22); + +} // namespace +} // namespace xla diff --git a/third_party/xla/xla/xla.proto b/third_party/xla/xla/xla.proto index a23060f7f1ee10..4a56ebe77e5e97 100644 --- a/third_party/xla/xla/xla.proto +++ b/third_party/xla/xla/xla.proto @@ -94,6 +94,7 @@ message DebugOptions { DYNAMIC_SLICE_FUSION = 9; CONVOLUTION = 11; HOST_EXECUTE = 12; + COLLECTIVES_KERNEL = 13; reserved "DYNAMIC_SLICE_COPY_FUSION"; reserved 10;