feat: TorchTRT Annotation Layer for Cuda generated kernels

[bowang007]

Description

This PR introduces torch_tensorrt.annotation, an experimental module for registering hand-written CUDA C++ kernels as both PyTorch custom ops (for eager execution) and TensorRT Quick Deployable Plugins with AOT support (for torch_tensorrt.compile).

Usage

  import torch, torch_tensorrt
  import torch_tensorrt.annotation as tta                                                                                                                           
   
  CU = """                                                                                                                                                          
  extern "C" __global__ void my_sigmoid(const float* x, int n, float* y) {
      int i = blockIdx.x * blockDim.x + threadIdx.x;
      if (i < n) y[i] = 1.0f / (1.0f + __expf(-x[i]));                                                                                                              
  }
  """                                                                                                                                                               
                  
  tta.auto_cuda_kernel_plugin(                                                                                                                                      
      "ann_ex::sigmoid",
      tta.KernelSpec(                                                                                                                                               
          kernel_source=CU, kernel_name="my_sigmoid",
          inputs=[tta.InputDecl("x")],                                                                                                                              
          outputs=[tta.OutputDecl("y", shape=tta.SameAs(0))],
          extras=[tta.Numel("x")],                                                                                                                                  
          geometry=tta.Elementwise(block=(256,), layout="flat"),                                                                                                    
      ),
  )   

After this call, torch.ops.ann_ex.sigmoid is available in eager and is embedded as a TensorRT plugin during torch_tensorrt.compile. The meta function, eager
launch, AOT implementation, and PyTorch schema are all derived from the KernelSpec.

API Surface

The module exposes two primary entry points, layered by declarativeness:

auto_cuda_kernel_plugin is the recommended default. The caller supplies a KernelSpec dataclass describing the kernel's inputs, outputs (with a shape relation such
as SameAs or ReduceDims), scalar extras (Numel, DimSize), and launch geometry (Elementwise or Reduction). The framework derives the meta function, eager CUDA
launch, TensorRT AOT implementation, and PyTorch schema. This path covers pointwise kernels (1-D flat or N-D grid launches), reductions (with optional keepdim),
multi-input kernels, and scalar (non-tensor) kernel arguments via ScalarInput.

manual_cuda_kernel_plugin is the lower-level alternative for kernels outside the declarative DSL — shape-changing outputs, multi-output kernels, or non-standard
launch geometries. The caller provides eager_fn and aot_fn directly; the decorator still registers the PyTorch op, TRT plugin, AOT implementation, and converter
in a single call.

A Custom(fn=...) geometry is also available for callers who want the declarative path's schema/meta derivation but need to hand-write the TRT KernelLaunchParams.

Type of change

• New feature (non-breaking change which adds functionality)

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