This is an archive of the discontinued LLVM Phabricator instance.

[mlir][nvvm] Generalize wmma ops to handle more types and shapes
ClosedPublic

Authored by ThomasRaoux on Oct 27 2021, 11:01 PM.

Details

Summary

wmma intrinsics have a large number of combinations, ideally we want to be able to target all the different variants. To avoid a combinatorial explosion in the number of mlir op we use attributes to represent the different variation of load/store/mma ops.
We also can generate with tablegen helpers to know which combination are available.
Using this we can avoid having too hardcode a path for specific shapes and can support more types.
This patch also adds boiler plates for tf32 op support.

Diff Detail

Event Timeline

ThomasRaoux created this revision.Oct 27 2021, 11:01 PM
ThomasRaoux requested review of this revision.Oct 27 2021, 11:01 PM

cleanup useless include

mravishankar added inline comments.Oct 29 2021, 2:05 PM
mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
154–161

Nit typo : s/instanciate/instantiate

163–172

Correct me if I am wrong, but this reads like defining a type and op semantics. Wouldnt it be possible to have then as new Type and Op?

ThomasRaoux added inline comments.Oct 29 2021, 2:10 PM
mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
163–172

This does hold all the parameters attached to the matrix/mma intrinsic if that's what you mean.
Do you mean creating a new opaque type in the dialect? Since this is LLVM dialect I want something 1:1 to llvm IR so I don't want to add a new type that would require extra logic during translation to llvm IR.
Maybe I misunderstood the question.

mravishankar requested changes to this revision.Oct 29 2021, 2:34 PM
mravishankar added inline comments.
mlir/test/Conversion/GPUToNVVM/wmma-ops-to-nvvm.mlir
26

Is the type here correct? the instruction is taking two operands and returning one. The type in the modified code is just taking a (ptr) -> (struct). Is there a reason for the change?

This revision now requires changes to proceed.Oct 29 2021, 2:34 PM
ThomasRaoux added inline comments.Oct 29 2021, 2:37 PM
mlir/test/Conversion/GPUToNVVM/wmma-ops-to-nvvm.mlir
26

The previous version was using anytype for the stride, since it has to be i32 I changed it in the declaration and we don't need to print it explicitly.

ThomasRaoux added inline comments.Oct 29 2021, 2:56 PM
mlir/test/Conversion/GPUToNVVM/wmma-ops-to-nvvm.mlir
26

That's overall how it is done for rest of the ops I see in LLVM dialect, if there is no ambiguity on the type to use it is not printed out.

ThomasRaoux marked an inline comment as done.Oct 29 2021, 3:00 PM

Overall, it mostly looks like a refactoring, but a lot of tablegen ninja code . If possible could we move most of it into C++ and use the attributes directly to get what you want?

mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
163–172

Ok, I see your point. I was thinking of some better way of representing this, but here is better to mirror LLVM

275

THis is very hard to read. COuld we just add a C++ method to do whatever this is doing?

402

Could we move this to c++ file. Its hard to read here.

439

THis isnt the format of the new instructions? Probably needs to be restated

Fix descriptions

mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
275

The problem is that because of how getIntrinsicID is my source of truth is in tablegen so I need to have some function generated from tablegen to query the different operation available. Unless we can dynamically generate the intrinsicID annd move getIntrinsicID I don't see an obvious way to avoid some sort of tablegen for this.

402

I don't think it is possible because it needs to instantiate the intrinsic ID. Unfortunately it is not possible to build the intrinsic iD dynamically from C++ and since there are over 100 intrinsics it would be a lot of code to write and would be error prone.

439

Good catch, fixed it.

mravishankar accepted this revision.Oct 29 2021, 5:08 PM

After looking at the generated code, this looks much better! Nice cleanup! Its hard to read the tblgen, but presumably nobody has to.

mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
549

Restate this cause nvvm.wmma.m*n*k.mma is not the instruction after this change.

558

Restate this cause nvvm.wmma.m*n*k.store is not the instruction after this change.

This revision is now accepted and ready to land.Oct 29 2021, 5:08 PM
bondhugula requested changes to this revision.Oct 29 2021, 5:42 PM

These are mostly requests for documentation and code comments.

mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
163–172

Can you have a doc comment for these individual classes? I see a high-level top-level comment but individual doc comments are important for posterity.

174–183

Doc comment here.

275

Can you add a comment capturing this note?

289–350

Again, without additional comments, these are going to be hard to read and maintain.

355–363

Doc comment here again.

412–422

We need to be documenting what these methods return.

mlir/lib/Conversion/GPUToNVVM/WmmaOpsToNvvm.cpp
38–39

Nit: -> when an unimplemented

mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp
137

Doc comment.

This revision now requires changes to proceed.Oct 29 2021, 5:42 PM

Add extra doc comments

ThomasRaoux marked 2 inline comments as done.

Add extra comments

ThomasRaoux marked 6 inline comments as done.Oct 31 2021, 10:39 PM
ThomasRaoux added inline comments.
mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
163–172

Added doc here and below.

275

Documented this line 401

mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp
137

there is already a doc comment on the declaration in the header file. I believe we usually don't duplicate those?

bondhugula added inline comments.Nov 1 2021, 8:43 AM
mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp
137

Oh, sorry, I missed that. We shouldn't duplicate it then.

This revision was not accepted when it landed; it landed in state Needs Review.Nov 1 2021, 10:29 AM
This revision was automatically updated to reflect the committed changes.