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

[xunsongh]

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 (#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)

[harryPorter828]

Hi @xunsongh,

At Graphcore we are looking to make changes to Kineto to support an IPU profiling backend.

We want to add another activity - e.g. IPUActivityAPI but we don't want to compile time link this to Kineto/PyTorch. i.e. We would like to register IPUActivityAPI (or perhaps IPUActivityProfiler) at run-time with Kineto.
We had a few questions relating to this RFC:

1. We understand that this RFC is a plan to add an XPUActivityAPI within CuptiActivityProfiler - extending CuptiActivityProfiler to support the XPU with additional XPU specific members and paths. This will be an addition to the existing CUPTI and Roctracer code paths. On its own, this change will add XPU support, but does not provide a generic way to add new profiling backends as you described in [RFC] Make the kineto extendable for other runtime than CUDA #647. Is our understanding correct?
2. What was the main difficulty with [RFC] Make the kineto extendable for other runtime than CUDA #647 and was there any investigation into how this could be broken into smaller changes?
3. You state that "To my knowledge, Kineto is planning to separate CUPTI and Roctracer from the main profiler." Is there any more information about this to give regarding details and timeline? Will this introduce a system to register profiling backends?

Thanks in advance for taking the time to go over these questions and for any answers you can provide.

[xunsongh]

Hi @harryPorter828 ,
Thank you for your attention to this RFC and for taking the time to read and respond. However, after an online discussion in April, our XPU team has decided to follow the advice of the Kineto developers. We will adopt the form of a Kineto plugin, registering the XPU profiler as a relatively independent child profiler to the existing main profiler of Kineto. This approach avoids making changes to the main profiler of Kineto and does not involve directly pushing XPU-related integration code into the Kineto code repository. Instead, this code will be placed in PyTorch. Moreover, this plugin-based implementation is not directly related to what I previously proposed and what you mentioned in #647; they can even be considered two separate topics.

As developers on the XPU side, our primary goal is to enable PyTorch XPU users to utilize a profiler supported by Kineto without changing their usage habits, meaning with minimal changes to the front-end code. To achieve this, we are willing to first adapt to all existing designs and code logic of Kineto, reaching this goal in a way that is acceptable to all parties. In the future, we are open to participating in the structural reconstruction of Kineto, continuing to develop it into a more open, general, and user-friendly library. #647 was merely a conceptual idea towards this goal, and through extensive practice over this period, we have become aware of various issues with the original plan.

As for the current this RFC, it has been rejected by Meta's meeting, so we do not intend to discuss or explain it further. Once we complete and submit our design based on the Kineto plugin, I will close this RFC to avoid any misunderstandings.

Thank you.

[harryPorter828]

Hi @xunsongh,

Thank you for your response.

Would you be able to elaborate on the Kineto plugin solution (which is different from #647) that you have mentioned? It was our understanding that a plugin type architecture does not yet exist in Kineto or as part of PyTorch profiler. Do you have a timescale for submitting the design for this new approach and will it be in the form of an RFC?

Regarding your point for 'continuing to develop it [Kineto] into a more open, general and user-friendly library', we would like to make it clear that we have similar goals and would be open to collaborating on these changes in the long run. However, we appreciate that adhering to existing Kineto designs may make sense for now.