Ordering async copies and sizing dynamic shared memory in the CUDA target

[FeathBow]

Thanks for pointing me to #265. Here's how I'd like to close the gap from #264, where finalised TIR doesn't order its async copies or say how much dynamic shared memory it needs. I've followed the layout in #265's current code, and the reproducer below ran on main before that refactor.

What's missing

loop_outputs.py shows both problems. It copies a row into shared memory at offset 0 with T.copy_async, then reads it back five lines later with T.copy, with no commit or wait in between. The host entry builds the launch without dynamic_smem, so the kernel gets 0 bytes, and the launcher never raises the limit above 48 KiB. The async copy kernels that do run, async_sync, split_stage, the rings in test_staged_ring.py and the GEMM in #264, add the ordering by hand, and the GEMM sets dynamic_smem=512 itself.

You can see this without a GPU. On main before #265, the driver below prints block = [32, 1, 1] dynamic_smem = 0, and the generated kernel calls copy_async and then copy on the same view. The build then fails with identifier "smem_base" is undefined. smem_base is declared where the kernel first uses it, which here is inside the first loop, so the second loop can't see it. I'd fix that as well.

Driver and output

import importlib.util
import sys

import tilefoundry
from tilefoundry.dsl import Topology
from tilefoundry.ir.core.module import Module
from tilefoundry.ir.tir.prim_function import PrimFunction
from tilefoundry.passes.transforms.host_entry import insert_default_host_entry

spec = importlib.util.spec_from_file_location("finalized", sys.argv[1])
loaded = importlib.util.module_from_spec(spec)
spec.loader.exec_module(loaded)
kernel = next(value for value in vars(loaded).values() if isinstance(value, PrimFunction))
module = Module(
    name="finalized",
    functions=(kernel,),
    entry=kernel.name,
    target=kernel.target,
    topologies=(Topology("cta", 1), Topology("thread", 32)),
)
launch = insert_default_host_entry(module).entry_function().body.body[0]
print("block =", [arg.value for arg in launch.args[4:7]], "dynamic_smem =", launch.callable.dynamic_smem)
tilefoundry.compile(module, target=kernel.target)

$ python drive.py tests/fixtures/schedule/tir/loop_outputs.py
block = [32, 1, 1] dynamic_smem = 0
...
device_0.cu(49): error: identifier "smem_base" is undefined
...
RuntimeError: link_modules: cmake-build failed (rc=2)

This is the part of the generated device_0.cu that reads the copy.

tilefoundry::ops::copy_async(value_4, value_1_6);
auto value_2_7 = tilefoundry::ops::tensor_view(out_tensor.data(), /* … */);
auto window_8 = tilefoundry::ops::tensor_view((value_2_7.data() + /* … */));
tilefoundry::ops::copy(value_1_6, window_8);

What I'd like to do

1. Size shared memory in the CUDA host entry. insert_default_host_entry sets dynamic_smem to the bytes the callee's integer smem views reach, including padding and every ring stage. It runs at BEFORE_CODEGEN, so this also covers TIR written by hand. The verifier checks that a launch covers its callee's views and fits the target, and the launcher calls cudaFuncSetAttribute above 48 KiB.
2. Add a CUDA ModulePass at PassPoint.AFTER_LOWER, so schedule finalize runs it right after ConvertHIRToTIR. Before a read of bytes that a cp.async wrote, the threads that issued the copy commit their group and wait for it with T.cp_async_wait(n), counting only their own groups, and then a T.sync over the issuing and reading threads makes the bytes visible. Before a write reuses smem bytes that other threads have read, the pass adds a T.sync. It works this out from what the TIR already holds, which is the operand effects, smem byte ranges, issuing mesh and stage rings. It follows each thread's commits and waits through branches and loops, and reuses any sync the author already wrote.
3. Order T.copy_async_bulk with mbarriers in the same pass, and T.copy_async_tensor once it has an emitter.

How I'd test it

I'd take the kernels that order themselves by hand, remove their syncs, run the pass, and check that their existing tests still pass. The finalised cp.async fixtures under tests/fixtures/schedule/tir/ should compile and match their plain programs on the GPU they target.

