[mlir][spirv] Fix encoding of cooperative matrix type to match SPIRV spec

[ThomasRaoux]

Summary

Scope, rows and columns need to be encoded in a separate constant operation.

[antiagainst]

Nice!

mlir/lib/Dialect/SPIRV/Serialization/Deserializer.cpp

Deserializer::processCooperativeMatrixType(ArrayRef<uint32_t> operands) {
 
   Type elementTy = getType(operands[1]);
   if (!elementTy) {
     return emitError(unknownLoc,
                      "OpTypeCooperativeMatrix references undefined <id> ")
            << operands[1];
   }
 
-  auto scope = spirv::symbolizeScope(operands[2]);
+  auto scope = spirv::symbolizeScope(getConstantInt(operands[2]).getInt());
   if (!scope) {
     return emitError(unknownLoc,
                      "OpTypeCooperativeMatrix references undefined scope <id> ")
            << operands[2];
   }
 
-  unsigned rows = operands[3];
-  unsigned columns = operands[4];
+  unsigned rows = getConstantInt(operands[3]).getInt();
+  unsigned columns = getConstantInt(operands[4]).getInt();
 
   typeMap[operands[0]] = spirv::CooperativeMatrixNVType::get(
       elementTy, scope.getValue(), rows, columns);
   return success();
 }
 
 LogicalResult
 Deserializer::processRuntimeArrayType(ArrayRef<uint32_t> operands) {

mlir/lib/Dialect/SPIRV/Serialization/Serializer.cpp

Serializer::prepareBasicType(Location loc, Type type, uint32_t resultID,
   if (auto cooperativeMatrixType =
           type.dyn_cast<spirv::CooperativeMatrixNVType>()) {
     uint32_t elementTypeID = 0;
     if (failed(processType(loc, cooperativeMatrixType.getElementType(),
                            elementTypeID))) {
       return failure();
     }
     typeEnum = spirv::Opcode::OpTypeCooperativeMatrixNV;
+    auto getConstantOp = [&](uint32_t id) {
+      auto attr = IntegerAttr::get(IntegerType::get(32, type.getContext()), id);
+      return prepareConstantInt(loc, attr);
+    };
     operands.push_back(elementTypeID);
-    operands.push_back(static_cast<uint32_t>(cooperativeMatrixType.getScope()));
-    operands.push_back(cooperativeMatrixType.getRows());
-    operands.push_back(cooperativeMatrixType.getColumns());
+    operands.push_back(
+        getConstantOp(static_cast<uint32_t>(cooperativeMatrixType.getScope())));
+    operands.push_back(getConstantOp(cooperativeMatrixType.getRows()));
+    operands.push_back(getConstantOp(cooperativeMatrixType.getColumns()));
     return success();
   }
 
   if (auto matrixType = type.dyn_cast<spirv::MatrixType>()) {
     uint32_t elementTypeID = 0;
     if (failed(processType(loc, matrixType.getElementType(), elementTypeID))) {
       return failure();
     }