How to make Tensor static and only call kernel by pointer?

[peterbjorgensen]

I am trying to turn Python cute dsl kernels into TensorRT QDP Plugins. It restricts the kind of parameters the PTX can have to (input_pointers..., scalar_arguments..., output_pointers...). I have it working for a simple "add_one_kernel" (see the program listing below if interested).

Instead of writing the kernel with pointer arguments, I would like to write it with some StaticTensorPointer class that automatically turns the layout etc. into compile-time variables and only the pointer ends up as a PTX parameter. I guess it's possible with a custom DynamicExpression, but I find the documentation a bit vague on this. Any ideas how to achieve this?

Another question: Is it possible to get the shared memory needed for a compiled kernel? I need it for the kernel launch parameters of the kernel in the plugin.

import tensorrt as trt
import numpy as np
from quack.cute_dsl_utils import torch2cute_dtype_map
import tensorrt.plugin as trtp
import torch
import torch_tensorrt
import cutlass.cute as cute

from quack.cute_dsl_utils import torch2cute_dtype_map

trt_logger = trt.Logger(trt.Logger.VERBOSE)

STATIC_SHAPE = (16, 16)
STATIC_STRIDE = (1, 16)

@cute.kernel
def add_one_kernel(x: cute.Pointer, y: cute.Pointer):
    STATIC_LAYOUT = cute.make_layout(shape=STATIC_SHAPE, stride=STATIC_STRIDE)
    tensor_x = cute.make_tensor(x, STATIC_LAYOUT)
    tensor_y = cute.make_tensor(y, STATIC_LAYOUT)
    tidx, _, _ = cute.arch.thread_idx()
    bidx, _, _ = cute.arch.block_idx()
    idx = bidx * cute.arch.block_dim()[0] + tidx
    if idx < cute.size(tensor_x):
        tensor_y[idx] = tensor_x[idx] + 1

@cute.jit
def apply_add_one_kernel(x: cute.Pointer, y: cute.Pointer):
    BLOCK_SIZE = 256
    num_blocks = (np.prod(STATIC_SHAPE).item() + BLOCK_SIZE - 1) // BLOCK_SIZE

    add_one_kernel(x, y).launch(
        grid = [num_blocks, 1, 1],
        block = [BLOCK_SIZE, 1, 1],
    )


@torch.library.custom_op("my::add_one", mutates_args=())  # type: ignore[misc]
def add_one(x: torch.Tensor) -> torch.Tensor:
    assert x.is_cuda
    y = torch.empty_like(x)
    input_dtype = torch2cute_dtype_map[x.dtype]
    y_ptr = cute.runtime.make_ptr(input_dtype, y.data_ptr(), cute.AddressSpace.gmem, assumed_align=32)
    x_ptr = cute.runtime.make_ptr(input_dtype, x.data_ptr(), cute.AddressSpace.gmem, assumed_align=32)

    compiled = cute.compile(apply_add_one_kernel, x_ptr, y_ptr, options="--keep-ptx")
    compiled(x_ptr, y_ptr)

    return y


@torch.library.register_fake("my::add_one")
def _(X: torch.Tensor) -> torch.Tensor:
    return X

@trtp.register("my::add_one")
def add_plugin_desc(X: trtp.TensorDesc) -> tuple[trtp.TensorDesc]:
    # Output has the same shape and dtype as the input.
    return X.like()


@trtp.aot_impl("my::add_one")
def add_plugin_aot_impl(
    X: trtp.TensorDesc, outputs: tuple[trtp.TensorDesc], tactic: int
) -> tuple[
    str | bytes, str | bytes, trtp.KernelLaunchParams, trtp.SymExprs
]:
    # Choose the pointer type based on the input dtype.
    if X.dtype == trt.float32:
        x_dtype = torch2cute_dtype_map[torch.float32]
    elif X.dtype == trt.float16:
        x_dtype = torch2cute_dtype_map[torch.float16]
    elif X.dtype == trt.bfloat16:
        x_dtype = torch2cute_dtype_map[torch.bfloat16]
    else:
        raise ValueError("Unexpected dtype")
    input_dtype = x_dtype
    y_ptr = cute.runtime.make_ptr(input_dtype, 32, cute.AddressSpace.gmem, assumed_align=32)
    x_ptr = cute.runtime.make_ptr(input_dtype, 64, cute.AddressSpace.gmem, assumed_align=32)

    compiled_kernel = cute.compile(apply_add_one_kernel, x_ptr, y_ptr, options="--keep-ptx")


    # Build symbolic launch parameters.
    BLOCK_SIZE = 256
    num_blocks = (np.prod(STATIC_SHAPE).item() + BLOCK_SIZE - 1) // BLOCK_SIZE
    launch_params = trtp.KernelLaunchParams()
    launch_params.grid_x = num_blocks  # number of thread blocks
    launch_params.block_x = BLOCK_SIZE  # threads per block
    launch_params.shared_mem = 0  # bytes of shared mem (I have no idea how to get this)

    # ``extra_args`` are scalar arguments appended to the kernel's argument list at
    # launch. Here ``n_elements`` is passed as a 32-bit symbolic integer so TRT
    # evaluates it from the actual tensor size at runtime.
    extra_args = trtp.SymIntExprs(0)

    kernel_name_ptx = list(compiled_kernel.kernel_info.keys())[0]
    return (
        kernel_name_ptx,  # kernel function name in PTX
        compiled_kernel.artifacts.PTX,  # PTX source — embedded in TRT engine
        launch_params,
        extra_args,
    )


