[RFC] Adding XeVM dialect to MLIR
TLDR;
We need a new dialect, XeVM, like NVVM or ROCDL, to supplement LLVM dialect and expose Intel Xe architecture specific features.
Background and Motivation
Intel is building an upstream MLIR lowering pipeline for Intel GPU to access Intel-specific hardware features. MLIR already has XeGPU dialect that exposes those features such as 2D block memory operation, mma and etc. But upstream still lacks a lowering pipeline from XeGPU dialect.
Other GPU vendors like NVidia and AMD lower their respective dialects, nvgpu and amdgpu, to LLVM dialect along with target specific dialects like nvvm and rocdl. The additional dialects help hide intrinsic calls that don’t fit well with MLIR passes as intrinsics calls are generic function calls instead of proper ops. And the dialects define target attributes that are required for module serialization under GPU pipeline.
By contrast, Intel needs to rely on SPIR-V dialect as lowering target instead of LLVM as upstream does not have a dialect like nvvm or rocdl for Intel GPU.
Challenges
We have encountered two major issues while developing a GPU pipeline for Intel GPU targeting SPIR-V dialect.
First, SPIR-V dialect is not integration friendly. The lowering pipeline does not provide flexible integration with popular GPU AI frameworks. For example, Triton and OpenXLA are yet to accept any downstream extensions to MLIR and any extra code for Intel GPU support needs to be added to the respective project repository. In addition, their lowering pipeline is implemented around LLVM dialect. Updating all such pipelines requires major refactoring and doing it in a clean way to accommodate SPIR-V dialect is a significant challenge as Intel does not own those projects. Also, SPIR-V dialect lowering path is less mature than that of LLVM in upstream. For example, vector to LLVM conversion pass covers more ops compared to vector to SPIR-V conversion pass.
Second, SPIR-V dialect is not future proof for Intel Xe Architecture. Intel currently exposes all advanced Xe architecture features through SPIR-V by adding SPIR-V extensions but that is not a permanent choice. However, explicit lowering through SPIR-V dialect signifies SPIR-V as a permanent layer in upstream lowering for Xe architecture. It forces SPIR-V as the architectural interface between MLIR and lower-level code generation stack. Using SPIR-V dialect as a way to access advanced features is not future proof and should be avoided.
Solution
To address the above challenges, we propose XeVM dialect.
XeVM dialect is an MLIR leaf dialect extending LLVM dialect with operation unique to Intel GPU in the same way as dialects like nvvm and rocdl. With XeVM dialect, framework integration for Intel GPU can be done easily as it fits well into pipelines designed around targeting LLVM.
XeVM dialect exposes all HW features for Xe architecture, and it does it in a way without relying on SPIR-V specification. This removes the role of SPIR-V serving as an architectural interface for Intel GPU and makes any code targeting XeVM future proof.
XeVM dialect will initially use LLVM’s SPIR-V backend, now an official LLVM backend, for module serialization, but this is an implementation detail subject to change depending on how Xe features are exposed in the future.
We want to highlight that XeVM dialect is not an optional but essential component needed for enabling Intel Xe architecture upstream which is the only reasonable way to integrate Intel GPU support for downstream projects.
Scope of XeVM and upstreaming plan
The scope of XeVM dialect is to represent all Xe architecture features that do not have an LLVM counterpart. Those include
-
Advanced memory operations such as load/store/prefetch operations on 2D tiles
-
Advanced arithmetic operations such as MMA
-
Barrier operations
-
Access to hardware indices
-
Other features
Upstreaming is planned in multiple steps. The following is a tentative list of future PRs and items covered in each step.
-
XeVM dialect definition and dialect op tests: Definition includes memory, mma and barrier operations along with XeVM target attribute.
-
XeVM attach target transform pass
-
XeVM to LLVM conversion pass
-
XeVM target for gpu kernel (gpu.object) creation
-
XeVM dialect integration tests using GPU compilation pipeline with XeVM target.
-
XeVM operations for hardware indices and GPU to XeVM conversion pass: XeVM lowering will have a temporary dependency on convert-gpu-to-llvm-spv pass till this point
Once XeVM upstreaming is complete, we will add XeGPU to XeVM lowering.
Related older proposal
Intel created an RFC to add GEN dialect and presented the proposal last year at the MLIR open design forum that eventually got rejected. XeVM may appear similar to GEN but they are designed for different goals.
They may appear similar as they are both based on LLVM and can solve framework integration issues. But there are major differences.
The primary goal of GEN was to expose select assembly instructions while still tied to SPIR-V specification.
In contrast, XeVM has a goal of freeing Intel GPU from using SPIR-V as a hard architecture interface. Making all investment targeting XeVM future proof. And XeVM has a broader scope of operation support. It is not tied to just select assembly instructions and intends to expose all hardware features for Intel Xe architecture.
GEN had to rely on external software, Khronos LLVM SPIR-V translator, for serialization.
In contrast, XeVM is not tied to a specific architecture interface and can use the DeFacto LLVM backend of choice for other Intel tools such as SYCL. The current backend is SPIR-V backend, but due to XeVM’s flexibility, it can move to other backends and stay in sync with the rest of Intel GPU eco system even if any such change happens in the future.
Link to initial PR for dialect definition: https://github.com/llvm/llvm-project/pull/144811
Link to POC with integration tests: GitHub - silee2/llvm-project at xevmWorkspace
