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