[Backport to 18] Implement SPV_KHR_bfloat16 extension (KhronosGroup#3099)

The extension add translation from LLVM's bfloat type to OpTypeFloat
%width% 16 %fp encoding% BFloat16KHR

Mangling follows LLVM's rules for the type.

Spec PR: KhronosGroup/SPIRV-Registry#323

---------

Signed-off-by: Sidorov, Dmitry <[email protected]>
Co-authored-by: Aziz, Michael <[email protected]>

Author: 2 people

include/LLVMSPIRVExtensions.inc

@@ -75,3 +75,4 @@ EXT(SPV_INTEL_maximum_registers)
 EXT(SPV_INTEL_bindless_images)
 EXT(SPV_INTEL_2d_block_io)
 EXT(SPV_INTEL_subgroup_matrix_multiply_accumulate)
+EXT(SPV_KHR_bfloat16)

lib/SPIRV/Mangler/ManglingUtils.cpp

@@ -28,6 +28,7 @@ static const char *PrimitiveNames[PRIMITIVE_NUM] = {
     "half",
     "float",
     "double",
+    "__bf16",
     "void",
     "...",
     "image1d_ro_t",
@@ -105,6 +106,7 @@ const char *MangledTypes[PRIMITIVE_NUM] = {
     "Dh",                                // HALF
     "f",                                 // FLOAT
     "d",                                 // DOUBLE
+    "u6__bf16",                          // __BF16
     "v",                                 // VOID
     "z",                                 // VarArg
     "14ocl_image1d_ro",                  // PRIMITIVE_IMAGE1D_RO_T
@@ -197,6 +199,7 @@ static const SPIRversion PrimitiveSupportedVersions[PRIMITIVE_NUM] = {
     SPIR12, // HALF
     SPIR12, // FLOAT
     SPIR12, // DOUBLE
+    SPIR12, // __BF16
     SPIR12, // VOID
     SPIR12, // VarArg
     SPIR12, // PRIMITIVE_IMAGE1D_RO_T

lib/SPIRV/Mangler/ParameterType.h

@@ -45,6 +45,7 @@ enum TypePrimitiveEnum {
   PRIMITIVE_HALF,
   PRIMITIVE_FLOAT,
   PRIMITIVE_DOUBLE,
+  PRIMITIVE_BFLOAT,
   PRIMITIVE_VOID,
   PRIMITIVE_VAR_ARG,
   PRIMITIVE_STRUCT_FIRST,

lib/SPIRV/SPIRVReader.cpp

@@ -313,6 +313,8 @@ std::optional<uint64_t> SPIRVToLLVM::getAlignment(SPIRVValue *V) {
 Type *SPIRVToLLVM::transFPType(SPIRVType *T) {
   switch (T->getFloatBitWidth()) {
   case 16:
+    if (T->isTypeFloat(16, FPEncodingBFloat16KHR))
+      return Type::getBFloatTy(*Context);
     return Type::getHalfTy(*Context);
   case 32:
     return Type::getFloatTy(*Context);
@@ -1521,7 +1523,9 @@ Value *SPIRVToLLVM::transValueWithoutDecoration(SPIRVValue *BV, Function *F,
       const llvm::fltSemantics *FS = nullptr;
       switch (BT->getFloatBitWidth()) {
       case 16:
-        FS = &APFloat::IEEEhalf();
+        FS =
+            (BT->isTypeFloat(16, FPEncodingBFloat16KHR) ? &APFloat::BFloat()
+                                                        : &APFloat::IEEEhalf());
         break;
       case 32:
         FS = &APFloat::IEEEsingle();

lib/SPIRV/SPIRVUtil.cpp

@@ -1330,6 +1330,8 @@ static SPIR::RefParamType transTypeDesc(Type *Ty,
     return SPIR::RefParamType(new SPIR::PrimitiveType(SPIR::PRIMITIVE_FLOAT));
   if (Ty->isDoubleTy())
     return SPIR::RefParamType(new SPIR::PrimitiveType(SPIR::PRIMITIVE_DOUBLE));
+  if (Ty->isBFloatTy())
+    return SPIR::RefParamType(new SPIR::PrimitiveType(SPIR::PRIMITIVE_BFLOAT));
   if (auto *VecTy = dyn_cast<FixedVectorType>(Ty)) {
     return SPIR::RefParamType(new SPIR::VectorType(
         transTypeDesc(VecTy->getElementType(), Info), VecTy->getNumElements()));

lib/SPIRV/SPIRVWriter.cpp

@@ -397,6 +397,16 @@ SPIRVType *LLVMToSPIRVBase::transType(Type *T) {
     }
   }
 
+  if (T->isBFloatTy()) {
+    BM->getErrorLog().checkError(
+        BM->isAllowedToUseExtension(ExtensionID::SPV_KHR_bfloat16),
+        SPIRVEC_RequiresExtension,
+        "SPV_KHR_bfloat16\n"
+        "NOTE: LLVM module contains bfloat type, translation of which "
+        "requires this extension");
+    return mapType(T, BM->addFloatType(16, FPEncodingBFloat16KHR));
+  }
+
   if (T->isFloatingPointTy())
     return mapType(T, BM->addFloatType(T->getPrimitiveSizeInBits()));
 

lib/SPIRV/libSPIRV/SPIRVEnum.h

@@ -221,6 +221,9 @@ template <> inline void SPIRVMap<SPIRVCapabilityKind, SPIRVCapVec>::init() {
                {CapabilityCooperativeMatrixKHR});
   ADD_VEC_INIT(internal::CapabilityCooperativeMatrixCheckedInstructionsINTEL,
                {CapabilityCooperativeMatrixKHR});
+  ADD_VEC_INIT(CapabilityBFloat16DotProductKHR, {CapabilityBFloat16TypeKHR});
+  ADD_VEC_INIT(CapabilityBFloat16CooperativeMatrixKHR,
+               {CapabilityBFloat16TypeKHR, CapabilityCooperativeMatrixKHR});
 }
 
 template <> inline void SPIRVMap<SPIRVExecutionModelKind, SPIRVCapVec>::init() {

lib/SPIRV/libSPIRV/SPIRVInstruction.h

@@ -693,6 +693,18 @@ class SPIRVBinary : public SPIRVInstTemplateBase {
       return static_cast<SPIRVWord>(VersionNumber::SPIRV_1_4);
     return static_cast<SPIRVWord>(VersionNumber::SPIRV_1_0);
   }
+  SPIRVCapVec getRequiredCapability() const override {
+    if (OpCode == OpDot) {
+      const SPIRVType *OpTy = getValueType(Ops[0]);
+      if (OpTy && OpTy->isTypeVector()) {
+        OpTy = OpTy->getVectorComponentType();
+        if (OpTy && OpTy->isTypeFloat(16, FPEncodingBFloat16KHR)) {
+          return getVec(CapabilityBFloat16DotProductKHR);
+        }
+      }
+    }
+    return SPIRVInstruction::getRequiredCapability();
+  }
 };
 
 template <Op OC>

lib/SPIRV/libSPIRV/SPIRVModule.cpp

@@ -254,7 +254,8 @@ class SPIRVModuleImpl : public SPIRVModule {
   template <class T> T *addType(T *Ty);
   SPIRVTypeArray *addArrayType(SPIRVType *, SPIRVConstant *) override;
   SPIRVTypeBool *addBoolType() override;
-  SPIRVTypeFloat *addFloatType(unsigned BitWidth) override;
+  SPIRVTypeFloat *addFloatType(unsigned BitWidth,
+                               unsigned FloatingPointEncoding) override;
   SPIRVTypeFunction *addFunctionType(SPIRVType *,
                                      const std::vector<SPIRVType *> &) override;
   SPIRVTypeInt *addIntegerType(unsigned BitWidth) override;
@@ -568,7 +569,8 @@ class SPIRVModuleImpl : public SPIRVModule {
   SPIRVTypeBool *BoolTy;
   SPIRVTypeVoid *VoidTy;
   SmallDenseMap<unsigned, SPIRVTypeInt *, 4> IntTypeMap;
-  SmallDenseMap<unsigned, SPIRVTypeFloat *, 4> FloatTypeMap;
+  SmallDenseMap<std::pair<unsigned, unsigned>, SPIRVTypeFloat *, 4>
+      FloatTypeMap;
   std::map<unsigned, SPIRVConstant *> LiteralMap;
   std::vector<SPIRVExtInst *> DebugInstVec;
   std::vector<SPIRVExtInst *> AuxDataInstVec;
@@ -993,12 +995,14 @@ SPIRVTypeInt *SPIRVModuleImpl::addIntegerType(unsigned BitWidth) {
   return addType(Ty);
 }
 
-SPIRVTypeFloat *SPIRVModuleImpl::addFloatType(unsigned BitWidth) {
-  auto Loc = FloatTypeMap.find(BitWidth);
+SPIRVTypeFloat *SPIRVModuleImpl::addFloatType(unsigned BitWidth,
+                                              unsigned FloatingPointEncoding) {
+  auto Desc = std::make_pair(BitWidth, FloatingPointEncoding);
+  auto Loc = FloatTypeMap.find(Desc);
   if (Loc != FloatTypeMap.end())
     return Loc->second;
-  auto *Ty = new SPIRVTypeFloat(this, getId(), BitWidth);
-  FloatTypeMap[BitWidth] = Ty;
+  auto *Ty = new SPIRVTypeFloat(this, getId(), BitWidth, FloatingPointEncoding);
+  FloatTypeMap[Desc] = Ty;
   return addType(Ty);
 }
 

lib/SPIRV/libSPIRV/SPIRVModule.h

@@ -242,7 +242,7 @@ class SPIRVModule {
   // Type creation functions
   virtual SPIRVTypeArray *addArrayType(SPIRVType *, SPIRVConstant *) = 0;
   virtual SPIRVTypeBool *addBoolType() = 0;
-  virtual SPIRVTypeFloat *addFloatType(unsigned) = 0;
+  virtual SPIRVTypeFloat *addFloatType(unsigned, unsigned = FPEncodingMax) = 0;
   virtual SPIRVTypeFunction *
   addFunctionType(SPIRVType *, const std::vector<SPIRVType *> &) = 0;
   virtual SPIRVTypeImage *addImageType(SPIRVType *,

lib/SPIRV/libSPIRV/SPIRVNameMapEnum.h

@@ -640,6 +640,9 @@ template <> inline void SPIRVMap<Capability, std::string>::init() {
   add(CapabilityFPMaxErrorINTEL, "FPMaxErrorINTEL");
   add(CapabilityCacheControlsINTEL, "CacheControlsINTEL");
   add(CapabilityRegisterLimitsINTEL, "RegisterLimitsINTEL");
+  add(CapabilityBFloat16TypeKHR, "BFloat16TypeKHR");
+  add(CapabilityBFloat16DotProductKHR, "BFloat16DotProductKHR");
+  add(CapabilityBFloat16CooperativeMatrixKHR, "BFloat16CooperativeMatrixKHR");
   add(CapabilityArithmeticFenceEXT, "ArithmeticFenceEXT");
   add(CapabilitySubgroup2DBlockIOINTEL, "Subgroup2DBlockIOINTEL");
   add(CapabilitySubgroup2DBlockTransformINTEL, "Subgroup2DBlockTransformINTEL");

lib/SPIRV/libSPIRV/SPIRVType.cpp

@@ -162,8 +162,15 @@ bool SPIRVType::isTypeComposite() const {
          isTypeJointMatrixINTEL() || isTypeCooperativeMatrixKHR();
 }
 
-bool SPIRVType::isTypeFloat(unsigned Bits) const {
-  return isType<SPIRVTypeFloat>(this, Bits);
+bool SPIRVType::isTypeFloat(unsigned Bits,
+                            unsigned FloatingPointEncoding) const {
+  if (!isType<SPIRVTypeFloat>(this))
+    return false;
+  if (Bits == 0)
+    return true;
+  const auto *ThisFloat = static_cast<const SPIRVTypeFloat *>(this);
+  return ThisFloat->getBitWidth() == Bits &&
+         ThisFloat->getFloatingPointEncoding() == FloatingPointEncoding;
 }
 
 bool SPIRVType::isTypeOCLImage() const {

lib/SPIRV/libSPIRV/SPIRVType.h

@@ -84,7 +84,8 @@ class SPIRVType : public SPIRVEntry {
   bool isTypeEvent() const;
   bool isTypeDeviceEvent() const;
   bool isTypeReserveId() const;
-  bool isTypeFloat(unsigned Bits = 0) const;
+  bool isTypeFloat(unsigned Bits = 0,
+                   unsigned FloatingPointEncoding = FPEncodingMax) const;
   bool isTypeImage() const;
   bool isTypeOCLImage() const;
   bool isTypePipe() const;
@@ -201,16 +202,31 @@ class SPIRVTypeFloat : public SPIRVType {
 public:
   static const Op OC = OpTypeFloat;
   // Complete constructor
-  SPIRVTypeFloat(SPIRVModule *M, SPIRVId TheId, unsigned TheBitWidth)
-      : SPIRVType(M, 3, OC, TheId), BitWidth(TheBitWidth) {}
+  SPIRVTypeFloat(SPIRVModule *M, SPIRVId TheId, unsigned TheBitWidth,
+                 unsigned TheFloatingPointEncoding)
+      : SPIRVType(M, 3 + (TheFloatingPointEncoding != FPEncodingMax), OC,
+                  TheId),
+        BitWidth(TheBitWidth), FloatingPointEncoding(TheFloatingPointEncoding) {
+  }
   // Incomplete constructor
-  SPIRVTypeFloat() : SPIRVType(OC), BitWidth(0) {}
+  SPIRVTypeFloat()
+      : SPIRVType(OC), BitWidth(0), FloatingPointEncoding(FPEncodingMax) {}
 
   unsigned getBitWidth() const { return BitWidth; }
 
+  unsigned getFloatingPointEncoding() const { return FloatingPointEncoding; }
+
+  std::optional<ExtensionID> getRequiredExtension() const override {
+    if (isTypeFloat(16, FPEncodingBFloat16KHR))
+      return ExtensionID::SPV_KHR_bfloat16;
+    return {};
+  }
+
   SPIRVCapVec getRequiredCapability() const override {
     SPIRVCapVec CV;
-    if (isTypeFloat(16)) {
+    if (isTypeFloat(16, FPEncodingBFloat16KHR)) {
+      CV.push_back(CapabilityBFloat16TypeKHR);
+    } else if (isTypeFloat(16)) {
       CV.push_back(CapabilityFloat16Buffer);
       auto Extensions = getModule()->getSourceExtension();
       if (std::any_of(Extensions.begin(), Extensions.end(),
@@ -222,14 +238,34 @@ class SPIRVTypeFloat : public SPIRVType {
   }
 
 protected:
-  _SPIRV_DEF_ENCDEC2(Id, BitWidth)
+  void encode(spv_ostream &O) const override {
+    assert(WordCount == 3 || WordCount == 4);
+    auto Encoder = getEncoder(O);
+    Encoder << Id << BitWidth;
+    if (WordCount == 4)
+      Encoder << FloatingPointEncoding;
+  }
+
+  void decode(std::istream &I) override {
+    assert(WordCount == 3 || WordCount == 4);
+    auto Decoder = getDecoder(I);
+    Decoder >> Id >> BitWidth;
+    if (WordCount == 4)
+      Decoder >> FloatingPointEncoding;
+  }
+
   void validate() const override {
     SPIRVEntry::validate();
     assert(BitWidth >= 16 && BitWidth <= 64 && "Invalid bit width");
+    assert(
+        (FloatingPointEncoding == FPEncodingMax ||
+         (BitWidth == 16 && FloatingPointEncoding == FPEncodingBFloat16KHR)) &&
+        "Invalid floating point encoding");
   }
 
 private:
   unsigned BitWidth; // Bit width
+  unsigned FloatingPointEncoding;
 };
 
 class SPIRVTypePointer : public SPIRVType {
@@ -1142,7 +1178,10 @@ class SPIRVTypeCooperativeMatrixKHR : public SPIRVType {
     return ExtensionID::SPV_KHR_cooperative_matrix;
   }
   SPIRVCapVec getRequiredCapability() const override {
-    return getVec(CapabilityCooperativeMatrixKHR);
+    auto CV = getVec(CapabilityCooperativeMatrixKHR);
+    if (CompType->isTypeFloat(16, FPEncodingBFloat16KHR))
+      CV.push_back(CapabilityBFloat16CooperativeMatrixKHR);
+    return CV;
   }
 
   SPIRVType *getCompType() const { return CompType; }

test/extensions/KHR/SPV_KHR_bfloat16/bfloat16.ll

@@ -0,0 +1,36 @@
+; RUN: llvm-as %s -o %t.bc
+; RUN: llvm-spirv %t.bc --spirv-ext=+SPV_KHR_bfloat16 -o %t.spv
+; RUN: llvm-spirv %t.spv -to-text -o - | FileCheck %s --check-prefix=CHECK-SPIRV
+
+; RUN: llvm-spirv -r %t.spv -o %t.rev.bc
+; RUN: llvm-dis %t.rev.bc -o - | FileCheck %s --check-prefix=CHECK-LLVM
+
+; RUN: not llvm-spirv %t.bc 2>&1 | FileCheck %s --check-prefix=CHECK-ERROR
+
+; CHECK-ERROR: RequiresExtension: Feature requires the following SPIR-V extension:
+; CHECK-ERROR-NEXT: SPV_KHR_bfloat16
+; CHECK-ERROR-NEXT: NOTE: LLVM module contains bfloat type, translation of which
+; CHECK-ERROR-SAME: requires this extension
+
+source_filename = "bfloat16.cpp"
+target datalayout = "e-i64:64-v16:16-v24:32-v32:32-v48:64-v96:128-v192:256-v256:256-v512:512-v1024:1024-n8:16:32:64"
+target triple = "spirv64-unknown-unknown"
+
+; CHECK-SPIRV-DAG: Capability BFloat16TypeKHR
+; CHECK-SPIRV-DAG: Extension "SPV_KHR_bfloat16"
+; CHECK-SPIRV: 4 TypeFloat [[BFLOAT:[0-9]+]] 16 0
+; CHECK-SPIRV: 4 TypeVector [[#]] [[BFLOAT]] 2
+
+; CHECK-LLVM: [[ADDR1:]] = alloca bfloat
+; CHECK-LLVM: [[ADDR2:]] = alloca <2 x bfloat>
+; CHECK-LLVM: [[DATA1:]] = load bfloat, ptr [[ADDR1]]
+; CHECK-LLVM: [[DATA2:]] = load <2 x bfloat>, ptr [[ADDR2]]
+
+define spir_kernel void @test() {
+entry:
+  %addr1 = alloca bfloat
+  %addr2 = alloca <2 x bfloat>
+  %data1 = load bfloat, ptr %addr1
+  %data2 = load <2 x bfloat>, ptr %addr2
+  ret void
+}

test/extensions/KHR/SPV_KHR_bfloat16/bfloat16_dot.ll

@@ -0,0 +1,39 @@
+; RUN: llvm-as %s -o %t.bc
+; RUN: llvm-spirv %t.bc --spirv-ext=+SPV_KHR_bfloat16 -o %t.spv
+; RUN: llvm-spirv %t.spv -to-text -o - | FileCheck %s --check-prefix=CHECK-SPIRV
+
+; RUN: llvm-spirv -r %t.spv -o %t.rev.bc
+; RUN: llvm-dis %t.rev.bc -o - | FileCheck %s --check-prefix=CHECK-LLVM
+
+source_filename = "bfloat16_dot.cpp"
+target datalayout = "e-i64:64-v16:16-v24:32-v32:32-v48:64-v96:128-v192:256-v256:256-v512:512-v1024:1024-n8:16:32:64"
+target triple = "spirv64-unknown-unknown"
+
+; CHECK-SPIRV-DAG: Capability BFloat16TypeKHR
+; CHECK-SPIRV-DAG: Capability BFloat16DotProductKHR
+; CHECK-SPIRV-DAG: Extension "SPV_KHR_bfloat16"
+; CHECK-SPIRV: 4 TypeFloat [[BFLOAT:[0-9]+]] 16 0
+; CHECK-SPIRV: 4 TypeVector [[#]] [[BFLOAT]] 2
+; CHECK-SPIRV: Dot
+
+; CHECK-LLVM: %addrA = alloca <2 x bfloat>
+; CHECK-LLVM: %addrB = alloca <2 x bfloat>
+; CHECK-LLVM: %dataA = load <2 x bfloat>, ptr %addrA
+; CHECK-LLVM: %dataB = load <2 x bfloat>, ptr %addrB
+; CHECK-LLVM: %call = call spir_func bfloat @_Z3dotDv2_u6__bf16S_(<2 x bfloat> %dataA, <2 x bfloat> %dataB)
+
+declare spir_func bfloat @_Z3dotDv2_u6__bf16Dv2_S_(<2 x bfloat>, <2 x bfloat>)
+
+define spir_kernel void @test() {
+entry:
+  %addrA = alloca <2 x bfloat>
+  %addrB = alloca <2 x bfloat>
+  %dataA = load <2 x bfloat>, ptr %addrA
+  %dataB = load <2 x bfloat>, ptr %addrB
+  %call = call spir_func bfloat @_Z3dotDv2_u6__bf16Dv2_S_(<2 x bfloat> %dataA, <2 x bfloat> %dataB)
+  ret void
+}
+
+!opencl.ocl.version = !{!7}
+
+!7 = !{i32 2, i32 0}

test/extensions/KHR/SPV_KHR_bfloat16/cooperative_matrix_bfloat16.ll

@@ -0,0 +1,34 @@
+; RUN: llvm-as < %s -o %t.bc
+; RUN: llvm-spirv %t.bc --spirv-ext=+SPV_KHR_cooperative_matrix,+SPV_KHR_bfloat16 -o %t.spv
+; RUN: llvm-spirv %t.spv -to-text -o %t.spt
+; RUN: FileCheck < %t.spt %s --check-prefix=CHECK-SPIRV
+
+; RUN: llvm-spirv --spirv-ext=+SPV_KHR_cooperative_matrix,+SPV_KHR_bfloat16 -r %t.spv -o %t.rev.bc
+; RUN: llvm-dis < %t.rev.bc | FileCheck %s --check-prefix=CHECK-LLVM
+
+; CHECK-SPIRV-DAG: Capability CooperativeMatrixKHR
+; CHECK-SPIRV-DAG: Capability BFloat16TypeKHR
+; CHECK-SPIRV-DAG: Capability BFloat16CooperativeMatrixKHR
+; CHECK-SPIRV-DAG: Extension "SPV_KHR_cooperative_matrix"
+; CHECK-SPIRV-DAG: Extension "SPV_KHR_bfloat16"
+
+; CHECK-SPIRV-DAG: 4 TypeFloat [[#BFloatTy:]] 16 0
+; CHECK-SPIRV-DAG: TypeInt [[#Int32Ty:]] 32 0
+; CHECK-SPIRV-DAG: Constant [[#Int32Ty]] [[#Const12:]] 12
+; CHECK-SPIRV-DAG: Constant [[#Int32Ty]] [[#Const3:]] 3
+; CHECK-SPIRV-DAG: Constant [[#Int32Ty]] [[#Const2:]] 2
+; CHECK-SPIRV-DAG: TypeCooperativeMatrixKHR [[#MatTy:]] [[#BFloatTy]] [[#Const3]] [[#Const12]] [[#Const12]] [[#Const2]]
+; CHECK-SPIRV-DAG: Constant [[#BFloatTy]] [[#]] 16256
+; CHECK-SPIRV: CompositeConstruct [[#MatTy]]
+
+; CHECK-LLVM: call spir_func target("spirv.CooperativeMatrixKHR", bfloat, 3, 12, 12, 2) @_Z26__spirv_CompositeConstructu6__bf16(bfloat 0xR3F80)
+
+target datalayout = "e-i64:64-v16:16-v24:32-v32:32-v48:64-v96:128-v192:256-v256:256-v512:512-v1024:1024-n8:16:32:64"
+target triple = "spir64-unknown-unknown"
+
+declare spir_func target("spirv.CooperativeMatrixKHR", bfloat, 3, 12, 12, 2) @_Z26__spirv_CompositeConstructu6__bf16(bfloat)
+
+define spir_kernel void @test() {
+  %mat = call spir_func target("spirv.CooperativeMatrixKHR", bfloat, 3, 12, 12, 2) @_Z26__spirv_CompositeConstructu6__bf16(bfloat 1.0)
+  ret void
+}