🚀 The feature, motivation and pitch

Summary

This RFC adds kernel-level profiling support for PrivateUse1 backends through a Kineto plugin (PU1PTI). It delivers a vendor-facing C API (PrivateUse1TracingApi), dynamic plugin loading so vendors can integrate without patching Kineto, an OpenReg reference implementation, and vendor-facing documentation.

Motivation

OpenReg is PyTorch's reference implementation for out-of-tree hardware backends. It demonstrates device model, operator dispatch, memory management, streams, autograd etc. But if a vendor wants kernel-level profiling (individual kernels, runtime/driver events, memory operations, CPU-device correlation), there is no reference demonstrating how an out-of-tree accelerator can integrate with Kineto's profiling infrastructure.

Today, the only references for Kineto profiling integration are CUPTI (CUDA), XPUPTI (XPU), and AIUPTI (AIU). All three are tightly coupled to their respective vendor libraries, compiled into the Kineto source tree via #ifdef, and not designed as integration guides for external vendors.

At present, PrivateUse1 backends lack:

• No IActivityProfiler implementation for PrivateUse1. Activity types (PRIVATEUSE1_RUNTIME, PRIVATEUSE1_DRIVER, CONCURRENT_KERNEL, GPU_MEMCPY, GPU_MEMSET) are defined in Kineto but nothing populates them.

• No vendor-facing API. XPU/AIU use proprietary vendor APIs (XPUPTI, AIUPTI) directly — there is no generic, documented C API for PrivateUse1.

• No dynamic loading mechanism. XPU/AIU use compile-time #ifdef registration in init.cpp; vendors must modify Kineto source to add a new backend.

• No documentation for the IActivityProfiler contract, buffer protocol, or correlation flow from an external vendor's perspective.

• No reference implementation. OpenReg only implements ProfilerStubs for fallback mode (KINETO_PRIVATEUSE1_FALLBACK), providing operator-level timing only.

Impact: vendors must reverse-engineer Kineto internals, leading to high integration cost, fragile implementations, and limited adoption.

This proposal addresses the gap by introducing a vendor-facing C API, a Kineto plugin that implements IActivityProfiler for PrivateUse1, dynamic plugin loading, and OpenReg as the reference implementation. Critically, this enables any out-of-tree accelerator vendor to integrate kernel-level profiling with Kineto without patching Kineto itself — vendors ship their own .so implementing the stable C API, and the plugin loader discovers it at runtime. This makes PrivateUse1 profiling truly seamless and out-of-tree, matching the same zero-patch philosophy that PrivateUse1 provides for device registration, operator dispatch, and memory management.

Goals:

1. PU1PTI plugin: A minimal IActivityProfiler implementation in Kineto that produces CONCURRENT_KERNEL, GPU_MEMCPY, GPU_MEMSET, PRIVATEUSE1_RUNTIME, and PRIVATEUSE1_DRIVER activities. Registered via registerProfilerFactory().

2. Vendor-facing C API (PrivateUse1TracingApi.h): A small, stable, ABI-versioned C API that vendors implement in their .so. Functions: enable/disable, buffer callbacks, correlation push/pop, flush, record iteration, version check.

3. Dynamic plugin loading: Kineto discovers and loads vendor .so at runtime (KINETO_PLUGIN_PATH or standard paths) via dlopen/dlsym. No Kineto source changes per vendor.

4. OpenReg reference implementation: Minimal working implementation of PrivateUse1TracingApi in OpenReg demonstrating how vendors integrate kernel-level profiling.

5. Documentation: Vendor guide covering C API, buffer protocol, correlation, plugin lifecycle, and troubleshooting.

Non-goals: Production-grade plugin features, performance metrics (FLOPs, occupancy), advanced analysis (roofline, SASS), automatic profiler detection API, vendor-specific optimizations.

Proposed Implementation

Detailed design document: openreg_profiling_proposal contains the full implementation specification.

Architecture Overview

torch.profiler.profile(activities=[PrivateUse1])
        │
        ▼
Kineto Core (ActivityProfilerProxy → Controller → CuptiActivityProfiler)
        ├── addChildActivityProfiler()
        ▼
PrivateUse1ActivityProfiler (implements IActivityProfiler)
        │
        ├── PrivateUse1ActivityProfilerSession (start/stop/processTrace)
        ├── PrivateUse1ActivityApi (C++ wrapper, calls vendor C API)
        └── PrivateUse1PluginLoader (discovers & loads vendor .so)
        │
        ▼
Vendor .so (e.g. libopenreg_tracing.so)
        │
        ├── Implements PrivateUse1TracingApi C functions
        ├── Kernel launch tracking
        ├── Runtime call tracking
        └── Memory operation tracking

The implementation has five components:

1. PrivateUse1TracingApi.h (C API header)

Defines the contract between Kineto and vendor implementations. Equivalent to what cupti.h is for CUDA, but open, Kineto-defined, and simplified:

