Materialise Memory::SMEM placement in emitted kernels (and the four things blocking cp.async)
Nobody has claimed this yet.
- Dominant language
- Rust
- Stars
- 14
- Forks
- 2
- Avg merge
- 12h 42m
- Merged PRs (30d)
- 61
Description
Summary
transfer(x, Memory::SMEM) is checked and priced but never becomes code. The emitted kernel reads every operand with ld.global and has zero .shared. So the memory algebra's SMEM predictions are claims about code we do not emit, and cp.async — which the A100 measurement showed we would need to make a "streamed" transfer real — has nothing to attach to.
Found while scoring a pre-registered prediction on a 2×A100 (vx-review#27, reverted in 7e4550f0). The prediction said an A100 fill into SMEM is a hardware copy engine because cp.async exists from sm_80. The ISA fact is true; the inference was wrong, because we never emit cp.async, and we never emit shared memory at all.
Evidence
The shipped FlashAttention kernel, tests/backend/pass/flash_attention_disaggregated.vx, compiled with --emit-llvm:
shared 0
cp.async 0
ld.global many
local_depot 1 <- spilling to local/stack memory
grep for workgroup attributions, address-space handling, or anything SMEM-related in src/dialect/VxLowering.cpp returns nothing.
The four blockers, worst first
1. Nothing places data in shared memory (the real one)
There is no shared-memory allocation, no gpu.func workgroup attribution, no address-space-3 memref anywhere in the lowering. Memory::SMEM is a pure analysis construct: capacity-checked by E6010, priced by derived_transfer_cost, and dropped before codegen.
This is the actual feature, and it is worth doing on its own merits regardless of cp.async. The flash kernel re-reads K and V from global memory on every query iteration — the textbook case for staging a tile in shared memory, and now measurable on an A100.
2. The device dialect allow-list excludes nvgpu
isDeviceLowerableDialect in src/dialect/VxLowering.cpp:961 permits arith|cf|gpu|math|memref|scf. A kernel containing nvgpu.device_async_copy is classified not-device-ready and silently dropped from GPU compilation — it falls back rather than failing loudly, which is its own problem. One line.
3. The pass pipeline has no convert-nvgpu-to-nvvm
The NVPTX sequence is nvvm-attach-target → convert-gpu-to-nvvm → arith → math → gpu-to-llvm → reconcile-unrealized-casts → gpu-module-to-binary. Even if nvgpu ops got in, they would never lower and gpu-module-to-binary would fail. One line.
4. cp.async only pays inside a software-pipelined loop
This is what makes it real work rather than plumbing. cp.async is asynchronous — issue, cp.async.commit_group, cp.async.wait_group. Issued and immediately waited on it is no faster than a load/store pair. It pays when tile k+1 is prefetched while tile k is being computed on, which means restructuring the kernel into a pipelined form. It also constrains granularity to 4, 8 or 16 bytes per thread.
Why this matters beyond performance
crossing: streamed (vx-review#26) is currently unfalsifiable by construction. Vx cannot emit a streamed route, so sequenced is correct for all Vx-generated code — not as an empirical result but as a fact about the compiler. The memory algebra was pricing a hardware capability when it only ever needs to price the code we emit.
That also means hiraditya/Vx.1#26's composition rule cannot be properly tested until (1) lands. Every measurement supporting it is a measurement of load-then-store, because that is all anyone was in a position to measure.
Order of work
(2) and (3) are trivial and useless alone. The order is:
- Materialise SMEM placement — a real shared-memory buffer in the kernel, with the tile copied in and read from there. Verify against the A100: the flash kernel should stop re-reading K and V from global.
- Open the allow-list and the pipeline for
nvgpu— and make a rejected kernel fail loudly rather than silently falling back. - Emit
nvgpu.device_async_copyfor the fill, still synchronous. - Software-pipeline the loop so the copy overlaps compute, which is the only version that is faster.
Only after (4) does crossing: streamed describe anything Vx does, and only then can vx-review#26's rule be tested rather than assumed.
Done when
- a tile placed in
Memory::SMEMappears as.sharedin the emitted PTX - the flash kernel reads K and V from shared memory rather than re-reading global
- a kernel using an unsupported dialect fails with a diagnostic instead of silently falling back
- measured on an A100: shared-memory version against the current global-only one
-
crossing: streamedbecomes testable, and vx-review#26 is re-scored against it
Contributor guide
No contributing guide indexed for this repository
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Research direction
Start with src/dialect/VxLowering.cpp, especially isDeviceLowerableDialect at line 961 and the NVPTX pass sequence described in the issue. Inspect tests/backend/pass/flash_attention_disaggregated.vx and its emitted LLVM/PTX to establish the current global-only behavior. Done means SMEM appears as .shared, K and V use it, unsupported dialects fail diagnostically, and the A100 comparison and streamed-crossing validation are recorded.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- rust
- Domain
- compilers, performance
- Issue type
- Feature
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Quiet
- Clarity
- Mostly clear
- Newbie friendliness
- 35/100