Skip to content

Cherry-pick openxla/xla#44428: [ROCm] packed bf16 atomic add for scatter/segment_sum #987

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
magaonka-amd wants to merge 1 commit into from
Closed

Cherry-pick openxla/xla#44428: [ROCm] packed bf16 atomic add for scatter/segment_sum #987

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
81 changes: 69 additions & 12 deletions xla/codegen/emitters/transforms/atomic_rmw_utils.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -84,6 +84,16 @@ std::optional<ml::AtomicBinOp> 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<arith::ExtFOp>()) {
return ext.getIn();
}
return value;
}

} // namespace

// Returns atomic op modifier and the atomic bin op kind.
Expand All @@ -92,9 +102,8 @@ std::optional<std::pair<Value, ml::AtomicBinOp>> 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<f32> atomic store.
Expand All @@ -103,19 +112,67 @@ std::optional<std::pair<Value, ml::AtomicBinOp>> 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.<binop> %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<arith::TruncFOp>();
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);
}

Expand Down
40 changes: 40 additions & 0 deletions xla/codegen/emitters/transforms/tests/lower_tensors.mlir
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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<index, 32 : i32>>} {
func.func private @add(%arg0: f32, %arg1: f32) -> f32 {
%sum = arith.addf %arg0, %arg1 : f32
Expand Down Expand Up @@ -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
Expand Down
Loading