[RFC] Support XPU Backend With PTI-sdk in Kineto

Motivation

The goal is to enable Kineto to support XPU pathways without altering its current design and with minimal maintenance costs. This will allow Kineto's profiler, designed and implemented for CUDA devices, to work with PyTorch running on XPU devices without the need for additional extensions.

Background

i. Built-in XPU in PyTorch

We are in the process of upstreaming the necessary code to provide direct support for a built-in XPU backend with various features in PyTorch. Soon, PyTorch will be able to compile and run on XPU devices without requiring additional extensions.

ii. About Extending Kineto

We previously entertained the idea of extending Kineto's main profiler to support any backend (https://github.com/pytorch/kineto/issues/647). However, after a period of research and experimentation, we believe that for supporting the XPU backend alone, this approach is overly comprehensive and not suitable as an immediate first step.

iii. About PTI-sdk Library

We have developed an independent PTI library called PTI-sdk, which we will soon release as a public product. This library can be installed separately on XPU devices and its usage and API semantics are highly similar to some parts of the existing CUPTI.

Design and Implementation

Our design is based on the current design of Kineto and is highly similar to how Kineto integrates with CUPTI. In our design, we use a separate XPUActivityApi to support all interactions between the main Kineto profiler and PTI-sdk, a separate XPUActivityBuffer to obtain raw data collected by PTI-sdk, and define and implement methods for handling various kinds of XPU activities within CuptiActivityProfiler, which ultimately feed into various data structures derived from XPUActivity.

In terms of implementation, XPUActivityApi will be implemented as a separate class, estimated to add about 500 lines of code across two new files; XPUActivityBuffer and XPUActivity along with their derived classes will also be implemented as separate classes, following Kineto's existing interface specifications and semantic assumptions, estimated to add about 400 lines of code across three new files; necessary modifications to the existing implementation of CuptiActivityProfiler will be made to add functions for the XPU pathway, estimated to affect two existing files and add about 100 lines of code.

Further Maintenance

i. Effort Estimation

In the plan, the internal maintenance of XPUActivityApi, XPUActivityBuffer, and XPUActivity along with their derived classes can largely be undertaken by XPU engineers, with Kineto's core development and maintenance remaining unaffected by their status. However, some modifications at the abstraction level may require alignment between both parties. Maintenance of CuptiActivityProfiler may involve multiple parties, but the maintenance of the XPU path can largely reference the CUPTI path, thus reducing maintenance efforts.

ii. Impact of Refactoring

To my knowledge, Kineto is planning to separate CUPTI and Roctracer from the main profiler. Our changes, due to their association with the existing CuptiActivityProfiler, may add some refactoring burden, but this is expected to be minimal. In my vision, when Kineto refactors CuptiActivityProfiler, it could handle the XPU path using the same approach as for the CUDA path.

About Testing

Planning to ensure functional correctness within the XPU pathway by implementing a series of unit tests.

Timeline

With the expected release of PTI-sdk in April, we plan to submit a pull request for code review at the beginning of April and make things ready for merging in the middle of April.

Additional Information

Associated Kineto PRs: (To be pasted here)
Associated PyTorch PRs: (To be pasted here)

Start by reviewing the existing CuptiActivityProfiler and its CUPTI integration, then compare the proposed XPUActivityApi, XPUActivityBuffer, and XPUActivity abstractions with those interfaces. The issue proposes new classes across five files and changes to two existing files, with unit tests covering functional correctness of the XPU pathway as the completion criterion.