torch_tensorrt.dynamo.conversion.plugins.generate_plugin_converter(
    "my::add_one",
    supports_dynamic_shapes=False,
    requires_output_allocator=False,
    use_aot_if_available=True,
)


class MyModel(torch.nn.Module):
    def __init__(self):
        super().__init__()

    def forward(self, X: torch.Tensor) -> torch.Tensor:
        res = torch.ops.my.add_one.default(X)
        return res


if __name__ == "__main__":


    my_model = MyModel().to("cuda").eval()
    m = torch.full(STATIC_SHAPE, 2, device="cuda", dtype=torch.float32)

    assert my_model(m)[0][0] == 3.0
    assert my_model(m)[1][1] == 3.0

    with torch_tensorrt.logging.debug():
        trt_inputs = [m]
        model_trt = torch_tensorrt.compile(
            my_model,
            inputs=trt_inputs,
            min_block_size=1,
        )
        print("Model compiled successfully!")
        print("Running inference with compiled model...")
        with torch.no_grad():
            for i in range(10):
                res = model_trt(m)
                assert torch.allclose(res, my_model(m)), "Results do not match!"

        program = torch_tensorrt.dynamo.trace(
            my_model,
            arg_inputs=trt_inputs,
        )
        engine_bytes: bytes = torch_tensorrt.dynamo.convert_exported_program_to_serialized_trt_engine(
            program,
            arg_inputs=trt_inputs,
            require_full_compilation=True,
            immutable_weights=True,
            enabled_precisions={torch.float16, torch.bfloat16, torch.float32},
            use_explicit_typing=False,
            optimization_level=5,
            version_compatible=True,
        )
    with open("cutlass_example_plugin.engine", "wb") as f:
        f.write(engine_bytes)

    print("Inference successful!")

[sylvesterkaczmarek]

For the pointer-only ABI, I think your current pattern is already the simplest one: keep cute.Pointer as the kernel argument and reconstruct the Tensor using a static layout inside the kernel.

Since the shape/stride are compile-time constants, they do not need to become PTX parameters. I wouldn't introduce a custom DynamicExpression unless you specifically need a reusable abstraction across many layouts.

For the shared-memory question, current CUTLASS has:

cutlass.utils.get_kernel_smem_size(kernel)

It returns the total static shared-memory allocation for the kernel. The important detail is that it must be called inside a @cute.jit body after the kernel's .launch() has been traced/registered.

SmemAllocator also calculates its total usage automatically. For your exact add_one_kernel, shared_mem = 0 is correct because that kernel does not allocate any shared memory, but get_kernel_smem_size() is the API I'd use once the real plugin kernel does.

[Hosseinghorbani0]

The key distinction seems to be between making the tensor layout static at the CuTe DSL level and keeping the underlying data pointer dynamic at the kernel ABI level.

If the goal is to generate a TensorRT QDP Plugin with only pointers and scalar arguments in the PTX parameter list, I would expect the tensor's shape/layout metadata to be represented as compile-time expressions while the base address remains a runtime parameter.

For the shared-memory question, I think the important part is whether the compiled kernel's shared-memory requirement is exposed through the generated kernel metadata rather than being derived from the runtime tensor object. If the allocation size depends only on static layout/tile parameters, it should be possible to determine it during compilation.

Would the recommended approach be to define a custom DynamicExpression whose value is the runtime pointer while its associated layout remains a static CuTe expression? Or is there already an existing CuTe DSL pattern for representing a statically-shaped tensor backed by a runtime pointer?

Also, for TensorRT integration, is there a supported way to query the generated kernel's required dynamic shared-memory size after compilation, rather than duplicating the shared-memory calculation in the plugin implementation?