From 94b5fd2f03637b107031ec3eff177df1b98f12a6 Mon Sep 17 00:00:00 2001 From: Yaoyao Ding Date: Thu, 2 Jul 2026 21:10:25 +0000 Subject: [PATCH] [Codegen] Fix shared-memory under-allocation for kernels with a shared workspace Kernels that request a shared workspace (tcgen05.alloc, reduce, scan, shuffle) could under-request dynamic shared memory, causing an illegal memory access at runtime (observed in examples/blackwell_matmul/matmul_v2.py). Root cause was two-fold: 1. smem_alloc_ctx.finalize() placed the workspace at a 128-aligned offset but computed dynamic_smem_bytes from the (possibly unaligned) high-water mark, so the alignment padding between the arena top and the workspace start was never reserved. The workspace tail then spilled past the requested dynamic shared memory. Now the total is sized from the aligned workspace offset. 2. The barrier allocation in mbarrier_alloc_ctx.finalize() (added in #154) bypassed the shared-memory allocator, manually bumping maximum_allocated by the raw barrier byte count (unaligned) and leaving the region unregistered in the free list. This left the high-water mark unaligned, exposing bug (1). Reverted to allocating the barriers through allocate_shared_tensor(), which keeps the allocator accounting self-consistent. This is a byte-for-byte no-op for kernels without a shared workspace above the barriers (e.g. the Hopper matmul examples). Verified matmul_v2 now runs correctly and tests/kernels/matmul/test_matmul_v2.py passes. Co-Authored-By: Claude Opus 4.8 (1M context) Signed-off-by: Yaoyao Ding --- .../tilus/backends/contexts/mbarrier_alloc_ctx.py | 15 ++++----------- python/tilus/backends/contexts/smem_alloc_ctx.py | 6 +++++- 2 files changed, 9 insertions(+), 12 deletions(-) diff --git a/python/tilus/backends/contexts/mbarrier_alloc_ctx.py b/python/tilus/backends/contexts/mbarrier_alloc_ctx.py index b5a4a8ca..73c5ab8b 100644 --- a/python/tilus/backends/contexts/mbarrier_alloc_ctx.py +++ b/python/tilus/backends/contexts/mbarrier_alloc_ctx.py @@ -26,6 +26,8 @@ from tilus.hidet.ir.primitives.cuda.smem import dynamic_shared_memory from tilus.hidet.ir.primitives.cuda.sync import syncthreads from tilus.hidet.ir.primitives.cuda.vars import threadIdx +from tilus.ir.layout import ops +from tilus.ir.tensor import SharedTensor class BarrierAllocContext(BaseEmitContext): @@ -44,17 +46,8 @@ def finalize(self): # No barriers to allocate return - # Place barriers past the smem high-water mark rather than into a slot that - # was freed by FreeShared. The smem allocator only tracks the static - # post-emit free list; at runtime, data tensors that get freed *after* the - # main loop are still alive while the barriers are in use, so reusing their - # slot would cause TMA writes to clobber the barrier state. - smem_alloc_ctx = self.contexts.smem_alloc_ctx - nbytes = num_barriers * uint64.nbytes - alignment = 128 - aligned_hwm = (smem_alloc_ctx.smem_allocator.maximum_allocated + alignment - 1) // alignment * alignment - virtual_smem_addr = aligned_hwm - smem_alloc_ctx.smem_allocator.maximum_allocated = aligned_hwm + nbytes + tensor = SharedTensor(dtype=uint64, shape=(num_barriers,), optional_layout=ops.shared_row_major(num_barriers)) + virtual_smem_addr = self.contexts.smem_alloc_ctx.allocate_shared_tensor(tensor, nbytes=tensor.storage_nbytes) sb = StmtBuilder() sb.declare( v=self.barrier_addr, diff --git a/python/tilus/backends/contexts/smem_alloc_ctx.py b/python/tilus/backends/contexts/smem_alloc_ctx.py index 9cb24768..981fbcf9 100644 --- a/python/tilus/backends/contexts/smem_alloc_ctx.py +++ b/python/tilus/backends/contexts/smem_alloc_ctx.py @@ -155,7 +155,11 @@ def finalize(self): # define the shared workspace variable if needed if self.shared_workspace_var is not None: workspace_offset = (maximum_allocated + 127) // 128 * 128 # align to 128 bytes - maximum_allocated += self.shared_workspace_bytes + # Size the arena from the aligned workspace offset, not the (possibly + # unaligned) high-water mark: otherwise the alignment padding between + # `maximum_allocated` and `workspace_offset` is not reserved and the + # workspace tail spills past the requested dynamic shared memory. + maximum_allocated = workspace_offset + self.shared_workspace_bytes sb = StmtBuilder() sb.declare(self.shared_workspace_var, init=dynamic_shared_memory(workspace_offset, dtype=uint8)) self.kernel_prepend(sb.finish())