example: using nvrtc kernel for aot plugin

[bowang007]

Description

Please include a summary of the change and which issue is fixed. Please also include relevant motivation and context. List any dependencies that are required for this change.

Fixes # (issue)

Type of change

Please delete options that are not relevant and/or add your own.

• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds functionality)
• Breaking change (fix or feature that would cause existing functionality to not work as expected)
• This change requires a documentation update

Checklist:

• My code follows the style guidelines of this project (You can use the linters)
• I have performed a self-review of my own code
• I have commented my code, particularly in hard-to-understand areas and hacks
• I have made corresponding changes to the documentation
• I have added tests to verify my fix or my feature
• New and existing unit tests pass locally with my changes
• I have added the relevant labels to my PR in so that relevant reviewers are notified

[narendasan]

Can't we use the autogenerated registration for this?

[bowang007]

we could do that if we isolate the registration part generation

[narendasan]

Lets use the autogenerated version,

[narendasan]

@bowang007 any updates here?

[github-actions]

There are some changes that do not conform to Python style guidelines:

--- /home/runner/work/TensorRT/TensorRT/examples/dynamo/nvrtc_aot_plugin.py	2025-12-10 23:20:21.730652+00:00
+++ /home/runner/work/TensorRT/TensorRT/examples/dynamo/nvrtc_aot_plugin.py	2025-12-10 23:20:54.395887+00:00
@@ -78,10 +78,11 @@
# ============================================================================
# 3. Register Custom Op in PyTorch
# ============================================================================
# We register the custom operation with PyTorch so it can be used in models.
# The 'mutates_args=()' argument tells PyTorch this op is functional (doesn't modify inputs in-place).
+

@torch.library.custom_op("pointwise_sigmoid_ops::pointwise_sigmoid", mutates_args=())  # type: ignore[misc]
def pointwise_sigmoid(X: torch.Tensor) -> torch.Tensor:
    """
    Implementation of the custom op for PyTorch eager execution.
@@ -126,10 +127,11 @@
# 4. Register Fake Implementation (Meta Kernel)
# ============================================================================
# The fake implementation is crucial for TorchDynamo. It tells the compiler
# about the output shape and data type without actually running the kernel.
# This is used during the tracing phase.
+

@torch.library.register_fake("pointwise_sigmoid_ops::pointwise_sigmoid")
def _(input: torch.Tensor) -> torch.Tensor:
    """Fake implementation for TorchDynamo tracing of base operation."""
    return torch.empty_like(input)
@@ -175,11 +177,11 @@
) -> Tuple[
    Union[str, bytes], Union[str, bytes], trtp.KernelLaunchParams, trtp.SymExprs
]:
    # Get the PTX code from our pre-compiled module
    compiled_kernel = _module.code.decode("utf-8")
-    
+
    # Calculate grid and block dimensions based on input shape
    N = input.shape_expr.numel()
    launch_params = trtp.KernelLaunchParams()
    block = 256
    launch_params.grid_x = trtp.cdiv(N, block)
@@ -214,10 +216,11 @@


# ============================================================================
# 7. Test the Model
# ============================================================================
+

class PointwiseSigmoidModel_WithTRTWrapper(torch.nn.Module):
    """
    Test model that uses the TRT wrapper with custom_op() registration.
    """
@@ -245,11 +248,11 @@
            inputs=trt_inputs,
            enabled_precisions={torch.float32},
            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(input)
                assert torch.allclose(

[narendasan]

LGTM once you switch to using the generated registration.

[narendasan]

Also lint your code