Add ComponentTypeInterpretation for joint matrix type

It specifies how to interpret 'Component Type' when components
of a joint matrix are storages for values of different types,
for example float for TF32, unsigned short for bfloat16.

Spec update:
intel/llvm#8175

Signed-off-by: Sidorov, Dmitry <[email protected]>

Author: MrSidims

lib/SPIRV/SPIRVReader.cpp

@@ -445,9 +445,28 @@ Type *SPIRVToLLVM::transType(SPIRVType *T, bool UseTPT) {
                                        (unsigned)S};
     if (auto *Use = MT->getUse())
       Params.push_back(static_cast<SPIRVConstant *>(Use)->getZExtIntValue());
+    std::string ComponentTypeName;
+    auto *CTI = MT->getCTI();
+    if (!CTI) {
+      ComponentTypeName = transTypeToOCLTypeName(MT->getCompType());
+    } else {
+      switch (static_cast<SPIRVConstant *>(CTI)->getZExtIntValue()) {
+      case internal::InternalJointMatrixCTI::TF32:
+        ComponentTypeName = "tf32";
+        break;
+      case internal::InternalJointMatrixCTI::Bfloat16:
+        ComponentTypeName = "bfloat16";
+        break;
+      case internal::InternalJointMatrixCTI::PackedInt2:
+      case internal::InternalJointMatrixCTI::PackedInt4:
+        // Do nothing just now
+        break;
+      default:
+        llvm_unreachable("Unexpected joint matrix component type");
+      }
+    }
     return mapType(T, getSPIRVType(internal::OpTypeJointMatrixINTEL,
-                                   transTypeToOCLTypeName(MT->getCompType()),
-                                   Params, !UseTPT));
+                                   ComponentTypeName, Params, !UseTPT));
   }
   case OpTypeForwardPointer: {
     SPIRVTypeForwardPointer *FP =

lib/SPIRV/SPIRVWriter.cpp

@@ -626,6 +626,20 @@ SPIRVType *LLVMToSPIRVBase::transPointerType(SPIRVType *ET, unsigned AddrSpc) {
 // simply not true.
 SPIRVType *LLVMToSPIRVBase::transSPIRVJointMatrixINTELType(
     SmallVector<std::string, 8> Postfixes) {
+  auto ParseInteger = [this](StringRef Postfix) -> ConstantInt * {
+    unsigned long long N = 0;
+    if (consumeUnsignedInteger(Postfix, 10, N)) {
+      BM->getErrorLog().checkError(
+          false, SPIRVEC_InvalidLlvmModule,
+          "TypeJointMatrixINTEL expects integer parameters");
+      return 0;
+    }
+    return getUInt32(M, N);
+  };
+  std::vector<SPIRVValue *> Args;
+  for (size_t I = 1; I != Postfixes.size(); ++I)
+    Args.emplace_back(transConstant(ParseInteger(Postfixes[I])));
+
   Type *ElemTy = nullptr;
   StringRef Ty{Postfixes[0]};
   auto NumBits = llvm::StringSwitch<unsigned>(Ty)
@@ -634,32 +648,27 @@ SPIRVType *LLVMToSPIRVBase::transSPIRVJointMatrixINTELType(
                      .Case("int", 32)
                      .Case("long", 64)
                      .Default(0);
-  if (NumBits)
+  if (NumBits) {
     ElemTy = IntegerType::get(M->getContext(), NumBits);
-  else if (Ty == "half")
+  } else if (Ty == "half") {
     ElemTy = Type::getHalfTy(M->getContext());
-  else if (Ty == "float")
+  } else if (Ty == "float") {
     ElemTy = Type::getFloatTy(M->getContext());
-  else if (Ty == "double")
+  } else if (Ty == "double") {
     ElemTy = Type::getDoubleTy(M->getContext());
-  else if (Ty == "bfloat16")
+  } else if (Ty == "bfloat16") {
     ElemTy = Type::getInt16Ty(M->getContext());
-  else
+    auto *CTI = transConstant(getUInt32(M, static_cast<uint64_t>(
+            internal::InternalJointMatrixCTI::Bfloat16)));
+    Args.push_back(CTI);
+  } else if (Ty == "tf32") {
+    ElemTy = Type::getFloatTy(M->getContext());
+    auto *CTI = transConstant(getUInt32(M, static_cast<uint64_t>(
+            internal::InternalJointMatrixCTI::TF32)));
+    Args.push_back(CTI);
+  } else {
     llvm_unreachable("Unexpected type for matrix!");
-
-  auto ParseInteger = [this](StringRef Postfix) -> ConstantInt * {
-    unsigned long long N = 0;
-    if (consumeUnsignedInteger(Postfix, 10, N)) {
-      BM->getErrorLog().checkError(
-          false, SPIRVEC_InvalidLlvmModule,
-          "TypeJointMatrixINTEL expects integer parameters");
-      return 0;
-    }
-    return getUInt32(M, N);
-  };
-  std::vector<SPIRVValue *> Args;
-  for (size_t I = 1; I != Postfixes.size(); ++I)
-    Args.emplace_back(transConstant(ParseInteger(Postfixes[I])));
+  }
   return BM->addJointMatrixINTELType(transType(ElemTy), Args);
 }
 

lib/SPIRV/libSPIRV/SPIRVType.h

@@ -1088,6 +1088,7 @@ class SPIRVTypeJointMatrixINTEL : public SPIRVType {
   SPIRVValue *getLayout() const { return Args[2]; }
   SPIRVValue *getScope() const { return Args[3]; }
   SPIRVValue *getUse() const { return Args.size() > 4 ? Args[4] : nullptr; }
+  SPIRVValue *getCTI() const { return Args.size() > 5 ? Args[5] : nullptr; }
 };
 
 } // namespace SPIRV

lib/SPIRV/libSPIRV/spirv_internal.hpp

@@ -94,6 +94,13 @@ enum InternalJointMatrixLayout {
 
 enum InternalJointMatrixUse { MatrixA = 0, MatrixB = 1, Accumulator = 2 };
 
+enum InternalJointMatrixCTI {
+  TF32 = 0,
+  Bfloat16 = 1,
+  PackedInt2 = 2,
+  PackedInt4 = 3
+};
+
 enum InternalBuiltIn {
   IBuiltInSubDeviceIDINTEL = 6135,
   IBuiltInGlobalHWThreadIDINTEL = 6136,