From 28f387922980494f2de2888b16dcef3a6a22131d Mon Sep 17 00:00:00 2001 From: Feathbow Date: Wed, 21 Oct 2026 23:34:29 +0100 Subject: [PATCH 1/2] fix(codegen): raise when the device refuses a kernel launch --- docs/spec/codegen.md | 4 +- docs/spec/runtime.md | 5 +++ src/tilefoundry/codegen/cuda/module.py | 13 ++++++- .../cuda/templates/cuda_preamble.cu.j2 | 2 + tests/ops/tir/test_launch.py | 38 ++++++++++++++++++- 5 files changed, 58 insertions(+), 4 deletions(-) diff --git a/docs/spec/codegen.md b/docs/spec/codegen.md index 5201a0975..64c4ee483 100644 --- a/docs/spec/codegen.md +++ b/docs/spec/codegen.md @@ -44,7 +44,9 @@ flowchart LR - **Host / device boundary.** A host `LinkableModule` MUST NOT reference CUDA or CuTe symbols or types. A CUDA `LinkableModule` owns the kernels and their C-ABI launchers. The host module invokes device code only through that - C-ABI launcher. + C-ABI launcher. The launcher MUST read the launch status right after the + launch and, when the device refuses the launch, throw naming the kernel and + the CUDA error, so a refused launch never returns as though the kernel ran. The target-specific generator behavior — how a CPU vs CUDA function emits, program-shape / dynamic-CTA accessors, and the `ShardLayout` runtime mapping — diff --git a/docs/spec/runtime.md b/docs/spec/runtime.md index c64ad7f14..6790e129a 100644 --- a/docs/spec/runtime.md +++ b/docs/spec/runtime.md @@ -393,6 +393,11 @@ Launch geometry (grid / block extents) is embedded into the generated host entry or supplied by an authored `launch(...)`; it is never carried as metadata past codegen. +A launch the device refuses raises `RuntimeError` from the call, naming the +kernel and the CUDA error; the call never returns as though the kernel ran. A +fault while the kernel runs is asynchronous and surfaces at the next +synchronizing CUDA call. + ### 1.6 `compile.py` diff --git a/src/tilefoundry/codegen/cuda/module.py b/src/tilefoundry/codegen/cuda/module.py index a7c86d2ac..cfd2ba611 100644 --- a/src/tilefoundry/codegen/cuda/module.py +++ b/src/tilefoundry/codegen/cuda/module.py @@ -75,7 +75,10 @@ def _define_kernel(fn: ModuleFunction, kernel: CallableSignature, ctx: CudaCodeg def _define_launcher( launcher: CallableSignature, kernel: CallableSignature, ctx: CudaCodegenContext ) -> None: - """The ``extern "C"`` entry: the geometry it was handed, spent on one launch.""" + """The ``extern "C"`` entry: the geometry it was handed, spent on one launch. + + A launch the device refuses throws, so the caller never reads outputs the kernel did not write. + """ ctx.emit(f'extern "C" void {launcher.name}({ctx.parameters(launcher, exported=True)}) {{') ctx.indent() _tell_the_block(launcher, ctx) @@ -83,6 +86,14 @@ def _define_launcher( ctx.emit(f"dim3 grid({grid_x}, 1, 1);") ctx.emit(f"dim3 block({block_x}, 1, 1);") ctx.emit(f"{kernel.name}<<>>({kernel_arguments(kernel, ctx)});") + ctx.emit("if (cudaError_t status = cudaGetLastError(); status != cudaSuccess) {") + ctx.indent() + ctx.emit( + f'throw std::runtime_error(std::string("tilefoundry: launch of {kernel.name} failed: ") ' + "+ cudaGetErrorString(status));" + ) + ctx.dedent() + ctx.emit("}") ctx.dedent() ctx.emit("}") diff --git a/src/tilefoundry/codegen/cuda/templates/cuda_preamble.cu.j2 b/src/tilefoundry/codegen/cuda/templates/cuda_preamble.cu.j2 index deec05438..2ae5bd8eb 100644 --- a/src/tilefoundry/codegen/cuda/templates/cuda_preamble.cu.j2 +++ b/src/tilefoundry/codegen/cuda/templates/cuda_preamble.cu.j2 @@ -2,6 +2,8 @@ // Which levels this program names: everything from here to the finest. #define TILEFOUNDRY_PROGRAM_TOPOLOGY {{ program_topology }} #include +#include +#include namespace tilefoundry { {%- for dim in program_dims %} diff --git a/tests/ops/tir/test_launch.py b/tests/ops/tir/test_launch.py index a6f99f56c..77effa3e1 100644 --- a/tests/ops/tir/test_launch.py +++ b/tests/ops/tir/test_launch.py @@ -1,10 +1,12 @@ -"""Verify host-to-device launch argument contracts.""" +"""Verify host-to-device launch contracts.""" from __future__ import annotations import pytest +import torch -from tilefoundry.dsl import Tensor +import tilefoundry +from tilefoundry.dsl import * from tilefoundry.ir.core import Constant, Var, VerifyError from tilefoundry.ir.tir.launch import Launch from tilefoundry.ir.tir.prim_function import PrimFunction @@ -41,3 +43,35 @@ def test_launch_rejects_forwarded_type_that_differs_from_device_param() -> None: with pytest.raises(VerifyError, match=r"forwarded arg\[0\] type"): verify_module((host, device)) + + +_CUDA = CudaTarget("nvidia.h200_sxm") + + +@module(entry="over_asking_host", target=_CUDA, topologies=(Topology("thread", 128),)) +class OverAskingLaunch: + """One kernel launched with 1 MiB of dynamic shared memory. + + No CUDA device grants a block that much, so every device refuses the launch. + """ + + @prim_func(target=_CUDA) + def over_asking_device(a: Tensor[(4, 32), "f32"]): + with Mesh((Topology("thread", 128),), layout=(4, 32), names=("w", "t")) as m: + view = T.tensor_view(T.ptr_of(a), layout=((4 @ m.w, 32 @ m.t), (32, 1))) + reg = T.alloc_tensor( + tensor_type=Tensor[(4, 32), "f32", ((4 @ m.w, 32 @ m.t), (32, 1)), "rmem"] + ) + T.copy(view, reg) + T.copy(reg, view) + + @prim_func(target=CpuTarget()) + def over_asking_host(a: Tensor[(4, 32), "f32"]): + launch(over_asking_device, a, grid=(1, 1, 1), block=(128, 1, 1), dynamic_smem=1048576) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +def test_a_launch_the_device_refuses_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")) From 5f5de7f23e282f8aacdd019e75fdfd5563d82d51 Mon Sep 17 00:00:00 2001 From: Feathbow Date: Thu, 22 Oct 2026 07:46:12 +0100 Subject: [PATCH 2/2] fix(codegen): return the launch error to the host entry --- docs/spec/codegen.md | 6 ++++-- src/tilefoundry/codegen/cpu/module.py | 15 ++++++++++++--- src/tilefoundry/codegen/cuda/module.py | 17 ++++++----------- .../codegen/cuda/templates/cuda_preamble.cu.j2 | 2 -- tests/codegen/test_device_dispatch.py | 2 +- tests/codegen/test_host_multi_launch.py | 2 +- 6 files changed, 24 insertions(+), 20 deletions(-) diff --git a/docs/spec/codegen.md b/docs/spec/codegen.md index 64c4ee483..cd2a2a79a 100644 --- a/docs/spec/codegen.md +++ b/docs/spec/codegen.md @@ -45,8 +45,10 @@ flowchart LR CuTe symbols or types. A CUDA `LinkableModule` owns the kernels and their C-ABI launchers. The host module invokes device code only through that C-ABI launcher. The launcher MUST read the launch status right after the - launch and, when the device refuses the launch, throw naming the kernel and - the CUDA error, so a refused launch never returns as though the kernel ran. + launch and return the CUDA error text when the device refuses the launch, or + null otherwise. The host entry MUST raise when a launcher returns an error, + so a refused launch never returns as though the kernel ran and no exception + crosses the C ABI. The target-specific generator behavior — how a CPU vs CUDA function emits, program-shape / dynamic-CTA accessors, and the `ShardLayout` runtime mapping — diff --git a/src/tilefoundry/codegen/cpu/module.py b/src/tilefoundry/codegen/cpu/module.py index 96c417606..a8fb15f04 100644 --- a/src/tilefoundry/codegen/cpu/module.py +++ b/src/tilefoundry/codegen/cpu/module.py @@ -195,7 +195,10 @@ def _forwarded(entry: PrimFunction, launch: Evaluate, signature: CallableSignatu def _one_launch( entry: PrimFunction, launch: Evaluate, callee: PrimFunction, ctx: CpuCodegenContext ) -> list[str]: - """The lines one ``Launch`` becomes: the checks it owes, then the call.""" + """The lines one ``Launch`` becomes: the checks it owes, then the call. + + The launcher returns the CUDA error text of a launch the device refused, and the call raises it. + """ _reject_unsupported_config(launch.callable) launcher = ctx.signature_of(callee) forwarded = _forwarded(entry, launch, launcher) @@ -204,7 +207,13 @@ def _one_launch( _placement_line(var.name, param.type.storage) for var, param in zip(forwarded, launcher.params) ] - lines.append(f"{launcher.name}({ctx.arguments(bound, ctx.codegen_context.targets[id(callee)])});") + call = f"{launcher.name}({ctx.arguments(bound, ctx.codegen_context.targets[id(callee)])})" + lines.append(f"if (const char* error = {call}) {{") + lines.append( + f' TVM_FFI_THROW(RuntimeError) << "tilefoundry: launch of {callee.name} failed: "' + " << error;" + ) + lines.append("}") return lines @@ -215,7 +224,7 @@ def _declare_callee(callee: PrimFunction, ctx: CpuCodegenContext) -> str: it cannot come to disagree about what the symbol takes. """ launcher = ctx.signature_of(callee) - return f'extern "C" void {launcher.name}({ctx.parameters(launcher, ctx.codegen_context.targets[id(callee)], exported=True)});' + return f'extern "C" const char* {launcher.name}({ctx.parameters(launcher, ctx.codegen_context.targets[id(callee)], exported=True)});' def _host_entry( diff --git a/src/tilefoundry/codegen/cuda/module.py b/src/tilefoundry/codegen/cuda/module.py index cfd2ba611..4f8a1ab6e 100644 --- a/src/tilefoundry/codegen/cuda/module.py +++ b/src/tilefoundry/codegen/cuda/module.py @@ -43,7 +43,7 @@ def _emit_function(fn: ModuleFunction, ctx: CudaCodegenContext) -> LinkableFunct kernel = kernel_signature(launcher, ctx.codegen_context.symbols.by_function[id(fn)].symbol) return LinkableFunction( name=fn.name, - declaration=f'extern "C" void {launcher.name}({ctx.parameters(launcher, exported=True)});', + declaration=f'extern "C" const char* {launcher.name}({ctx.parameters(launcher, exported=True)});', definition=ctx.capture(lambda scope: _define(fn, launcher, kernel, scope)), ) @@ -77,23 +77,18 @@ def _define_launcher( ) -> None: """The ``extern "C"`` entry: the geometry it was handed, spent on one launch. - A launch the device refuses throws, so the caller never reads outputs the kernel did not write. + 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. """ - ctx.emit(f'extern "C" void {launcher.name}({ctx.parameters(launcher, exported=True)}) {{') + 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"{kernel.name}<<>>({kernel_arguments(kernel, ctx)});") - ctx.emit("if (cudaError_t status = cudaGetLastError(); status != cudaSuccess) {") - ctx.indent() - ctx.emit( - f'throw std::runtime_error(std::string("tilefoundry: launch of {kernel.name} failed: ") ' - "+ cudaGetErrorString(status));" - ) - ctx.dedent() - ctx.emit("}") + ctx.emit("cudaError_t status = cudaGetLastError();") + ctx.emit("return status == cudaSuccess ? nullptr : cudaGetErrorString(status);") ctx.dedent() ctx.emit("}") diff --git a/src/tilefoundry/codegen/cuda/templates/cuda_preamble.cu.j2 b/src/tilefoundry/codegen/cuda/templates/cuda_preamble.cu.j2 index 2ae5bd8eb..deec05438 100644 --- a/src/tilefoundry/codegen/cuda/templates/cuda_preamble.cu.j2 +++ b/src/tilefoundry/codegen/cuda/templates/cuda_preamble.cu.j2 @@ -2,8 +2,6 @@ // Which levels this program names: everything from here to the finest. #define TILEFOUNDRY_PROGRAM_TOPOLOGY {{ program_topology }} #include -#include -#include namespace tilefoundry { {%- for dim in program_dims %} diff --git a/tests/codegen/test_device_dispatch.py b/tests/codegen/test_device_dispatch.py index a2075a868..592e5037a 100644 --- a/tests/codegen/test_device_dispatch.py +++ b/tests/codegen/test_device_dispatch.py @@ -54,7 +54,7 @@ def test_a_prototype_becomes_one_kernel_and_one_shim() -> None: assert source.count("__global__") == 1 assert ( source.count( - f'extern "C" void {launcher_symbol(mangle(TirSquare, TirSquare.square_device))}(' + f'extern "C" const char* {launcher_symbol(mangle(TirSquare, TirSquare.square_device))}(' ) == 2 ) diff --git a/tests/codegen/test_host_multi_launch.py b/tests/codegen/test_host_multi_launch.py index eb0b57cf9..1f632d747 100644 --- a/tests/codegen/test_host_multi_launch.py +++ b/tests/codegen/test_host_multi_launch.py @@ -195,7 +195,7 @@ def test_same_device_function_can_be_launched_twice() -> None: source = _source((device,), (_launch(device, a, block=32), _launch(device, a, block=64)), (a,)) mod, _ = _module((device,), (), ()) symbol = launcher_symbol(mangle(mod, device)) - assert source.count(f'extern "C" void {symbol}') == 1 + assert source.count(f'extern "C" const char* {symbol}') == 1 assert source.count(f"{symbol}(a.data_ptr()") == 2