[XeGPU] uArch definition (PR 1/N)

[mshahneo]

Update:

After the first iteration of review comment, the current PR is update in the following way:

• The main pivot for this iteration is the utilities. These utilities are expose via a range of Interfaces (e.g. TileOpInterface, MatrixOpInterface). A specific instruction for a specific Architecture need to implement theses utilities.
• The uArch information is now directly embedded in the utility implementation. We are not using .yaml files. However, I haven't removed the .yaml file yet, since, they can still be utilized if choses to do so in the future.
• Also, last but not least, the current version provides an end-to-end prototype of how uArch can be used:
  -- Full uArch definition, utilities and all (at least for one instruction for now, DPAS).
  -- Definition for 2 uArchs (PVC, BMG).
  -- A pass to attach the uArch name to the module attribute using DLTI.
  -- Shows the usage of the uArch interface to verify XeGPU ops.
  -- Few test cases.

P.S. I still kept some old code intentionally, if we want to go back or use some of it. I'll clean up once we agree on the process.

================================

Hi All,

This is the initial PR for uArch definition. The primary purpose for this PR is to have a discussion and finalize the design. So please think of it as an RFC as well.

First target we want to achieve here is to finalize the skeleton:

• Skeleton of the YAML file, what information we want in there
• Skeleton of Data structures we want to have (e.g., necessary to create utilities to support the need)

Due to this purpose, the main focus point of this PR are the following files:

• mlir/lib/Dialect/XeGPU/Utils/intel_gpu_pvc.yaml : holds the YAML definition of uArch (PVC is used for this use case)
• mlir/include/mlir/Dialect/XeGPU/Utils/uArch.h : holds the base uArch struct that can be re-used by specialized uArchs such as PVC.
• mlir/include/mlir/Dialect/XeGPU/Utils/IntelGpuPVC.h : holds the PVC-specific structs.

mlir/lib/Dialect/XeGPU/Utils/IntelGpuPVC.cpp is added as placeholder for YAML mapping. This file is intentionally not complete. We want to reach a consensus on the YAML, and C++ structures before filling this one in.

Please provide your valuable feedback as we try to finalize this together.

[Garra1980]

this one should be on the top I guess

[mshahneo]

Ah, yes, I think I auto-added by co-pilot suggestion. Will remove it :)

[adam-smnk]

A general question on the structure - what's the benefit of maintaining both yaml and C++ definitions?
Do we expect sth to reuse yaml separately from C++? Can the C++ bindings be autogenerated or will it be a lot of manual boilerplate?

[adam-smnk]

Opaque string mapping doesn't feel user friendly. Especially as it seems to be wrapping numerical bit width that could be a numerical key directly - I assume it's some limitation due to using (conversion from) yaml.

matrix_size is even worse offender in that regard. These wrappers are hard to use without yaml definitions which again makes me wonder if we really need to split it into two separate files.

[mshahneo]

Opaque string mapping doesn't feel user friendly. Especially as it seems to be wrapping numerical bit width that could be a numerical key directly - I assume it's some limitation due to using (conversion from) yaml.

You are right, the limitation is due to the yaml mapping.

matrix_size is even worse offender in that regard. These wrappers are hard to use without yaml definitions which again makes me wonder if we really need to split it into two separate files.

Sorry, the matrix size is a mistake. The value should be a vector of uint. But yes, in general you are right, the structs and and yaml to some extent have to be used in conjunction.

[mshahneo]

A general question on the structure - what's the benefit of maintaining both yaml and C++ definitions? Do we expect sth to reuse yaml separately from C++? Can the C++ bindings be autogenerated or will it be a lot of manual boilerplate?

The C++ structure actually comes out of a necessity. To use yaml mapping utlity in LLVM, we have to have a C++ object mapped to it, hence the structures. And since we need the C++ structures anyway, I wanted to make some of them to be re-usable across uArchs.

Do we expect sth to reuse yaml separately from C++

Not really, not in this context anyway. But we wanted to keep the base structs in a such a way that they can be used in a non-LLVM cases, but that's more of hope not necessasity.

Can the C++ bindings be autogenerated or will it be a lot of manual boilerplate?

I thought about this, but it takes me back to tablegen. Should we use tablegen?

[adam-smnk]

A general question on the structure - what's the benefit of maintaining both yaml and C++ definitions? Do we expect sth to reuse yaml separately from C++? Can the C++ bindings be autogenerated or will it be a lot of manual boilerplate?

