RFC: Global TensorRT engine profiling controls for opaque AOTInductor packages

[apbose]

Summary

Add process-global TensorRT engine profiling controls that are reachable through torch.ops.tensorrt. This makes execution profiling usable for TensorRT engines deserialized inside an opaque AOTInductor package, where callers do not have a TRTEngine handle.

This RFC addresses #4679.

Motivation

Today, profiling is enabled through methods on an individual TorchBind TRTEngine object:

engine.enable_profiling()
engine.set_profile_format("perfetto")

That works when Python owns the engine object. It does not work for an AOTInductor package:

Python caller
    |
    v
Opaque AOTI runner
    |
    +-- hidden TensorRT engine A
    +-- hidden TensorRT engine B
    `-- hidden TensorRT engine C

No engine handle is returned to Python.

The packaged artifact is also the deployment form whose performance we want to measure. The proposed API therefore configures profiling at process scope, while preserving the existing per-engine API for directly accessible engines.

Proposed API

torch.ops.tensorrt.set_profile_path_prefix("/tmp/profiles")
torch.ops.tensorrt.set_profile_format("perfetto")  # or "trex"
torch.ops.tensorrt.set_profile_execution(True)

outputs = aoti_runner(*inputs)

torch.ops.tensorrt.set_profile_execution(False)

Matching getters expose the active process-wide configuration:

torch.ops.tensorrt.get_profile_execution()
torch.ops.tensorrt.get_profile_format()
torch.ops.tensorrt.get_profile_path_prefix()

The recommended order is path, format, then enable, so an engine cannot begin profiling with an unintended path or format. The operations remain independent; calling only set_profile_execution(True) uses the defaults (perfetto and the system temporary directory).

Scope

The new API intentionally applies one profiling configuration to all TensorRT engines in the process. It does not replace the existing per-engine methods.

Accessible TRTEngine       -> existing per-engine methods
Opaque AOTI engines        -> new process-global torch operators

Independent configuration of individual hidden AOTI engines is a non-goal for this change. That would require stable serialized engine identifiers, a runtime registry and conflict rules between global and per-engine settings.

Design

Configuration state

runtime.cpp owns one coherent configuration:

struct GlobalProfilingConfig {
  bool enabled;
  std::string profile_format;
  std::string profile_path_prefix;
  uint64_t generation;
};

It also owns:

• global_profiling_mutex, protecting reads and writes of the complete configuration;
• global_profiling_generation, an atomic process-wide configuration version.

Every setter updates one field under the mutex and increments the generation. The generation is a monotonically increasing revision number, not a profiling stage.

set path    -> generation 0 -> 1
set format  -> generation 1 -> 2
enable      -> generation 2 -> 3
disable     -> generation 3 -> 4

Per-engine state

Each TRTEngine stores:

uint64_t applied_global_profiling_generation = 0;

This is the last process-global configuration revision applied by that engine. It is deliberately separate from the shared atomic generation:

Global revision = 3

Engine A applied revision = 3  -> current
Engine B applied revision = 1  -> stale
Engine C applied revision = 0  -> stale

Engines synchronize lazily on their next invocation. This avoids maintaining a registry or eagerly finding every engine hidden inside loaded packages.

End-to-end flow

%%{init: {"flowchart": {"minNodeWidth": 280, "wrappingWidth": 250, "padding": 32, "rankSpacing": 65}}}%%
flowchart TD
    U["<br/>Python calls global profiling setters<br/>path · format · enabled<br/>"] --> G["<br/>Lock global mutex<br/>Update GlobalProfilingConfig<br/>"]
    G --> V["<br/>Increment atomic<br/>global configuration generation<br/>"]
    V --> A["<br/>AOTI runner invokes<br/>a hidden TRTEngine<br/>"]
    A --> L["<br/>Lock this engine's<br/>per-engine mutex<br/>"]
    L --> C{"<br/>Global generation equals<br/>this engine's applied generation?<br/>"}
    C -- Yes --> F["<br/>Fast path<br/>Keep current engine state<br/>"]
    C -- No --> S["<br/>Lock global mutex briefly<br/>Copy coherent config snapshot<br/>"]
    S --> P["<br/>Apply configuration to this engine<br/>Set path and format<br/>Enable or disable profiler<br/>"]
    P --> E["<br/>Record the applied generation<br/>on this engine<br/>"]
    F --> X["<br/>Obtain a stable TensorRT<br/>IExecutionContext<br/>"]
    E --> X
    X --> R["<br/>Execute the TensorRT engine<br/>while its mutex remains held<br/>"]
    R --> T["<br/>Write Perfetto or TREX<br/>profiling traces<br/>"]
    T --> O["<br/>Return outputs<br/>RAII unlocks the engine mutex<br/>"]

Loading

The corresponding execution ordering is:

execute_engine(inputs, engine)
    |
    +-- normalize input devices
    +-- lock engine->mu
    +-- engine->sync_global_profiling_state()
    +-- perform lazy NCCL binding when required
    +-- obtain a stable IExecutionContext pointer
    +-- configure inputs, outputs and workspace
    +-- enqueue TensorRT execution
    +-- collect outputs and emit profiling traces
    `-- return; engine mutex unlocks through RAII

Concurrency model

The two locks and the atomic counter have different responsibilities:

atomic global generation
    -> lock-free fast check on every engine invocation

global profiling mutex
    -> coherent snapshot of enabled/format/path when settings changed

per-engine mutex
    -> protects that engine's profiler and TensorRT execution context

The engine mutex must be acquired before synchronization because enabling profiling attaches a profiler to the execution context, while disabling profiling may invalidate that context. It remains held through execution so another thread cannot use the same context while it is being changed. Different engines have different mutexes and can continue executing concurrently.

Trace output

When profiling is enabled, existing engine profiling machinery emits the applicable traces, including:

• engine execution profile;
• input packing profile;
• output packing profile;
• enqueue profile.

The existing TRTEngineProfiler and Perfetto/TREX implementations are reused. This proposal changes reachability and synchronization, not the profiling data format itself.

Validation

Two levels of coverage are included:

1. A C++ runtime unit test sets path, format and enabled state; checks the coherent snapshot, public getters and generation; and restores the original process-global state.
2. A Python AOTI end-to-end test compiles and packages a TensorRT model, loads the opaque runner, enables profiling without obtaining an engine handle, executes it and verifies that engine, input, output and enqueue trace files are produced. It restores the original global settings in finally.

Follow-up considerations

• A global dump_engine_timing() entry point could expose accumulated per-engine/per-phase counters from opaque packages.
• If users need independent configuration of hidden engines, a separate RFC should define stable engine identifiers and global-versus-local precedence.
• The process-global nature of these operators should be prominent in API documentation; names prefixed with global_ could make the scope even more explicit.