Add support for SPV_INTEL_joint_matrix extension (KhronosGroup#1165)

The spec is available at intel/llvm#4373

Signed-off-by: Alexey Sotkin <[email protected]>

Author: AlexeySotkin

include/LLVMSPIRVExtensions.inc

@@ -46,3 +46,4 @@ EXT(SPV_INTEL_debug_module)
 EXT(SPV_INTEL_runtime_aligned)
 EXT(SPV_INTEL_arithmetic_fence)
 EXT(SPV_INTEL_bfloat16_conversion)
+EXT(SPV_INTEL_joint_matrix)

lib/SPIRV/CMakeLists.txt

@@ -40,6 +40,7 @@ add_llvm_library(LLVMSPIRVLib
     Analysis
     BitWriter
     Core
+    Demangle
     IRReader
     Linker
     Support

lib/SPIRV/SPIRVInternal.h

@@ -265,6 +265,22 @@ SPIRVMap<SPIRVExtInstSetKind, std::string, SPIRVExtSetShortName>::init() {
 typedef SPIRVMap<SPIRVExtInstSetKind, std::string, SPIRVExtSetShortName>
     SPIRVExtSetShortNameMap;
 
+template <>
+inline void SPIRVMap<internal::InternalJointMatrixLayout, std::string>::init() {
+  add(internal::RowMajor, "matrix.rowmajor");
+  add(internal::ColumnMajor, "matrix.columnmajor");
+  add(internal::PackedA, "matrix.packed.a");
+  add(internal::PackedB, "matrix.packed.b");
+}
+typedef SPIRVMap<internal::InternalJointMatrixLayout, std::string>
+    SPIRVMatrixLayoutMap;
+
+template <> inline void SPIRVMap<spv::Scope, std::string>::init() {
+  add(ScopeWorkgroup, "scope.workgroup");
+  add(ScopeSubgroup, "scope.subgroup");
+}
+typedef SPIRVMap<spv::Scope, std::string> SPIRVMatrixScopeMap;
+
 #define SPIR_MD_COMPILER_OPTIONS "opencl.compiler.options"
 #define SPIR_MD_KERNEL_ARG_ADDR_SPACE "kernel_arg_addr_space"
 #define SPIR_MD_KERNEL_ARG_ACCESS_QUAL "kernel_arg_access_qual"
@@ -312,6 +328,7 @@ const static char ConstantSampler[] = "ConstantSampler";
 const static char PipeStorage[] = "PipeStorage";
 const static char ConstantPipeStorage[] = "ConstantPipeStorage";
 const static char VmeImageINTEL[] = "VmeImageINTEL";
+const static char JointMatrixINTEL[] = "JointMatrixINTEL";
 } // namespace kSPIRVTypeName
 
 namespace kSPR2TypeName {
@@ -1054,6 +1071,8 @@ bool postProcessBuiltinsReturningStruct(Module *M, bool IsCpp = false);
 
 bool postProcessBuiltinsWithArrayArguments(Module *M, bool IsCpp = false);
 
+template <typename T>
+MetadataAsValue *map2MDString(LLVMContext &C, SPIRVValue *V);
 } // namespace SPIRV
 
 #endif // SPIRV_SPIRVINTERNAL_H

lib/SPIRV/SPIRVReader.cpp

@@ -443,6 +443,23 @@ Type *SPIRVToLLVM::transType(SPIRVType *T, bool IsClassMember) {
                                             SPIRAddressSpace::SPIRAS_Global));
   }
 