The C++ structure actually comes out of a necessity. To use yaml mapping utlity in LLVM, we have to have a C++ object mapped to it, hence the structures. And since we need the C++ structures anyway, I wanted to make some of them to be re-usable across uArchs.

Do we expect sth to reuse yaml separately from C++

Not really, not in this context anyway. But we wanted to keep the base structs in a such a way that they can be used in a non-LLVM cases, but that's more of hope not necessasity.

It feels like the overhead of this approach in its complexity and possible maintenance cost is not really justified. On its own it seems fine but these C++ bindings are not great. 😅
Overall, I'd say yaml is not really a first-class citizen in MLIR so it'd face extra scrutiny during upstreaming.

Can the C++ bindings be autogenerated or will it be a lot of manual boilerplate?

I thought about this, but it takes me back to tablegen. Should we use tablegen?

It's worth a try if what we need easily translates to tablegen entries and structure. That is, if you find yourself at any point fighting against tablegen or in need to introduce custom "hacks", I would not necessarily double down on it.

I would suggest starting with user interfaces to flesh out what's really needed. Then it should become clearer where to store all the uarch details e.g., embed shared memory size directly in getShmemSize() or create a separate struct or pull it from a constexpr dictionary that contains everything etc.

Next quick test would be to see how well it all holds up for another target like battlemage.

[mshahneo]

A general question on the structure - what's the benefit of maintaining both yaml and C++ definitions? Do we expect sth to reuse yaml separately from C++? Can the C++ bindings be autogenerated or will it be a lot of manual boilerplate?

The C++ structure actually comes out of a necessity. To use yaml mapping utlity in LLVM, we have to have a C++ object mapped to it, hence the structures. And since we need the C++ structures anyway, I wanted to make some of them to be re-usable across uArchs.

Do we expect sth to reuse yaml separately from C++

Not really, not in this context anyway. But we wanted to keep the base structs in a such a way that they can be used in a non-LLVM cases, but that's more of hope not necessasity.

It feels like the overhead of this approach in its complexity and possible maintenance cost is not really justified. On its own it seems fine but these C++ bindings are not great. 😅 Overall, I'd say yaml is not really a first-class citizen in MLIR so it'd face extra scrutiny during upstreaming.

Can the C++ bindings be autogenerated or will it be a lot of manual boilerplate?

I thought about this, but it takes me back to tablegen. Should we use tablegen?

It's worth a try if what we need easily translates to tablegen entries and structure. That is, if you find yourself at any point fighting against tablegen or in need to introduce custom "hacks", I would not necessarily double down on it.

I would suggest starting with user interfaces to flesh out what's really needed. Then it should become clearer where to store all the uarch details e.g., embed shared memory size directly in getShmemSize() or create a separate struct or pull it from a constexpr dictionary that contains everything etc.

Next quick test would be to see how well it all holds up for another target like battlemage.

Thanks, Adam. I think that's a good idea. Let me try to collect the uArch info battlemage. I looked into them during the design, but more details may help us finding the approach that would be easy to maintain.

[mshahneo]

Hi Adam (@adam-smnk), Igor (@Garra1980 ),

Modified the uarch definition based on our discussion.
In this version the Pivot is the utilities which necessary instructions have to implement.
I also moved the information to C++ (as part of get functions directly instead of in yaml).

Please let me know what you think. I only added implementation for DPAS. Will add Load/store/prefetch if we agree on the process.

Thanks a lot in advance :).

P.S. I still kept some old code intentionally, if we want to go back or use some of it. I'll clean up once we agree on the process.

[mshahneo]

@silee2 , @nbpatel , @chencha3 , @charithaintc , @Jianhui-Li, @akroviakov
please let me know if you have any comments/suggestions.
Thank you so much in advance :)

[Garra1980]

also cc @tkarna

[tkarna]

Do you think xegpu could at some point support the Data Layout and Target Information (DLTI) dialect for querying the uarch info?
FYI @rolfmorel

[rolfmorel]

Thanks for adding me, @tkarna ! Agreed, it would be great to collaborate on ways to expose this target info through DLTI.

I think @tkarna's work on performant schedules for Xe - which have plentiful magic values derivable from hw info - should be a good exemplar of a user of Xe's DLTI. Showing how to automatically derive those values through DLTI would be a great success indicator.

I had a quick attempt at what exposing this data through DLTI might look like (here using the transform dialect though you can also do the same queries from C++):

