diff --git a/targets/nvidia/tilefoundry_nvidia/cuda/codegen/context.py b/targets/nvidia/tilefoundry_nvidia/cuda/codegen/context.py index d72aa68c..bcf50110 100644 --- a/targets/nvidia/tilefoundry_nvidia/cuda/codegen/context.py +++ b/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: diff --git a/targets/nvidia/tilefoundry_nvidia/cuda/codegen/module.py b/targets/nvidia/tilefoundry_nvidia/cuda/codegen/module.py index 0c4386df..edbbd27f 100644 --- a/targets/nvidia/tilefoundry_nvidia/cuda/codegen/module.py +++ b/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,11 +72,18 @@ 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( @@ -82,8 +91,10 @@ def _define_launcher( ) -> 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() @@ -91,6 +102,15 @@ def _define_launcher( 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}<<>>({kernel_arguments(kernel, ctx)});") ctx.emit("cudaError_t status = cudaGetLastError();") ctx.emit("return status == cudaSuccess ? nullptr : cudaGetErrorString(status);") diff --git a/targets/nvidia/tilefoundry_nvidia/cuda/codegen/tir/prim_function.py b/targets/nvidia/tilefoundry_nvidia/cuda/codegen/tir/prim_function.py index 52ddfbf4..bcbf5211 100644 --- a/targets/nvidia/tilefoundry_nvidia/cuda/codegen/tir/prim_function.py +++ b/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 diff --git a/targets/nvidia/tilefoundry_nvidia/cuda/docs/spec/codegen.md b/targets/nvidia/tilefoundry_nvidia/cuda/docs/spec/codegen.md index 50373c11..60d84a1b 100644 --- a/targets/nvidia/tilefoundry_nvidia/cuda/docs/spec/codegen.md +++ b/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: `. -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,7 +38,7 @@ 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: ... ``` @@ -43,9 +46,11 @@ class CudaCodegenContext(EmitContext): - 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. diff --git a/targets/nvidia/tilefoundry_nvidia/cuda/docs/spec/passes.md b/targets/nvidia/tilefoundry_nvidia/cuda/docs/spec/passes.md index ec7efaf2..707264d3 100644 --- a/targets/nvidia/tilefoundry_nvidia/cuda/docs/spec/passes.md +++ b/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. diff --git a/targets/nvidia/tilefoundry_nvidia/cuda/passes/transforms/host_entry.py b/targets/nvidia/tilefoundry_nvidia/cuda/passes/transforms/host_entry.py index b2129b44..7e8bd726 100644 --- a/targets/nvidia/tilefoundry_nvidia/cuda/passes/transforms/host_entry.py +++ b/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, diff --git a/tests/nvidia/cuda/codegen/test_shared_memory_base.py b/tests/nvidia/cuda/codegen/test_shared_memory_base.py new file mode 100644 index 00000000..a3c9a622 --- /dev/null +++ b/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) diff --git a/tests/nvidia/cuda/fixtures/codegen/two_loop_stage.py b/tests/nvidia/cuda/fixtures/codegen/two_loop_stage.py new file mode 100644 index 00000000..ac3da3ce --- /dev/null +++ b/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) diff --git a/tests/nvidia/cuda/fixtures/ops/wide_shared_stage.py b/tests/nvidia/cuda/fixtures/ops/wide_shared_stage.py new file mode 100644 index 00000000..ed5c3f4f --- /dev/null +++ b/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) diff --git a/tests/nvidia/cuda/fixtures/passes/device_entries_without_host.py b/tests/nvidia/cuda/fixtures/passes/device_entries_without_host.py index 2731ad78..67d4f641 100644 --- a/tests/nvidia/cuda/fixtures/passes/device_entries_without_host.py +++ b/tests/nvidia/cuda/fixtures/passes/device_entries_without_host.py @@ -1,4 +1,4 @@ -"""Device entries with no host entry: a single kernel and a dispatch prototype.""" +"""Device entries with no host entry: single kernels and dispatch prototypes.""" from __future__ import annotations @@ -51,3 +51,59 @@ def large(x: Tensor[(_S,), "f32"]): view = T.tensor_view(T.ptr_of(x), layout=_rows(128), shape=(128,)) T.copy(view, view) T.sync(thread) + + +@module(entry="stage_twice", target=_CUDA, topologies=(Topology("thread", 128),)) +class _SharedRing: + """A kernel declaring two slots of 512 bytes in shared memory; the second ends at byte 1024.""" + + @prim_func(target=_CUDA) + def stage_twice(x: Tensor[(128,), "f32"]): + with Mesh((Topology("thread", 128),), Layout((128,), (1,))) as _thread: + _slots = ( + T.tensor_view(0, dtype="f32", storage="smem", layout=_rows(128), shape=(128,)), + T.tensor_view(512, dtype="f32", storage="smem", layout=_rows(128), shape=(128,)), + ) + + +@module(entry="stage", target=_CUDA, topologies=(Topology("thread", 128),)) +class _StagedPrototype: + """A prototype whose variants declare different slots; the farther ends at byte 2560.""" + + @prim_func(target=_CUDA) + def stage(x: Tensor[(_S,), "f32"]): + pass + + @stage.specialize(RangePattern("S", 1, 127)) + def near(x: Tensor[(_S,), "f32"]): + with Mesh((Topology("thread", 128),), Layout((128,), (1,))) as _thread: + _slot = T.tensor_view( + 0, dtype="f32", storage="smem", layout=_rows(128), shape=(128,) + ) + + @stage.specialize(RangePattern("S", 128, 255)) + def far(x: Tensor[(_S,), "f32"]): + with Mesh((Topology("thread", 128),), Layout((128,), (1,))) as _thread: + _slot = T.tensor_view( + 2048, dtype="f32", storage="smem", layout=_rows(128), shape=(128,) + ) + + +@module( + entry="stage_halves", + target=_CUDA, + topologies=(Topology("cta", 2), Topology("thread", 128)), +) +class _CtaHalves: + """Two CTAs, each declaring its 128 KiB half of a (65536,) f32 view in shared memory.""" + + @prim_func(target=_CUDA) + def stage_halves(x: Tensor[(128,), "f32"]): + with Mesh((Topology("cta", 2),), Layout((2,), (1,))) as cta: + _halves = T.tensor_view( + 0, + dtype="f32", + storage="smem", + layout=ShardLayout(Layout((65536,), (1,)), (S(0),), cta), + shape=(65536,), + ) diff --git a/tests/nvidia/cuda/ops/tir/test_launch.py b/tests/nvidia/cuda/ops/tir/test_launch.py index 8003e3f6..ab8d630e 100644 --- a/tests/nvidia/cuda/ops/tir/test_launch.py +++ b/tests/nvidia/cuda/ops/tir/test_launch.py @@ -8,6 +8,7 @@ import torch import tilefoundry +from tests.nvidia.cuda.ops.tir.targets import CUDA_TARGETS, retarget from tilefoundry.cli.source import load_authored_ir from tilefoundry.dsl import * from tilefoundry.ir.core import Constant, Var, VerifyError @@ -59,3 +60,17 @@ def test_refused_launch_raises() -> None: runtime = tilefoundry.compile(OverAskingLaunch, target=_CUDA) with pytest.raises(RuntimeError, match=r"tilefoundry: launch of \w+ failed: "): runtime(torch.ones(4, 32, device="cuda")) + + +@pytest.mark.parametrize("target", CUDA_TARGETS) +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +def test_launch_lifts_the_dynamic_shared_memory_default(target: CudaTarget) -> None: + """An 80 KiB stage runs once its launch asks for 80 KiB and lifts CUDA's 48 KiB default.""" + WideSharedStage = load_authored_ir(f"{OPS / 'wide_shared_stage.py'}:WideSharedStage") + runtime = tilefoundry.compile(retarget(WideSharedStage, target), target=target) + torch.manual_seed(5) + source = torch.randn(128, 160, dtype=torch.float32, device="cuda") + out = torch.empty_like(source) + runtime(source, out) + torch.cuda.synchronize() + torch.testing.assert_close(out, source, rtol=0, atol=0) diff --git a/tests/nvidia/cuda/passes/test_host_entry.py b/tests/nvidia/cuda/passes/test_host_entry.py index bdc455e1..5033544a 100644 --- a/tests/nvidia/cuda/passes/test_host_entry.py +++ b/tests/nvidia/cuda/passes/test_host_entry.py @@ -39,12 +39,15 @@ class HostEntryCase(Case): class DeviceEntry(HostEntryCase): - """A device entry with no host entry, launched on one block of 128 threads.""" + """A device entry with no host entry, launched on blocks of 128 threads.""" + + ctas = 1 def guard(self, state, problems): problems.run("gains a host entry that launches it", self.gains_a_host_entry, state) problems.run("leaves every recorded target alone", self.leaves_targets_alone, state) problems.run("reads its geometry", self.reads_its_geometry, state) + problems.run("asks for the shared memory it reaches", self.asks_for_its_reach, state) def gains_a_host_entry(self, authored): lowered = insert_default_host_entry(authored) @@ -60,10 +63,14 @@ def leaves_targets_alone(self, authored): def reads_its_geometry(self, authored): assert _geometry(_launch_of(insert_default_host_entry(authored))) == ( - (1, 1, 1), + (self.ctas, 1, 1), (128, 1, 1), ) + def asks_for_its_reach(self, authored): + launch = _launch_of(insert_default_host_entry(authored)) + assert launch.callable.dynamic_smem == self.dynamic_smem + @case class LoneKernel(DeviceEntry): @@ -71,6 +78,7 @@ class LoneKernel(DeviceEntry): path: Path = DEVICE_ENTRIES name: str | None = "_OneKernel" + dynamic_smem: int = 0 @case @@ -82,6 +90,35 @@ class Prototype(DeviceEntry): path: Path = DEVICE_ENTRIES name: str | None = "_Prototype" + dynamic_smem: int = 0 + + +@case +class SharedRing(DeviceEntry): + """Two slots of 512 bytes at offsets 0 and 512, so the launch asks for 1024.""" + + path: Path = DEVICE_ENTRIES + name: str | None = "_SharedRing" + dynamic_smem: int = 1024 + + +@case +class StagedPrototype(DeviceEntry): + """One launch serves every variant, so it asks for the farthest any of them reaches.""" + + path: Path = DEVICE_ENTRIES + name: str | None = "_StagedPrototype" + dynamic_smem: int = 2560 + + +@case +class CtaHalves(DeviceEntry): + """Shared memory belongs to a CTA, so a view split across two CTAs asks for one half.""" + + path: Path = DEVICE_ENTRIES + name: str | None = "_CtaHalves" + ctas: int = 2 + dynamic_smem: int = 131072 @case