I'll start once #265 has landed. Let me know if you'd like any of this done differently.

[zhen8838]

Thanks @FeathBow for the careful write-up and the reproducer. This is exactly the gap #264 exposed, and the plan fits the extension points from #265. Some direction before you start:

1. smem_base scope. Codegen should only consume the result of memory allocation; it should not decide where a shared-memory base lives. So before changing where smem_base is declared, please check that the allocation attributes each smem allocation to the right scope (region and lifetime). If the staged buffer in loop_outputs.py is recorded as owned by the first loop, that accounting is the bug, and the declaration point follows once it is fixed.

2. dynamic_smem. Shared memory is allocated against the device's smem capacity from the target facts, so a kernel may use all of it by default. The host entry can therefore default dynamic_smem to that target fact, and the launcher raises the limit with cudaFuncSetAttribute when it is above 48 KiB. The host entry does not need to compute a size from the views.

3. The ordering pass, and 4. the mbarriers for bulk copies. A CUDA ModulePass is the right shape. Insert it in the middle of the lowering to TIR, where the producers, the consumers and each buffer's stage are still known. The mbarriers for bulk copies can use the same protocol. Two papers are worth reading first:

• Tawa (CGO 2026): asynchronous references; a tile program split into producer and consumer warps that communicate through full/empty barriers.
• Twill (OSDI 2026): software pipelining and warp specialization solved together as one constraint problem.

This is also where we want the compiler to go next:

1. Timing analysis for async pipelines in schedule finalize: from each instruction's unit and cycles, the buffer count S and the dependence distances, bound the initiation interval (resource bound, recurrence bound, and S × II ≥ the latency from writing a buffer to reading it). An agent can then choose S, and see whether TMA or the MMA is the bottleneck, without a GPU.
2. A Tawa-style warp-specialized pipeline skeleton derived from the sequential TIR, which already carries the mesh scopes, buffers=S and slot = k % S: producer/consumer roles, full/empty barriers and phases, checked (pairing, phase, release count) before anything runs.
3. Twill-style resource-constrained scheduling: the II and a (stage, cycle) for each instruction under unit occupancy and register/smem capacity, with a prologue, a steady state and an epilogue, including a resource turn in which two consumer warpgroups take the tensor core in turn.

5. Tests. The most direct check: run schedule finalize on a scheduled HIR fixture, compile the result, run it on the GPU and compare it with the plain program. If that runs, the pass did its job.

6. Expectations. TileFoundry's codegen is still far from complete. In TileOps, an agent currently transcribes the incomplete TIR into a TileLang implementation. So expect to find more gaps like smem_base along the way; each one makes a good small PR.

[zhen8838]

Thanks again @FeathBow. #265 has landed, so here is an update before you start.

Code layout and packaging. #265 and #268 moved the CUDA target into its own plugin package, tilefoundry-nvidia, under targets/nvidia. Core loads each backend through the tilefoundry.targets entry point, and the plugin registers its own operations and passes. So you can develop this work in that directory, or even out of tree as a separate package. See docs/develop.md §6 and the target loading contract.

A first step. I suggest you do the ordering in two parts:

1. During lowering, after shared memory is allocated, mark the program: which buffers the async copies write, which threads issue them, and which stage each buffer is in.
2. When the program is lowered to TIR, read those marks and generate the program with its synchronization: commit and wait for cp.async, then T.sync.

The analysis then stays where this information still exists, and the lowering to TIR only translates the marks.

If you have questions, ask here at any time.

[FeathBow]

A slightly late update from me. I've read Tawa and Twill, and along the way had a look at the Exo-GPU and Trainium papers.

Your two steps make sense to me. To get the order right I need to know which bytes each async copy writes, which threads issue it, and which reads and later writes touch those bytes. Once shared memory is allocated the storage answers are there, and with the issue plan they give the actual accesses and threads. So I'd work it all out at that point, record it as marks, and let the lowering turn the marks into commits, waits and syncs. A few small fixes come first so the e2e test you suggested can actually run.

[FeathBow]

Roughly, the steps look like this.