module attribute { #xevm.target<arch = "intel_gpu_pvc", ... = ?any other flags that vary across hardware?> } {
    ...
}
module {transform.with_named_sequence} {
    transform.named_sequence @__main__(%mod: !transform.any_op) {
        %number_of_eus_per_xe_core = transform.dlti.query ["xe_core","#eus"] at %mod  // #xevm.target answers 8
        %number_of_threads_per_eu = transform.dlti.query ["xe_core","eu","#threads"] at %mod  // #xevm.target answers [8, 4]
        
        %shared_memory_size = transform.dlti.query ["xe","shared_memory","size"] at %mod  // #xevm.target answers 524288
        %shared_memory_alignment = transform.dlti.query ["xe","shared_memory","alighment"] at %mod  // #xevm.target answers 64

        %grf_bitwidth = transform.dlti.query ["xe","grf","bitwidth"] at %mod  // #xevm.target answers 512
        %grf_modes = transform.dlti.query ["xe","grf","modes"] at %mod  // #xevm.target answers ["small", "large"]
        %grf_count_per_thread = transform.dlti.query ["xe","grf","mode","small","per_thread"] at %mod  // #xevm.target answers 128

        %dpas_opcode = transform.dlti.query ["xe","inst","dpas","opcode"] at %mod  // #xevm.target answers 0x83
        %dpas_systolic_depth = transform.dlti.query ["xe","inst","dpas","systolic_depth"] at %mod  // #xevm.target answers [8]
        %dpas_repeat_count = transform.dlti.query ["xe","inst","dpas","repeat_count"] at %mod  // #xevm.target answers 1,2,3,4,5,6,7,8 as multiple associations for single result handle OR [1, 8] OR #dlti.map<"min" = 1, "max" = 8> OR #dlti.range<1,2,...,8>/#tune.range<1,2,...,8> which would also respond to "min" and "max" queries,
        %dpas_ops_per_channel = transform.dlti.query ["xe","inst","dpas","ops_per_channel"] at %mod  // #xevm.target answers #dlti.map<19: 1, 16: 2, 8: 4, 4: 8, 2:8>
        
        %dpas_supported_types = transform.dlti.query ["xe","inst","dpas","supported_types"] at %mod  // #xevm.target answers #xevm.dpas_types so that ...
        %dpas_supported_types_B_wildcard = transform.dlti.query ["xe","inst","dpas","supported_types", f16, f16, f16, #dlti._] at %mod  // #xevm.dpas_types (yielded by #xevm.target which answered prefix) answers f16
        %dpas_supported_types_Dst_wildcard = transform.dlti.query ["xe","inst","dpas","supported_types", #dlti._, f16, f16, f16] at %mod  // #xevm.dpas_types (yielded by #xevm.target which answered prefix) answers f16,f32 (i.e. two type params associated to the handle)
        %dpas_supported_types_Dst_and_Acc_wildcard:2 = transform.dlti.query ["xe","inst","dpas","supported_types", #dlti._, #dlti._, f16, f16] at %mod  // #xevm.dpas_types (yielded by #xevm.target which answered prefix) answers f16,f32,f16,f32 for first result and f16,f16,f32,f32 for second result
        
        ... // transform ops would take the above Values as parameters
    }
}

That is, there's an #xevm.target attribute (or #xegpu.target) that handles most queries directly though it can also delegate to attributes like #xevm.dpas_types (in this case constructed by xevm.target) if so desired. This means there should be a XeTargetAttr implementing the DLTIQueryInterface's single query(..) method. We can discuss what's the best mechanism for implementing this method.

Not every query above is (easily) implemented with current upstream DLTI. I am working on a PR (oriented towards DLTI attributes that respond to queries by querying for target info and applying cost models) that will enable all the above that I will push as a draft this week.

[mshahneo]

@tkarna , yes, we would like to.
Current plan was to utilize DLTI only to query/expose the target triple.

As @rolfmorel pointed out, the current DLTI is too verbose and simple, @adam-smnk & I spoke about using DLTI directly, the issue was mapping all the hardware info in the DLTI attribute was making the IR way too verbose and big (e.g., think the .yaml file in the atttibute list :().

Thanks @rolfmorel for the detailed example. I remember you said that you envision that a query can essentially take a C++ function as key. That's why in the current iteration, I tried to make the utilities a first class citizen, this way, any ops that implement these interfaces can be queried via DLTI:

%dpas_supported_M_sizes = transform.dlti.query ["xe","inst","dpas","get_suppported_M", "bf16"] at %mod // #xevm.target answers the possible M sizes for type bf16...

This would essentially call that specific uArch's implementation of getSupportedM() function.

I'd like to discuss more with you all on this.

[adam-smnk]

I like the general direction 👍
The overall concept to have common descriptors (the uArch.h) which can have platform specific implementations (PVC, battlemage variants etc.) is great.

I haven't paid too close attention to exact implementation yet. However, it is a bit hard to really judge these helper APIs and their exact abstraction right now. I think it might be best to just start with a small prototype as soon as possible and just iterate on the go.

@chencha3 I think the on-going XeGPU verifier relaxation could great opportunity to run more hands-on experiments with this infrastructure. Perhaps a small uArch-aware validation pass for XeGPU or XeVM ops?

@tkarna If you could share what key hardware properties you already use for tuning or which further info could help you in that, it'd be super valuable feedback to help steer this design.

[adam-smnk]

What exactly is the meaning of no_of_dims.
In this case, if dims.size() is different than no_of_dims then what happens?

[mshahneo]

You are right, we could just remove no_of_dims and get it from dims.size().
Maybe remove the class itself. Let me see.

[tkarna]

@tkarna If you could share what key hardware properties you already use for tuning or which further info could help you in that, it'd be super valuable feedback to help steer this design.

In my preliminary tests I've use the following parameters for 4k matmul benchmark.

Fixed parameters:

• DPAS tile sizes (n, m, k)

Tunable parameters:

• WG tile (n, m); SG tile (n, m); K tile size
• Prefetch tile sizes for A and B
• Block (load) tile sizes for A and B

Implemented constrains: tile size must not exceed the parent tile size and divide parent tile evenly. SG tiling must not exceed max number of threads (32). Prefetch tile size must be compatible with the SG level tread configuration (coop prefetching).

There are other hardware constraints I have not explicitly expressed yet, for example maximum/supported load tile size and VNNI packing constraints for B. Some configs do spill registers and run slower.

[mshahneo]

I like the general direction 👍 The overall concept to have common descriptors (the uArch.h) which can have platform specific implementations (PVC, battlemage variants etc.) is great.

I haven't paid too close attention to exact implementation yet. However, it is a bit hard to really judge these helper APIs and their exact abstraction right now. I think it might be best to just start with a small prototype as soon as possible and just iterate on the go.

@chencha3 I think the on-going XeGPU verifier relaxation could great opportunity to run more hands-on experiments with this infrastructure. Perhaps a small uArch-aware validation pass for XeGPU or XeVM ops?

@tkarna If you could share what key hardware properties you already use for tuning or which further info could help you in that, it'd be super valuable feedback to help steer this design.

Thank you so much, @adam-smnk for the feedback :-) .
I am working on the prototype, I should be able push the changes today.

[mshahneo]

Hi @adam-smnk , @rolfmorel , @tkarna , @Jianhui-Li, @silee2, @chencha3 , @charithaintc,

I added modified the PR to add a prototype. Please take a look at your conveniences.

[Garra1980]

Am I correct that this file is supposed to replace current https://github.com/llvm/llvm-project/blob/main/mlir/include/mlir/Dialect/XeGPU/IR/XeGPUTargetInfo.h?

[mshahneo]

Potentially, yes.

[Garra1980]

I guess this piece can be hidden inside uArch utilities so client just say smth like
auto res = verify(op, "dpas");
?

[mshahneo]

The issue with this is that it requires the utility to have all the parameters for of the op.
Alternatively, the op would have to implement the same interface so that the uArch utility can get the parameters.

[chencha3]

It seems the test checks are missing here.

[mshahneo]

Ah, sorry Chao, this test is supposed to be removed. The other 2 tests covers the mechanism. I pushed it initially forgot to remove it.

[chencha3]

I want to raise the following discussion points:

1. NV and AMD are putting nvvm-attach-target and rocdl-attach-target under the GPU dialect. Not sure whether it is better for us to follow their way too.
2. They are attached to GPUModuleOp, a finer granularity than us. It may be a requirement of GPU dialect. A potential benefit is that they may be able to support the case that a ModuleOp containing multiple GPUModuleOp, with each GPUMoudleOp map to different GPUs from different vendors.
3. They used NVVMTargetAttr and ROCDLTargetAttr to store the info. Now, we have XeVMTargetAttr, can we reuse it here too? Considering we only capture the device name here, which is also captured by XeVMTargetAttr?

[chencha3]

Sang Ik created this PR: llvm#147372. is it similar to the pass here?

[akroviakov]

I agree, this appears similar to the xevm's attach target pass that is yet to be merged. After the merge, the chip string of the xevm target should be appropriate to query uarch info:

    auto targetDeviceArch =
        mlir::xegpu::uArch::uArchMap::instance().get(targetDeviceNameStr);

[rolfmorel]

Do the XeVM targets have a sensible notion of a LLVM triple along side that of chip? How about (LLVM) features? If so, it would probably make sense to have a #xevm.target attribute, a la the proposed above, which implements the LLVMTargetAttrInterface from this WIP upstream PR:
https://github.com/llvm/llvm-project/pull/145899/files#diff-6c2503d165a7390c955c3c4fa4a76fd1991633b5ec597a1b1fd92731e1f3684dR572

I am in the process of making #nvvm.target and #rocdl.target attributes implement this interface (in addition to the generic #llvm.target attr in that PR). If sensible, the XeVM dialect should probably (try to) mirror this.

[rolfmorel]

Given this in https://github.com/llvm/llvm-project/pull/147372/files#diff-3da7b8747032bb581f119fd49037d06706443c3c5db56d4d28cc74053a9b754dR284 in the Add xevm-attach-target transform pass PR, I take it there's a sensible triple.

I will have a look at making xevm.target also implement my interface.

[mshahneo]

Thanks Chao, Rolf, Artem.

You are right once XeVM target attribute gets upstreamed, we should use it. This pass is here just to show the use case.

[rolfmorel]

👍

On that note, I found that the #xevm.target attr itself already got merged: https://github.com/llvm/llvm-project/blob/b9d7513bf134febe72c05a04ff20f87191d7213a/mlir/include/mlir/Dialect/LLVMIR/XeVMOps.td#L521

We can get cracking!

[Jianhui-Li]

why opcode ?

[mshahneo]

I added some extra fields in the structs in case they become useful in the future, they are not useful in our current use-case. I can remove them if we don't think they are necessary.

[Jianhui-Li]

Tile is overloaded term. It looks to me just a shape. Does there any pre-defined data type you can reuse?

[mshahneo]

Sorry, Jianhui. I am not using it in the current version of the implementation. I kept it if we want to use it in some fashion later.

The current implementation uses std::vector.

[Garra1980]

@mshahneo Looks like we should remove older version of the interface and keep the one we want to proceed with

[Jianhui-Li]

why not just use DiscreteTile, if this is not used that often.

[mshahneo]

Kept two different ways to specify the range.

[Jianhui-Li]

It seems that DiscreteTile represents shapes combined from lists of sizes. How about name them ShapeCombination or something more specific to our problem and less general?

[mshahneo]

That makes sense. Current implementation actually uses std::vector.

[Jianhui-Li]

The term "restriction" is also very generic. Please consider making it specific. Validate the instance belongs to a combination?

[mshahneo]

Restriction is supposed to be model all restriction.
And a common validate() could just apply all the restrictions. However, we can always go back to specific validations. I'll think about this.

[Jianhui-Li]

I am a bit confused. I thought you will read the information from dlti, not hard-coded inside the code.

[adam-smnk]

These hardware properties are static so I'd say they should be hardcoded somewhere.
Whether you can extract them from LLVM, inside DLTI query function, or as a separate helper like this (or IMEX's target info) is matter of what's most convenient.

I don't think we should be dumping whole hardware specification into IR as DLTI attributes.

[Jianhui-Li]

We should discuss where to check the uArch info. The uArch info could be checked in the lowering passes, and the op definition should check general constraints (like a.m = c.m).

[mshahneo]

Actually the way the uArch info is set up. Once the target is attached, it can be queried from everywhere (pass and dialect).

[Jianhui-Li]

how is this file being loaded?

[adam-smnk]

That's an older artifact, we shouldn't have yaml at all.

[Jianhui-Li]

General: I don't think the yaml should write following the terminology in the internal uArch spec for PVC (like simd, systolic_depth, repeat_count, vnni ...). Instead, it should written following the SPIRV extnesion (or the OCL C function) (no mentioning of simd, using m/n/k, packing).

[mshahneo]

Hi Jianhui,

We are not using the YAML in the current version.