diff --git a/xla/codegen/emitters/transforms/atomic_rmw_utils.cc b/xla/codegen/emitters/transforms/atomic_rmw_utils.cc index 23388b8047289..ebfed15099286 100644 --- a/xla/codegen/emitters/transforms/atomic_rmw_utils.cc +++ b/xla/codegen/emitters/transforms/atomic_rmw_utils.cc @@ -84,6 +84,16 @@ std::optional GetAtomicBinOp(Operation* modifier_op, .Default([](Operation* op) { return std::nullopt; }); } +// Looks through an arith.extf widening cast and returns the narrower source +// value. Low-precision floating-point reductions (e.g. bf16) are computed in a +// wider type, so the atomic modifier appears in the body as extf(modifier). +Value LookThroughExtF(Value value) { + if (auto ext = value.getDefiningOp()) { + return ext.getIn(); + } + return value; +} + } // namespace // Returns atomic op modifier and the atomic bin op kind. @@ -92,9 +102,8 @@ std::optional> GetAtomicModifierParameters( Type element_type = op.getInput().getType().getElementType(); auto& operations = op.getBody()->getOperations(); auto terminator = op.getBody()->getTerminator(); - if (operations.size() > 2) { - return std::nullopt; - } + Value block_arg = op.getBody()->getArgument(0); + // If the body contains only the terminator, then it is an atomic store. if (operations.size() == 1) { // TODO(b/336367145): Support complex atomic store. @@ -103,19 +112,67 @@ std::optional> GetAtomicModifierParameters( } return std::nullopt; } - // Match the kind of the atomic op. - // TODO(rocm): Match bf16 ops - mlir::Operation* modifier_op = &operations.front(); + + // Simple case: a single binary modifier op followed by the terminator, + // operating directly on the atomic element type. + if (operations.size() == 2) { + mlir::Operation* modifier_op = &operations.front(); + auto kind = GetAtomicBinOp(modifier_op, element_type); + if (!kind.has_value()) { + return std::nullopt; + } + // Find the modifier arg that does not match the argument of `atomic_rmw` + // body. + Value modifier_arg = modifier_op->getOperand(0) == block_arg + ? modifier_op->getOperand(1) + : modifier_op->getOperand(0); + return std::make_pair(modifier_arg, *kind); + } + + // Widened low-precision case (e.g. bf16): the reduction is computed in a + // wider type, so the body looks like: + // %0 = arith.extf %current : bf16 to f32 + // %1 = arith.extf %modifier : bf16 to f32 + // %2 = arith. %0, %1 : f32 + // %3 = arith.truncf %2 : f32 to bf16 + // xla.yield %3 : bf16 + // + // bf16 has no native arithmetic on the relevant targets, so + // FloatNormalization (driven by GpuFloatSupport, which reports bf16 + // add/mul/etc. as unsupported) legitimately wraps the combiner in f32 + // conversions. We do NOT undo that normalization globally; we only look + // through it *here*, where we are about to emit a hardware atomic (e.g. + // global_atomic_pk_add_bf16) that performs the very same widen-add-round + // internally. Recovering the narrow (bf16) modifier lets the lowering use + // that packed atomic instead of a slow compare-and-swap loop, without + // claiming a native bf16 arithmetic instruction. + auto trunc_op = terminator->getOperand(0).getDefiningOp(); + if (!trunc_op || trunc_op.getType() != element_type) { + return std::nullopt; + } + mlir::Operation* modifier_op = trunc_op.getIn().getDefiningOp(); + if (!modifier_op || modifier_op->getNumOperands() != 2) { + return std::nullopt; + } auto kind = GetAtomicBinOp(modifier_op, element_type); if (!kind.has_value()) { return std::nullopt; } - // Find the modifier arg that does not match the argument of `atomic_rmw` - // body. - Value block_arg = op.getBody()->getArgument(0); - Value modifier_arg = modifier_op->getOperand(0) == block_arg - ? modifier_op->getOperand(1) - : modifier_op->getOperand(0); + Value lhs = LookThroughExtF(modifier_op->getOperand(0)); + Value rhs = LookThroughExtF(modifier_op->getOperand(1)); + Value modifier_arg; + if (lhs == block_arg) { + modifier_arg = rhs; + } else if (rhs == block_arg) { + modifier_arg = lhs; + } else { + return std::nullopt; + } + // The recovered modifier must have the atomic element type for the downstream + // direct-atomic emission to be valid. + if (modifier_arg.getType() != element_type) { + return std::nullopt; + } return std::make_pair(modifier_arg, *kind); } diff --git a/xla/codegen/emitters/transforms/tests/lower_tensors.mlir b/xla/codegen/emitters/transforms/tests/lower_tensors.mlir index 4469755829bc5..7d39ef8758ecd 100644 --- a/xla/codegen/emitters/transforms/tests/lower_tensors.mlir +++ b/xla/codegen/emitters/transforms/tests/lower_tensors.mlir @@ -30,6 +30,10 @@ // RUN: -xla-lower-tensors="gpu_device_info='rocm_compute_capability {gcn_arch_name: \"gfx90a:sramecc+:xnack\"}'" \ // RUN: | FileCheck %s --check-prefix=CHECK-GFX90A-MI200 +// RUN: emitters_opt %s --allow-unregistered-dialect -split-input-file \ +// RUN: -xla-lower-tensors="gpu_device_info='rocm_compute_capability {gcn_arch_name: \"gfx942:sramecc+:xnack\"}'" \ +// RUN: | FileCheck %s --check-prefix=CHECK-GFX942-MI300 + module attributes {dlti.dl_spec = #dlti.dl_spec<#dlti.dl_entry>} { func.func private @add(%arg0: f32, %arg1: f32) -> f32 { %sum = arith.addf %arg0, %arg1 : f32 @@ -764,6 +768,42 @@ func.func @direct_atomic_rmw_fadd_bf16(%in: tensor<8xbf16>, // ----- +// The scatter/reduce combiner for bf16 is computed in f32, so the atomic_rmw +// body is the widened pattern extf -> addf(f32) -> truncf. The modifier is a +// genuine bf16 value (a constant would have its extf folded to f32). On MI300 +// this must lower to a packed `atomicrmw fadd <2 x bf16>` rather than a slow +// compare-and-swap loop. +func.func @direct_atomic_rmw_fadd_bf16_widened(%in: tensor<8xbf16>, + %i: index, %arg: bf16) -> (tensor<8xbf16>) { + %ret = xla.atomic_rmw %in[%i] : tensor<8xbf16> { + ^bb0(%current : bf16): + %ext0 = arith.extf %current : bf16 to f32 + %ext1 = arith.extf %arg : bf16 to f32 + %add = arith.addf %ext0, %ext1 : f32 + %res = arith.truncf %add : f32 to bf16 + xla.yield %res : bf16 + } + return %ret : tensor<8xbf16> +} +// CHECK-GFX942-MI300-LABEL: @direct_atomic_rmw_fadd_bf16_widened +// CHECK-GFX942-MI300-DAG: %[[C_NEG4:.*]] = llvm.mlir.constant(-4 : i64) : i64 +// CHECK-GFX942-MI300-DAG: %[[C2:.*]] = llvm.mlir.constant(2 : i32) : i32 +// CHECK-GFX942-MI300-DAG: %[[C8:.*]] = llvm.mlir.constant(8 : i32) : i32 +// CHECK-GFX942-MI300: %[[ADDR:.*]] = llvm.getelementptr +// CHECK-GFX942-MI300: %[[ADDR_INT:.*]] = llvm.ptrtoint %[[ADDR]] +// CHECK-GFX942-MI300: %[[ADDR_MASKED:.*]] = llvm.and %[[ADDR_INT]], %[[C_NEG4]] +// CHECK-GFX942-MI300: %[[ADDR_TRUNC:.*]] = llvm.trunc %[[ADDR_INT]] +// CHECK-GFX942-MI300: %[[OFFSET:.*]] = llvm.and %[[ADDR_TRUNC]], %[[C2]] +// CHECK-GFX942-MI300: %[[SHIFT:.*]] = llvm.mul %[[OFFSET]], %[[C8]] +// CHECK-GFX942-MI300: %[[VAL_INT:.*]] = llvm.bitcast %{{.*}} : bf16 to i16 +// CHECK-GFX942-MI300: %[[VAL_WIDE:.*]] = llvm.zext %[[VAL_INT]] : i16 to i32 +// CHECK-GFX942-MI300: %[[VAL_SHIFT:.*]] = llvm.shl %[[VAL_WIDE]], %[[SHIFT]] +// CHECK-GFX942-MI300: %[[ADDR_PTR:.*]] = llvm.inttoptr %[[ADDR_MASKED]] +// CHECK-GFX942-MI300: %[[VAL:.*]] = llvm.bitcast %[[VAL_SHIFT]] : i32 to vector<2xbf16> +// CHECK-GFX942-MI300: llvm.atomicrmw fadd %[[ADDR_PTR]], %[[VAL]] syncscope("agent-one-as") monotonic + +// ----- + func.func @direct_atomic_rmw_fadd_f64(%in: tensor<8xf64>, %i: index) -> (tensor<8xf64>) { %c2 = arith.constant 2.0 : f64