+  case internal::OpTypeJointMatrixINTEL: {
+    auto *MT = static_cast<SPIRVTypeJointMatrixINTEL *>(T);
+    auto R = static_cast<SPIRVConstant *>(MT->getRows())->getZExtIntValue();
+    auto C = static_cast<SPIRVConstant *>(MT->getColumns())->getZExtIntValue();
+    std::stringstream SS;
+    SS << kSPIRVTypeName::PostfixDelim;
+    SS << transTypeToOCLTypeName(MT->getCompType());
+    auto L = static_cast<SPIRVConstant *>(MT->getLayout())->getZExtIntValue();
+    auto S = static_cast<SPIRVConstant *>(MT->getScope())->getZExtIntValue();
+    SS << kSPIRVTypeName::PostfixDelim << R << kSPIRVTypeName::PostfixDelim << C
+       << kSPIRVTypeName::PostfixDelim << L << kSPIRVTypeName::PostfixDelim
+       << S;
+    std::string Name =
+        getSPIRVTypeName(kSPIRVTypeName::JointMatrixINTEL, SS.str());
+    return mapType(T, getOrCreateOpaquePtrType(M, Name));
+  }
+
   default: {
     auto OC = T->getOpCode();
     if (isOpaqueGenericTypeOpCode(OC) || isSubgroupAvcINTELTypeOpCode(OC))
@@ -3097,7 +3114,7 @@ Instruction *SPIRVToLLVM::transSPIRVBuiltinFromInst(SPIRVInstruction *BI,
   bool AddRetTypePostfix = false;
   if (OC == OpImageQuerySizeLod || OC == OpImageQuerySize ||
       OC == OpImageRead || OC == OpSubgroupImageBlockReadINTEL ||
-      OC == OpSubgroupBlockReadINTEL)
+      OC == OpSubgroupBlockReadINTEL || OC == internal::OpJointMatrixLoadINTEL)
     AddRetTypePostfix = true;
 
   bool IsRetSigned = false;

lib/SPIRV/SPIRVRegularizeLLVM.cpp

@@ -41,6 +41,8 @@
 #include "SPIRVInternal.h"
 #include "libSPIRV/SPIRVDebug.h"
 
+#include "llvm/ADT/StringExtras.h" // llvm::isDigit
+#include "llvm/Demangle/Demangle.h"
 #include "llvm/IR/InstVisitor.h"
 #include "llvm/IR/Instructions.h"
 #include "llvm/IR/Operator.h"
@@ -104,7 +106,7 @@ class SPIRVRegularizeLLVMBase {
   void buildUMulWithOverflowFunc(Function *UMulFunc);
 
   static std::string lowerLLVMIntrinsicName(IntrinsicInst *II);
-
+  void adaptStructTypes(StructType *ST);
   static char ID;
 
 private:
@@ -291,6 +293,58 @@ void SPIRVRegularizeLLVMBase::lowerUMulWithOverflow(
   UMulIntrinsic->setCalledFunction(UMulFunc);
 }
 
+void SPIRVRegularizeLLVMBase::adaptStructTypes(StructType *ST) {
+  if (!ST->hasName())
+    return;
+  StringRef STName = ST->getName();
+  STName.consume_front("struct.");
+  StringRef MangledName = STName.substr(0, STName.find('.'));
+
+  // Demangle the name of a template struct and parse the template
+  // parameters which look like:
+  // <signed char, 2ul, 2ul, (spv::MatrixLayout)0, (spv::Scope)3>
+  // The result should look like SPIR-V friendly LLVM IR:
+  // %spirv.JointMatrixINTEL._char_2_2_0_3
+  if (MangledName.startswith("_ZTSN5__spv24__spirv_JointMatrixINTEL")) {
+    std::string DemangledName = llvm::demangle(MangledName.str());
+    StringRef Name(DemangledName);
+    Name = Name.slice(Name.find('<') + 1, Name.rfind('>'));
+    std::stringstream SPVName;
+    // Name = signed char, 2ul, 2ul, (spv::MatrixLayout)0, (spv::Scope)3
+    auto P = Name.split(", ");
+    // P.first = "signed char
+    // P.second = "2ul, 2ul, (spv::MatrixLayout)0, (spv::Scope)3"
+    StringRef ElemType = P.first;
+    // remove possile qualifiers, like "const" or "signed"
+    ElemType.consume_back(" const");
+    size_t Space = ElemType.rfind(' ');
+    if (Space != StringRef::npos)
+      ElemType = ElemType.substr(Space + 1);
+    P = P.second.split(", ");
+    // P.first = "2ul"
+    // P.second = "2ul, (spv::MatrixLayout)0, (spv::Scope)3"
+    StringRef Rows = P.first.take_while(llvm::isDigit);
+    P = P.second.split(", ");
+    // P.first = "2ul"
+    // P.second = "(spv::MatrixLayout)0, (spv::Scope)3"
+    StringRef Cols = P.first.take_while(llvm::isDigit);
+    P = P.second.split(", ");
+    // P.first = "(spv::MatrixLayout)0"
+    // P.second = "(spv::Scope)3"
+    StringRef Layout = P.first.substr(P.first.rfind(')') + 1);
+    StringRef Scope = P.second.substr(P.second.rfind(')') + 1);
+
+    SPVName << kSPIRVTypeName::PrefixAndDelim
+            << kSPIRVTypeName::JointMatrixINTEL << kSPIRVTypeName::Delimiter
+            << kSPIRVTypeName::PostfixDelim << ElemType.str()
+            << kSPIRVTypeName::PostfixDelim << Rows.str()
+            << kSPIRVTypeName::PostfixDelim << Cols.str()
+            << kSPIRVTypeName::PostfixDelim << Layout.str()
+            << kSPIRVTypeName::PostfixDelim << Scope.str();
+    ST->setName(SPVName.str());
+  }
+}
+
 bool SPIRVRegularizeLLVMBase::runRegularizeLLVM(Module &Module) {
   M = &Module;
   Ctx = &M->getContext();
@@ -430,6 +484,9 @@ bool SPIRVRegularizeLLVMBase::regularize() {
     }
   }
 
+  for (StructType *ST : M->getIdentifiedStructTypes())
+    adaptStructTypes(ST);
+
   if (SPIRVDbgSaveRegularizedModule)
     saveLLVMModule(M, RegularizedModuleTmpFile);
   return true;

lib/SPIRV/SPIRVUtil.cpp

@@ -52,6 +52,7 @@
 #include "llvm/Bitcode/BitcodeWriter.h"
 #include "llvm/IR/IRBuilder.h"
 #include "llvm/IR/IntrinsicInst.h"
+#include "llvm/IR/Metadata.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/ErrorHandling.h"
@@ -148,7 +149,17 @@ std::string mapLLVMTypeToOCLType(const Type *Ty, bool Signed) {
     Ss << mapLLVMTypeToOCLType(EleTy, Signed) << Size;
     return Ss.str();
   }
-  return "invalid_type";
+  // It is expected that `Ty` can be mapped to `ReturnType` from "Optional
+  // Postfixes for SPIR-V Builtin Function Names" section of
+  // SPIRVRepresentationInLLVM.rst document (aka SPIRV-friendly IR).
+  // If `Ty` is not a scalar or vector type mentioned in the document (return
+  // value of some SPIR-V instructions may be represented as pointer to a struct
+  // in LLVM IR) we can mangle the type.
+  BuiltinFuncMangleInfo MangleInfo;
+  std::string MangledName =
+      mangleBuiltin("", const_cast<Type *>(Ty), &MangleInfo);
+  // Remove "_Z0"(3 characters) from the front of the name
+  return MangledName.erase(0, 3);
 }
 
 std::string mapSPIRVTypeToOCLType(SPIRVType *Ty, bool Signed) {
@@ -2122,4 +2133,16 @@ std::string getSPIRVFriendlyIRFunctionName(const std::string &UniqName,
   return mangleBuiltin(UniqName, ArgTys, &MangleInfo);
 }
 
+template <typename T>
+MetadataAsValue *map2MDString(LLVMContext &C, SPIRVValue *V) {
+  if (V->getOpCode() != OpConstant)
+    return nullptr;
+  uint64_t Const = static_cast<SPIRVConstant *>(V)->getZExtIntValue();
+  std::string Str = SPIRVMap<T, std::string>::map(static_cast<T>(Const));
+  return MetadataAsValue::get(C, MDString::get(C, Str));
+}
+template MetadataAsValue *
+map2MDString<internal::InternalJointMatrixLayout>(LLVMContext &, SPIRVValue *);
+template MetadataAsValue *map2MDString<spv::Scope>(LLVMContext &, SPIRVValue *);
+
 } // namespace SPIRV

lib/SPIRV/SPIRVWriter.cpp

@@ -561,7 +561,38 @@ SPIRVType *LLVMToSPIRVBase::transSPIRVOpaqueType(Type *T) {
     return mapType(T, BM->addQueueType());
   else if (TN == kSPIRVTypeName::PipeStorage)
     return mapType(T, BM->addPipeStorageType());
-  else
+  else if (TN == kSPIRVTypeName::JointMatrixINTEL) {
+    Type *ElemTy = nullptr;
+    StringRef Ty{Postfixes[0]};
+    auto NumBits = llvm::StringSwitch<unsigned>(Ty)
+                       .Case("char", 8)
+                       .Case("short", 16)
+                       .Case("int", 32)
+                       .Case("long", 64)
+                       .Default(0);
+    if (NumBits)
+      ElemTy = IntegerType::get(M->getContext(), NumBits);
+    else if (Ty == "half")
+      ElemTy = Type::getHalfTy(M->getContext());
+    else if (Ty == "float")
+      ElemTy = Type::getFloatTy(M->getContext());
+    else if (Ty == "double")
+      ElemTy = Type::getDoubleTy(M->getContext());
+    else
+      llvm_unreachable("Unexpected type for matrix!");
+
+    auto ParseInteger = [this](StringRef Postfix) -> ConstantInt * {
+      unsigned long long N = 0;
+      consumeUnsignedInteger(Postfix, 10, N);
+      return getUInt32(M, N);
+    };
+    SPIRVValue *Rows = transConstant(ParseInteger(Postfixes[1]));
+    SPIRVValue *Columns = transConstant(ParseInteger(Postfixes[2]));
+    SPIRVValue *Layout = transConstant(ParseInteger(Postfixes[3]));
+    SPIRVValue *Scope = transConstant(ParseInteger(Postfixes[4]));
+    return mapType(T, BM->addJointMatrixINTELType(transType(ElemTy), Rows,
+                                                  Columns, Layout, Scope));
+  } else
     return mapType(T,
                    BM->addOpaqueGenericType(SPIRVOpaqueTypeOpCodeMap::map(TN)));
 }

lib/SPIRV/libSPIRV/SPIRVInstruction.h

@@ -3251,6 +3251,28 @@ class SPIRVBfloat16ConversionINTELInstBase : public SPIRVUnaryInst<OC> {
 _SPIRV_OP(ConvertFToBF16INTEL)
 _SPIRV_OP(ConvertBF16ToFINTEL)
 #undef _SPIRV_OP
+
+class SPIRVJointMatrixINTELInstBase : public SPIRVInstTemplateBase {
+protected:
+  llvm::Optional<ExtensionID> getRequiredExtension() const override {
+    return ExtensionID::SPV_INTEL_joint_matrix;
+  }
+};
+
+class SPIRVJointMatrixINTELInst : public SPIRVJointMatrixINTELInstBase {
+  SPIRVCapVec getRequiredCapability() const override {
+    return getVec(internal::CapabilityJointMatrixINTEL);
+  }
+};
+
+#define _SPIRV_OP(x, ...)                                                      \
+  typedef SPIRVInstTemplate<SPIRVJointMatrixINTELInst, internal::Op##x##INTEL, \
+                            __VA_ARGS__>                                       \
+      SPIRV##x##INTEL;
+_SPIRV_OP(JointMatrixLoad, true, 6, true)
+_SPIRV_OP(JointMatrixStore, false, 5, true)
+_SPIRV_OP(JointMatrixMad, true, 7)
+#undef _SPIRV_OP
 } // namespace SPIRV
 
 #endif // SPIRV_LIBSPIRV_SPIRVINSTRUCTION_H

lib/SPIRV/libSPIRV/SPIRVModule.cpp

@@ -244,6 +244,9 @@ class SPIRVModuleImpl : public SPIRVModule {
   SPIRVEntry *addTypeStructContinuedINTEL(unsigned NumMembers) override;
   void closeStructType(SPIRVTypeStruct *T, bool) override;
   SPIRVTypeVector *addVectorType(SPIRVType *, SPIRVWord) override;
+  SPIRVTypeJointMatrixINTEL *addJointMatrixINTELType(SPIRVType *, SPIRVValue *,
+                                                     SPIRVValue *, SPIRVValue *,
+                                                     SPIRVValue *) override;
   SPIRVType *addOpaqueGenericType(Op) override;
   SPIRVTypeDeviceEvent *addDeviceEventType() override;
   SPIRVTypeQueue *addQueueType() override;
@@ -897,6 +900,14 @@ SPIRVTypeVector *SPIRVModuleImpl::addVectorType(SPIRVType *CompType,
                                                 SPIRVWord CompCount) {
   return addType(new SPIRVTypeVector(this, getId(), CompType, CompCount));
 }
+
+SPIRVTypeJointMatrixINTEL *SPIRVModuleImpl::addJointMatrixINTELType(
+    SPIRVType *CompType, SPIRVValue *Rows, SPIRVValue *Columns,
+    SPIRVValue *Layout, SPIRVValue *Scope) {
+  return addType(new SPIRVTypeJointMatrixINTEL(this, getId(), CompType, Rows,
+                                               Columns, Layout, Scope));
+}
+
 SPIRVType *SPIRVModuleImpl::addOpaqueGenericType(Op TheOpCode) {
   return addType(new SPIRVTypeOpaqueGeneric(TheOpCode, this, getId()));
 }

lib/SPIRV/libSPIRV/SPIRVModule.h

@@ -95,6 +95,7 @@ class SPIRVAsmINTEL;
 class SPIRVAsmCallINTEL;
 class SPIRVTypeBufferSurfaceINTEL;
 class SPIRVTypeTokenINTEL;
+class SPIRVTypeJointMatrixINTEL;
 
 typedef SPIRVBasicBlock SPIRVLabel;
 struct SPIRVTypeImageDescriptor;
@@ -242,6 +243,9 @@ class SPIRVModule {
   virtual SPIRVEntry *addTypeStructContinuedINTEL(unsigned NumMembers) = 0;
   virtual void closeStructType(SPIRVTypeStruct *, bool) = 0;
   virtual SPIRVTypeVector *addVectorType(SPIRVType *, SPIRVWord) = 0;
+  virtual SPIRVTypeJointMatrixINTEL *
+  addJointMatrixINTELType(SPIRVType *, SPIRVValue *, SPIRVValue *, SPIRVValue *,
+                          SPIRVValue *) = 0;
   virtual SPIRVTypeVoid *addVoidType() = 0;
   virtual SPIRVType *addOpaqueGenericType(Op) = 0;
   virtual SPIRVTypeDeviceEvent *addDeviceEventType() = 0;