mirror of
https://github.com/tile-ai/tilelang.git
synced 2026-10-24 07:04:00 +08:00
Xuyang LiandLeiWang1999
c5e5307621
* [BugFix] Return previous value for atomic min/max/add (scalar + vectorized) `atomic_min`/`atomic_max` (scalar) and `atomic_addx2`/`atomic_addx4` (vectorized) typed their call with a non-void return under `return_prev=True` but always selected the void `*_elem_op`, so codegen emitted an empty binding (`prev = ;`) and nvcc failed. Select the `*_ret_elem_op` variant under `return_prev`, as `atomic_add` already does: - min/max: the ret ops, codegen handlers, and device templates (AtomicMaxRet/AtomicMinRet) already existed. - addx2/addx4 (float32): register the ret ops (builtin.cc/.h) and add codegen handlers emitting AtomicAddx2Ret/AtomicAddx4Ret (already in the device headers) as expressions; the return is typed as the packed 2-/4-lane vector so the store matches. - addx2 (fp16/bf16): the ret store LHS is `uint1`; codegen wraps the half2/bf162 result with `tl::to_uint1`. - addx4 (fp16/bf16): add AtomicAddx4Ret(half_t*)/(bfloat16_t*) device overloads that run two AtomicAddx2Ret halves and return them packed as uint2 (a fp16x4 is stored as uint2). The 16-bit x4 path is per-pair atomic (two half2 atomics), not whole-quad single-atomic -- hardware has no vector atomic wider than `.f16x2` for 16-bit types. This matches the existing fp32-x4 scalar fallback, and the docstring only promises "quad-width atomic addition". Adds tests across scalar/x2/x4 and float32/fp16/bf16. Fixes #2667 * minor fix for missing conditiions --------- Co-authored-by: LeiWang1999 <leiwang1999@outlook.com>