vx-lang / vx-lang/Vx

Materialise Memory::SMEM placement in emitted kernels (and the four things blocking cp.async)

Open
#352 2 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

codegen enhancement runtime
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-targetconvert-gpu-to-nvvm → arith → math → gpu-to-llvmreconcile-unrealized-castsgpu-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:

  1. 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.
  2. Open the allow-list and the pipeline for nvgpu — and make a rejected kernel fail loudly rather than silently falling back.
  3. Emit nvgpu.device_async_copy for the fill, still synchronous.
  4. 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::SMEM appears as .shared in 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: streamed becomes testable, and vx-review#26 is re-scored against it

Contributor guide

No contributing guide indexed for this repository

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.