[RFC] MergeSharedMemoryAllocations unions shared buffers across thread-divergent branches, aliasing live data

[Rachmanino]

Summary

MergeSharedMemoryAllocations treats the then and else branches of every IfThenElse as separate lifetime scopes, so shared-memory buffers used exclusively in opposite branches get packed as a union instead of a sum.

That is correct for a branch decided per CTA (e.g. T.sm_specialize_scope, which branches on blockIdx), because shared memory is private to each CTA and a given CTA only ever executes one side.

It is not correct for a branch decided per thread. A threadIdx-dependent condition diverges inside a single CTA, which has exactly one shared-memory arena, and the two branches run concurrently. Unioning their buffers gives two simultaneously-live buffers the same base address.

The result is silent shared-memory corruption. This pass is part of the shared CUDA pipeline (tilelang/cuda/pipeline.py), so it is not limited to the distributed paths.

Reproduction

import re
import tilelang
import tilelang.language as T
from tilelang import tvm

SIZE = 32768


def _apply(func):
    target = tvm.target.Target({"kind": "cuda", "arch": "sm_90a"})
    mod = tvm.IRModule.from_expr(func.with_attr("global_symbol", "main"))
    with target:
        mod = tvm.tirx.transform.BindTarget(target)(mod)
        mod = tilelang.transform.PlanAndUpdateBufferAllocationLocation()(mod)
        mod = tilelang.transform.HoistGlobalBufferAllocations()(mod)
        mod = tilelang.transform.LowerOpaqueBlock()(mod)
        mod = tilelang.transform.FlattenBuffer()(mod)
        mod = tilelang.transform.MergeSharedMemoryAllocations(enable_aggressive_merge=True)(mod)
    return mod.script(show_meta=False)


@T.prim_func
def thread_varying(A: T.Tensor((1,), "uint8"), B: T.Tensor((1,), "uint8")):
    with T.Kernel(1, threads=128):
        tx = T.get_thread_binding()
        X = T.alloc_shared((SIZE,), "uint8")
        Y = T.alloc_shared((SIZE,), "uint8")
        if tx < 64:          # decided per THREAD -> both branches run at once
            X[0] = A[0]
            B[0] = X[0]
        else:
            Y[0] = A[0]
            B[0] = Y[0]


print(_apply(thread_varying))

Observed

X and Y are both placed at offset 0 of a 32 KB arena:

bx = T.launch_thread("blockIdx.x", 1)
buf_dyn_shmem = T.alloc_buffer((32768,), "uint8", scope="shared.dyn")
X: T.handle("uint8", "shared.dyn") = T.handle_add_byte_offset(buf_dyn_shmem.data, 0)
Y: T.handle("uint8", "shared.dyn") = T.handle_add_byte_offset(buf_dyn_shmem.data, 0)
tx = T.launch_thread("threadIdx.x", 128)

One CTA (blockIdx.x extent 1), 128 threads, one 32 KB arena, and two logically distinct buffers sharing byte 0. With threads=128 and tx < 64, warps 0–1 execute the then branch while warps 2–3 execute the else branch. Different warps are scheduled independently, so both branches are genuinely in flight at the same time and the two buffers clobber each other. (Even a single diverging warp can have both paths resident under Volta+ independent thread scheduling.)

Expected

The sum, i.e. what a control kernel with both buffers simultaneously live produces, and what upstream TileLang produces:

buf_dyn_shmem = T.alloc_buffer((65536,), "uint8", scope="shared.dyn")
X: ... = T.handle_add_byte_offset(buf_dyn_shmem.data, 0)
Y: ... = T.handle_add_byte_offset(buf_dyn_shmem.data, 32768)

Substituting bx < 4 under T.Kernel(8, ...) for tx < 64 gives the same 32 KB union — and there it is correct, because each CTA takes only one branch and has its own arena.

case	arena	X offset	Y offset	correct?
both buffers live (control)	65536	0	32768	yes
if bx < 4 (per CTA)	32768	0	0	yes — union is the intended win
if tx < 64 (per thread)	32768	0	0	no — aliases live data

Root cause

src/transform/merge_shared_memory_allocations.cc:411

void VisitStmt_(const IfThenElseNode *op) final {
  this->VisitExpr(op->condition);
  VisitBranchScope(op->then_case);
  if (op->else_case.defined()) {
    const Stmt &else_case = op->else_case.value();
    VisitBranchScope(else_case);
  }
}

Upstream TileLang (550e25d) has:

void VisitStmt_(const IfThenElseNode *op) final { VisitNewScope(op); }

which puts the whole if in one lifetime scope and is therefore conservative. The per-branch split was introduced in e9ee575, and its own comment names sm_specialize as the motivation — but the code applies to every IfThenElse, with no check on what kind of condition it is.

Note the trigger condition and the bug condition are the same: the union only happens when each buffer is used exclusively inside one branch, which is exactly what makes the liveness intervals disjoint. So this is not an unlucky corner case; it is the intended optimization applied to the wrong class of branch.

Two existing mechanisms limit the blast radius but do not cover the general case:

• VisitWarpSpecializationBody bypasses VisitStmt_(IfThenElseNode) for the kWarpSpecializationScope partition if, visiting both branches in one scope. Kernels using T.ws() are therefore safe.
• SharedMemoryConflictCollector adds conflicts only between tma_store sources and other buffers touched inside a warp-specialization scope.

Neither helps a hand-written if tid < N { ... } else { ... }, which is the realistic way to hit this.

Suggested fix

Take the per-branch path only when the condition is provably CTA-uniform, and otherwise fall back to upstream's VisitNewScope(op).

Prefer a positive allowlist over "does the condition mention threadIdx":

• uniform: blockIdx.* iter vars, PrimFunc params / symbolic shape vars, constants, and pure arithmetic over those
• not provably uniform: threadIdx.* iter vars, any BufferLoad, any unrecognised Call or Var

A negative check would miss indirection such as tmp[0] = threadIdx.x; if (tmp[0] < 64). Misclassifying in the conservative direction only costs shared memory (it reproduces upstream behaviour), whereas misclassifying the other way is silent corruption — so the analysis should default to "not uniform".

The pass already visits the thread_extent AttrStmts and can distinguish blockIdx.* from threadIdx.* via the IterVar thread tag, so collecting the two var sets is cheap.

sm_specialize branches on blockIdx-derived values and stays classified uniform, so the optimization the change was written for is preserved. Worth confirming empirically that the dynamic shared-memory usage of examples/distributed/allgather_gemm/example_allgather_gemm_specialized.py and examples/distributed/gemm_allreduce/example_gemm_ar_specialized.py does not regress.

Suggested regression tests

In testing/python/transform/test_tilelang_transform_merge_shared_memory.py:

• thread-varying condition with branch-exclusive buffers → arena must be the sum and the offsets must differ
• blockIdx-varying condition → arena may be the union, pinning the intended optimization in place

Caveat on fixing this

The fix returns shared memory that kernels were (unsoundly) saving via thread-varying unions. A kernel that currently fits only because of this could start failing the shared-memory limit at compile time. That is the right trade — a compile error instead of silent corruption — but it is a visible behaviour change.

Environment

• branch release/v0.0.2-tilelang-550e25d (PR [Release] TileScale 0.0.2 on upstream TileLang 550e25d #61), reproduced at 42374143
• upstream baseline TileLang 550e25d493a93729cb087e4ecb587c19028d3cea
• repro is target-only (sm_90a) and needs no GPU; confirmed on B200 / CUDA 13.2