diff --git a/docsrc/py_api/kernels.rst b/docsrc/py_api/kernels.rst index a6eda8ff4b..2b00f93ab8 100644 --- a/docsrc/py_api/kernels.rst +++ b/docsrc/py_api/kernels.rst @@ -17,19 +17,22 @@ torch_tensorrt.kernels Overview -------- -The ``kernels`` module registers NVRTC-compiled CUDA C++ kernels as -TensorRT Quick Deployable Plugins. Tensor-only declarative kernels use -Ahead-of-Time (AOT) plugin launches when available; kernels with -``ScalarInput`` compile through TensorRT's QDP JIT path because QDP AOT -extra arguments currently support symbolic integer expressions, not -arbitrary runtime floats. - -A single function — :func:`cuda_kernel_op` — handles both the declarative -case (drive everything from a :class:`KernelSpec` dataclass) and the -override case (supply ``meta_fn`` / ``eager_fn`` / ``aot_fn`` / ``schema`` -keyword arguments when the declarative DSL doesn't cover your kernel). -:func:`ptx_op` is a parallel entry point for kernels that are already -compiled to PTX bytes. +The ``kernels`` module registers custom kernels — CUDA C++ compiled with +NVRTC, or cuTile — as TensorRT Quick Deployable Plugins. Tensor-only +declarative kernels use Ahead-of-Time (AOT) plugin launches when +available; kernels with ``ScalarInput`` compile through TensorRT's QDP JIT +path because QDP AOT extra arguments currently support symbolic integer +expressions, not arbitrary runtime floats. + +Three entry points share one registration funnel: + +* :func:`cuda_kernel_op` handles both the declarative case (drive + everything from a :class:`KernelSpec` dataclass) and the override case + (supply ``meta_fn`` / ``eager_fn`` / ``aot_fn`` / ``schema`` keyword + arguments when the declarative DSL doesn't cover your kernel). +* :func:`ptx_op` registers kernels that are already compiled to PTX bytes. +* :func:`cutile_op` registers a ``@ct.kernel`` cuTile program, compiling it + to PTX and deriving the AOT launch for you. Entry points ------------ @@ -117,18 +120,113 @@ Pre-compiled PTX entry point .. autofunction:: ptx_op +cuTile entry point +------------------ + +.. autofunction:: cutile_op + +:func:`cutile_op` is the cuTile analogue of :func:`cuda_kernel_op`. It +compiles the kernel once with ``cuda.tile.compilation.export_kernel``, then +registers the PyTorch custom op, the TRT plugin descriptor, the AOT impl +embedding the PTX, and the Torch-TensorRT converter:: + + import cuda.tile as ct + import tensorrt.plugin as trtp + import torch + import torch_tensorrt.kernels as ttk + + TILE = 128 + + @ct.kernel + def add_one_kernel(x, out, tile_size: ct.Constant[int]): + pid = ct.bid(0) + ct.store(out, index=(pid,), + tile=ct.load(x, index=(pid,), shape=(tile_size,)) + 1.0) + + def add_one_meta(X: torch.Tensor) -> torch.Tensor: + return torch.empty_like(X) + + ttk.cutile_op( + "my::add_one", + kernel=add_one_kernel, + signature={"x": "fp32", "out": "fp32"}, + meta_fn=add_one_meta, + grid=lambda inputs, outputs: ( + trtp.cdiv(inputs[0].shape_expr.numel(), TILE), + ), + constants={"tile_size": TILE}, + ) + +``signature`` lists the kernel's array parameters in declaration order — +inputs first, then outputs — mapped to their element type: a +:class:`torch.dtype`, a dtype name like ``"float32"``, or a short alias +like ``"fp32"``. ``constants`` supplies the ``ct.Constant`` values baked +into the compiled symbol. ``grid`` receives ``trtp.TensorDesc`` objects, so use +``.shape_expr`` to stay symbolic and keep one engine valid across shapes. + +Because the launch is built from symbolic shape expressions, ``cutile_op`` +supports dynamic shapes by default. Pass ``eager_fn`` to also give the op +a CUDA implementation outside TensorRT, or ``aot_fn`` to replace the +derived launch entirely. Threads-per-block comes from the ``.reqntid`` the +compiled kernel declares — cuTile vectorizes, so this is often below the +tile size, and it is a hard requirement rather than a hint. + +.. note:: + + cuTile groups each array's parameters as ``(ptr, extents..., strides...)`` + in kernel-declaration order, which is *not* the + ``(input_ptrs..., extra_args..., output_ptrs...)`` order TensorRT's AOT + launcher uses. ``cutile_op`` permutes the compiled PTX's ``.entry`` + parameter list and supplies the matching extents and strides as AOT extra + arguments. A mismatch here does not fail loudly — the kernel would read + whatever TensorRT placed in each slot — so the parameter count of the + compiled PTX is checked against the signature at registration time and a + disagreement raises :class:`RuntimeError`. + + ``ndim`` (default 1) is the rank each array is compiled for; a rank-1 + array's extent is the tensor's element count, which is what a kernel + written against a flattened view expects and what lets one registration + accept any input shape. Pass ``ndim=`` or a ``"[rank]"`` signature + entry for kernels that index multi-dimensional tiles. + + A ``cutile_op`` registration compiles a single PTX for the given dtypes + and ``constants``. Inputs or outputs whose dtypes differ from the + compiled ones are detected during conversion: the op is left out of the + engine and runs in PyTorch, with a warning naming the mismatch. Register + a second op if you need a second dtype — multi-config autotuning is not + yet supported. + + ``tileiras``, the cuTile compiler, ships with ``cuda-tile`` but is not on + ``PATH`` by default; add the package's bin directory (e.g. + ``/nvidia/cu13/bin``) before registering cuTile kernels. + + ``tileiras`` emits the PTX ISA of the toolkit it was built against, which + can be newer than the installed driver loads. Nothing catches that on its + own: TensorRT builds the engine, and at inference the plugin logs + ``onShapeChange status -1`` while ``enqueue`` still returns, so the model + silently produces wrong numbers. :func:`cutile_op` therefore offers the + compiled PTX to the driver at registration. If it is refused, the + ``.version`` header is lowered to what the driver accepts, the result is + re-checked, and a warning names both versions; if no header makes it + loadable, registration raises. Aligning the driver with the cuda-tile + toolchain removes the step; ``max_ptx_version=`` pins the header manually. + Kernel signature convention --------------------------- -All entry points assume the ``__global__`` kernel takes its arguments in -the fixed order:: +All entry points assume the kernel takes its arguments in the fixed +order:: (input_ptrs..., extras..., output_ptrs...) -Pointers are ``void*`` cast to the appropriate element type. Extras -follow the order declared in :attr:`KernelSpec.extras` for the -declarative path, or the order your ``aot_fn`` builds for the override -path. +This matches the order TensorRT passes tensor pointers and AOT extra +arguments. In a CUDA C++ ``__global__`` kernel, pointers are ``void*`` +cast to the appropriate element type; a cuTile kernel declares its arrays +in ``(inputs..., outputs...)`` order and ``cutile_op`` rewrites the +compiled PTX into the layout above. Extras follow the order declared in +:attr:`KernelSpec.extras` for the declarative path, the extents and strides +:func:`cutile_op` derives from the signature, or the order your ``aot_fn`` +builds for the override path. Error behavior -------------- diff --git a/docsrc/tutorials/extensibility/plugins/index.rst b/docsrc/tutorials/extensibility/plugins/index.rst index ef381c0e38..9449e6d4e1 100644 --- a/docsrc/tutorials/extensibility/plugins/index.rst +++ b/docsrc/tutorials/extensibility/plugins/index.rst @@ -1,8 +1,8 @@ Plugins ======= -Register custom CUDA and Triton kernels as TensorRT plugins — from -auto-generated Python plugins to AOT-compiled C++ plugins for use +Register custom CUDA, Triton, and cuTile kernels as TensorRT plugins — +from auto-generated Python plugins to AOT-compiled C++ plugins for use in serialized engines. .. toctree:: @@ -15,3 +15,4 @@ in serialized engines. Example: Custom Kernels with NVRTC in TensorRT AOT Plugins <../../_rendered_examples/dynamo/nvrtc_aot_plugin> Example: Auto-derived CUDA Kernel Plugins via cuda_kernel_op <../../_rendered_examples/dynamo/cuda_kernel_op> Example: Pre-compiled PTX Kernels via ptx_op <../../_rendered_examples/dynamo/ptx_op> + Example: cuTile Kernel AOT Plugins via cutile_op <../../_rendered_examples/dynamo/cutile_op> diff --git a/docsrc/tutorials/extensibility/plugins/plugins.rst b/docsrc/tutorials/extensibility/plugins/plugins.rst index 1bfaa22be7..a6ddb8fd7a 100644 --- a/docsrc/tutorials/extensibility/plugins/plugins.rst +++ b/docsrc/tutorials/extensibility/plugins/plugins.rst @@ -27,6 +27,10 @@ depending on your kernel language and performance requirements: - CUDA C++ via NVRTC - Pre-compiled PTX embedded in engine - :ref:`nvrtc_aot_plugin` + * - QDP declarative (AOT) + - cuTile + - Pre-compiled PTX embedded in engine + - :ref:`cutile_op` * - Manual (legacy) - Triton / any - JIT callback into Python at runtime @@ -120,6 +124,7 @@ For complete end-to-end examples see: * :ref:`auto_generate_plugins` — Triton kernel, QDP JIT plugin * :ref:`aot_plugin` — Triton kernel, QDP AOT plugin (pre-compiled PTX, no Python overhead at runtime) * :ref:`nvrtc_aot_plugin` — CUDA C++ kernel compiled with NVRTC, QDP AOT plugin +* :ref:`cutile_op` — cuTile kernel, QDP AOT plugin registered in one ``cutile_op`` call * :ref:`custom_kernel_plugins` — manual plugin + converter registration (legacy approach) ---- diff --git a/docsrc/user_guide/compilation/unsupported_ops.rst b/docsrc/user_guide/compilation/unsupported_ops.rst index 14bac94aee..b3848f9466 100644 --- a/docsrc/user_guide/compilation/unsupported_ops.rst +++ b/docsrc/user_guide/compilation/unsupported_ops.rst @@ -275,6 +275,12 @@ FAQ :ref:`aot_plugin` for an end-to-end example of compiling a Triton kernel as a TRT plugin for use in a serialized engine. +**"I have a cuTile kernel. Can I use it in a serialized TRT engine?"** + + Yes — ``torch_tensorrt.kernels.cutile_op`` compiles a ``@ct.kernel`` program + ahead of time and registers the op, plugin, and converter in one call, with + the PTX embedded in the engine. See :ref:`cutile_op`. + **"Operator X is listed as supported but my model still falls back"** The converter may only support specific overloads or dtype combinations. Check: diff --git a/examples/dynamo/cutile_op.py b/examples/dynamo/cutile_op.py new file mode 100644 index 0000000000..21402f61bf --- /dev/null +++ b/examples/dynamo/cutile_op.py @@ -0,0 +1,134 @@ +""" +.. _cutile_op: + +Register a cuTile kernel as an AOT QDP plugin via ``torch_tensorrt.kernels.cutile_op`` +======================================================================================= + +``cutile_op`` is the cuTile counterpart of :ref:`cuda_kernel_op`, and the +declarative form of the hand-written AOT plugin in :ref:`aot_plugin`: it takes +a ``@ct.kernel`` program, compiles it ahead of time, and registers the PyTorch +custom op, the TensorRT plugin descriptor, the AOT impl (with the compiled PTX +embedded in the engine), and the Torch-TensorRT converter — in one call. + +The piece ``cutile_op`` exists to handle is the calling convention. cuTile +expands every array parameter into ``(ptr, extents..., strides...)``, grouped +per array in declaration order. TensorRT's AOT plugin launcher instead passes +``(input_ptrs..., extra_args..., output_ptrs...)``. Those two orders do not +agree, and a mismatch does not fail — the kernel reads whatever landed in each +slot and returns plausible-looking garbage. ``cutile_op`` permutes the compiled +PTX's parameter list and supplies the matching extents and strides as AOT extra +arguments, so the kernel binds what it expects. +""" + +import argparse + +import cuda.tile as ct +import tensorrt.plugin as trtp +import torch + +import torch_tensorrt +import torch_tensorrt.kernels as ttk + +parser = argparse.ArgumentParser() +parser.add_argument("--min_block_size", type=int, default=1) +ARGS, _ = parser.parse_known_args() + +# %% +# Step 1: Define the cuTile kernel (pure cuTile, nothing TensorRT-specific) +# ------------------------------------------------------------------------- +# +# Array parameters come first — inputs then outputs — followed by the +# ``ct.Constant`` parameters, which are baked into the compiled symbol. + +TILE_SIZE = 128 + + +@ct.kernel +def add_one_kernel(x, out, tile_size: ct.Constant[int]): + pid = ct.bid(0) + tile = ct.load(x, index=(pid,), shape=(tile_size,)) + ct.store(out, index=(pid,), tile=tile + 1.0) + + +# %% +# Step 2: Describe the op and register it with a single ``cutile_op`` call +# ------------------------------------------------------------------------ +# +# * ``signature`` — the array parameters and their element types, in +# declaration order. ``ndim`` defaults to 1, matching a kernel written +# against a flattened view, so a rank-1 array's extent is the tensor's +# element count and the op accepts any input shape. +# * ``constants`` — the ``ct.Constant`` values to compile for. +# * ``grid`` — the launch grid in tiles, computed from ``trtp.TensorDesc`` +# inputs. Using ``.shape_expr`` keeps it symbolic so one engine covers a +# range of shapes. +# * ``meta_fn`` — shape/dtype inference for FakeTensors (the Torch schema is +# inferred from its type hints). +# * ``eager_fn`` — optional; lets ``torch.ops.my.add_one`` also run outside +# TensorRT. + + +def add_one_meta(X: torch.Tensor) -> torch.Tensor: + return torch.empty_like(X) + + +def add_one_eager(X: torch.Tensor) -> torch.Tensor: + Y = torch.empty_like(X) + flat_x = X.contiguous().reshape(-1) + flat_y = Y.reshape(-1) + ct.launch( + torch.cuda.current_stream().cuda_stream, + (ct.cdiv(flat_x.numel(), TILE_SIZE), 1, 1), + add_one_kernel, + (flat_x, flat_y, TILE_SIZE), + ) + return Y + + +ttk.cutile_op( + "my::add_one", + kernel=add_one_kernel, + signature={"x": "fp32", "out": "fp32"}, + meta_fn=add_one_meta, + grid=lambda inputs, outputs: (trtp.cdiv(inputs[0].shape_expr.numel(), TILE_SIZE),), + constants={"tile_size": TILE_SIZE}, + eager_fn=add_one_eager, + supports_dynamic_shapes=True, +) + + +# %% +# Step 3: Use it — the op lowers to the AOT QDP plugin inside the engine +# ---------------------------------------------------------------------- + + +class AddOne(torch.nn.Module): + def forward(self, x: torch.Tensor) -> torch.Tensor: + return torch.ops.my.add_one(x) + + +if __name__ == "__main__": + x = torch.randn(4, 256, device="cuda", dtype=torch.float32) + model = AddOne().cuda().eval() + + ref = x + 1 + assert torch.allclose(model(x), ref), "eager path mismatch" + print("eager path OK") + + trt_model = torch_tensorrt.compile( + model, + inputs=[x], + min_block_size=ARGS.min_block_size, + ) + print("engine compiled with the AOT QDP plugin") + + trt_out = trt_model(x) + if torch.allclose(trt_out, ref, atol=1e-5): + print("cutile_op AOT QDP plugin ran correctly under Torch-TensorRT") + else: + print( + "WARNING: TRT output did not match. Check that the CUDA driver " + "supports the kernel's PTX ISA (a toolkit newer than the driver " + "can require capping via max_ptx_version); registration, compile " + "and engine build succeeded." + ) diff --git a/py/torch_tensorrt/kernels/__init__.py b/py/torch_tensorrt/kernels/__init__.py index d35e895402..d856c37836 100644 --- a/py/torch_tensorrt/kernels/__init__.py +++ b/py/torch_tensorrt/kernels/__init__.py @@ -22,6 +22,12 @@ Useful when the PTX comes from an external compiler (Triton, a cached NVRTC output, etc.). +``cutile_op`` — register a ``@ct.kernel`` cuTile program. Compiles the kernel + to PTX for you, reorders its parameters from cuTile's per-array + ``(ptr, extents..., strides...)`` layout into TensorRT's launch order, and + derives the AOT launch — so you only supply the array ``signature``, the + ``ct.Constant`` values, a ``grid``, and a ``meta_fn``. + Minimal example — declarative ``cuda_kernel_op``:: import torch, torch_tensorrt @@ -65,7 +71,7 @@ SameAs, ScalarInput, ) -from torch_tensorrt.kernels._ops import cuda_kernel_op, ptx_op +from torch_tensorrt.kernels._ops import cuda_kernel_op, cutile_op, ptx_op __all__ = [ "Custom", @@ -80,5 +86,6 @@ "SameAs", "ScalarInput", "cuda_kernel_op", + "cutile_op", "ptx_op", ] diff --git a/py/torch_tensorrt/kernels/_cutile.py b/py/torch_tensorrt/kernels/_cutile.py new file mode 100644 index 0000000000..ff7767481d --- /dev/null +++ b/py/torch_tensorrt/kernels/_cutile.py @@ -0,0 +1,923 @@ +"""cuTile backend for ``torch_tensorrt.kernels.cutile_op``. + +Compiles a ``@ct.kernel`` cuTile program ahead of time and reshapes the result +into what TensorRT's AOT Quick Deployable Plugin launcher expects. + +Pipeline +-------- +1. **Signature validation** — the ``signature`` names the kernel's array + parameters (inputs then outputs) and ``constants`` the ``ct.Constant[...]`` + parameters that follow them. Both are checked against the op's arity before + anything is compiled. +2. **CUBIN compilation** — ``cuda.tile.compilation.export_kernel`` builds a + CUBIN for a ``KernelSignature`` of ``ArrayConstraint`` / ``ConstantConstraint`` + entries derived from the signature. +3. **PTX extraction** — the CUBIN embeds its PTX in a debug section; it is + recovered from the raw ELF bytes. +4. **Parameter reordering** — the cuTile kernel ABI groups parameters per array + as ``(ptr, extents..., strides...)``, in declaration order. TRT's AOT + launcher passes ``(input_ptrs..., extra_args..., output_ptrs...)``. The + ``.entry`` parameter list is permuted so the two agree, and + :func:`build_extra_args` produces the matching extents / strides. + +Every step that could silently bind the wrong argument raises instead. A +misordered launch does not fail — the kernel reads whatever TensorRT happened +to place in those slots and returns plausible-looking garbage. +""" + +from __future__ import annotations + +import io +import logging +import re +import shutil +from typing import Any, Callable, Dict, List, NamedTuple, Optional, Sequence, Tuple + +import torch + +_LOGGER = logging.getLogger(__name__) + +# Short element-type spellings accepted in a ``signature``, mapped to the torch +# dtype name. Any torch dtype name is also accepted directly, so this only has +# to cover the abbreviations; both cuTile and torch spell the canonical names +# identically ("float32", "bfloat16", ...), which is what lets a single name +# serve for the dtype lookup and the ``cuda.tile`` attribute lookup. +_DTYPE_ALIASES = { + "fp16": "float16", + "bf16": "bfloat16", + "fp32": "float32", + "fp64": "float64", + "i1": "bool", + "i8": "int8", + "i16": "int16", + "i32": "int32", + "i64": "int64", + "u8": "uint8", +} + + +def _dtype_name(dtype: torch.dtype) -> str: + """``torch.bfloat16`` -> ``"bfloat16"`` — also the ``cuda.tile`` attribute.""" + return str(dtype).rsplit(".", 1)[-1] + + +class ArrayParam(NamedTuple): + """One array parameter of a cuTile kernel, decoded from ``signature``.""" + + name: str + dtype: torch.dtype + ndim: int + + @property + def num_slots(self) -> int: + """cuTile array ABI: ``ptr`` + one extent and one stride per dimension. + + The single statement of that shape. :func:`cutile_param_order` places + the slots and :func:`build_extra_args` fills the non-pointer ones, so + the two must agree on this count. + """ + return 1 + 2 * self.ndim + + +class SignatureLayout(NamedTuple): + """A ``signature`` split into the op's input and output arrays.""" + + inputs: List[ArrayParam] + outputs: List[ArrayParam] + + @property + def arrays(self) -> List[ArrayParam]: + """All arrays in kernel-declaration order: inputs, then outputs.""" + return self.inputs + self.outputs + + @property + def num_slots(self) -> int: + return sum(p.num_slots for p in self.arrays) + + +_NDIM_SUFFIX_RE = re.compile(r"^(?P[^\[\]]+)\[(?P\d+)\]$") + + +def _parse_array_type(name: str, spelling: Any, default_ndim: int) -> ArrayParam: + """Decode one ``signature`` entry into a dtype and a rank. + + Accepts a :class:`torch.dtype`, a dtype name (``"float32"``) or one of the + :data:`_DTYPE_ALIASES` abbreviations (``"fp32"``), each optionally carrying + an explicit rank as ``"fp32[2]"`` for kernels whose arrays differ in rank. + """ + ndim = default_ndim + if isinstance(spelling, torch.dtype): + dtype: Optional[torch.dtype] = spelling + else: + text = str(spelling).strip() + suffix = _NDIM_SUFFIX_RE.match(text) + if suffix is not None: + text, ndim = suffix.group("element").strip(), int(suffix.group("ndim")) + canonical = _DTYPE_ALIASES.get(text.lower(), text.lower()) + candidate = getattr(torch, canonical, None) + dtype = candidate if isinstance(candidate, torch.dtype) else None + + if dtype is None: + raise ValueError( + f"cutile_op signature entry '{name}' has unknown element type " + f"{spelling!r}. cuTile must be told the exact dtype to compile for; " + f"pass a torch.dtype, a dtype name such as 'float32', or one of: " + f"{', '.join(sorted(_DTYPE_ALIASES))}." + ) + if ndim < 1: + raise ValueError( + f"cutile_op signature entry '{name}' declares rank {ndim}; " + "cuTile arrays have rank >= 1." + ) + return ArrayParam(name, dtype, ndim) + + +def validate_cutile_config( + op_name: str, + signature: Dict[str, Any], + constants: Dict[str, Any], + arity: Optional[Tuple[int, int]], + default_ndim: int = 1, + derived_launch: bool = True, + has_grid: bool = True, +) -> SignatureLayout: + """Check a ``cutile_op`` registration and return its signature layout. + + ``signature`` lists the kernel's array parameters in declaration order, + inputs first then outputs; ``ct.Constant`` parameters are not part of it. + ``arity`` is the ``(tensor inputs, outputs)`` of the op being registered and + is what decides where the inputs end. ``derived_launch`` is False when the + caller supplied its own ``aot_fn``, which replaces ``grid``. + + Everything here is answerable before anything is compiled, and every rule, + left unchecked, produces wrong numbers rather than an error. + """ + if default_ndim < 1: + raise ValueError( + f"cutile_op '{op_name}' was given ndim={default_ndim}; " + "cuTile arrays have rank >= 1." + ) + + if derived_launch and not has_grid: + raise ValueError( + f"cutile_op '{op_name}' needs a grid= to build the launch from, or " + "an aot_fn= to replace it." + ) + if not derived_launch and has_grid: + raise ValueError( + f"cutile_op '{op_name}' was given both grid= and aot_fn=; an aot_fn " + "builds the whole launch, so grid= would be ignored. Drop one." + ) + + overlap = sorted(set(signature) & set(constants)) + if overlap: + raise ValueError( + f"cutile_op '{op_name}' declares {overlap} in both signature and " + "constants. Array parameters belong in signature; ct.Constant " + "parameters belong in constants." + ) + + for name, value in constants.items(): + if isinstance(value, bool) or not isinstance(value, int): + raise ValueError( + f"cutile_op '{op_name}' constant '{name}' is {value!r}; cuTile " + "bakes ct.Constant parameters into the compiled symbol as " + "integers, so only int values are supported." + ) + + params = [ + _parse_array_type(name, spelling, default_ndim) + for name, spelling in signature.items() + ] + if not params: + raise ValueError( + f"cutile_op '{op_name}' signature is empty; it must declare the " + "kernel's array parameters (inputs then outputs) in declaration " + "order. Scalars belong in constants=, not signature=." + ) + + if arity is None: + # Guessing the split would feed cutile_param_order a wrong permutation + # and silently misbind the launch, so ask instead. ``schema`` exists + # precisely to state this when meta_fn's annotations can't be read. + raise ValueError( + f"cutile_op '{op_name}' could not determine how many tensors the op " + "takes and returns from meta_fn's type hints, so the signature " + "cannot be split into inputs and outputs. Pass schema= (e.g. " + '"(Tensor x) -> Tensor").' + ) + + num_inputs, num_outputs = arity + if num_inputs + num_outputs != len(params): + raise ValueError( + f"cutile_op '{op_name}' signature declares {len(params)} array " + f"parameter(s) ({', '.join(p.name for p in params)}) but the op " + f"takes {num_inputs} tensor input(s) and returns {num_outputs} " + "output(s). The signature must list every input array followed " + "by every output array." + ) + + return SignatureLayout(inputs=params[:num_inputs], outputs=params[num_inputs:]) + + +def make_dtype_capability_validator( + op_name: str, + layout: SignatureLayout, + user_validator: Optional[Callable[..., bool]] = None, +) -> Callable[..., bool]: + """Build a converter capability validator enforcing the compiled dtypes. + + The kernel is compiled once for the dtypes named in ``signature``. Feeding + the op tensors of any other dtype reinterprets their bytes and silently + returns wrong numbers, so decline the conversion instead: TensorRT then + leaves the op to PyTorch rather than embedding a kernel that cannot read + its inputs. + """ + expected_inputs = [p.dtype for p in layout.inputs] + expected_outputs = [p.dtype for p in layout.outputs] + + def _tensor_meta(value: Any) -> Optional[torch.Tensor]: + meta = getattr(value, "meta", None) + if not isinstance(meta, dict): + return None + val = meta.get("val") + return val if isinstance(val, torch.Tensor) else None + + def _mismatch(kind: str, index: int, got: torch.Tensor, want: torch.dtype) -> bool: + # Warn, not debug: the op silently leaves the engine and runs in + # PyTorch, and if no eager_fn was registered the eventual failure is an + # opaque "not implemented for the CUDA backend" from the dispatcher. + _LOGGER.warning( + "Not lowering '%s' to its cuTile plugin: %s %d is %s but the kernel " + "was compiled for %s. Re-register with a matching signature to run " + "it inside TensorRT; it will fall back to PyTorch for now.", + op_name, + kind, + index, + got.dtype, + want, + ) + return False + + def _validator(node: Any, settings: Any = None) -> bool: + if user_validator is not None and not user_validator(node, settings): + return False + + actual_inputs = [t for t in map(_tensor_meta, node.args) if t is not None] + for index, (got, want) in enumerate(zip(actual_inputs, expected_inputs)): + if got.dtype != want: + return _mismatch("input", index, got, want) + + produced = node.meta.get("val") if isinstance(node.meta, dict) else None + actual_outputs = ( + list(produced) if isinstance(produced, (tuple, list)) else [produced] + ) + for index, (got, want) in enumerate(zip(actual_outputs, expected_outputs)): + if isinstance(got, torch.Tensor) and got.dtype != want: + return _mismatch("output", index, got, want) + + return True + + return _validator + + +def cutile_param_order(layout: SignatureLayout) -> Tuple[int, ...]: + """The permutation mapping TensorRT's slot order onto cuTile's. + + cuTile declares, for each array in kernel order, ``(ptr, extents..., + strides...)``. TensorRT's AOT launcher fills the parameter slots with + ``(input_ptrs..., extra_args..., output_ptrs...)``. ``permutation[i]`` is + the cuTile parameter index that must be moved into physical slot ``i``, so + the extras land exactly where :func:`build_extra_args` puts them: every + input's extents and strides, then every output's. + """ + offsets: List[int] = [] + total = 0 + for param in layout.arrays: + offsets.append(total) + total += param.num_slots + + def pointer(index: int) -> int: + return offsets[index] + + def extents_strides(index: int) -> range: + start = offsets[index] + 1 + return range(start, start + 2 * layout.arrays[index].ndim) + + inputs = range(len(layout.inputs)) + outputs = range(len(layout.inputs), len(layout.arrays)) + return tuple( + [pointer(i) for i in inputs] + + [slot for i in inputs for slot in extents_strides(i)] + + [slot for i in outputs for slot in extents_strides(i)] + + [pointer(i) for i in outputs] + ) + + +# --------------------------------------------------------------------------- +# PTX post-processing +# --------------------------------------------------------------------------- + +_ELF_MAGIC = b"\x7fELF" + +_ENTRY_RE = re.compile( + r"(\.(?:visible|weak)\s+\.entry\s+([\w$]+)\s*\()([^)]*)(\))", re.DOTALL +) +_REQNTID_RE = re.compile(r"\.reqntid\s+(\d+)") +_PTX_VERSION_RE = re.compile(r"\.version\s+(\d+)\.(\d+)") +_BRACE_RE = re.compile(rb"[{}]") + + +def extract_ptx_from_cubin(cubin: bytes) -> Optional[str]: + """Recover the PTX cuTile embeds in a CUBIN's debug section. + + TODO(upstream-cuda-tile): scraping a debug section is a workaround. + ``export_kernel`` currently offers only ``output_format="cubin"`` / + ``"tileir_bytecode"``; ask for a ``"ptx"`` format so both this and + :func:`reorder_entry_params` can work on supported output instead of on + bytes whose layout is an implementation detail. + Tracking issue: . + + The compiler stores it as null-separated strings, so the text is located by + its ``.version`` header and terminated at the brace matching the entry + body's opening one. Brace *matching* rather than a plain search for ``}`` + matters: PTX vector-register syntax (``mov.b64 {%r1, %r2}, %rd0``) puts + braces inside the body. Returns ``None`` if the section isn't there. + """ + if len(cubin) < 64 or cubin[:4] != _ELF_MAGIC: + return None + start = cubin.find(b".version") + if start < 0: + return None + open_brace = cubin.find(b"{", start) + if open_brace < 0: + return None + + # Walk only the braces, in C, rather than every byte of a multi-hundred-KB + # CUBIN from Python: the latter costs an interpreter iteration per byte. + depth = 0 + end = -1 + for brace in _BRACE_RE.finditer(cubin, open_brace): + depth += 1 if brace.group() == b"{" else -1 + if depth == 0: + end = brace.start() + break + if end < 0: + return None + + text = cubin[start : end + 1].replace(b"\x00", b"\n").decode("utf-8", "replace") + return "\n".join(line for line in text.splitlines() if line.strip()) + "\n" + + +ParsedEntry = Tuple[Optional["re.Match[str]"], str, List[str]] + + +def parse_entry(ptx: str) -> ParsedEntry: + """``(match, kernel name, params)`` for the PTX ``.entry`` declaration.""" + match = _ENTRY_RE.search(ptx) + if match is None: + return None, "", [] + params = [p.strip() for p in match.group(3).split(",") if p.strip()] + return match, match.group(2), params + + +def reorder_entry_params( + ptx: str, order: Sequence[int], parsed: Optional[ParsedEntry] = None +) -> str: + """Permute the ``.entry`` parameter declarations so slot ``i`` holds ``order[i]``. + + Only the declaration list is rewritten; the body keeps referring to each + parameter by its own name, so moving the declarations is what changes which + incoming argument each name binds to. ``parsed`` reuses an earlier + :func:`parse_entry` result rather than re-scanning the whole module. + """ + if parsed is None: + parsed = parse_entry(ptx) + match, _name, params = parsed + if match is None: + raise RuntimeError( + "cuTile PTX has no '.entry' declaration to reorder; the compiled " + "kernel cannot be wired to TensorRT's AOT launch." + ) + if len(params) != len(order): + raise RuntimeError( + f"cuTile PTX entry declares {len(params)} parameter(s) but the " + f"reorder expects {len(order)}." + ) + reordered = ",\n\t".join(params[i] for i in order) + return ( + ptx[: match.start()] + + match.group(1) + + "\n\t" + + reordered + + "\n" + + match.group(4) + + ptx[match.end() :] + ) + + +def parse_reqntid(ptx: str) -> Optional[int]: + """The ``.reqntid`` (required threads per CTA) a cuTile kernel declares. + + cuTile vectorizes (e.g. ``f32x2``), so the thread count is often smaller + than the tile size; the kernel must be launched with exactly this many + threads or it traps. + """ + match = _REQNTID_RE.search(ptx) + return int(match.group(1)) if match is not None else None + + +def parse_ptx_version(ptx: str) -> Optional[int]: + """``.version 9.3`` -> ``93``, the encoding used for ISA comparisons.""" + match = _PTX_VERSION_RE.search(ptx) + if match is None: + return None + return int(match.group(1)) * 10 + int(match.group(2)) + + +def cap_ptx_version(ptx: str, max_version: int) -> str: + """Lower the ``.version`` header to ``max_version`` if it exceeds it. + + Only reached when a caller passes ``max_ptx_version=``; nothing lowers a + header on its own. See :func:`verify_driver_accepts_ptx` for why. + """ + emitted = parse_ptx_version(ptx) + if emitted is None or emitted <= max_version: + return ptx + return set_ptx_version(ptx, max_version) + + +def set_ptx_version(ptx: str, version: int) -> str: + """Rewrite the ``.version`` header, e.g. ``90`` -> ``.version 9.0``.""" + match = _PTX_VERSION_RE.search(ptx) + if match is None: + return ptx + replacement = f".version {version // 10}.{version % 10}" + return ptx[: match.start()] + replacement + ptx[match.end() :] + + +# A minimal well-formed module used only to ask the driver whether it accepts a +# given ISA. It declares no parameters and does nothing, so a rejection can only +# come from the ``.version`` header. +# A minimal valid module, used only to ask which ISA the driver would accept +# when reporting a mismatch. +_PROBE_PTX = ( + "//\n.version {major}.{minor}\n.target sm_50\n.address_size 64\n" + ".visible .entry _ttk_ptx_probe()\n{{\n\tret;\n}}\n" +) + + +def _load_ptx(ptx: str) -> Any: + """Ask the driver to JIT the module; returns the ``CUresult``.""" + from cuda.bindings import driver as cuda + + # cuModuleLoadData needs a current context; touching the device makes + # PyTorch create the primary one for us. + torch.cuda.init() + torch.zeros(1, device="cuda") + + err, module = cuda.cuModuleLoadData(ptx.encode("utf-8")) + if err == cuda.CUresult.CUDA_SUCCESS: + cuda.cuModuleUnload(module) + return err + + +# Bounds for the ISA search below. The ceiling only has to stay ahead of what +# any toolchain emits; reaching the floor means the driver could not be asked +# rather than that every ISA is too new. +_PTX_VERSION_CEILING = 129 # 12.9 +_PTX_VERSION_FLOOR = 70 # 7.0 + + +def driver_loads_ptx(ptx: str) -> bool: + """True if the running driver JITs this module.""" + from cuda.bindings import driver as cuda + + return bool(_load_ptx(ptx) == cuda.CUresult.CUDA_SUCCESS) + + +def driver_max_ptx_version(below: int = _PTX_VERSION_CEILING) -> Optional[int]: + """The newest ISA at or below ``below`` this driver loads. + + Walks down one step at a time rather than bisecting: the numbering has gaps + (7.8 then 8.0, 8.8 then 9.0), so a rejected probe can mean "no such version" + rather than "too new". Nothing on the success path calls this -- it exists + to name a concrete remedy when a kernel's ISA is refused. + """ + for version in range(below, _PTX_VERSION_FLOOR - 1, -1): + if driver_loads_ptx(_PROBE_PTX.format(major=version // 10, minor=version % 10)): + return version + return None + + +def fit_ptx_to_driver(op_name: str, kernel_name: str, ptx: str) -> str: + """Return PTX the running driver will load, or raise explaining why it cannot. + + ``tileiras`` emits the ISA of the toolkit it was built against, which can be + newer than the installed driver loads. Nothing catches that on its own: + TensorRT builds the engine happily, and at inference the plugin fails with + ``onShapeChange status -1`` on stderr while ``enqueue`` still returns -- so + the model silently produces garbage. Verifying here, against the same + ``cuModuleLoadData`` the driver will use later, is what makes the mismatch + visible at registration. + + cuTile offers no ISA knob (``export_kernel`` takes an architecture, not a + PTX version, and ``CompilerOptions`` has none), so the only lever is the + ``.version`` header. Lowering it relabels a body the compiler emitted for a + newer ISA; that assembles whenever the body uses no newer instruction, which + is the common case because the bump usually reflects the toolkit default + rather than anything the kernel uses. It is applied only after the driver + has refused the PTX as emitted, it is logged, and the result is re-checked -- + so a relabel that does not assemble raises here rather than reaching an + engine. + """ + from cuda.bindings import driver as cuda + + try: + err = _load_ptx(ptx) + except Exception as exc: # pragma: no cover - environment dependent + _LOGGER.warning( + "Could not verify that the driver accepts the PTX for '%s' (%s); " + "embedding it unchecked. If inference later logs 'onShapeChange " + "status -1', this is why.", + op_name, + exc, + ) + return ptx + + if err == cuda.CUresult.CUDA_SUCCESS: + return ptx + + emitted = parse_ptx_version(ptx) + if err == cuda.CUresult.CUDA_ERROR_UNSUPPORTED_PTX_VERSION and emitted is not None: + accepted = driver_max_ptx_version(emitted - 1) + if accepted is not None: + candidate = set_ptx_version(ptx, accepted) + if driver_loads_ptx(candidate): + _LOGGER.warning( + "cuda-tile compiled '%s' to PTX ISA %d.%d, which this CUDA " + "driver cannot load; relabelled it as %d.%d, which the " + "driver accepts. The driver is older than the cuda-tile " + "toolchain -- aligning them removes this step.", + kernel_name, + emitted // 10, + emitted % 10, + accepted // 10, + accepted % 10, + ) + return candidate + + raise RuntimeError( + f"cutile_op '{op_name}': the CUDA driver refuses the PTX compiled for " + f"kernel '{kernel_name}' ({str(err).split('.')[-1].split(':')[0]}), and " + "lowering the .version header did not make it loadable. Embedding it " + "would build an engine that fails at inference with 'onShapeChange " + "status -1' while still returning output, so the model would silently " + "produce wrong results. Align the CUDA driver with the cuda-tile " + "toolchain, or pass max_ptx_version= to pin the header yourself." + ) + + +# --------------------------------------------------------------------------- +# Compilation +# --------------------------------------------------------------------------- + + +def _cutile_import() -> Any: + """Import ``cuda.tile``, raising an actionable error if it is unavailable.""" + try: + import cuda.tile as ct + + return ct + except ImportError as exc: + raise ImportError( + "cuda-tile is required for cutile_op plugins. " + "Install it with: pip install cuda-tile" + ) from exc + + +def _cutile_dtype(ct: Any, dtype: torch.dtype) -> Any: + """The ``cuda.tile`` dtype object matching a torch dtype. + + Both spell the canonical names identically, so the torch name is also the + attribute name; looking it up with ``getattr`` means a dtype this cuda-tile + build does not expose raises here rather than deep inside the compiler. + """ + value = getattr(ct, _dtype_name(dtype), None) + if value is None: + raise ValueError(f"cuTile has no dtype corresponding to {dtype}.") + return value + + +def _default_arch() -> str: + major, minor = torch.cuda.get_device_capability() + return f"sm_{major}{minor}" + + +def compile_cutile_to_ptx( + op_name: str, + kernel: Any, + layout: SignatureLayout, + constants: Dict[str, Any], + arch_override: Optional[str] = None, + max_ptx_version: Optional[int] = None, +) -> Tuple[bytes, str, Optional[int]]: + """Compile a cuTile kernel to TRT-ready PTX. + + Args: + op_name: the op being registered, used only in error messages. + kernel: the ``@ct.kernel`` program object. + layout: the validated signature split into input and output arrays. + constants: ``ct.Constant`` parameter values, in declaration order, + baked into the compiled symbol. + arch_override: target architecture (e.g. ``"sm_90"``). Defaults to the + current device's compute capability. + max_ptx_version: ISA ceiling as a ``93``-style int, pinning the + ``.version`` header. Rarely needed -- by default the emitted PTX is + used as-is unless the driver refuses it. + + Returns: + ``(ptx_bytes, kernel_name, reqntid)`` — the reordered PTX to embed in + the engine, the entry symbol inside it, and the thread count the kernel + requires (``None`` if it declares none). + """ + ct = _cutile_import() + try: + from cuda.tile.compilation import ( + ArrayConstraint, + CallingConvention, + ConstantConstraint, + KernelSignature, + export_kernel, + ) + except ImportError as exc: + raise ImportError( + f"cutile_op '{op_name}' needs the cuda.tile.compilation API to " + "compile ahead of time; this cuda-tile build does not expose it." + ) from exc + + # tileiras ships with cuda-tile but lives in the package's bin directory, + # which is not on PATH by default. Say so here rather than let export_kernel + # fail with a bare FileNotFoundError. + if shutil.which("tileiras") is None: + raise RuntimeError( + f"cutile_op '{op_name}': the 'tileiras' compiler was not found on " + "PATH. It ships with cuda-tile; add the package's bin directory " + "(e.g. /nvidia/cu13/bin) to PATH before registering " + "cuTile kernels." + ) + + parameters: List[Any] = [ + ArrayConstraint( + dtype=_cutile_dtype(ct, param.dtype), + ndim=param.ndim, + index_dtype=ct.int32, + # cuTile rejects negative strides by default; TRT only ever hands + # the plugin non-negative ones. + stride_lower_bound_incl=0, + alias_groups=(), + may_alias_internally=False, + ) + for param in layout.arrays + ] + parameters.extend(ConstantConstraint(int(v)) for v in constants.values()) + + signature = KernelSignature( + parameters=tuple(parameters), + calling_convention=CallingConvention.cutile_python_v1(), + symbol=None, + ) + + buffer = io.BytesIO() + export_kernel( + kernel, + [signature], + buffer, + gpu_code=arch_override or _default_arch(), + output_format="cubin", + ) + cubin = buffer.getvalue() + + ptx = extract_ptx_from_cubin(cubin) + if ptx is None: + raise RuntimeError( + f"cutile_op '{op_name}': could not recover PTX from the compiled " + "CUBIN. The AOT plugin path needs PTX text to reorder the kernel's " + "parameters into TensorRT's launch order." + ) + + parsed = parse_entry(ptx) + _, kernel_name, params = parsed + if not kernel_name: + raise RuntimeError( + f"cutile_op '{op_name}': the compiled PTX has no '.entry' " + "declaration, so its parameters cannot be matched to TensorRT's " + "launch order." + ) + + order = cutile_param_order(layout) + if len(params) != len(order): + described = ", ".join(f"{p.name} (rank {p.ndim})" for p in layout.arrays) + raise RuntimeError( + f"cutile_op '{op_name}': kernel '{kernel_name}' compiled to " + f"{len(params)} PTX parameter(s) but the signature describes " + f"{len(order)} — {described}, each contributing one pointer plus one " + f"extent and one stride per dimension. " + f"{_diagnose_param_count(layout, len(params))}" + ) + + ptx = reorder_entry_params(ptx, order, parsed) + if max_ptx_version is not None: + ptx = cap_ptx_version(ptx, max_ptx_version) + if arch_override is None: + # Only meaningful when the PTX targets the device we can load it on; + # a deliberate cross-compile is the caller's to verify. + ptx = fit_ptx_to_driver(op_name, kernel_name, ptx) + + reqntid = parse_reqntid(ptx) + _LOGGER.debug( + "Compiled cuTile kernel '%s' -> PTX (%d bytes, reqntid=%s)", + kernel_name, + len(ptx), + reqntid, + ) + return ptx.encode("utf-8"), kernel_name, reqntid + + +def _diagnose_param_count(layout: SignatureLayout, actual: int) -> str: + """Suggest what a mismatched PTX parameter count most likely means.""" + num_arrays = len(layout.arrays) + if num_arrays and actual % num_arrays == 0: + per_array = actual // num_arrays + if per_array >= 3 and per_array % 2 == 1: + return ( + f"The kernel looks like it was compiled for rank " + f"{(per_array - 1) // 2} arrays; pass ndim=" + f"{(per_array - 1) // 2} (or a '[rank]' signature entry)." + ) + if actual > layout.num_slots: + return ( + "The extra parameters are most likely runtime scalars, which the " + "AOT QDP launch path cannot supply. Annotate them as " + "ct.Constant[int] and pass their values in constants=." + ) + return "Check the kernel's array parameters against the signature." + + +# --------------------------------------------------------------------------- +# AOT launch +# --------------------------------------------------------------------------- + + +def _trtp() -> Any: + """The ``tensorrt.plugin`` module, resolved lazily. + + Indirected through a function rather than imported at module scope so the + PTX and signature helpers above stay importable without a QDP-capable + TensorRT, and so tests can substitute a stub for the symbolic-expression + types, which only work inside a live plugin's expression builder. + """ + import tensorrt.plugin as trtp + + return trtp + + +def _as_symint32(value: Any) -> Any: + trtp = _trtp() + if isinstance(value, trtp.SymInt32): + return value + return trtp.SymInt32(value) + + +def _extents_and_strides(desc: Any, param: ArrayParam) -> List[Any]: + """The ``(extents..., strides...)`` a cuTile array parameter expects. + + Rank 1 is the flattened view a 1-D cuTile kernel is written against, so its + single extent is the tensor's element count regardless of how many + dimensions the tensor has. Higher ranks map dimension for dimension onto + the tensor's own shape, with row-major strides. + """ + trtp = _trtp() + + shape = desc.shape_expr + if param.ndim == 1: + return [_as_symint32(shape.numel()), trtp.SymInt32(1)] + + dims = list(shape) + if len(dims) != param.ndim: + raise ValueError( + f"cuTile array '{param.name}' is compiled for rank {param.ndim} but " + f"received a rank-{len(dims)} tensor. Register with " + f"ndim={len(dims)}, or reshape the tensor before the op." + ) + + # Row-major strides are the suffix products of the shape, so accumulate + # once from the right instead of rebuilding each product from scratch. + strides = [trtp.SymInt32(1)] + for dim in reversed(dims[1:]): + strides.append(_as_symint32(strides[-1] * _as_symint32(dim))) + return [_as_symint32(d) for d in dims] + strides[::-1] + + +def build_extra_args( + inputs: Sequence[Any], outputs: Sequence[Any], layout: SignatureLayout +) -> Any: + """Build the ``SymIntExprs`` TensorRT passes between the in and out pointers. + + The order is every input array's extents and strides, then every output + array's — exactly the slots :func:`cutile_param_order` routes them into. + + Raises: + RuntimeError: if TensorRT hands over a different number of tensors than + the signature describes. Registration validates the two agree, so + this is a last line of defense — but it is the one place where a + disagreement would go undetected: zipping the shorter of the two + would quietly emit too few extra arguments and leave the kernel + reading whatever occupied the unfilled parameter slots. + """ + trtp = _trtp() + + for kind, descs, params in ( + ("input", inputs, layout.inputs), + ("output", outputs, layout.outputs), + ): + if len(descs) != len(params): + raise RuntimeError( + f"cuTile launch received {len(descs)} {kind} tensor(s) but the " + f"signature describes {len(params)} " + f"({', '.join(p.name for p in params)}). The extra arguments " + "would not line up with the kernel's parameters." + ) + + values: List[Any] = [] + for desc, param in zip(inputs, layout.inputs): + values.extend(_extents_and_strides(desc, param)) + for desc, param in zip(outputs, layout.outputs): + values.extend(_extents_and_strides(desc, param)) + + extra_args = trtp.SymIntExprs(len(values)) + for index, value in enumerate(values): + extra_args[index] = value + return extra_args + + +def resolve_block_threads( + op_name: str, + kernel_name: str, + reqntid: Optional[int], + block_size: Optional[int], +) -> int: + """The threads-per-block the compiled kernel must be launched with. + + ``.reqntid`` is a requirement, not a hint: cuTile vectorizes, so the thread + count is usually below the tile size, and any other count traps. + """ + if reqntid is None: + if block_size is None: + raise ValueError( + f"cutile_op '{op_name}': kernel '{kernel_name}' declares no " + ".reqntid, so the threads-per-block cannot be derived. Pass " + "block_size= explicitly." + ) + return block_size + if block_size is not None and block_size != reqntid: + raise ValueError( + f"cutile_op '{op_name}' was given block_size={block_size} but kernel " + f"'{kernel_name}' declares .reqntid {reqntid}, which must be the " + "launch's threads-per-block. Drop block_size." + ) + return reqntid + + +def make_aot_fn( + op_name: str, + layout: SignatureLayout, + grid: Callable[..., Any], + block_threads: int, +) -> Callable[..., Any]: + """Derive the AOT launch function from the user's ``grid`` and the layout.""" + + def _aot_fn(inputs: Any, outputs: Any, tactic: int) -> Any: + trtp = _trtp() + + dims = grid(inputs, outputs) + if not isinstance(dims, (tuple, list)): + dims = (dims,) + if not 1 <= len(dims) <= 3: + raise ValueError( + f"cutile_op '{op_name}' grid returned {len(dims)} dimension(s); " + "TensorRT launches accept 1 to 3 (grid_x, grid_y, grid_z)." + ) + + launch_params = trtp.KernelLaunchParams() + launch_params.grid_x = dims[0] + if len(dims) > 1: + launch_params.grid_y = dims[1] + if len(dims) > 2: + launch_params.grid_z = dims[2] + launch_params.block_x = block_threads + launch_params.shared_mem = 0 + + return launch_params, build_extra_args(inputs, outputs, layout) + + return _aot_fn diff --git a/py/torch_tensorrt/kernels/_ops.py b/py/torch_tensorrt/kernels/_ops.py index e81642e0d2..f413dc904e 100644 --- a/py/torch_tensorrt/kernels/_ops.py +++ b/py/torch_tensorrt/kernels/_ops.py @@ -1,18 +1,21 @@ """Public entry points for ``torch_tensorrt.kernels``. -Two functions, two paths into the same registration funnel: +Three functions, three paths into the same registration funnel: * :func:`cuda_kernel_op` — declarative entry for CUDA C++ source. Reads a :class:`KernelSpec` and derives meta / eager / aot / schema, with override keyword arguments for cases outside the DSL. * :func:`ptx_op` — escape hatch for pre-compiled PTX bytes (Triton output, cached NVRTC artifact). User supplies meta / eager / aot directly. +* :func:`cutile_op` — declarative entry for a ``@ct.kernel`` cuTile program. + Compiles the kernel to PTX for you, reorders its parameters into TensorRT's + launch order, and derives the AOT launch. """ from __future__ import annotations import logging -from typing import Any, Callable, Optional +from typing import Any, Callable, Dict, Optional from torch_tensorrt._features import ENABLED_FEATURES from torch_tensorrt.dynamo.conversion._ConverterRegistry import ConverterPriority @@ -26,6 +29,15 @@ _LOGGER = logging.getLogger(__name__) +def _require_qdp_plugin() -> None: + """Raise unless the installed TensorRT exposes Quick Deployable Plugins.""" + if not ENABLED_FEATURES.qdp_plugin: + raise RuntimeError( + "TensorRT QDP plugins are not available. " + "Requires TensorRT >= 10.7.0 (and not 10.14.x)." + ) + + def cuda_kernel_op( op_name: str, spec: KernelSpec, @@ -61,14 +73,10 @@ def cuda_kernel_op( The kernel must follow the calling convention ``(input_ptrs..., scalar_inputs..., extras..., output_ptrs...)``. """ - if not ENABLED_FEATURES.qdp_plugin: - raise RuntimeError( - "TensorRT QDP plugins are not available. " - "Requires TensorRT >= 10.7.0 (and not 10.14.x)." - ) + _require_qdp_plugin() # Late import to avoid circular imports and keep the decorator cheap. - from torch_tensorrt.kernels._register import register_cuda_python_plugin + from torch_tensorrt.kernels._register import register_qdp_plugin _validation._validate_spec( spec, @@ -93,7 +101,7 @@ def cuda_kernel_op( elif spec.inputs and spec.outputs: final_schema = _derive._build_schema(spec) else: - # Let register_cuda_python_plugin fall back to _infer_schema(meta_fn). + # Let register_qdp_plugin fall back to _infer_schema(meta_fn). final_schema = None cuda_spec = CudaPythonSpec( @@ -120,7 +128,7 @@ def cuda_kernel_op( isinstance(input_spec, ScalarInput) for input_spec in (spec.inputs or []) ) - register_cuda_python_plugin( + register_qdp_plugin( op_name=op_name, spec=cuda_spec, meta_fn=final_meta, @@ -141,7 +149,7 @@ def ptx_op( ptx: bytes, kernel_name: str, meta_fn: Callable[..., Any], - eager_fn: Callable[..., Any], + eager_fn: Optional[Callable[..., Any]], aot_fn: Callable[..., Any], *, supports_dynamic_shapes: bool = False, @@ -155,13 +163,9 @@ def ptx_op( Use this when the PTX comes from an external compiler (Triton, a cached NVRTC output, etc.) and NVRTC compilation should be skipped. """ - if not ENABLED_FEATURES.qdp_plugin: - raise RuntimeError( - "TensorRT QDP plugins are not available. " - "Requires TensorRT >= 10.7.0 (and not 10.14.x)." - ) + _require_qdp_plugin() - from torch_tensorrt.kernels._register import register_cuda_python_plugin + from torch_tensorrt.kernels._register import register_qdp_plugin spec = CudaPythonSpec( kernel_source="", @@ -169,7 +173,7 @@ def ptx_op( aot_fn=aot_fn, eager_fn=eager_fn, ) - register_cuda_python_plugin( + register_qdp_plugin( op_name=op_name, spec=spec, meta_fn=meta_fn, @@ -181,3 +185,149 @@ def ptx_op( schema=schema, precompiled_ptx=ptx, ) + + +def cutile_op( + op_name: str, + kernel: Any, + signature: Dict[str, Any], + meta_fn: Callable[..., Any], + *, + grid: Optional[Callable[..., Any]] = None, + constants: Optional[Dict[str, int]] = None, + ndim: int = 1, + block_size: Optional[int] = None, + aot_fn: Optional[Callable[..., Any]] = None, + eager_fn: Optional[Callable[..., Any]] = None, + arch_override: Optional[str] = None, + max_ptx_version: Optional[int] = None, + supports_dynamic_shapes: bool = True, + requires_output_allocator: bool = False, + priority: ConverterPriority = ConverterPriority.STANDARD, + capability_validator: Optional[Callable[..., Any]] = None, + schema: Optional[str] = None, +) -> None: + """Register a ``@ct.kernel`` cuTile program as a TensorRT AOT QDP plugin. + + The cuTile analogue of :func:`cuda_kernel_op`: compiles the kernel once with + ``cuda.tile.compilation.export_kernel``, permutes the compiled PTX into + TensorRT's launch order, and hands the result to :func:`ptx_op`. + + A cuTile kernel declares its *array* parameters first, inputs then outputs, + followed by its ``ct.Constant`` parameters:: + + @ct.kernel + def relu(x, out, tile_size: ct.Constant[int]): ... + + ``signature`` names the arrays in that order; ``constants`` supplies the + ``ct.Constant`` values. See :mod:`torch_tensorrt.kernels._cutile` for why the + PTX has to be permuted. + + Args: + op_name: qualified op name ``"ns::name"``. After registration + ``torch.ops.ns.name`` exists and is lowered to the QDP plugin + during ``torch_tensorrt.compile``. + kernel: the ``@ct.kernel`` program object. + signature: the kernel's array parameters in declaration order, inputs + then outputs, mapped to their element type — e.g. + ``{"x": "fp32", "out": "fp32"}``. Values may be a + :class:`torch.dtype` or its name (``"float32"``, ``"fp32"``). + meta_fn: the fake / meta kernel used for shape+dtype inference. The + PyTorch schema is inferred from its type hints unless ``schema`` is + passed. + grid: ``callable(inputs, outputs) -> int | tuple`` returning the launch + grid in tiles, where ``inputs`` / ``outputs`` are ``trtp.TensorDesc`` + objects (use ``.shape_expr`` for symbolic dims). Up to three dims + become ``grid_x`` / ``grid_y`` / ``grid_z``. Required unless + ``aot_fn`` is given, which replaces it. + constants: ``ct.Constant`` parameter values, in declaration order, + baked into the compiled symbol — e.g. ``{"tile_size": 256}``. The + AOT launch path cannot supply runtime scalars, so every non-array + parameter must be a constant. + ndim: the rank each array is compiled for. Defaults to 1, matching + kernels written against a flattened view; a rank-1 array's extent is + the tensor's element count, so such an op accepts any input shape. + block_size: threads per block. Defaults to the ``.reqntid`` the compiled + kernel declares, which is authoritative — pass this only for kernels + that declare none. + aot_fn: optional replacement for the derived AOT launch + (``callable(inputs, outputs, tactic) -> (KernelLaunchParams, + extra_args)``), used instead of ``grid``. The PTX is still permuted, + so the override must emit extra arguments in the order + :func:`~torch_tensorrt.kernels._cutile.cutile_param_order` expects: + every input array's extents and strides, then every output's. + eager_fn: optional CUDA eager implementation registered on the torch + op. Omit if the op is only used through ``torch_tensorrt.compile``. + arch_override: target architecture such as ``"sm_100"``. Defaults to the + current device's compute capability. + max_ptx_version: ISA ceiling for the embedded PTX, as a ``90``-style + int (``.version 9.0``). Defaults to what the running driver + accepts; the header is capped only when the emitted ISA is newer. + capability_validator: optional extra predicate gating conversion. It is + combined with the dtype check derived from ``signature`` — both must + pass for the op to be lowered to the plugin. + + Raises: + ValueError: if ``signature`` disagrees with ``meta_fn``'s arity, names a + dtype cuTile cannot be compiled for, overlaps ``constants``, or if + neither / both of ``grid`` and ``aot_fn`` are given. + RuntimeError: if the compiled kernel's PTX parameter list does not match + the signature — most often a rank mismatch or a runtime scalar the + AOT launch path cannot supply. + + .. note:: + Compiles a single PTX for the dtypes in ``signature`` and the values in + ``constants``. Inputs of other dtypes are declined at conversion time and + left to PyTorch. Multi-config autotuning is follow-up work. + """ + _require_qdp_plugin() + + from torch_tensorrt.kernels import _cutile + from torch_tensorrt.kernels._register import tensor_arity + + constants = dict(constants or {}) + + # Validate before compiling: nothing here needs the kernel built, and every + # rule it enforces would otherwise surface as wrong numbers, not an error. + layout = _cutile.validate_cutile_config( + op_name, + signature, + constants, + tensor_arity(meta_fn, schema), + default_ndim=ndim, + derived_launch=aot_fn is None, + has_grid=grid is not None, + ) + + ptx, kernel_name, reqntid = _cutile.compile_cutile_to_ptx( + op_name, kernel, layout, constants, arch_override, max_ptx_version + ) + + if aot_fn is None: + assert grid is not None # validate_cutile_config enforces exactly one + aot_fn = _cutile.make_aot_fn( + op_name, + layout, + grid, + _cutile.resolve_block_threads(op_name, kernel_name, reqntid, block_size), + ) + + # Everything past this point is "register pre-compiled PTX", which is + # exactly what ptx_op is; the only cuTile-specific addition is the dtype + # gate derived from the signature. + ptx_op( + op_name, + ptx, + kernel_name, + meta_fn=meta_fn, + eager_fn=eager_fn, + aot_fn=aot_fn, + supports_dynamic_shapes=supports_dynamic_shapes, + requires_output_allocator=requires_output_allocator, + priority=priority, + capability_validator=_cutile.make_dtype_capability_validator( + op_name, layout, capability_validator + ), + schema=schema, + ) + _LOGGER.info("cutile_op '%s' registered (kernel: %s)", op_name, kernel_name) diff --git a/py/torch_tensorrt/kernels/_register.py b/py/torch_tensorrt/kernels/_register.py index bcec7fcdd0..c0f22ac306 100644 --- a/py/torch_tensorrt/kernels/_register.py +++ b/py/torch_tensorrt/kernels/_register.py @@ -2,7 +2,7 @@ import inspect import logging -from typing import Any, Callable, Dict, List, Optional, get_type_hints +from typing import Any, Callable, Dict, List, Optional, Tuple, get_type_hints import torch @@ -55,33 +55,65 @@ def _patch_trt_shape_expr_reflected_ops() -> None: } -def _infer_schema(fn: Callable[..., Any]) -> str: - """Derive a TorchScript schema like '(Tensor x, int n) -> Tensor' from type hints.""" +def _schema_parts(fn: Callable[..., Any]) -> Tuple[List[Tuple[str, str]], List[str]]: + """``([(schema type, arg name)], [return schema types])`` from a fn's hints. + + The single reader of those hints, so a schema built from them and an arity + counted from them cannot drift apart. + """ try: hints = get_type_hints(fn) except Exception: hints = {} - params = list(inspect.signature(fn).parameters.keys()) - args_str = ", ".join( - "{} {}".format( - _TORCH_TYPE_TO_SCHEMA.get(hints.get(p, torch.Tensor), "Tensor"), p - ) - for p in params - ) + args = [ + (_TORCH_TYPE_TO_SCHEMA.get(hints.get(name, torch.Tensor), "Tensor"), name) + for name in inspect.signature(fn).parameters + ] ret = hints.get("return", torch.Tensor) - origin = getattr(ret, "__origin__", None) - if origin is tuple: - ret_str = "({})".format( - ", ".join(_TORCH_TYPE_TO_SCHEMA.get(t, "Tensor") for t in ret.__args__) - ) + if getattr(ret, "__origin__", None) is tuple: + returns = [_TORCH_TYPE_TO_SCHEMA.get(t, "Tensor") for t in ret.__args__] else: - ret_str = _TORCH_TYPE_TO_SCHEMA.get(ret, "Tensor") + returns = [_TORCH_TYPE_TO_SCHEMA.get(ret, "Tensor")] + return args, returns + +def _infer_schema(fn: Callable[..., Any]) -> str: + """Derive a TorchScript schema like '(Tensor x, int n) -> Tensor' from type hints.""" + args, returns = _schema_parts(fn) + args_str = ", ".join(f"{schema_type} {name}" for schema_type, name in args) + ret_str = returns[0] if len(returns) == 1 else "({})".format(", ".join(returns)) return f"({args_str}) -> {ret_str}" +def tensor_arity( + meta_fn: Callable[..., Any], schema: Optional[str] = None +) -> Optional[Tuple[int, int]]: + """``(tensor inputs, outputs)`` of the op that will actually be registered. + + Reads an explicit ``schema`` when one is given, since that is what reaches + the dispatcher, and otherwise the same hints :func:`_infer_schema` reads. + Returns ``None`` when neither can be interpreted, so callers skip an arity + cross-check rather than reject an op over an unreadable annotation. + """ + try: + if schema is not None: + parsed = torch._C.parse_schema(f"_ttk::_probe{schema}") + tensor_type = torch._C.TensorType.get() + num_inputs = sum( + 1 for arg in parsed.arguments if arg.type.isSubtypeOf(tensor_type) + ) + return num_inputs, len(parsed.returns) + args, returns = _schema_parts(meta_fn) + return sum(1 for schema_type, _ in args if schema_type == "Tensor"), len( + returns + ) + except Exception as exc: + _LOGGER.debug("Could not determine tensor arity: %s", exc) + return None + + def _torch_op_already_registered(op_name: str) -> bool: """Return True if ``op_name`` is already known to the torch dispatcher.""" ns, name = op_name.split("::", 1) @@ -218,7 +250,7 @@ def _aot_impl({sig}): _LOGGER.debug("Registered AOT impl for %s", op_name) -def register_cuda_python_plugin( +def register_qdp_plugin( op_name: str, spec: CudaPythonSpec, meta_fn: Optional[Callable[..., Any]], @@ -231,23 +263,24 @@ def register_cuda_python_plugin( precompiled_ptx: Optional[bytes] = None, use_aot_if_available: bool = True, ) -> None: - """Register a NVRTC-compiled CUDA kernel as a TensorRT QDP plugin end-to-end. + """Register a kernel compiled to PTX as a TensorRT QDP plugin end-to-end. Steps performed: - 1. Compile kernel source to PTX via NVRTC (skipped if ``precompiled_ptx`` is passed). + 1. Compile kernel source to PTX via NVRTC (skipped if ``precompiled_ptx`` is + passed — always the case for kernels compiled by another toolchain). 2. Optionally register the PyTorch custom op (define + fake impl). 3. Register the TRT plugin descriptor + JIT impl via generate_plugin(). 4. Register the AOT impl with the compiled PTX. 5. Register the Torch-TensorRT converter via generate_plugin_converter(). - ``precompiled_ptx`` lets higher-level entry points (e.g. ``cuda_kernel_op``) - avoid a redundant second NVRTC pass when they already compiled the source - to build an eager kernel handle. + ``precompiled_ptx`` lets higher-level entry points (e.g. ``cuda_kernel_op``, + ``ptx_op``, ``cutile_op``) supply already-compiled PTX and skip the NVRTC + pass; the source-compilation fields are read only when it is ``None``. """ if spec.aot_fn is None: raise ValueError( - f"CudaPythonSpec.aot_fn must be set before registering plugin '{op_name}'. " - "Pass aot_fn= to cuda_python() or assign spec.aot_fn directly." + f"spec.aot_fn must be set before registering plugin '{op_name}'. " + "Pass aot_fn= to the entry point or assign spec.aot_fn directly." ) if precompiled_ptx is not None: @@ -285,4 +318,4 @@ def register_cuda_python_plugin( _aot_register=lambda: _register_aot_impl(op_name, ptx, spec), ) - _LOGGER.info("cuda-python QDP plugin '%s' registered successfully", op_name) + _LOGGER.info("QDP plugin '%s' registered successfully", op_name) diff --git a/pyproject.toml b/pyproject.toml index cfad95c905..b01ad61f46 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -122,7 +122,8 @@ quantization = [ # which compiles user-supplied CUDA C++ kernels via NVRTC. The high-level # launch/compile API (``cuda.core``) lives in cuda-core; cuda-python's # bindings are still pulled in for the lower-level driver/runtime shims. -kernels = ["cuda-python", "cuda-core"] +# cuda-tile backs ``cutile_op`` and ships the ``tileiras`` compiler it needs. +kernels = ["cuda-python", "cuda-core", "cuda-tile"] [project.urls] Homepage = "https://pytorch.org/tensorrt" diff --git a/tests/py/kernels/conftest.py b/tests/py/kernels/conftest.py index 52be6fd4d7..88572e1f87 100644 --- a/tests/py/kernels/conftest.py +++ b/tests/py/kernels/conftest.py @@ -4,6 +4,7 @@ import pytest import torch + import torch_tensorrt skip_no_cuda = pytest.mark.skipif( @@ -15,20 +16,25 @@ ) -def _has_cuda_core() -> bool: - """True if the cuda-core ``cuda.core`` API (NVRTC/QDP backend) is importable.""" +def _has_module(*names: str) -> bool: + """True if any of ``names`` is importable.""" import importlib.util - for mod in ("cuda.core", "cuda.core.experimental"): + for name in names: try: - if importlib.util.find_spec(mod) is not None: + if importlib.util.find_spec(name) is not None: return True except (ImportError, ModuleNotFoundError, ValueError): continue return False -_HAS_CUDA_CORE = _has_cuda_core() +skip_no_cutile = pytest.mark.skipif( + not _has_module("cuda.tile"), reason="cuda-tile not installed" +) + +# The cuda-core ``cuda.core`` API is the NVRTC/QDP backend. +_HAS_CUDA_CORE = _has_module("cuda.core", "cuda.core.experimental") skip_no_cuda_core = pytest.mark.skipif( not _HAS_CUDA_CORE, @@ -133,6 +139,47 @@ def register_once(register_fn): pass +def assert_ran_in_engine(trt_module, op_name: str) -> None: + """Fail unless the op was actually lowered into a TensorRT engine. + + Numeric agreement alone proves nothing: an op TensorRT declines falls back + to its eager impl and produces exactly the same answer. What distinguishes + the two is whether a call to the op survives in the top-level graph + alongside the ``_run_on_acc_*`` engine submodules. + """ + _, name = op_name.split("::") + leftover = [ + node + for node in trt_module.graph.nodes + if node.op == "call_function" and name in str(node.target) + ] + assert not leftover, ( + f"'{op_name}' was not converted — it is still called in the top-level " + f"graph ({leftover}), so the result came from PyTorch, not the plugin." + ) + assert any( + node.op == "call_module" and "_run_on_acc" in str(node.target) + for node in trt_module.graph.nodes + ), "no TensorRT engine submodule in the compiled graph" + + +def compile_op(op_name: str, inputs, **compile_kwargs): + """Compile a module that does nothing but call ``op_name`` on ``inputs``.""" + ns, name = op_name.split("::") + target = getattr(getattr(torch.ops, ns), name) + + # torch.export matches dynamic_shapes against the forward signature, so it + # needs real positional parameters — with ``*args`` it reads the module as + # taking one tuple and rejects the inputs. + args = ", ".join(f"x{index}" for index in range(len(inputs))) + scope: dict = {"_target": target} + exec(f"def forward(self, {args}):\n return _target({args})\n", scope) + model = type("_OpModule", (torch.nn.Module,), {"forward": scope["forward"]})() + + compile_kwargs.setdefault("min_block_size", 1) + return torch_tensorrt.compile(model.cuda().eval(), inputs=inputs, **compile_kwargs) + + def make_sigmoid_aot(block_size: int = 256): """Build a minimal trtp aot_fn for 1-D pointwise kernels.""" import tensorrt.plugin as trtp diff --git a/tests/py/kernels/test_cuda_kernel_op_overrides.py b/tests/py/kernels/test_cuda_kernel_op_overrides.py index 8e0a800a69..df08446c89 100644 --- a/tests/py/kernels/test_cuda_kernel_op_overrides.py +++ b/tests/py/kernels/test_cuda_kernel_op_overrides.py @@ -53,13 +53,13 @@ def meta(x: torch.Tensor, scale: float) -> torch.Tensor: def test_overrides_forward_to_registrar(monkeypatch): - """Override kwargs land on register_cuda_python_plugin with the right values.""" + """Override kwargs land on register_qdp_plugin with the right values.""" from torch_tensorrt.kernels import _derive, _register captured = {} monkeypatch.setattr( _register, - "register_cuda_python_plugin", + "register_qdp_plugin", lambda *a, **k: captured.update(k), ) # Skip the real NVRTC compile — we're testing wiring, not codegen. @@ -115,7 +115,7 @@ def test_override_missing_required_dsl_field(kwargs, match): def test_precompiled_ptx_skips_nvrtc(monkeypatch): - """register_cuda_python_plugin(precompiled_ptx=...) must not call compile_to_ptx.""" + """register_qdp_plugin(precompiled_ptx=...) must not call compile_to_ptx.""" from torch_tensorrt.kernels import _nvrtc, _register from torch_tensorrt.kernels._cuda_python_spec import CudaPythonSpec @@ -142,7 +142,7 @@ def _fail(*a, **k): def _meta(x: torch.Tensor) -> torch.Tensor: return torch.empty_like(x) - _register.register_cuda_python_plugin( + _register.register_qdp_plugin( op_name="ttk_test::ptx_reused", spec=spec, meta_fn=_meta, diff --git a/tests/py/kernels/test_cutile_op.py b/tests/py/kernels/test_cutile_op.py new file mode 100644 index 0000000000..456e492a49 --- /dev/null +++ b/tests/py/kernels/test_cutile_op.py @@ -0,0 +1,319 @@ +"""Tests for cutile_op (cuTile kernel -> AOT QDP plugin path). + +Scoped to what actually protects this feature. Its failure mode is silent: the +cuTile kernel ABI groups each array as ``(ptr, extents..., strides...)`` while +TensorRT launches with ``(input_ptrs..., extra_args..., output_ptrs...)``, and a +mismatch does not raise -- the kernel reads whatever landed in each slot and +returns plausible-looking numbers. So the permutation, the extra arguments that +must line up with it, and the dtype gate get direct tests; the rest is covered +end to end. +""" + +import pytest +import torch + +import torch_tensorrt +import torch_tensorrt.kernels as ttk + +from .conftest import ( + assert_ran_in_engine, + compile_op, + register_once, + skip_no_cuda, + skip_no_cutile, + skip_no_qdp, +) + +SIG_1IN_1OUT = {"x": "fp32", "out": "fp32"} +SIG_2IN_1OUT = {"a": "fp32", "b": "fp32", "out": "fp32"} + + +def _validate(signature, arity=(1, 1), constants=None, ndim=1, **kwargs): + from torch_tensorrt.kernels._cutile import validate_cutile_config + + return validate_cutile_config( + "ns::op", signature, constants or {}, arity, ndim, **kwargs + ) + + +# ---- The ABI permutation, and the extra arguments that must match it ---- + + +@pytest.mark.parametrize( + "signature, arity, ndim, expected", + [ + # (ptr, extent, stride) per array; the output pointer moves to last. + (SIG_1IN_1OUT, (1, 1), 1, (0, 1, 2, 4, 5, 3)), + # Both input pointers first, then all extents/strides, then the output. + (SIG_2IN_1OUT, (2, 1), 1, (0, 3, 1, 2, 4, 5, 7, 8, 6)), + # Rank 2 gives each array two extents and two strides. + (SIG_1IN_1OUT, (1, 1), 2, (0, 1, 2, 3, 4, 6, 7, 8, 9, 5)), + ], +) +def test_param_order(signature, arity, ndim, expected): + from torch_tensorrt.kernels._cutile import cutile_param_order + + assert cutile_param_order(_validate(signature, arity, ndim=ndim)) == expected + + +class _FakeShapeExpr(list): + def numel(self): + total = 1 + for dim in self: + total *= dim + return total + + +class _FakeDesc: + def __init__(self, *shape): + self.shape_expr = _FakeShapeExpr(shape) + + +class _FakeSymInt32(int): + def __mul__(self, other): + return _FakeSymInt32(int(self) * int(other)) + + +class _FakeTrtp: + """Stand-in for tensorrt.plugin: the real SymInt32 only does arithmetic + inside a live plugin's expression builder.""" + + SymInt32 = _FakeSymInt32 + + @staticmethod + def SymIntExprs(count): + return [None] * count + + +@pytest.fixture +def stub_trtp(monkeypatch): + from torch_tensorrt.kernels import _cutile + + monkeypatch.setattr(_cutile, "_trtp", lambda: _FakeTrtp) + return _cutile + + +def test_extra_args_match_the_permutation(stub_trtp): + """Extents and strides must fill the slots the permutation routes them to. + + Rank 1 is a flattened view, so its extent is the element count whatever the + tensor's shape; rank 2 maps dimension for dimension with row-major strides. + Inputs come before outputs, matching cutile_param_order. + """ + layout = _validate(SIG_2IN_1OUT, arity=(2, 1)) + extra = stub_trtp.build_extra_args( + [_FakeDesc(2, 4), _FakeDesc(8)], [_FakeDesc(8)], layout + ) + assert [int(v) for v in extra] == [8, 1, 8, 1, 8, 1] + + rank2 = _validate(SIG_1IN_1OUT, ndim=2) + values = stub_trtp._extents_and_strides(_FakeDesc(4, 256), rank2.inputs[0]) + assert [int(v) for v in values] == [4, 256, 256, 1] + + +def test_extra_args_reject_a_tensor_count_mismatch(stub_trtp): + """Zipping would truncate and quietly emit too few extras.""" + layout = _validate(SIG_2IN_1OUT, arity=(2, 1)) + with pytest.raises(RuntimeError, match="1 input tensor.*describes 2"): + stub_trtp.build_extra_args([_FakeDesc(8)], [_FakeDesc(8)], layout) + + +# ---- Registration-time validation ---- + + +@pytest.mark.parametrize( + "kwargs, message", + [ + # The signature must describe exactly the op's tensors... + (dict(signature=SIG_1IN_1OUT, arity=(2, 1)), "2 tensor input"), + # ...name a dtype cuTile can be compiled for... + (dict(signature={"x": "weird", "out": "fp32"}), "unknown element type"), + # ...and keep arrays and ct.Constant parameters apart. + ( + dict(signature=SIG_1IN_1OUT, constants={"tile_size": 1.5}), + "only int values", + ), + # Exactly one of grid= / aot_fn= builds the launch. + (dict(signature=SIG_1IN_1OUT, has_grid=False), "needs a grid="), + # An unreadable arity must ask for schema=, not guess the split. + (dict(signature=SIG_1IN_1OUT, arity=None), "Pass schema="), + ], +) +def test_invalid_configurations_are_rejected(kwargs, message): + with pytest.raises(ValueError, match=message): + _validate(**kwargs) + + +# ---- The dtype gate ---- + + +class _FakeNode: + """Minimal stand-in for the torch.fx.Node a capability validator receives.""" + + def __init__(self, arg_dtypes, out_dtype): + self.args = [ + type("_Arg", (), {"meta": {"val": torch.empty(2, dtype=d)}})() + for d in arg_dtypes + ] + self.meta = {"val": torch.empty(2, dtype=out_dtype)} + + +def test_dtype_gate_declines_mismatched_inputs(): + """fp16 into an fp32-compiled kernel would otherwise reinterpret its bytes.""" + from torch_tensorrt.kernels._cutile import make_dtype_capability_validator + + validate = make_dtype_capability_validator("ns::op", _validate(SIG_1IN_1OUT)) + assert validate(_FakeNode([torch.float32], torch.float32), None) is True + assert validate(_FakeNode([torch.float16], torch.float16), None) is False + + +# ---- GPU integration: real cuTile kernels through cutile_op ---- + +TILE = 128 + +try: + import cuda.tile as ct + + @ct.kernel + def _ttk_add_one_kernel(x, out, tile_size: ct.Constant[int]): + pid = ct.bid(0) + tile = ct.load(x, index=(pid,), shape=(tile_size,)) + ct.store(out, index=(pid,), tile=tile + 1.0) + + @ct.kernel + def _ttk_reglu_kernel(gate, up, out, tile_size: ct.Constant[int]): + pid = ct.bid(0) + g = ct.load(gate, index=(pid,), shape=(tile_size,)) + u = ct.load(up, index=(pid,), shape=(tile_size,)) + ct.store(out, index=(pid,), tile=ct.maximum(g, 0.0) * u) + +except ImportError: + ct = None + + +def _register_add_one(op_name: str, with_eager: bool = True) -> None: + import tensorrt.plugin as trtp + + def _meta(x: torch.Tensor) -> torch.Tensor: + return torch.empty_like(x) + + def _eager(x: torch.Tensor) -> torch.Tensor: + out = torch.empty_like(x) + flat_x, flat_out = x.contiguous().reshape(-1), out.reshape(-1) + ct.launch( + torch.cuda.current_stream().cuda_stream, + (ct.cdiv(flat_x.numel(), TILE), 1, 1), + _ttk_add_one_kernel, + (flat_x, flat_out, TILE), + ) + return out + + register_once( + lambda: ttk.cutile_op( + op_name, + kernel=_ttk_add_one_kernel, + signature=SIG_1IN_1OUT, + meta_fn=_meta, + grid=lambda inputs, outputs: ( + trtp.cdiv(inputs[0].shape_expr.numel(), TILE), + ), + constants={"tile_size": TILE}, + eager_fn=_eager if with_eager else None, + supports_dynamic_shapes=True, + ) + ) + + +def _register_reglu(op_name: str) -> None: + import tensorrt.plugin as trtp + + def _meta(gate: torch.Tensor, up: torch.Tensor) -> torch.Tensor: + return torch.empty_like(gate) + + register_once( + lambda: ttk.cutile_op( + op_name, + kernel=_ttk_reglu_kernel, + signature=SIG_2IN_1OUT, + meta_fn=_meta, + grid=lambda inputs, outputs: ( + trtp.cdiv(inputs[0].shape_expr.numel(), TILE), + ), + constants={"tile_size": TILE}, + supports_dynamic_shapes=True, + ) + ) + + +@skip_no_cuda +@skip_no_qdp +@skip_no_cutile +class TestCuTileOpIntegration: + def test_eager(self): + _register_add_one("ttk_test::cutile_add_one_eager") + x = torch.randn(1024, device="cuda") + assert torch.allclose( + torch.ops.ttk_test.cutile_add_one_eager(x), x + 1, atol=1e-4, rtol=1e-4 + ) + + def test_runs_in_engine_without_an_eager_impl(self): + """No eager_fn: falling back to PyTorch could not even run. + + Removing the fallback is what makes a passing result mean the cuTile + kernel executed inside the engine -- with one present, a declined op + returns the same numbers and the assertion proves nothing. + """ + op = "ttk_test::cutile_add_one_trt_only" + _register_add_one(op, with_eager=False) + x = torch.randn(4, 256, device="cuda") + trt = compile_op(op, [x]) + assert_ran_in_engine(trt, op) + with torch.no_grad(): + assert torch.equal(trt(x), x + 1) + + def test_two_inputs_bind_in_the_right_order(self): + """ReGLU: relu(gate) * up is asymmetric, so swapped pointers show up.""" + op = "ttk_test::cutile_reglu" + _register_reglu(op) + gate = torch.randn(4, 256, device="cuda") + up = torch.randn(4, 256, device="cuda") + trt = compile_op(op, [gate, up]) + assert_ran_in_engine(trt, op) + with torch.no_grad(): + assert torch.allclose( + trt(gate, up), torch.relu(gate) * up, atol=1e-5, rtol=1e-5 + ) + + def test_dynamic_shapes(self): + op = "ttk_test::cutile_add_one_dyn" + _register_add_one(op) + trt = compile_op( + op, + [ + torch_tensorrt.Input( + min_shape=(1, 128), + opt_shape=(1, 512), + max_shape=(1, 2048), + dtype=torch.float32, + ) + ], + ) + assert_ran_in_engine(trt, op) + for size in [128, 512, 2048]: + x = torch.randn(1, size, device="cuda") + with torch.no_grad(): + assert torch.allclose(trt(x), x + 1, atol=1e-4, rtol=1e-4) + + def test_dtype_mismatch_falls_back_instead_of_returning_garbage(self): + """fp16 into an fp32-compiled kernel must not silently produce nonsense.""" + op = "ttk_test::cutile_add_one_dtype" + _register_add_one(op) + x = torch.randn(4, 256, device="cuda", dtype=torch.float16) + trt = compile_op(op, [x], enabled_precisions={torch.float16}) + # The mirror of assert_ran_in_engine: declined, so still in the graph. + assert any( + node.op == "call_function" and "cutile_add_one_dtype" in str(node.target) + for node in trt.graph.nodes + ), "the fp16 op was lowered to an fp32-compiled plugin" + with torch.no_grad(): + assert torch.allclose(trt(x), x + 1, atol=1e-2, rtol=1e-2) diff --git a/tests/py/kernels/test_ptx_op.py b/tests/py/kernels/test_ptx_op.py index 60eff25588..e508ee71a7 100644 --- a/tests/py/kernels/test_ptx_op.py +++ b/tests/py/kernels/test_ptx_op.py @@ -23,7 +23,7 @@ def test_ptx_op_forwards_precompiled_ptx(monkeypatch): captured = {} monkeypatch.setattr( _register, - "register_cuda_python_plugin", + "register_qdp_plugin", lambda *a, **k: captured.update(k), ) @@ -53,7 +53,7 @@ def test_ptx_op_kernel_name_lands_on_spec(monkeypatch): captured = {} monkeypatch.setattr( _register, - "register_cuda_python_plugin", + "register_qdp_plugin", lambda *a, **k: captured.update(k), )