content: """
  Using Custom Kernels with NVRTC in TensorRT AOT Plugins
  =======================================================
  
  This example demonstrates how to use the NVIDIA Runtime Compilation (NVRTC) library
  to compile custom CUDA kernels at runtime and integrate them into a TensorRT
  Ahead-Of-Time (AOT) plugin.
  
  This approach is powerful because it allows you to:
  1. Write raw CUDA C++ code for maximum performance.
  2. Compile it on-the-fly, adapting to the specific GPU architecture.
  3. Wrap it in a TensorRT plugin without writing a separate C++ plugin library.
  4. Integrate it seamlessly into Torch-TensorRT's compilation flow.
  
  The example performs a simple pointwise Sigmoid operation: f(x) = 1 / (1 + exp(-x)).
  """
  
  from typing import List, Tuple, Union
  
  import torch
  
  import torch_tensorrt
  
  # ============================================================================
  # 1. Define the CUDA Kernel Source
  # ============================================================================
  # We define the CUDA kernel source code as a Python string.
  # This code will be compiled by NVRTC.
  # Note that we use extern "C" to avoid name mangling, making it easier to
  # retrieve the kernel function by name later.
  
  cu_code = """
  // Simple pointwise Sigmoid kernel: f(x) = 1 / (1 + exp(-x))
  extern "C" __global__ void pointwise_sigmoid_kernel_nvrtc(const float* __restrict__ input,
                                                            const int size,
                                                            float* __restrict__ output) {
      const int idx = blockIdx.x * blockDim.x + threadIdx.x;
  
      if (idx < size) {
          const float x = input[idx];
          // use fast device intrinsic to avoid headers
          output[idx] = 1.0f / (1.0f + __expf(-x));
      }
  }
  """
  
  # ============================================================================
  # 2. Compile the Kernel using NVRTC (for eager mode)
  # ============================================================================
  # Before defining the Torch custom op, we compile the kernel so we can run it
  # in standard PyTorch (eager mode) for verification and testing.
  # We use the cuda-python library's NVRTC bindings.
  
  from cuda.core.experimental import Device as _CudaDevice
  from cuda.core.experimental import LaunchConfig as _LaunchConfig
  from cuda.core.experimental import Program as _CudaProgram
  from cuda.core.experimental import ProgramOptions as _CudaProgramOptions
  from cuda.core.experimental import launch as _cuda_launch
  
  # Initialize CUDA device and stream
  _cuda_device = _CudaDevice()
  _cuda_device.set_current()
  _cuda_stream = _cuda_device.create_stream()
  
  # Configure compilation options
  _program_options = _CudaProgramOptions(
      std="c++17",
      arch=f"sm_{_cuda_device.arch}",  # Target the current GPU architecture
      include_path=["/usr/local/cuda/include"],
  )
  
  # Create and compile the program
  _program = _CudaProgram(cu_code, code_type="c++", options=_program_options)
  _module = _program.compile("ptx", name_expressions=("pointwise_sigmoid_kernel_nvrtc",))
  _kernel = _module.get_kernel("pointwise_sigmoid_kernel_nvrtc")
  
  
  # ============================================================================
  # 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.
      This function launches the pre-compiled NVRTC kernel.
      """
      assert X.is_cuda, "Tensor must be on CUDA device."
      assert X.dtype == torch.float32, "For this test, expected float32 input"
  
      Y = torch.empty_like(X)
      N = int(X.numel())
  
      block = 256
      grid_x = max(1, (N + block - 1) // block)
      config = _LaunchConfig(grid=(grid_x), block=(block))
  
      # Helper class to wrap PyTorch's stream for cuda-python
      class _PyTorchStreamWrapper:
          def __init__(self, pt_stream):
              self.pt_stream = pt_stream
  
          def __cuda_stream__(self):
              stream_id = self.pt_stream.cuda_stream
              return (0, stream_id)
  
      pt_stream = torch.cuda.current_stream()
      s = _cuda_device.create_stream(_PyTorchStreamWrapper(pt_stream))
  
      # Launch kernel with raw pointers
      _cuda_launch(
          s,
          config,
          _kernel,
          X.data_ptr(),
          N,
          Y.data_ptr(),
      )
  
      return Y
  
  
  # ============================================================================
  # 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)
  
  
  # ============================================================================
  # 5. Define TensorRT AOT Plugin
  # ============================================================================
  # Now we define how this operation should be handled within TensorRT.
  # We use the torch_tensorrt plugin auto generation feature and the AOT implementation using NVRTC.
  
  import tensorrt.plugin as trtp
  from cuda.core.experimental import Device, LaunchConfig, Program, ProgramOptions
  
  torch_tensorrt.dynamo.conversion.plugins.generate_plugin(
      "pointwise_sigmoid_ops::pointwise_sigmoid"
  )
  
  
  # This is where the magic happens. We provide the compiled PTX code and
  # launch parameters to TensorRT. This code runs during engine building.
  @trtp.aot_impl("pointwise_sigmoid_ops::pointwise_sigmoid")
  def sigmoid_aot_nvrtc_impl(
      X: trtp.TensorDesc,
      outputs: Tuple[trtp.TensorDesc],
      tactic: int,
  ) -> 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 = X.shape_expr.numel()
      launch_params = trtp.KernelLaunchParams()
      block = 256
      launch_params.grid_x = trtp.cdiv(N, block)
      launch_params.block_x = block
      launch_params.shared_mem = 0
  
      # Pass the number of elements (N) as an extra argument to the kernel
      extra_args = trtp.SymIntExprs(1)
      extra_args[0] = trtp.SymInt32(N)
  
      # Return: kernel name, PTX code, launch params, kernel arguments
      return (
          "pointwise_sigmoid_kernel_nvrtc",
          compiled_kernel,
          launch_params,
          extra_args,
      )
  
  
  # ============================================================================
  # 6. Generate Plugin Converter
  # ============================================================================
  # This registers the mapping between the PyTorch custom op and the TensorRT plugin.
  # It tells Torch-TensorRT: "When you see 'pointwise_sigmoid_ops::pointwise_sigmoid',
  # replace it with the TensorRT plugin we just defined."
  
  torch_tensorrt.dynamo.conversion.plugins.generate_plugin_converter(
      "pointwise_sigmoid_ops::pointwise_sigmoid",
      supports_dynamic_shapes=True,
      requires_output_allocator=False,
  )
  
  
  # ============================================================================
  # 7. Test the Model
  # ============================================================================
  
  
  class PointwiseSigmoidModel_WithTRTWrapper(torch.nn.Module):
      """
      Test model that uses the TRT wrapper with custom_op() registration.
      """
  
      def forward(self, input: torch.Tensor) -> torch.Tensor:
          z = torch.ops.pointwise_sigmoid_ops.pointwise_sigmoid(input)
          return z
  
  
  if __name__ == "__main__":
      model = PointwiseSigmoidModel_WithTRTWrapper().to("cuda").eval()
      input = torch.randn(1, 1024, device="cuda", dtype=torch.float32)
  
      print("PyTorch baseline result:")
      print(torch.sigmoid(input))
  
      print("Custom Op eager result:")
      print(model(input))
  
      print("\nCompiling with Torch-TensorRT...")
      with torch_tensorrt.logging.debug():
          trt_inputs = [input]
          model_trt = torch_tensorrt.compile(
              model,
              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(
                      res, model(input), rtol=1e-2, atol=1e-2
                  ), "Results do not match!"
  
      print("Inference successful!")
  