tripy-compilation

nvtripy Compilation Pipeline

When to Use

• Compiling functions or modules with tp.compile()
• Defining runtime inputs with InputInfo and DimensionInputInfo
• Working with compiled Executable objects
• Configuring compilation options (optimization level, timing cache)
• Understanding the Trace → MLIR → TensorRT flow
• Using dynamic shapes with min/opt/max bounds
• Debugging compilation failures

Compilation Flow

User Function → Trace (graph) → MLIR (IR) → TensorRT (engine) → Executable

1. Trace: The function is called with tracer tensors to record the computation graph
2. MLIR: The trace graph is lowered to MLIR using the tensorrt dialect
3. TensorRT: MLIR is compiled to a TensorRT engine
4. Executable: The engine is wrapped in a callable Executable object

tp.compile() — The Main Entry Point

compiled_fn = tp.compile(
    func,                    # Function or Module to compile
    optimization_level=3,    # 0-5, higher = better runtime, longer compile
    args=[...],              # Positional arguments
    kwargs={...},            # Keyword arguments
)

Argument Types

The compiled Executable only accepts parameters that were InputInfo/DimensionInputInfo in the original compile() call.

InputInfo — Tensor Runtime Inputs

# Static shape
inp = tp.InputInfo(shape=(2, 4), dtype=tp.float32)
# shape_bounds: min=(2,4), opt=(2,4), max=(2,4)

# Dynamic dimensions (min, opt, max)
inp = tp.InputInfo(shape=((1, 2, 3), 4), dtype=tp.float32)
# First dim: min=1, opt=2, max=3; second dim: fixed at 4
# shape_bounds: min=(1,4), opt=(2,4), max=(3,4)

# Named dimensions (must be equal at runtime)
window = tp.NamedDimension("window", 3, 5, 7)
inp = tp.InputInfo(shape=(1, window, window), dtype=tp.float32)
# Both dims named "window" must have the same value at runtime

DimensionInputInfo — Scalar Dimension Inputs

For functions that take scalar shape values as parameters:

dim_info = tp.DimensionInputInfo(value_bounds=(1, 2, 4))
# min=1, opt=2, max=4

Used when a function parameter controls a reshape or dynamic shape operation.

Executable — Running Compiled Functions

# The executable's signature matches the InputInfo parameters
compiled_fn = tp.compile(add, args=[
    tp.InputInfo((2, 4), dtype=tp.float32),  # "a"
    tp.InputInfo((2, 4), dtype=tp.float32),  # "b"
])

# Call with evaluated tensors
a = tp.ones((2, 4), dtype=tp.float32).eval()
b = tp.ones((2, 4), dtype=tp.float32).eval()
result = compiled_fn(a, b)

Key Executable properties

• input_infos: Dict of parameter name → InputInfo
• stream: The CUDA stream used for execution
• __signature__: Compatible with inspect.signature() for introspection

Important: .eval() for inputs

Runtime inputs to compiled functions should be evaluated tensors (not lazy). Use .eval() to force evaluation before passing to the executable.

Compiling Modules

class MyModel(tp.Module):
    def __init__(self):
        super().__init__()
        self.linear = tp.Linear(3, 4)

    def forward(self, x):
        return self.linear(x)

model = MyModel()
# Load real weights before compiling
model.linear.weight = tp.Tensor(weight_data)
model.linear.bias = tp.Tensor(bias_data)

compiled_model = tp.compile(
    model,
    args=[tp.InputInfo(shape=(2, 3), dtype=tp.float32)],
)

When compiling a Module:

• The module's state_dict() entries are named for readable traces
• Weights become compile-time constants (baked into the engine)
• Only InputInfo arguments become runtime inputs

Dynamic Shapes

Basic dynamic dimensions

compiled_add = tp.compile(
    add,
    args=[
        tp.InputInfo(shape=((1, 2, 3), 2), dtype=tp.float32),
        tp.InputInfo(shape=((1, 2, 3), 2), dtype=tp.float32),
    ],
)

# Works for any first-dim size in [1, 3]:
small = compiled_add(tp.ones((1, 2)).eval(), tp.ones((1, 2)).eval())
big = compiled_add(tp.ones((3, 2)).eval(), tp.ones((3, 2)).eval())

Named dimensions for constraints

window_size = tp.NamedDimension("window_size", 3, 5, 7)
inp = tp.InputInfo((1, window_size, window_size), dtype=tp.float32)
# Both dimensions named "window_size" must be equal at runtime

Scalar dimension inputs

def dynamic_reshape(x, s):
    return tp.reshape(x, (-1, s))

compiled_reshape = tp.compile(
    dynamic_reshape,
    args=[
        tp.InputInfo(shape=(3, (2, 4, 6)), dtype=tp.float32),
        tp.DimensionInputInfo(value_bounds=(1, 2, 4)),
    ],
)

result = compiled_reshape(tp.ones((3, 4)).eval(), tp.DimensionSize(2))
assert result.shape == (6, 2)

Compilation Options

Optimization Level

Timing Cache

tp.config.timing_cache_file_path = "/path/to/cache"

The timing cache stores kernel profiling data across compilations, significantly speeding up repeated compilations with similar operations.

Compiler Internals

The MLIR compiler (nvtripy/backend/mlir/compiler.py) uses these options:

• --tensorrt-timing-cache-path: Path to timing cache
• --tensorrt-builder-opt-level: Optimization level (0-5)
• --force-entrypoints-return-allocs: Memory management
• --mlir-elide-elementsattrs-if-larger: Debug readability
• --tensorrt-layer-info-dir: TensorRT layer debug info

Function Requirements for tp.compile

The function passed to compile() must:

1. Be pure — no side effects (print, assert, file I/O)
2. Return Tensor(s) — only Tensor return types supported
3. No collection inputs — List[Tensor] or Dict[str, Tensor] will be frozen as constants
4. No variadic args — *args and **kwargs are frozen at compile time

Checklist

• InputInfo used for all runtime tensor inputs
• DimensionInputInfo used for scalar shape parameters
• Dynamic dimension bounds specified as (min, opt, max) tuples
• Module weights loaded before calling tp.compile()
• Function is pure (no side effects)
• Runtime inputs .eval()'d before passing to executable
• Timing cache configured for repeated compilations
• Optimization level appropriate for use case