[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