This is an archive of the discontinued LLVM Phabricator instance.

[MLIR][SPIRV] Add intel joint matrix ops
ClosedPublic

Authored by nirvedhmeshram on Aug 10 2022, 9:02 AM.

Diff Detail

Event Timeline

nirvedhmeshram created this revision.Aug 10 2022, 9:02 AM
nirvedhmeshram requested review of this revision.Aug 10 2022, 9:02 AM

Cool, thanks for adding this! A few comments inlined.

mlir/include/mlir/Dialect/SPIRV/IR/SPIRVBase.td
446

Nit: match the order of definition in the above.

1395

This should not imply Shader, which is the ~root capability for Vulkan API? I didn't find the spec saying it implies Shader.

1474

Nit: match the order of definitions above

4058

Nit: Alignment is a bit off here.

mlir/include/mlir/Dialect/SPIRV/IR/SPIRVJointMatrixOps.td
1

Nit: alignment is off here.

46

Now we have v1.6. Should make this as v1.6 here. Also applies for following ops.

67

This seems to be the doc for the type, not for JointMatrixLoadINTEL?

86

This does not match with the assemblyFormat anymore. With assemblyFormat we don't need to document this; so just remove it.

120

Please define the layout as an attribute like SPV_ScopeAttr, instead of using the magic SPV_Integer here. You can follow how SPV_ScopeAttr is defined and copy that for Matrix Layout (https://github.com/intel/llvm/blob/c2f8602c843d2ba87dc1c9057f1422f23a77ed66/sycl/doc/design/spirv-extensions/SPV_INTEL_joint_matrix.asciidoc#matrix-layout). Please make sure to put it outside of the autogen region in SPIRVBase.td. Defining it that way will give us better IR parsing/printing and allows wiring up capability checks.

It's a bit tricker given this SPV_MatrixLayoutAttr would need special treatment like SPV_Scope. Search where SPV_ScopeAttr are listed in SPIRVUtilsGen.cpp and make sure you add SPV_MatrixLayoutAttr there.

mlir/include/mlir/Dialect/SPIRV/IR/SPIRVTypes.h
441

Also getLayout?

mlir/lib/Dialect/SPIRV/IR/SPIRVDialect.cpp
351

It should be rows 'x' columns 'x' element-type ',' layout ',' scope?

352

Missing layout here? As said before we should define SPV_MatrixLayoutAttr like SPV_ScopeAttr. Then you can add the parsing rules for matrix layout like scope in the following. So the full type should look like !spv.jointmatrix<8x8xf32, RowMajor, Subgroup>.

mlir/lib/Dialect/SPIRV/IR/SPIRVOps.cpp
3924

Hmm, StorageBuffer and PhysicalStorageBuffer are actually in the Shader capability hierarchy and for Vulkan. Are we sure we want those here? I'd assume this extension is more for OpenCL/oneAPI, which I'd think we need Workgroup , CrossWorkGroup, and some other ones fitting OpenCL/oneAPI. (The spec is unclear about this particular part, which could be a good question to ask.)

mlir/lib/Target/SPIRV/Serialization/Serializer.cpp
612

Missing layout here.

mlir/test/Dialect/SPIRV/IR/joint-matrix-ops.mlir
1

This is the test for IR correctness. Please also add a test file under Target/SPIRV/ to test serialization and deserialization.

antiagainst requested changes to this revision.Aug 10 2022, 11:36 AM
This revision now requires changes to proceed.Aug 10 2022, 11:36 AM
antiagainst added inline comments.Aug 10 2022, 11:39 AM
mlir/include/mlir/Dialect/SPIRV/IR/SPIRVAttributes.td
67

Please also put a comment in the above to mention the URL of the spec.

antiagainst added inline comments.Aug 10 2022, 11:43 AM
mlir/test/Dialect/SPIRV/IR/joint-matrix-ops.mlir
53

Actually these element wise ops are not explicitly said to be allowed. I suspect they should. But good to double check with Intel folks if you can.

Addressing reviewer comments

nirvedhmeshram marked 12 inline comments as done.Aug 15 2022, 12:58 PM
nirvedhmeshram added inline comments.
mlir/lib/Dialect/SPIRV/IR/SPIRVOps.cpp
3924

I did not get a response from intel, https://github.com/intel/llvm/issues/6568 at the moment. I am only allowing Workgroup and CrossWorkGroup for now

mlir/test/Dialect/SPIRV/IR/joint-matrix-ops.mlir
53

Let me know if you want me to remove this support.

nirvedhmeshram marked an inline comment as done.
nirvedhmeshram marked 2 inline comments as done.
antiagainst requested changes to this revision.Aug 15 2022, 2:46 PM

Awesome, just some final nits!

mlir/include/mlir/Dialect/SPIRV/IR/SPIRVBase.td
3836

Nit: Like others, use SPV_ML_* here. (ML is the short for MatrixLayout.)

3836

Could you define these new enum attribute after L3969? Here, this is the section for autogenerated enum attributes; it will be removed once we run the script to refresh definitions.

mlir/lib/Dialect/SPIRV/IR/SPIRVOps.cpp
3924

Okay SG!

mlir/test/Dialect/SPIRV/IR/joint-matrix-ops.mlir
53

Let's wait for Intel's reply. Fine to keep it right now.

This revision now requires changes to proceed.Aug 15 2022, 2:46 PM
nirvedhmeshram marked 2 inline comments as done.
antiagainst accepted this revision.Aug 15 2022, 3:03 PM

Awesome! Do you have write access to llvm repo? Do you need me to land it?

This revision is now accepted and ready to land.Aug 15 2022, 3:03 PM

Awesome! Do you have write access to llvm repo? Do you need me to land it?

yes I do, will land it, thanks.

This revision was automatically updated to reflect the committed changes.