Skip to content

[Codegen] Fix shared-memory under-allocation for kernels #158

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
yaoyaoding merged 1 commit into from
Jul 5, 2026
Merged

[Codegen] Fix shared-memory under-allocation for kernels #158

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
15 changes: 4 additions & 11 deletions python/tilus/backends/contexts/mbarrier_alloc_ctx.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand All @@ -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,
Expand Down
6 changes: 5 additions & 1 deletion python/tilus/backends/contexts/smem_alloc_ctx.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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())
Expand Down
Loading