feat: add cutile_op for AOT QDP plugins from cuda tile kernels

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 ``"<dtype>[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.
+   ``<site-packages>/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
 --------------

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>

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

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:

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

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",
 ]