fix(codegen): raise when the device refuses a kernel launch

docs/spec/codegen.md

@@ -44,7 +44,11 @@ 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 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 —

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`
 
 <!-- generated: py-compile -->

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(

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)),
     )
 
@@ -75,14 +75,20 @@ 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."""
-    ctx.emit(f'extern "C" void {launcher.name}({ctx.parameters(launcher, exported=True)}) {{')
+    """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.
+    """
+    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}<<<grid, block, {smem}, nullptr>>>({kernel_arguments(kernel, ctx)});")
+    ctx.emit("cudaError_t status = cudaGetLastError();")
+    ctx.emit("return status == cudaSuccess ? nullptr : cudaGetErrorString(status);")
     ctx.dedent()
     ctx.emit("}")
 

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
     )

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
 
 

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"))