1. dynamic_smem, defaulting to the target's shared memory capacity, with the launcher calling cudaFuncSetAttribute above 48 KiB.
2. smem_base, once I've checked how the allocation scopes each smem buffer.
3. Two gaps on the sm_89 GEMM fixture, each as a small PR. Codegen rejects the T.ptr_of that finalize takes of the sharded accumulator, and the finalised copies aren't split over the warp's threads.
4. Marking, which I'd start as a CUDA ModulePass at BEFORE_LOWER (happy to move it into ConvertHIRToTIR if that suits you better). For each copy scheduled as T.copy_async(), it takes the bytes written and the issuing threads from the same StoragePlan and IssuePlan answers the lowering uses, and the stage from the lowering's slot rule for buffers=S, nested and dynamic loops included. Then it finds the reads of those bytes and any later write that overlaps them. If it can't pin down the threads or the bytes, it stops with an error.
5. Lowering. ConvertHIRToTIR commits at the group boundaries the marks give. Before a read, or before bytes get overwritten, the issuing threads wait for the async copies that must have landed by then, with cp_async_wait(n) counting the groups committed since. A T.sync over everyone involved then protects the read or the reuse.
6. Checks. The same PR checks the waits and barriers in the final TIR, a bit like Exo-GPU's abstract machine that tracks which threads have touched each element and who can see it, with a broken variant for each check. It also carries the e2e test.
7. The mbarriers for bulk copies come after that, since they need shared memory of their own.

There might be a couple of small optimisations along the way. The Trainium paper swaps the barrier at each loop iteration for completion counts worked out from the loop counters, and proves that, for the loops and dependences it supports, this adds no extra waiting. Our marks already say who writes and reads each stage, so where the pipeline order allows it we might get down to one barrier per stage. In the sm_89 GEMM the allocator also puts A and B at the same offset. With the fragments read synchronously as they are now, giving B its own 256 bytes might save two of the syncs that guard that reuse. Both need measuring.

Does this look reasonable to you? I'll follow up with the PRs one at a time.

[zhen8838]

Thanks @FeathBow for the detailed plan and for reading the papers. It looks reasonable to me. A few notes:

1. PR size. I did the reorg so that the code structure does not get in the way of development. Size and pace your PRs in the way that suits you. For example, changes inside one package can go in one PR.

2. Marking at BEFORE_LOWER. You can do it there now. Since #269, finalize receives the memory analysis result, so BEFORE_LOWER passes run after shared memory is allocated and before ConvertHIRToTIR. There are two limits:

• The pass must change the module in place. ConvertHIRToTIR checks that it gets the same module that was analyzed, so the marks have to be metadata on that module.
• StoragePlan and IssuePlan are built inside ConvertHIRToTIR, and a target pass receives only the module. A BEFORE_LOWER pass would have to build them again.

I did not think these through fully in the reorg, so change them as you need.

3. Synchronization and allocation. Automatic synchronization really has to be decided together with memory allocation: the allocator decides which bytes are reused, and that decides where the syncs go. Your A/B offset observation in the sm_89 GEMM is one example. For the first stage, let's keep it simple and treat them as two separate steps: allocate first, then insert the synchronization from that allocation. Later I hope we can study the combined problem together in more depth.

[FeathBow]

Before I start on the cp.async synchronisation, I'd like to check one change with you.

Right now schedule finalize writes no commit, wait or T.sync after a T.copy_async, so the copy that reads the stage can run before the data has landed. In the sm_89 GEMM the A and B tiles also share offset 0, so the copy into B can overwrite A while other threads are still reading it.

In my earlier reply I planned to mark the HIR at BEFORE_LOWER and have lowering turn the marks into synchronisation. Reading the lowering code, it also writes copies that have no HIR call of their own, such as the loop carry transfers and InsertSlice. Marks on the HIR would miss those, and filling them in would mean repeating the address reasoning on the HIR.

So I plan to add the synchronisation as a core pass that runs right after ConvertHIRToTIR. By then the allocation is already in the TIR, with each shared memory buffer a T.tensor_view at its byte offset, so the pass can work from the addresses directly and sees every access, including the ones lowering adds. It treats existing commits, waits and T.sync as already in place. Before the first conflicting access it adds a commit and a wait on the threads that issued the copy, and a T.sync in the nearest scope that both sides reach. Loops go through a fixed point, in the same way as Triton's Membar. The pass runs in schedule finalize only.

