fix(cuda): size a kernel's dynamic shared memory and declare its base once

targets/nvidia/tilefoundry_nvidia/cuda/codegen/context.py

@@ -9,7 +9,7 @@
 
 from __future__ import annotations
 
-from collections.abc import Iterable, Mapping
+from collections.abc import Callable, Iterable, Mapping
 
 from tilefoundry.codegen.context import EmitContext
 from tilefoundry.codegen.signature import CallableSignature, TensorSignature
@@ -95,17 +95,25 @@ def __init__(
         self._next_barrier_id = 1
         self.needs_grid_barrier_state = False
         """Set while emitting a grid barrier, which the module declares state for."""
-        self._has_smem_base = False
+        self._names_smem_base = False
 
-    def reset_smem_base(self) -> None:
-        """Start a kernel with no dynamic shared-memory base declaration."""
-        self._has_smem_base = False
+    def declare_smem_base(self, write: Callable[[CudaCodegenContext], None]) -> None:
+        """Run *write*, declaring the dynamic shared-memory base ahead of it if it names the base.
 
-    def smem_base(self) -> str:
-        """Return the byte-addressed dynamic shared-memory base, declaring it once."""
-        if not self._has_smem_base:
+        The base addresses the kernel's whole dynamic shared memory, so it is
+        declared once before the first statement, where every later scope can
+        reach it.
+        """
+        self._names_smem_base = False
+        body = self.capture(write)
+        if self._names_smem_base:
             self.emit("extern __shared__ unsigned char smem_base[];")
-            self._has_smem_base = True
+        if body:
+            self._lines.extend(body.split("\n"))
+
+    def smem_base(self) -> str:
+        """Name the byte-addressed dynamic shared-memory base."""
+        self._names_smem_base = True
         return "smem_base"
 
     def bind_extents(self, params: Iterable[TensorSignature]) -> None:

targets/nvidia/tilefoundry_nvidia/cuda/codegen/module.py

@@ -37,6 +37,8 @@
 
 Geometry = tuple[tuple[object, object, object], tuple[object, object, object]]
 
+_DEFAULT_DYNAMIC_SMEM = 48 * 1024
+
 
 def _emit_function(fn: ModuleFunction, ctx: CudaCodegenContext) -> LinkableFunction:
     """*fn* in both of its positions, written from the one convention it is called by.
@@ -70,27 +72,45 @@ def _define_kernel(fn: ModuleFunction, kernel: CallableSignature, ctx: CudaCodeg
     ctx.reset_barrier_ids()
     ctx.emit(f"__global__ void {kernel.name}({ctx.parameters(kernel)}) {{")
     ctx.indent()
+    ctx.declare_smem_base(lambda scope: _write_kernel_body(fn, kernel, scope))
+    ctx.dedent()
+    ctx.emit("}")
+
+
+def _write_kernel_body(
+    fn: ModuleFunction, kernel: CallableSignature, ctx: CudaCodegenContext
+) -> None:
+    """What the kernel was told, then what it was written to do."""
     if kernel.leading:
         ctx.emit(f"tilefoundry::program_meta({PROGRAM_META.name});")
     ctx.emit_node(fn)
-    ctx.dedent()
-    ctx.emit("}")
 
 
 def _define_launcher(
     launcher: CallableSignature, kernel: CallableSignature, ctx: CudaCodegenContext
 ) -> None:
     """The ``extern "C"`` entry: the geometry it was handed, spent on one launch.
 
-    It returns the CUDA error text when the device refuses the launch and null otherwise, so
-    the host raises without an exception crossing the C ABI.
+    CUDA grants a kernel 48 KiB of dynamic shared memory unless it opts in to more,
+    so a larger request first raises the kernel's limit to what the launch asks for.
+    It returns the CUDA error text when the device refuses either step and null
+    otherwise, so the host raises without an exception crossing the C ABI.
     """
     ctx.emit(f'extern "C" const char* {launcher.name}({ctx.parameters(launcher, exported=True)}) {{')
     ctx.indent()
     _tell_the_block(launcher, ctx)
     grid_x, block_x, smem = (signature.name for signature in launcher.trailing)
     ctx.emit(f"dim3 grid({grid_x}, 1, 1);")
     ctx.emit(f"dim3 block({block_x}, 1, 1);")
+    ctx.emit(f"if ({smem} > {_DEFAULT_DYNAMIC_SMEM}) {{")
+    ctx.indent()
+    ctx.emit(
+        f"cudaError_t raised = cudaFuncSetAttribute({kernel.name}, "
+        f"cudaFuncAttributeMaxDynamicSharedMemorySize, {smem});"
+    )
+    ctx.emit("if (raised != cudaSuccess) return cudaGetErrorString(raised);")
+    ctx.dedent()
+    ctx.emit("}")
     ctx.emit(f"{kernel.name}<<<grid, block, {smem}, nullptr>>>({kernel_arguments(kernel, ctx)});")
     ctx.emit("cudaError_t status = cudaGetLastError();")
     ctx.emit("return status == cudaSuccess ? nullptr : cudaGetErrorString(status);")

targets/nvidia/tilefoundry_nvidia/cuda/codegen/tir/prim_function.py

@@ -98,7 +98,6 @@ def _subject(fn: PrimFunction, pattern: RangePattern, ctx: CudaCodegenContext) -
 @register_codegen(CudaTarget, Role.EMIT, PrimFunction)
 def _emit(fn: PrimFunction, ctx: CudaCodegenContext) -> None:
     """Write what runs inside one ``__global__``: its buffers, then its statements."""
-    ctx.reset_smem_base()
     if fn.variants:
         _dispatch(fn, ctx)
         return

targets/nvidia/tilefoundry_nvidia/cuda/docs/spec/codegen.md

@@ -12,10 +12,13 @@ body does not implement multi-device execution. Skeletons retain signatures,
 parameter views, specialization dispatch and launchers, replacing statements
 with `// TODO: <fn>`.
 
-The device unit owns kernels and C-ABI launchers. A launcher reads launch
-status immediately and returns the CUDA error text or null. Host entries
-raise on an error; no exception crosses the C ABI and a refused launch never
-returns as though the kernel ran. A host unit MUST NOT reference CUDA or CuTe
+The device unit owns kernels and C-ABI launchers. A launcher asked for more
+than 48 KiB of dynamic shared memory MUST first raise its kernel's
+`cudaFuncAttributeMaxDynamicSharedMemorySize` to that size. A launcher reads
+launch status immediately and returns the CUDA error text or null, including
+when the device refuses the raised limit. Host entries raise on an error; no
+exception crosses the C ABI and a refused launch never returns as though the
+kernel ran. A host unit MUST NOT reference CUDA or CuTe
 types.
 
 CUDA registers the host writer for shared `Launch` at
@@ -35,17 +38,19 @@ class CudaCodegenContext(EmitContext):
     def bind_extents(self, params) -> None: ...
     def reset_barrier_ids(self) -> None: ...
     def alloc_barrier_id(self) -> int: ...
-    def reset_smem_base(self) -> None: ...
+    def declare_smem_base(self, write: Callable[[CudaCodegenContext], None]) -> None: ...
     def smem_base(self) -> str: ...
     def dtype_to_cpp(self, dtype_name: str) -> str: ...
 ```
 
 - constraints:
   - `dynamic_extents` is derived from parameter types by `bind_extents`.
     `launches` records settled launch geometry keyed by `id(fn)`.
-  - `smem_base` declares one byte-addressed dynamic shared-memory base per
-    kernel. Numeric `TensorView` addresses add byte offsets before conversion
-    to CuTe shared-memory pointers.
+  - Each kernel declares one byte-addressed dynamic shared-memory base, ahead
+    of its first statement, when its body names the base through `smem_base`.
+    `declare_smem_base` writes the body first so every scope in the kernel
+    reaches that one declaration. Numeric `TensorView` addresses add byte
+    offsets before conversion to CuTe shared-memory pointers.
   - Handlers obtain target type spellings and hardware counters through the
     context rather than rediscovering them from IR.
 

targets/nvidia/tilefoundry_nvidia/cuda/docs/spec/passes.md

@@ -66,5 +66,11 @@ def insert_default_host_entry(module: Module) -> Module:
   - The launch geometry MUST be settled here and written into the `Launch`, so
     nothing downstream derives it again. A prototype states no body of its own,
     so its variants state its geometry and they MUST agree.
+  - The synthesized TIR entry's `Launch` MUST ask for the dynamic shared memory
+    its device function reaches: the farthest byte any integer-address
+    shared-memory `TensorView` of the function or of its variants spans, each
+    view sized as a staged slot reserves it
+    ([analysis §1.2.2](../../../../../../docs/spec/analysis.md#122-memory)), and 0
+    when there is none. An authored host entry keeps the size its author wrote.
   - It MUST reject ambiguous device-function sets, a non-entry CPU function, or
     a launch-provided dynamic CTA extent.

targets/nvidia/tilefoundry_nvidia/cuda/passes/transforms/host_entry.py

@@ -2,25 +2,28 @@
 
 The entry is synthesized, never retargeted: which target a function records is
 what its author wrote, and a pass that rewrites it decides in the wrong place.
-Launch geometry is settled here and written into the ``Launch``, so nothing
-downstream derives it a second time. A launch-provided (dynamic) grid is not
-supported by the synthesized entry.
+Launch geometry, and the dynamic shared memory a kernel reaches, are settled
+here and written into the ``Launch``, so nothing downstream derives them a
+second time. A launch-provided (dynamic) grid is not supported by the
+synthesized entry.
 """
 
 from __future__ import annotations
 
 from dataclasses import replace
 
-from tilefoundry.ir.core import Var
+from tilefoundry.analysis.allocation import storage_bytes
+from tilefoundry.ir.core import Call, Constant, Expr, Var
 from tilefoundry.ir.core.module import Module
 from tilefoundry.ir.hir.function import Function
 from tilefoundry.ir.module_utils import list_call_paths
 from tilefoundry.ir.tir.launch import launch_call
+from tilefoundry.ir.tir.memory import TensorView
 from tilefoundry.ir.tir.prim_function import PrimFunction
 from tilefoundry.ir.tir.stmts import Sequential
-from tilefoundry.ir.types import StorageKind, TensorType, TupleType, UnitType
+from tilefoundry.ir.types import StorageKind, TensorType, Topology, TupleType, UnitType
 from tilefoundry.ir.types.layout import Layout
-from tilefoundry.ir.visitor import ExprVisitor
+from tilefoundry.ir.visitor import ExprVisitor, StmtExprMutator
 from tilefoundry.passes.pass_base import ModulePass
 from tilefoundry.passes.pass_manager import PassPoint, register_pass
 from tilefoundry.target import CpuTarget
@@ -88,6 +91,43 @@ def _launch_config_of(
     return (grid, 1, 1), (block, 1, 1)
 
 
+class _SharedMemoryReach(StmtExprMutator):
+    """The farthest byte any integer-address shared-memory view reaches.
+
+    It visits every expression a statement embeds and returns each unchanged.
+    An integer ``TensorView`` address is a byte offset from the dynamic
+    shared-memory base, and each view spans what the allocator reserves for it
+    within one unit of the level that owns shared memory.
+    """
+
+    def __init__(self, level: str | None, topologies: tuple[Topology, ...]) -> None:
+        self.level = level
+        self.topologies = topologies
+        self.end = 0
+
+    def visit_Call(self, call: Call, ctx=None) -> Expr:
+        address = call.args[0] if isinstance(call.target, TensorView) else None
+        if isinstance(address, Constant) and type(address.value) is int:
+            span = storage_bytes(call.type, topology_level=self.level, topologies=self.topologies)
+            self.end = max(self.end, address.value + span)
+        return self._expr_generic_visit(call, ctx)
+
+
+def _dynamic_shared_memory(module: Module, device_fn: PrimFunction) -> int:
+    """Bytes of dynamic shared memory *device_fn* and its variants reach in one owner unit.
+
+    A split at the level that owns shared memory, or a coarser one, gives each
+    unit its own share; a finer split shares that unit's storage.
+    """
+    owner = device_fn.target.storage_owners.owner.get(StorageKind.SMEM)
+    topologies = module.effective_topologies()
+    level = owner if any(topology.name == owner for topology in topologies) else None
+    reach = _SharedMemoryReach(level, topologies)
+    for fn in (device_fn, *device_fn.variants):
+        reach.visit(fn.body)
+    return reach.end
+
+
 def insert_default_host_entry(module: Module) -> Module:
     """Keep CPU entries or synthesize entries according to the module's IR kind."""
     entry = module.entry_function()
@@ -107,7 +147,9 @@ def entry_for_one_device_function(module: Module) -> Module:
     device_fn = _device_callee(module)
     entry_params = tuple(Var(type=p.type, name=p.name) for p in device_fn.params)
     grid, block = _launch_config_of(module)
-    launch = launch_call(device_fn, entry_params, grid, block)
+    launch = launch_call(
+        device_fn, entry_params, grid, block, dynamic_smem=_dynamic_shared_memory(module, device_fn)
+    )
     name = _fresh_entry_name({fn.name for fn in module.functions})
     entry = PrimFunction(
         name=name,

tests/nvidia/cuda/codegen/test_shared_memory_base.py

@@ -0,0 +1,31 @@
+"""The dynamic shared-memory base, declared once where every scope of a kernel reaches it."""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+import pytest
+import torch
+
+import tilefoundry
+import tilefoundry_nvidia.cuda.codegen  # noqa: F401 -- trigger emitter autodiscovery
+from tests.nvidia.cuda.ops.tir.targets import CUDA_TARGETS, retarget
+from tilefoundry.cli.source import load_authored_ir
+from tilefoundry_nvidia import CudaTarget
+
+CODEGEN = Path(__file__).parents[1] / "fixtures" / "codegen"
+
+
+@pytest.mark.parametrize("target", CUDA_TARGETS)
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA")
+def test_second_loop_reaches_the_shared_memory_base(target: CudaTarget) -> None:
+    """Both loops stage ``x`` through offset 0, the second outside the block of the first."""
+    TwoLoopStage = load_authored_ir(f"{CODEGEN / 'two_loop_stage.py'}:TwoLoopStage")
+    runtime = tilefoundry.compile(retarget(TwoLoopStage, target), target=target)
+    source = torch.arange(8, dtype=torch.bfloat16, device="cuda").view(2, 4)
+    first = torch.full_like(source, -1.0)
+    second = torch.full_like(source, -1.0)
+    runtime(source, first, second)
+    torch.cuda.synchronize()
+    assert torch.equal(first, source)
+    assert torch.equal(second, source)

tests/nvidia/cuda/fixtures/codegen/two_loop_stage.py

@@ -0,0 +1,56 @@
+"""``loop_outputs``'s two staging loops, run by one thread with the wait each read back needs."""
+
+from __future__ import annotations
+
+from tilefoundry.dsl import *
+from tilefoundry_nvidia import CudaTarget
+
+_CUDA = CudaTarget("nvidia.h200_sxm")
+
+
+@module(
+    entry="two_loops",
+    target=_CUDA,
+    topologies=(Topology("cta", 1), Topology("thread", 1)),
+)
+class TwoLoopStage:
+    """Each loop stages rows of ``x`` through its own view of shared-memory offset 0.
+
+    One thread issues every copy and reads it back, so each byte has one writer. The
+    second loop names the dynamic shared-memory base outside the block of the first.
+    """
+
+    @prim_func(target=_CUDA)
+    def two_loops(
+        x: Tensor[(2, 4), "bf16"], first: Tensor[(2, 4), "bf16"], second: Tensor[(2, 4), "bf16"]
+    ):
+        with Mesh((Topology("cta", 1),), layout=(1,), names=("d0",)) as _cta:
+            with Mesh((Topology("thread", 1),), layout=(1,), names=("d0",)) as _thread:
+                for row in range(0, 2, 1):
+                    source = T.tensor_view(
+                        T.ptr_of(x[row:row + 1, 0:0 + 1]), layout=(1, 4), shape=(1, 4)
+                    )
+                    landing = T.tensor_view(
+                        0, dtype="bf16", storage="smem", layout=(1, 4), shape=(1, 4)
+                    )
+                    T.copy_async(source, landing)
+                    T.cp_async_commit()
+                    T.cp_async_wait(n=0)
+                    window = T.tensor_view(
+                        T.ptr_of(first[row:row + 1, 0:0 + 4]), layout=(1, 4), shape=(1, 4)
+                    )
+                    T.copy(landing, window)
+                for row_2 in range(0, 2, 1):
+                    source_2 = T.tensor_view(
+                        T.ptr_of(x[row_2:row_2 + 1, 0:0 + 1]), layout=(1, 4), shape=(1, 4)
+                    )
+                    landing_2 = T.tensor_view(
+                        0, dtype="bf16", storage="smem", layout=(1, 4), shape=(1, 4)
+                    )
+                    T.copy_async(source_2, landing_2)
+                    T.cp_async_commit()
+                    T.cp_async_wait(n=0)
+                    window_2 = T.tensor_view(
+                        T.ptr_of(second[row_2:row_2 + 1, 0:0 + 4]), layout=(1, 4), shape=(1, 4)
+                    )
+                    T.copy(landing_2, window_2)

tests/nvidia/cuda/fixtures/ops/wide_shared_stage.py

@@ -0,0 +1,33 @@
+"""A tile staged through more dynamic shared memory than CUDA grants without opting in."""
+
+from __future__ import annotations
+
+from tilefoundry.dsl import *
+from tilefoundry_nvidia import CudaTarget
+
+_CUDA = CudaTarget("nvidia.h200_sxm")
+
+
+def _rows(mesh) -> ShardLayout:
+    """One contiguous 160-element row per thread."""
+    return ShardLayout(Layout((128, 160), (160, 1)), (S(0),), mesh)
+
+
+@module(entry="wide_stage", target=_CUDA, topologies=(Topology("thread", 128),))
+class WideSharedStage:
+    """One block stages a (128, 160) f32 tile, 80 KiB, through shared memory and back.
+
+    The module has no host entry, so the synthesized one decides what the launch asks for.
+    """
+
+    @prim_func(target=_CUDA)
+    def wide_stage(x: Tensor[(128, 160), "f32"], out: Tensor[(128, 160), "f32"]):
+        with Mesh((Topology("thread", 128),), Layout((128,), (1,))) as thread:
+            source = T.tensor_view(T.ptr_of(x), layout=_rows(thread))
+            stage = T.tensor_view(
+                0, dtype="f32", storage="smem", layout=_rows(thread), shape=(128, 160)
+            )
+            target = T.tensor_view(T.ptr_of(out), layout=_rows(thread))
+            T.copy(source, stage)
+            T.sync(thread)
+            T.copy(stage, target)