SIGN IN SIGN UP

[Bugfix] Carry memory_order through vectorized atomic_add (#2924)

A tile-region `T.atomic_add(..., memory_order=...)` lost its ordering
whenever the add auto-vectorized: the vectorizer rebuilt the wide op from
only `{dst, src}`, so `AtomicAddx2`/`AtomicAddx4` were emitted without the
order argument and fell back to their `relaxed` default. The identical call
honored the order when it lowered scalar-wise, so whether the requested
ordering survived depended only on whether the per-thread extent was
divisible by the atomic vector width.

Carry the trailing operands through the wide reconstruction. Codegen and
the device helpers already accept the order (`args.size() > 2` guard, and
`int memory_order` parameters that branch to ordered PTX), and
`atomic_addx2/x4_elem_op` are declared with `num_inputs = 3` -- the
vectorizer was simply under-filling the op. The order is copied as-is
rather than visited, since it is a lane-invariant scalar that must not be
broadcast across the vector lanes.

Carrying the order also wakes up inline PTX that was previously
unreachable: with the order always defaulted to relaxed, the non-relaxed
branches were constant-folded away before ptxas ever saw them. Three
problems surface once they are live, all verified with ptxas 13.3:

1. `atom` with vector operands requires sm_90, so the ordered fp16 pair
   failed to assemble on sm_80/sm_86/sm_89 -- and fp16 vectorizes to x2 on
   every target, so that is every ordered fp16 region atomic_add. Use the
   packed `f16x2` form instead of `.v2.f16`: it encodes the same pairwise
   add but is available from sm_60, so the wide op stays a single atomic
   instruction on every supported target rather than needing an arch split.
   sm_89 SASS confirms a native ATOMG.E.ADD.F16x2.RN.STRONG.GPU. This also
   drops the per-half pack/unpack, since a half2 reinterprets directly to
   the 32-bit operand.

2. The bf16 v2 asm was missing the `.noftz` modifier that ptxas requires
   for `.bf16`, so `atom.{release,acquire,acq_rel}.gpu.global.add.v2.bf16`
   failed to assemble on every arch, sm_90 included. The scalar bf16
   helper next to it already spells it correctly.

3. bf16 has no packed equivalent below sm_90 -- `bf16x2` and `.v2.bf16`
   both require sm_90, as does scalar `atom.add.bf16` -- so unlike fp16 it
   genuinely cannot express the ordering in one instruction there. Guard
   the sm_90 path and emulate the ordering below it with device-scope
   fences around a relaxed bf16x2 atomicAdd. `__threadfence()` lowers to
   `fence.sc.gpu`, the same device scope as the `.gpu` qualifier on the
   ordered PTX, and is two-way where release and acquire need only one-way,
   so this is conservative rather than lossy. `AtomicAddx2Ret` takes the
   same guard, since it shares the helper. Nothing reachable instantiates
   the bf16 Ret overloads with a non-relaxed order today, so that path has
   no test; it is guarded to keep the file consistent and to stay correct
   once region-form `return_prev=True` is implemented.

Add a regression test in fp16 and bf16. It disables the kernel cache,
since the key does not cover the native library by default and would
otherwise serve stale pre-fix source, and it runs the kernel so the
ordered wide path is checked for correctness and not just for the emitted
argument. The bf16 case covers the pre-sm_90 fenced branch on Ampere/Ada
and the sm_90 instruction on Hopper.

Also correct the stale memory_order sentence in atomic_max and
atomic_min, which forward the order through the shared region lowering
just like atomic_add does.

Fixes #2688

Signed-off-by: Srijan Keshri <srijankeshri007@gmail.com>
S
Srijan Keshri committed
9f336776eef3b0cfafc0b08ce5248040ae1a26c3
Parent: f2ee1dd
Committed by GitHub <noreply@github.com> on 8/20/2026, 5:34:00 AM