Each Op that touches shared memory will state when its effects complete, either on return or as a cp.async group. An Op that states nothing counts as unknown, and finalize refuses to order a later access after an unknown one. A wgmma that reads a ring stage which a later copy overwrites is refused until wgmma has its own wait.

To test it I plan to use the run you suggested. It runs schedule finalize, compiles the result, runs it on the GPU and compares it with the plain program. This takes the place of the static check I mentioned earlier, because one analysis is enough to place the barriers and a second model would repeat the same rules. Triton, TVM and TileLang also rely on one analysis for their barriers.

Does this suit us?

[zhen8838]

Thanks @FeathBow for checking before you start, and for reading the lowering this closely. Before we pick a pass, I'd like to lay out how we see the problem, because I think it is larger than placing synchronisation points in the TIR.

1. Where this work sits

flowchart LR
    H["Scheduled HIR<br/>written by an agent<br/>op=T.copy_async(), buffers=S"]
    C["check_program"]
    M["analyze(memory)<br/>liveness + allocation<br/>(synchronous model)"]
    L["ConvertHIRToTIR<br/>(hir2tir)"]
    T["TIR"]
    G["CUDA codegen"]
    H --> C --> M --> L --> T --> G

Loading

• The agent writes the schedule. The scheduled HIR already says which copies are asynchronous (tf.schedule(..., op=T.copy_async())), which level each buffer lives in, and how many stages it has (buffers=S). It does not write commits, waits, barriers or addresses. The HIR is SSA: the instruction that reads a copied tile uses the copy's value, so the order between them is a data dependence.
• The memory analysis allocates conservatively. Liveness gives each value an interval from its definition event to its last use, and assumes every access happens at its own event (analysis §lifetimes). A staged value is live across its whole loop. Values whose intervals do not overlap may share bytes, which is why A and B share offset 0 in the sm_89 GEMM. The analysis knows nothing about asynchronous completion.
• Lowering is where the races appear. ConvertHIRToTIR turns SSA values into shared memory buffers that are written in place. The copy completes later than it is issued, bytes are reused between values, and a ring stage is reused every S iterations.

So the work is to rewrite the lowering so that it produces asynchronous TIR. A pass that patches the TIR after lowering cannot do all of it, for the reason in the next section.

2. Asynchronous TIR is not sequential TIR plus barriers

Below, simplified, is what finalize writes today for a staged GEMM (left), and the warp-specialised form we eventually want (right). The right side is pseudo-TIR; wait, arrive and arrive_on_completion stand for mbarrier operations.

Today: every thread runs every statement

Warp specialised: producers and consumers

