Questions about joint matrix extension

[nirvedhmeshram]

Hello, I am trying to add support for joint matrix as specified in this doc in MLIR
https://github.com/intel/llvm/blob/c2f8602c843d2ba87dc1c9057f1422f23a77ed66/sycl/doc/design/spirv-extensions/SPV_INTEL_joint_matrix.asciidoc
I did have some questions about the ops

1. The spec does not mention what storage class the pointer can support, we were thinking of allowing, Workgroup and CrossWorkgroup, is that right and should any other be allowed?

2. If you look at the cooperativematrix side https://htmlpreview.github.io/?https://github.com/KhronosGroup/SPIRV-Registry/blob/master/extensions/NV/SPV_NV_cooperative_matrix.html

in "3.32.13 Arithmetic Instructions"

there is "Allow cooperative matrix types for the following instructions (if the component type is appropriate): OpSNegate, OpFNegate, OpIAdd, OpFAdd, OpISub, OpFSub, OpFDiv, OpSDiv, and OpUDiv. Allow cooperative matrix types for OpMatrixTimesScalar."
should this true for the joint matrix ops too?

Here is a WIP patch https://reviews.llvm.org/D131586 , if any contributors have feedback

cc @MrSidims @jcranmer-intel (tagging you as I see you guys actively working on this on the open source side)

[MrSidims]

First of all thanks for working on this in MLIR!

Tagging @dkhaldi who might be interested as well.

UniformConstant and Generic are looking to be legit and commonly used Storage Classes for pointer to a joint matrix as well. I would suggest to allow all Storage Classes/address spaces from OpenCL programming model (for now).

Regarding arithmetic instructions: unfortunately at the moment only Mad arithmetic instructions are allowed as a part of joint matrix extension. But I do believe, that it will be changed in a future.

[dkhaldi]

Thank you @nirvedhmeshram for starting this,

The mad arithmetic instruction is the one added explicitly because it is the one available on tensor hardware currently. The other arithmetic instructions will be done using vector/scalar units. Besides vector instructions, there are other instructions that will be done on element basis like the ones needed for quantization. That's why, as a general solution, we suggest this can be done using element indexing.

So, the other arithmetic instructions are done via element wise indexing using OpVectorExtractDynamic and OpVectorInsertDynamic. OpJointMatrixWorkItemLengthINTEL is used to get number of elements from the joint matrix that are owned by a work item.

Having said that, if you think there is a subset of operations that must be explicitly added for optimizations purposes, we can consider adding them as well.

[nirvedhmeshram]

Than you for the responses @MrSidims and @dkhaldi , that makes sense and is very helpful and I will add the MLIR ops appropriately. Also adding @antiagainst to the conversation who is working with me on adding this support in MLIR.

[nirvedhmeshram]

@MrSidims @dkhaldi , would it be possible to share some SPIR-V code (e.g for a matrix multiplication) that uses joint matrix. This would help a lot to validate that we are doing the right thing in our MLIR based code generation.

[MrSidims]

Some test examples can be found here: https://github.com/KhronosGroup/SPIRV-LLVM-Translator/tree/main/test/transcoding/SPV_INTEL_joint_matrix and here: https://github.com/intel/llvm-test-suite/tree/intel/SYCL/Matrix
Not sure, how to upload files via github, but here what I did just now to obtain SPIR-V modules. Here is a test for Mad operation:
https://github.com/intel/llvm-test-suite/blob/intel/SYCL/Matrix/joint_matrix_su_int8.cpp
Call:
clang++ -fsycl -fsycl-device-only test.cpp - it will generate bc file
llvm-spirv --spirv-ext=+all test.bc - get .spv file, llvm-spirv binary can be found after this repo building
llvm-spirv --to-text test.spv - optionally get test SPIR-V representation

there is another option of clang++ to generate SPIR-V module right away, but I do it this way in case if I want to check both LLVM IR and SPIR-V

[antiagainst]

Also it seems spirv-tools does not recognize the extension yet. Do you have patches somewhere to add support to spirv-tools? It would make inspecting the generated spirv blob much easier.

[nirvedhmeshram]

Yes spirv-tools would be very helpful, I was able to use the -to-text feature and generate a human readable version of the binary. However, this MLIR generated binary crashes in level zero without any error message, @MrSidims @dkhaldi and others do you spot any mistakes in it?
https://gist.github.com/nirvedhmeshram/f2013e4ebe11fba5dd3572322f3d02d7

