[XeGPU] uArch definition (PR 1/N)

[mshahneo]

Update 2:

After another round of review update the uArch definition in the following way:

• Add common 5 APIs for all the interfaces.
• Make OpInterfaces more specific (e.g., 2D, 1D Block IO have seperate interfaces).
• Make the design more specific to Intel hardware.
• Remove the information we don't use, i.e., make the design simpler.

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.

[Jianhui-Li]

can we wrap these code as an XeGPU utility function? say getXeGPUChipStr()

[chencha3]

+1, I think it is worth to have a getXeArch() to hide defining PVCuArch and BMGuArch instantiation at the beginning. It sounds not scalable.

[chencha3]

It should be IntelGpuXe2.h here.

[chencha3]

Do Load and Store belong to different unit?

[chencha3]

I am not sure whether it is ok and necessary to expose the opcode here.

[chencha3]

nit: a format issue? Happy to see the restriction abstraction.

[chencha3]

nit: num_components? what is the major use of it?

[chencha3]

can XeCoreInfo be used by other archs?

[chencha3]

what is the string used for in alingment? the type name?

[chencha3]

nit: maybe we can remove PVCuArch namespace if it is not hard required.

[chencha3]

nit: format issue.
Is it also available for store2d, which doesn't have array length support in, e.g., pvc

[chencha3]

+1, I think it is worth to have a getXeArch() to hide defining PVCuArch and BMGuArch instantiation at the beginning. It sounds not scalable.