Adding a New Operation to nvtripy

When to Use

• Adding a new mathematical, tensor, or neural network operation
• Creating a new frontend function that maps to TensorRT/MLIR ops
• Extending the op registry with unary, binary, or custom operations

Architecture Overview

Operations in nvtripy follow a Frontend → Trace → MLIR pipeline:

1. Trace Op (nvtripy/trace/ops/): Defines the computational graph node — rank inference, dtype inference, and MLIR code generation.
2. Frontend Op (nvtripy/frontend/ops/): The public API function — exports, constraints, docstring, and bridges to the trace op via create_op().
3. Registration: Both __init__.py files must be updated so the op is discoverable.

Procedure

Step 1: Create the Trace Operation