Skip to content

[AMD] Guard degenerate axis info in atomic RMW lowering - #993

Open
WhatGhost wants to merge 1 commit into
flagos-ai:mainfrom
WhatGhost:amd-fix-atomic-axisinfo
Open

[AMD] Guard degenerate axis info in atomic RMW lowering#993
WhatGhost wants to merge 1 commit into
flagos-ai:mainfrom
WhatGhost:amd-fix-atomic-axisinfo

Conversation

@WhatGhost

@WhatGhost WhatGhost commented Aug 14, 2026

Copy link
Copy Markdown
Contributor

Problem

Any tensor-indexed atomic on a shared-memory pointer crashes the compiler on the AMD backend with a segfault, e.g.

buf = tle.gpu.alloc([BLOCK], dtype=tl.float32, scope=tle.gpu.smem, ...)
ptrs = tle.gpu.local_ptr(buf, (idx, ))
tl.atomic_add(ptrs, val)   # Segmentation fault during ConvertTritonAMDGPUToLLVM

Root cause

mentioned in #972

AtomicRMWOpConversion dereferences the axis-info entry of its pointer operand and indexes the contiguity vector without checking either:

unsigned contigWithinLanes =
    axisAnalysisPass.getAxisInfo(ptr)->getContiguity(threadOrder.front());

tle.local_pointers has no AxisInfoVisitor, so its result carries a degenerate rank-0 entry and the indexing walks off an empty vector. The value is only used to refine enableIntraWaveReduce, which is restricted to CDNA3/CDNA4, so on RDNA the crash happens while computing something that is never used.

Fix

Check the entry before using it and fall back to disabling intra-wave reduce
when the contiguity is unavailable.

Validation

Compiler behaviour, same input IR, only this patch differing:

before after
tle_local_pointers_to_llvm.mlir (new case) Segmentation fault (exit 139) PASS
  • Added lit coverage for tt.atomic_rmw on tle.local_pointers, checking it lowers to llvm.atomicrmw fadd ... : !llvm.ptr<3>.
  • Conversion/amd: 33/35 pass; the two failures (async_ops_to_llvm_gfx1250, cluster_load) are pre-existing CHECK mismatches unrelated to this path.
  • No CDNA regression: full lowering output of tritongpu_to_llvm.mlir for arch=gfx942 and arch=gfx950 is byte-identical before and after (matching md5), so the intra-wave reduce decision is untouched wherever axis info is well-formed.
  • test_core.py -k atomic on gfx1201: 832 passed, 132 skipped, 0 failed.
  • Runtime check on gfx1201: 128 threads contending 8 LDS slots each end up at exactly 16, lowered to ds_add_f32, so the atomics are real and not lost updates.

Co-authored-by: Cursor <cursoragent@cursor.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant