Fix zero-like GPU launches, architecture-aware caching, and JIT API

Author: whjthu

src/ninetoothed/__init__.py

@@ -1,6 +1,7 @@
 from ninetoothed.aot import aot
 from ninetoothed.build import build
 from ninetoothed.compiler import lower
+from ninetoothed.compiler.jit import jit
 from ninetoothed.dtype import (
     bfloat16,
     float16,
@@ -17,7 +18,6 @@
 )
 from ninetoothed.eval import _eval as eval
 from ninetoothed.eval import _subs as subs
-from ninetoothed.jit import jit
 from ninetoothed.make import make
 from ninetoothed.symbol import Symbol, block_size
 from ninetoothed.tensor import Tensor

src/ninetoothed/backends/emitters/cuda.py

@@ -288,6 +288,9 @@ def render_module(self, context: ModuleRenderContext) -> str:
 ) {{
     constexpr int threads = {threads};
     int64_t blocks = {blocks_expr};
+    if (blocks <= 0) {{
+        return static_cast<int>(cudaSuccess);
+    }}
     {kernel.kernel_name}_kernel<<<static_cast<unsigned int>(blocks), threads, 0, stream>>>(
         {args}
     );

src/ninetoothed/backends/materializers/cuda.py

@@ -96,16 +96,21 @@ def _cuda_wrapper(function, abi, tensor_specs):
     from ninetoothed.compiler.runtime import (
         KernelLaunchError,
         _bound_values,
+        _empty_launch,
         _first_output,
         _public_values,
     )
 
     spec_by_name = {spec.name: spec for spec in tensor_specs}
 
     def launch(*args, **kwargs):
+        public = _public_values(abi, args, kwargs, specs=tensor_specs)
+
+        if _empty_launch(abi, public):
+            return _first_output(abi, public)
+
         import torch
 
-        public = _public_values(abi, args, kwargs, specs=tensor_specs)
         values, keepalive = _bound_values(
             abi,
             public,

src/ninetoothed/compiler/__init__.py

@@ -11,15 +11,14 @@
     make,
     resolve_target,
 )
-from ninetoothed.compiler.jit import JIT, jit
+from ninetoothed.compiler.jit import jit
 from ninetoothed.compiler.runtime import load_built_artifact
 
 __all__ = [
     "Compiler",
     "Compilation",
     "CompileRequest",
     "DEFAULT_COMPILER",
-    "JIT",
     "aot",
     "compile_kernel",
     "jit",

src/ninetoothed/compiler/cache.py

@@ -49,7 +49,7 @@ def compilation_cache_key(compilation) -> str:
             "backend_options": request.backend_options,
             "pipeline": request.pipeline,
             "pass_options": request.pass_options,
-            "architecture": _architecture(),
+            "architecture": _architecture(compilation),
             "versions": _compiler_versions(),
         }
     )
@@ -191,13 +191,51 @@ def _compiler_versions() -> Mapping[str, str]:
     return result
 
 
-def _architecture() -> Mapping[str, str | None]:
-    return {
+def _architecture(compilation) -> Mapping[str, Any]:
+    backend = compilation.artifact.backend.value
+    backend_options = dict(compilation.request.backend_options or {})
+    architecture = {
         "machine": platform.machine(),
         "cuda_arch": os.environ.get("TORCH_CUDA_ARCH_LIST"),
         "cuda_visible_devices": os.environ.get("CUDA_VISIBLE_DEVICES"),
     }
 
+    if backend == "cuda":
+        target = str(backend_options.get("arch", "native"))
+        architecture["cuda_target"] = target
+
+        if target != "native":
+            return architecture
+
+    if backend in {"cuda", "tilelang", "triton"}:
+        architecture.update(_runtime_cuda_architecture())
+    return architecture
+
+
+def _runtime_cuda_architecture() -> Mapping[str, Any]:
+    import torch
+
+    if not torch.cuda.is_available():
+        return {
+            "cuda_current_capability": None,
+            "cuda_visible_capabilities": (),
+        }
+
+    def capability(device: int) -> str:
+        major, minor = torch.cuda.get_device_capability(device)
+
+        return f"sm_{major}{minor}"
+
+    current = capability(torch.cuda.current_device())
+    visible = tuple(
+        dict.fromkeys(capability(device) for device in range(torch.cuda.device_count()))
+    )
+
+    return {
+        "cuda_current_capability": current,
+        "cuda_visible_capabilities": visible,
+    }
+
 
 def _json_value(value: Any) -> Any:
     if value is None or isinstance(value, (bool, int, float, str)):

src/ninetoothed/compiler/jit.py

@@ -24,65 +24,22 @@ def jit(
         )
 
     def wrapper(application):
-        return JIT(
-            application,
-            caller=caller,
-            backend=backend,
-            kernel_name=kernel_name,
-            num_warps=num_warps,
-            num_stages=num_stages,
-            max_num_configs=max_num_configs,
-            pipeline=pipeline,
-            pass_options=pass_options,
-            backend_options=backend_options,
-        )()
-
-    return wrapper if func is None else wrapper(func)
-
-
-class JIT:
-    """Deferred materialization object used by :func:`jit`."""
-
-    def __init__(
-        self,
-        func,
-        *,
-        backend,
-        caller,
-        kernel_name,
-        num_warps,
-        num_stages,
-        max_num_configs,
-        pipeline,
-        pass_options,
-        backend_options=None,
-    ):
-        self.func = func
-        self._caller = caller
-        self._backend = backend
-        self._kernel_name = kernel_name or func.__name__
-        self._num_warps = num_warps
-        self._num_stages = num_stages
-        self._max_num_configs = max_num_configs
-        self._pipeline = pipeline
-        self._pass_options = pass_options
-        self._backend_options = dict(backend_options or {})
-
-    def __call__(self):
         return DEFAULT_COMPILER.materialize(
             CompileRequest(
-                application=self.func,
-                backend=self._backend,
-                caller=self._caller,
-                kernel_name=self._kernel_name,
-                num_warps=self._num_warps,
-                num_stages=self._num_stages,
-                max_num_configs=self._max_num_configs,
-                pipeline=self._pipeline,
-                pass_options=self._pass_options,
-                backend_options=self._backend_options,
+                application=application,
+                backend=backend,
+                caller=caller,
+                kernel_name=kernel_name or application.__name__,
+                num_warps=num_warps,
+                num_stages=num_stages,
+                max_num_configs=max_num_configs,
+                pipeline=pipeline,
+                pass_options=pass_options,
+                backend_options=backend_options,
             )
         )
 
+    return wrapper if func is None else wrapper(func)
+
 
-__all__ = ["JIT", "jit"]
+__all__ = ["jit"]

src/ninetoothed/compiler/runtime.py

@@ -276,6 +276,10 @@ def _runtime_wrapper(
 ):
     def launch(*args, **kwargs):
         public = _public_values(abi, args, kwargs, specs=specs)
+
+        if _empty_launch(abi, public):
+            return _first_output(abi, public)
+
         values, keepalive = _bound_values(abi, public, scalar_mode="value")
         bindings = abi.kernel_args
 

src/ninetoothed/jit.py

@@ -206,6 +206,7 @@ def test_optimization_metadata_contains_only_materialized_choices(self):
     def test_generic_cuda_launch_matches_emitted_thread_count(self):
         artifact = emit(_add_kernel(), "cuda")
         assert "constexpr int threads = 256;" in artifact.primary_source
+        assert "if (blocks <= 0)" in artifact.primary_source
         assert artifact.metadata["launch_block"] == ("256",)
 
     def test_artifact_can_write_all_sources(self, tmp_path):

tests/test_backend_registry.py

@@ -68,6 +68,38 @@ def test_aot_built_artifact_can_be_reloaded(backend, device, tmp_path):
     )
 
 
+@pytest.mark.parametrize("device", get_available_devices())
+@pytest.mark.parametrize("mode", ("jit", "aot"))
+def test_cuda_empty_tensor_is_a_no_op(mode, device, tmp_path):
+    tensors = tuple(Tensor(1, dtype=ninetoothed.float32) for _ in range(3))
+    compilation = DEFAULT_COMPILER.compile(
+        CompileRequest(
+            arrangement=_arrangement,
+            application=_application,
+            tensors=tensors,
+            backend="cuda",
+            kernel_name=f"empty_cuda_{mode}",
+        )
+    )
+    handle = DEFAULT_COMPILER.materialize(
+        compilation,
+        output_dir=tmp_path if mode == "aot" else None,
+        mode=mode,
+    )
+    input = torch.empty(0, device=device)
+    other = torch.empty_like(input)
+    output = torch.empty_like(input)
+
+    assert handle(input, other, output) is output
+    torch.cuda.synchronize()
+
+    if mode == "aot":
+        reloaded = load_built_artifact(handle._built_artifact)
+        reloaded_output = torch.empty_like(input)
+        assert reloaded(input, other, reloaded_output) is reloaded_output
+        torch.cuda.synchronize()
+
+
 _RELOAD_IN_SUBPROCESS = """
 import pickle
 import sys