[nirvedhmeshram]

Also, @MrSidims I tried the example above with a source build of intel llvm, I am running into this include issue, any thought how to overcome that

/home/nirvedh/llvm-project-shared/build/bin/clang++ -fsycl -fsycl-device-only llvm-test-suite/SYCL/Matrix/joint_matrix_su_int8.cpp
llvm-test-suite/SYCL/Matrix/joint_matrix_su_int8.cpp:15:10: fatal error: 'sycl/sycl.hpp' file not found
#include <sycl/sycl.hpp>
         ^~~~~~~~~~~~~~~
1 error generated.

[nirvedhmeshram]

Nevermind, I was able to generate the sycl example using docker,
https://gist.github.com/nirvedhmeshram/718038362d585c9d7b69cc06ec0ed5d1
now wondering what I am doing different

[nirvedhmeshram]

One major difference is that the sycl example uses a instrinsic call while the mlir generated code uses a Mad arithmetic instruction. Is this supported in the level zero driver?

[MrSidims]

Also it seems spirv-tools does not recognize the extension yet. Do you have patches somewhere to add support to spirv-tools? It would make inspecting the generated spirv blob much easier.

We don't have them yet, but we will certainly add them once we have the final (or almost final) version of the spec.

However, this MLIR generated binary crashes in level zero without any error message, @MrSidims @dkhaldi and others do you spot any mistakes in it?

llvm-spirv translator seems to be also OK with it in terms of generating SPIR-V binary out of it and consuming the SPIR-V module (translating it to LLVM IR). But there is an issue I see, I'll write it in the next bullet. Yet crashing without any error is definitely bad, @PawelJurek could you please help finding some one from GPU SPIR-V consumer side to take a look?

One major difference is that the sycl example uses a instrinsic call while the mlir generated code uses a Mad arithmetic instruction. Is this supported in the level zero driver?

Yes, L0 should know all 4 signed/unsinged Mad instructions (Mad, SUMad, USMad, UUMad).
Yet I see an issue in the SPIR-V module you have sent. Note these lines:
60 Name 22 "_Z29__spirv_JointMatrixSUMadINTELIahiLm8ELm32ELm8ELN5__spv12MatrixLayoutE0ELS1_3ELS1_0ELNS0_5Scope4FlagE3EEPNS0_24__spirv_JointMatrixINTELIT1_XT2_EXT4_EXT7_EXT8_EEEPNS4_IT_XT2_EXT3_EXT5_EXT8_EEEPNS4_IT0_XT3_EXT4_EXT6_EXT8_EEES7_S3_"
...
5 Function 16 22 0 21
...
8 FunctionCall 16 397 22 442 462 401 15

So as you see, here we have a function call to a function, that is suspiciously named __spirv_JointMatrixSUMadINTEL. It shouldn't be a function, instead it should be a SPIR-V instruction JointMatrixSUMadINTEL and the last lines of the module should be probably looking like:

8 JointMatrixLoadINTEL 20 462 461 457 15 15 14 
7 JointMatrixSUMadINTEL 16 22 442 462 401 15 

Yet L0 shouldn't silently crash.

file not found

That is quite unexpected. Let me check, if I'm doing something special to compile SYCL programs, meanwhile, could you please check, if replacing sycl/sycl.hpp with CL/sycl.hpp helps?

[nirvedhmeshram]

Thanks @MrSidims , I am looking for problems in the MLIR generated code here https://gist.github.com/nirvedhmeshram/f2013e4ebe11fba5dd3572322f3d02d7
that does not have the function call you mentioned but thats the one that crashes silently.

 could you please check, if replacing sycl/sycl.hpp with CL/sycl.hpp helps?

In the docker container I got this

test.cpp:15:10: fatal error: 'cl/sycl.hpp' file not found
#include <cl/sycl.hpp>

note that the sycl/sycl.hpp does work but that gives the function call. I did try to run that one and it gave this error

root@bfcedd6d3a57:clang++ -fsycl test.cpp
root@bfcedd6d3a57:~/sycl_workspace# ./a.out 
terminate called after throwing an instance of 'sycl::_V1::runtime_error'
  what():  joint matrix is not supported on host device. -33 (PI_ERROR_INVALID_DEVICE)
Aborted (core dumped)