• pu1TracingEnable(activityKind) / pu1TracingDisable(activityKind) — activity control
• pu1TracingSetCallbacks(requestFn, completeFn) — buffer management (producer-consumer pattern; completeFn omits ctx/streamId for simplicity)
• pu1PushExternalCorrelationId(id) / pu1PopExternalCorrelationId() — single correlation stack (unlike CUDA's dual CUSTOM0/CUSTOM1)
• pu1FlushAll() / pu1GetNextRecord(buffer, size, &record) — data retrieval
• pu1GetVersion() — ABI compatibility
• pu1ActivityRecord — one unified struct (kind, start_time, end_time, device_id, stream_id, correlation_id, name, metadata) for all activity types. Intentionally simpler than CUPTI's many per-type structs.

2. PU1PTI Kineto Plugin

Implements IActivityProfiler and IActivityProfilerSession inside Kineto. Vendors do not implement these C++ interfaces — PU1PTI does it once and delegates to the vendor's C API. Components:

• PrivateUse1ActivityProfiler — factory, returns sessions
• PrivateUse1ActivityProfilerSession — start/stop, processTrace, correlation
• PrivateUse1ActivityApi — calls vendor C functions via function pointers from loaded .so
• PrivateUse1ActivityHandlers — converts pu1ActivityRecord → Kineto GenericTraceActivity
• PrivateUse1ActivityBuffer — RAII buffer wrapper

3. Dynamic Plugin Loader

PrivateUse1PluginLoader discovers and loads vendor .so at Kineto initialization:

• Discovery: KINETO_PLUGIN_PATH (colon-separated) → standard paths (/usr/lib/kineto/plugins/, /usr/local/lib/kineto/plugins/, ~/.local/lib/kineto/plugins/)
• Loading: dlopen → dlsym("kineto_register_privateuse1_plugin") → pu1GetVersion() ABI check → call registration
• Error handling: env unset → no session (zero overhead); dlopen/dlsym failure → log warning, return nullptr; ABI mismatch → log error, return nullptr
• Vendor .so loaded once, remains loaded for process lifetime (no hot-reload)
• Registration in init.cpp under #ifdef HAS_PRIVATEUSE1_PROFILER

Why not use REGISTER_PRIVATEUSE1_PROFILER (PR #172154) only? That macro provides process-local registration (consistent with other PrivateUse1 components) but requires the vendor to implement the full IActivityProfiler C++ interface, depend on Kineto's C++ ABI, and be linked into the process. The dynamic loader + C API gives vendors a simpler, ABI-stable surface without Kineto build dependency.

4. OpenReg Reference Implementation

Minimal working implementation of the C API in OpenReg:

• openreg_tracing.h / tracing.cpp — implements all C API functions
• Kernel launch hooks — records kernel name, timestamps, correlation
• Runtime hooks — records runtime API calls
• Memory hooks — records memcpy/memset
• Builds as .so, loadable by the plugin loader

OpenReg is a reference for the contract and integration pattern. It does not need to run on every vendor's hardware. Other PrivateUse1 vendors (e.g. custom NPUs, OEM stacks, in-house accelerators) reimplement the same C API on their own stack.

5. Documentation

Vendor integration guide covering: C API contract, buffer protocol, correlation mechanism, plugin lifecycle, implementation steps, dynamic loading, troubleshooting, and OpenReg walkthrough.

Drawbacks

• Extra indirection layer (PU1PTI + C API) between Kineto and vendor — negligible runtime cost but more code to maintain in Kineto
• C API must be designed correctly upfront; changes require ABI version bump
• Dynamic loading adds complexity (dlopen, ABI versioning, error handling)
• Single-process scope initially; no multi-process plugin coordination

Alternatives

• Vendors implement IActivityProfiler directly (PR #172154 path): Lower Kineto maintenance but each vendor re-implements boilerplate, depends on Kineto C++ ABI, and must be linked into the process.
• Compile-time registration (like XPU/AIU): Simpler loading but requires vendor code in Kineto tree; not suitable for out-of-tree backends.
• PyTorch-level profiling only (ProfilerStubs fallback): Already exists but limited to operator-level timing; no kernel visibility.

Additional context

• Related PR: pytorch/pytorch#172154 — adds REGISTER_PRIVATEUSE1_PROFILER macro for process-local registration of PrivateUse1 profilers with Kineto. This RFC builds on that by adding the C API layer, dynamic loading, and reference implementation.
• Kineto XPU plugin: third_party/kineto/libkineto/src/plugin/xpupti/
• Kineto AIU plugin: third_party/kineto/libkineto/src/plugin/aiupti/
• IActivityProfiler interface: third_party/kineto/libkineto/include/IActivityProfiler.h

cc @robieta @chaekit @guotuofeng @guyang3532 @dzhulgakov @davidberard98 @briancoutinho @sraikund16 @sanrise @mwootton @divyanshk @jiannanWang @scotts @ryanzhang22 @mikaylagawarecki @malfet