Align translation of OpCooperativeMatrixLengthKHR to match the spec (#2964)

`SPV_KHR_cooperative_matrix` extension defines that the only argument
accepted in this instruction is `Matrix Type <id>`, not the pointer to
an actual matrix.

This resolves #2963

Original commit:
KhronosGroup/SPIRV-LLVM-Translator@197a800558787c9

Author: vmaksimo

llvm-spirv/lib/SPIRV/SPIRVReader.cpp

@@ -3627,8 +3627,7 @@ Instruction *SPIRVToLLVM::transBuiltinFromInst(const std::string &FuncName,
       Func->addFnAttr(Attribute::Convergent);
   }
   CallInst *Call;
-  if (OC == OpCooperativeMatrixLengthKHR &&
-      Ops[0]->getOpCode() == OpTypeCooperativeMatrixKHR) {
+  if (OC == OpCooperativeMatrixLengthKHR) {
     // OpCooperativeMatrixLengthKHR needs special handling as its operand is
     // a Type instead of a Value.
     llvm::Type *MatTy = transType(reinterpret_cast<SPIRVType *>(Ops[0]));

llvm-spirv/lib/SPIRV/SPIRVWriter.cpp

@@ -6787,6 +6787,10 @@ LLVMToSPIRVBase::transBuiltinToInstWithoutDecoration(Op OC, CallInst *CI,
                                       transValue(CI->getArgOperand(2), BB), BB);
     return BM->addStoreInst(transValue(CI->getArgOperand(0), BB), V, {}, BB);
   }
+  case OpCooperativeMatrixLengthKHR: {
+    return BM->addCooperativeMatrixLengthKHRInst(
+        transScavengedType(CI), transType(CI->getArgOperand(0)->getType()), BB);
+  }
   case OpGroupNonUniformShuffleDown: {
     Function *F = CI->getCalledFunction();
     if (F->arg_size() && F->getArg(0)->hasStructRetAttr()) {

llvm-spirv/lib/SPIRV/libSPIRV/SPIRVModule.cpp

@@ -279,6 +279,9 @@ class SPIRVModuleImpl : public SPIRVModule {
   SPIRVTypeTaskSequenceINTEL *addTaskSequenceINTELType() override;
   SPIRVInstruction *addTaskSequenceGetINTELInst(SPIRVType *, SPIRVValue *,
                                                 SPIRVBasicBlock *) override;
+  SPIRVInstruction *
+  addCooperativeMatrixLengthKHRInst(SPIRVType *, SPIRVType *,
+                                    SPIRVBasicBlock *) override;
   SPIRVType *addOpaqueGenericType(Op) override;
   SPIRVTypeDeviceEvent *addDeviceEventType() override;
   SPIRVTypeQueue *addQueueType() override;
@@ -1094,6 +1097,14 @@ SPIRVInstruction *SPIRVModuleImpl::addTaskSequenceGetINTELInst(
       BB);
 }
 
+SPIRVInstruction *SPIRVModuleImpl::addCooperativeMatrixLengthKHRInst(
+    SPIRVType *RetTy, SPIRVType *MatTy, SPIRVBasicBlock *BB) {
+  return addInstruction(
+      SPIRVInstTemplateBase::create(OpCooperativeMatrixLengthKHR, RetTy,
+                                    getId(), getVec(MatTy->getId()), BB, this),
+      BB);
+}
+
 SPIRVType *SPIRVModuleImpl::addOpaqueGenericType(Op TheOpCode) {
   return addType(new SPIRVTypeOpaqueGeneric(TheOpCode, this, getId()));
 }

llvm-spirv/lib/SPIRV/libSPIRV/SPIRVModule.h

@@ -272,6 +272,9 @@ class SPIRVModule {
   virtual SPIRVTypeTaskSequenceINTEL *addTaskSequenceINTELType() = 0;
   virtual SPIRVInstruction *
   addTaskSequenceGetINTELInst(SPIRVType *, SPIRVValue *, SPIRVBasicBlock *) = 0;
+  virtual SPIRVInstruction *
+  addCooperativeMatrixLengthKHRInst(SPIRVType *, SPIRVType *,
+                                    SPIRVBasicBlock *) = 0;
   virtual SPIRVTypeVoid *addVoidType() = 0;
   virtual SPIRVType *addOpaqueGenericType(Op) = 0;
   virtual SPIRVTypeDeviceEvent *addDeviceEventType() = 0;

llvm-spirv/test/extensions/INTEL/SPV_INTEL_joint_matrix/cooperative_matrix_checked.ll

@@ -32,8 +32,7 @@
 ; CHECK-SPIRV-DAG: TypeCooperativeMatrixKHR [[#MatTy3:]] [[#Int8Ty]] [[#Const2]] [[#Const48]] [[#Const12]] [[#Const1]]
 ; CHECK-SPIRV: CooperativeMatrixConstructCheckedINTEL [[#MatTy1]]
 ; CHECK-SPIRV: CooperativeMatrixLoadCheckedINTEL [[#MatTy2]] [[#Load1:]]
-; TODO: Pass Matrix Type Id instead of Matrix Id to CooperativeMatrixLengthKHR.
-; CHECK-SPIRV: CooperativeMatrixLengthKHR [[#Int32Ty]] [[#]] [[#Load1]]
+; CHECK-SPIRV: CooperativeMatrixLengthKHR [[#Int32Ty]] [[#]] [[#MatTy2]]
 ; CHECK-SPIRV: CooperativeMatrixLoadCheckedINTEL [[#MatTy3]]
 ; CHECK-SPIRV: CooperativeMatrixMulAddKHR [[#MatTy1]]
 ; CHECK-SPIRV: CooperativeMatrixStoreCheckedINTEL

llvm-spirv/test/extensions/INTEL/SPV_INTEL_joint_matrix/cooperative_matrix_prefetch.ll

@@ -32,8 +32,7 @@
 ; CHECK-SPIRV-DAG: TypeCooperativeMatrixKHR [[#MatTy3:]] [[#Int8Ty]] [[#Const2]] [[#Const48]] [[#Const12]] [[#Const1]]
 ; CHECK-SPIRV: CompositeConstruct [[#MatTy1]]
 ; CHECK-SPIRV: CooperativeMatrixLoadKHR [[#MatTy2]] [[#Load1:]]
-; TODO: Pass Matrix Type Id instead of Matrix Id to CooperativeMatrixLengthKHR.
-; CHECK-SPIRV: CooperativeMatrixLengthKHR [[#Int32Ty]] [[#]] [[#Load1]]
+; CHECK-SPIRV: CooperativeMatrixLengthKHR [[#Int32Ty]] [[#]] [[#MatTy2]]
 ; CHECK-SPIRV: CooperativeMatrixPrefetchINTEL
 ; CHECK-SPIRV: CooperativeMatrixLoadKHR [[#MatTy3]]
 ; CHECK-SPIRV: CooperativeMatrixMulAddKHR [[#MatTy1]]

llvm-spirv/test/extensions/KHR/SPV_KHR_cooperative_matrix/cooperative_matrix.ll

@@ -30,8 +30,7 @@
 ; CHECK-SPIRV-DAG: TypeCooperativeMatrixKHR [[#MatTy3:]] [[#Int8Ty]] [[#Const2]] [[#Const48]] [[#Const12]] [[#Const1]]
 ; CHECK-SPIRV: CompositeConstruct [[#MatTy1]]
 ; CHECK-SPIRV: CooperativeMatrixLoadKHR [[#MatTy2]] [[#Load1:]]
-; TODO: Pass Matrix Type Id instead of Matrix Id to CooperativeMatrixLengthKHR.
-; CHECK-SPIRV: CooperativeMatrixLengthKHR [[#Int32Ty]] [[#]] [[#Load1]]
+; CHECK-SPIRV: CooperativeMatrixLengthKHR [[#Int32Ty]] [[#]] [[#MatTy2]]
 ; CHECK-SPIRV: CooperativeMatrixLoadKHR [[#MatTy3]]
 ; CHECK-SPIRV: CooperativeMatrixMulAddKHR [[#MatTy1]]
 ; CHECK-SPIRV: CooperativeMatrixStoreKHR