with Mesh(threads) as scope: for k in range(0, K, BK): s = (k // BK) % S T.copy_async(a[:, k:k+BK], a_stages[s]) T.copy_async(b[k:k+BK, :], b_stages[s]) # nothing orders the reads below # after the copies above T.copy(a_stages[s], a_frag) T.copy(b_stages[s], b_frag) T.tiled_mma(acc, a_frag, b_frag)

# full[S], empty[S]: mbarriers in shared memory with Mesh(threads[producer]) as producer: for k in range(0, K, BK): s = (k // BK) % S phase = (k // BK // S) % 2 wait(empty[s], phase ^ 1) # stage s released T.copy_async(a[:, k:k+BK], a_stages[s]) T.copy_async(b[k:k+BK, :], b_stages[s]) arrive_on_completion(full[s]) with Mesh(threads[consumer]) as consumer: for k in range(0, K, BK): s = (k // BK) % S phase = (k // BK // S) % 2 wait(full[s], phase) # stage s landed T.copy(a_stages[s], a_frag) T.copy(b_stages[s], b_frag) arrive(empty[s]) # stage s reusable T.tiled_mma(acc, a_frag, b_frag)

The smallest correct change keeps the left structure and adds cp_async_commit, cp_async_wait(0) and T.sync before each read and before each reuse. It is correct, but each iteration waits for its own copies, so nothing overlaps. Getting from the left to the right changes the structure of the program:

• One mesh becomes two role meshes, and the loop is written once for each role.
• The producer runs up to S stages ahead. The empty barriers bound how far.
• The barriers are new state: 2·S mbarriers that live in shared memory and are initialised before both loops.
• The phase of each wait comes from k and S.
• The consumer can release a stage as soon as its fragments are in registers, before the MMA. To know this, the lowering has to know when each read completes.

Only the lowering decides which mesh runs which statements, how many loops there are, and where the barriers live. That is why we see this as a rewrite of the lowering rather than a pass after it.

3. How to implement this stage

There are two ways, and both seem reasonable to us.

Option A: an asynchronous analysis, then a lowering that consumes it

flowchart LR
    M["analyze(memory)"] --> A["async analysis (new)<br/>per HIR event: bytes, threads,<br/>read/write, completion"]
    A --> P["async plan<br/>roles, commits, waits,<br/>barriers, keyed by HIR events"]
    P --> L["lowering in the CUDA target<br/>translates the plan"]

Loading

• A new analysis runs after allocation. For each event on the liveness timeline it records the access: bytes from the allocated offsets, threads from the operands' ShardLayout, read or write, and when the access completes (on return, as a cp.async group, or unknown).
• From these it builds a plan: which role each statement belongs to, and where the commits, waits and barriers go. The lowering only translates the plan, in the same way that it consumes StoragePlan today.
• Our suggestion for the implementation: rewrite the lowering in the CUDA target, tilefoundry-nvidia. Start from ConvertHIRToTIR and change it to consume the plan and write the asynchronous structure. The analysis can stay generic; only the lowering is CUDA-specific. A target cannot provide its own lowering yet, because finalize calls ConvertHIRToTIR directly, so we would add a hook for it.
• Your concern was that marks on the HIR miss the copies the lowering adds. We checked: every one of them has an event on the HIR liveness timeline, and whether it exists is decided by the allocation's alias answers.

Copy the lowering adds	HIR event	Emitted when
Loop carry, initial value	loop entry	no CARRY_INIT link
Loop carry, yielded value	loop backedge	no CARRY_YIELD link
InsertSlice, whole base and update window	the InsertSlice call	no INSERT_SLICE / WINDOW link
Unaliased Reshard, as its default Copy	the Reshard call	the reshard is not an alias
Zeros / FullLike fill	the call	always

The liveness events are in liveness.py.

• Advantage: the plan has the shape the future work needs (section 4), so a better producer can replace this analysis later without touching the lowering.
• Cost: the plan has to be able to express roles and restructuring, not only barrier positions, so it needs more design up front.

Option B: do it inside the lowering, CUDA-specific

• The lowering already holds everything this needs: the HIR and its SSA dependences, the allocation through StoragePlan, the issuing threads through IssuePlan, the copies it adds itself, and the stage rule for buffers=S. It can analyse while it lowers.
• Since cp.async groups and mbarriers are CUDA facts, this lowering also lives in tilefoundry-nvidia, behind the same hook as in Option A.
• Advantage: fastest to a running pipeline.
• Cost: analysis and emission are mixed, so moving to an async-aware allocation later means rewriting this part.

What both options must respect

1. Insert, do not move. The allocation assumes each access happens at its own event. A copy issued earlier than its HIR position can write bytes that another value still holds. Running ahead, as the producer does above, must come from the stages the schedule declares (buffers=S), not from moving issues.
2. An asynchronous read is live until its wait. The allocation ends a value's lifetime at its last issue. If a wgmma or a TMA store still reads those bytes, a wait must come before anything else writes them. An Op that cannot wait is refused, as you proposed.
3. Barrier state needs shared memory that the allocation has not reserved. cp.async groups and T.sync need none; mbarriers do. Before the warp-specialised form, the barriers have to become values the allocation can place, or the allocation has to reserve their space.

Because the lowering knows S and the stage rule, wait counts and phases follow from them directly. We do not need a fixed point over TIR addresses.

4. Where we want to end up

flowchart LR
    H["Scheduled HIR"] --> M["memory + async analysis<br/>async-aware lifetimes"]
    M --> R["resource-constrained solve<br/>II, (stage, cycle) per op,<br/>offsets, roles"]
    R --> P["async plan"]
    P --> L["async lowering<br/>(built in this stage)"]

Loading

• The memory analysis models asynchrony. A lifetime runs from issue to completion, and the allocator knows that reusing bytes forces a wait. In the sm_89 GEMM, it would then see that giving B its own 256 bytes removes the waits that guard the A/B reuse, which is the trade-off you pointed out.
• The solve is under resource constraints, in the spirit of Twill: functional units, shared memory and register capacity, and barrier count, with the initiation interval, the (stage, cycle) of each operation, the offsets and the roles decided together.
• The lowering from this stage is reused. If the solve produces the same kind of plan, the asynchronous lowering you build now turns a better plan into better code without changes. This is the main reason we lean towards keeping a plan between analysis and lowering, even in Option B's form.

5. Smaller points

• Op completion. We agree that each Op should state when its effects complete. TIR Ops already register their access relations in access_relation.py; the completion statement can sit next to them, so that bytes and threads come from one place.
• Tests. The end-to-end run is the right test. A race does not always show in one run, so please also run it under compute-sanitizer --tool racecheck and --tool synccheck. That gives an independent check without a second model.

Which option suits you? We are happy with either for this stage, as long as the plan, or something that can become one, stays between the analysis and the code it produces.

The details still need discussion: the format of the plan, how the roles are chosen, and where the barriers are stored, among others. This is not a simple problem, and it will probably take several rounds. Let's settle them here one at a time.

[zhen8838]

Thanks @FeathBow. An update on the hook from section 3: #277 has landed, so the lowering pass can now be customized by a target.

• finalize, compile and scaffold now run one pass list, default_pipeline. It obtains ConvertHIRToTIR through Target.resolve, so a target can supply its own.

• To do that, bind a subclass in tilefoundry-nvidia:

  @bind(CudaTarget, ConvertHIRToTIR)
  class CudaConvertHIRToTIR(ConvertHIRToTIR):
      visitor = CudaLowering  # replace only the visitor, or override run() to replace the whole pass

• The pass takes only the Target and CodegenOptions at construction (passes §7.1). It reads everything else, including the memory analysis, from the module it receives.

#277 also replaces the dynamic shared memory scan from #271. The host entry is now inserted before the lowering, while the memory analysis is still on the module, and its Launch takes the analyzed shared-memory peak as dynamic_smem (host_entry.py). Handwritten TIR states its own host entry. The rest of #271 stays: the launcher still raises the limit above 48 KiB, and each kernel still declares smem_base once.

So the next step is yours: override this lowering pass in the CUDA target so that it generates asynchronous TIR, with whichever option from section 3 suits you.

[FeathBow]

Thanks, the hook makes sense. I'll go with Option A on top of #277, so the analysis and the emission stay separate and a better producer can swap in for the analysis later.

For the first version, each Op declares its completion next to its access relation, as you suggested. The analysis sticks to the smallest correct change from section 2, with a commit, a wait and a barrier before each read and each reuse, and it refuses anything it can't prove safe. I'll check this holds up as I go, and test it end to end on the GPU on the supported cases that have a plain reference, under racecheck and synccheck as well.

Once that works, I'd like to compare placements in shared memory by the synchronisation they need. Reusing bytes can add ordering that separate storage wouldn't need. I'd keep that ordering explicit, and for each candidate placement recompute the affected dependences and regenerate the synchronisation, then weigh up the synchronisation against the shared memory it takes. It's a small piece of what you describe in section 4, where the memory analysis models asynchrony. I'd run the analysis over all the schedule fixtures and time the ones that run on the GPU.

Roles and mbarriers can wait, and I'd like to sort out where their barriers live when we get there.

[zhen8838]

Thanks @FeathBow, the plan sounds good to us: Option A on top of #277, starting from the smallest correct form and refusing what the analysis cannot prove safe, then comparing placements by the synchronisation they need.

This is not a simple problem, and it will likely take a few rounds before the pieces settle. We trust you with it. Please keep us posted here as you go, and raise anything that blocks you or needs a decision from us, so that we can work through it together.