SPIRVReader: handle direct types with CooperativeMatrixLengthKHR (#2695)

Translation of the attached test would currently fail due to the SPIRVReader attempting to process the `%matTy` operand as a regular value instead of a type. `OpCooperativeMatrixLengthKHR` seems to be pretty unique in taking an additional type operand beyond the result type, so special-case it in the reader. The translator currently accepts a non-type operand for `OpCooperativeMatrixLengthKHR` too, even though that's not within the specification; see various TODOs in the existing SPV_KHR_cooperative_matrix tests. Leave that relaxation in place, by only translating the operand as a type when it is an `OpTypeCooperativeMatrixKHR`. Original commit: KhronosGroup/SPIRV-LLVM-Translator@2b5f15d871aa39b
intel · Sep 6, 2024 · e992fe6 · e992fe6
1 parent d191784
commit e992fe6
Show file tree

Hide file tree

Showing 3 changed files with 37 additions and 2 deletions.
diff --git a/llvm-spirv/lib/SPIRV/SPIRVReader.cpp b/llvm-spirv/lib/SPIRV/SPIRVReader.cpp
@@ -3270,8 +3270,16 @@ Instruction *SPIRVToLLVM::transBuiltinFromInst(const std::string &FuncName,
         OC == OpControlBarrier)
       Func->addFnAttr(Attribute::Convergent);
   }
-  auto *Call =
-      CallInst::Create(Func, transValue(Ops, BB->getParent(), BB), "", BB);
+  CallInst *Call;
+  if (BI->getOpCode() == OpCooperativeMatrixLengthKHR &&
+      Ops[0]->getOpCode() == OpTypeCooperativeMatrixKHR) {
+    // OpCooperativeMatrixLengthKHR needs special handling as its operand is
+    // a Type instead of a Value.
+    llvm::Type *MatTy = transType(reinterpret_cast<SPIRVType *>(Ops[0]));
+    Call = CallInst::Create(Func, Constant::getNullValue(MatTy), "", BB);
+  } else {
+    Call = CallInst::Create(Func, transValue(Ops, BB->getParent(), BB), "", BB);
+  }
   setName(Call, BI);
   setAttrByCalledFunc(Call);
   SPIRVDBG(spvdbgs() << "[transInstToBuiltinCall] " << *BI << " -> ";

diff --git a/llvm-spirv/lib/SPIRV/libSPIRV/SPIRVInstruction.cpp b/llvm-spirv/lib/SPIRV/libSPIRV/SPIRVInstruction.cpp
@@ -146,6 +146,8 @@ SPIRVInstruction::getOperandTypes(const std::vector<SPIRVValue *> &Ops) {
     SPIRVType *Ty = nullptr;
     if (I->getOpCode() == OpFunction)
       Ty = reinterpret_cast<SPIRVFunction *>(I)->getFunctionType();
+    else if (I->getOpCode() == OpTypeCooperativeMatrixKHR)
+      Ty = reinterpret_cast<SPIRVType *>(I);
     else
       Ty = I->getType();
 

diff --git a/llvm-spirv/test/extensions/KHR/SPV_KHR_cooperative_matrix/length.spvasm b/llvm-spirv/test/extensions/KHR/SPV_KHR_cooperative_matrix/length.spvasm
@@ -0,0 +1,25 @@
+; RUN: spirv-as --target-env spv1.0 -o %t.spv %s
+; RUN: spirv-val %t.spv
+; RUN: llvm-spirv -r -o - %t.spv | llvm-dis | FileCheck %s
+
+OpCapability Addresses
+OpCapability Kernel
+OpCapability CooperativeMatrixKHR
+OpExtension "SPV_KHR_cooperative_matrix"
+OpMemoryModel Physical64 OpenCL
+OpEntryPoint Kernel %1 "testCoopMat"
+%void = OpTypeVoid
+%float = OpTypeFloat 32
+%fnTy = OpTypeFunction %void
+%uint = OpTypeInt 32 0
+%uint_3 = OpConstant %uint 3
+%uint_0 = OpConstant %uint 0
+%uint_8 = OpConstant %uint 8
+%matTy = OpTypeCooperativeMatrixKHR %float %uint_3 %uint_8 %uint_8 %uint_0
+%1 = OpFunction %void None %fnTy
+%2 = OpLabel
+%3 = OpCooperativeMatrixLengthKHR %uint %matTy
+OpReturn
+OpFunctionEnd
+
+; CHECK: call spir_func i32 @_Z34__spirv_CooperativeMatrixLengthKHRPU3AS143__spirv_CooperativeMatrixKHR__float_3_8_8_0(target("spirv.CooperativeMatrixKHR", float, 3, 8, 8, 0) zeroinitializer)