Page MenuHomePhabricator

[mlir][gpu] Add op to create MMA constant matrix
ClosedPublic

Authored by ThomasRaoux on Jun 7 2021, 10:33 PM.

Details

Summary

This allow creating a matrix with all elements set to a given value. This is needed to be able to implement a simple dot op.

Diff Detail

Event Timeline

ThomasRaoux created this revision.Jun 7 2021, 10:33 PM
ThomasRaoux requested review of this revision.Jun 7 2021, 10:33 PM
Herald added a project: Restricted Project. · View Herald TranscriptJun 7 2021, 10:33 PM
navdeepkk requested changes to this revision.Jun 9 2021, 9:27 AM

This is a great addition. We can bring in a scaling op also which scales mmaMatrix by a certain value. Maybe I can take that up.

mlir/include/mlir/Dialect/GPU/GPUOps.td
1043–1044

This should be changed/dropped I think.

1056

This doesn't match with the valid types of mmaMatrixType.

mlir/lib/Conversion/GPUToNVVM/WmmaOpsToNvvm.cpp
351–352

This comment needs to be updated.

This revision now requires changes to proceed.Jun 9 2021, 9:27 AM

Address review comments

ThomasRaoux marked 2 inline comments as done.Jun 9 2021, 9:44 AM
ThomasRaoux added inline comments.
mlir/include/mlir/Dialect/GPU/GPUOps.td
1043–1044

Changed it to only mention gpu.subgroup_mma_compute

This is a great addition. We can bring in a scaling op also which scales mmaMatrix by a certain value. Maybe I can take that up.

It would be nice to be able to handle most of the element-wise ops, ideally we should re-use the std ops but it looks like this would require infrastructure changes to MLIR (https://llvm.discourse.group/t/using-gpu-type-with-standard-ops/3542/2). The best short term solution is probably to add an op taking an attribute like GPU_AllReduceOperationAttr. This is a bit hacky but that would allow us to be able to generate interesting code using the mma ops.

navdeepkk accepted this revision.Jun 10 2021, 8:12 AM

This is a great addition. We can bring in a scaling op also which scales mmaMatrix by a certain value. Maybe I can take that up.

It would be nice to be able to handle most of the element-wise ops, ideally we should re-use the std ops but it looks like this would require infrastructure changes to MLIR (https://llvm.discourse.group/t/using-gpu-type-with-standard-ops/3542/2). The best short term solution is probably to add an op taking an attribute like GPU_AllReduceOperationAttr. This is a bit hacky but that would allow us to be able to generate interesting code using the mma ops.

Okay, That would allow us to use the same op and define the semantics of the pointwise op as we like. I will be happy to take this up in the near future. Let me know what you think.

This revision is now accepted and ready to land.Jun 10 2021, 8:12 AM
This revision was landed with ongoing or failed builds.Jun 10 2021, 8:34 AM
This revision was automatically updated to reflect the committed changes.

This is a great addition. We can bring in a scaling op also which scales mmaMatrix by a certain value. Maybe I can take that up.

It would be nice to be able to handle most of the element-wise ops, ideally we should re-use the std ops but it looks like this would require infrastructure changes to MLIR (https://llvm.discourse.group/t/using-gpu-type-with-standard-ops/3542/2). The best short term solution is probably to add an op taking an attribute like GPU_AllReduceOperationAttr. This is a bit hacky but that would allow us to be able to generate interesting code using the mma ops.

Okay, That would allow us to use the same op and define the semantics of the pointwise op as we like. I will be happy to take this up in the near future. Let me know what you think.

Yes that would be great. Feel free to pick it up. I'll sync up with you when I get close to needing it to make sure our timelines match but right now I can live with what exists now. My next step will most likely be adding transpose support (equivalent to the .col layout in wmma intrinsics).