[Backport to 14] Implement SPV_INTEL_bfloat16_arithmetic (KhronosGroup#3290) (KhronosGroup#3320) (KhronosGroup#3340)

The extension relaxes rules for bf16 type allowing to use it in some
arithmetic operations.

Spec is available here:
intel/llvm#18352

Co-authered by: Michael Aziz <michael.aziz@intel.com>

---------

Signed-off-by: Sidorov, Dmitry <dmitry.sidorov@intel.com>

Author: Dmitry Sidorov

include/LLVMSPIRVExtensions.inc

@@ -72,3 +72,4 @@ EXT(SPV_INTEL_bindless_images)
 EXT(SPV_INTEL_2d_block_io)
 EXT(SPV_INTEL_subgroup_matrix_multiply_accumulate)
 EXT(SPV_KHR_bfloat16)
+EXT(SPV_INTEL_bfloat16_arithmetic)

lib/SPIRV/SPIRVUtil.cpp

@@ -618,6 +618,11 @@ ParamType lastFuncParamType(StringRef MangledName) {
   char Mangled = Copy.back();
   std::string Mangled2 = Copy.substr(Copy.size() - 2);
 
+  std::string Mangled6 = Copy.substr(Copy.size() - 6);
+  if (Mangled6 == "__bf16") {
+    return ParamType::FLOAT;
+  }
+
   if (isMangledTypeFP(Mangled) || isMangledTypeHalf(Mangled2)) {
     return ParamType::FLOAT;
   } else if (isMangledTypeUnsigned(Mangled)) {
@@ -1847,6 +1852,9 @@ bool checkTypeForSPIRVExtendedInstLowering(IntrinsicInst *II, SPIRVModule *BM) {
       NumElems = VecTy->getNumElements();
       Ty = VecTy->getElementType();
     }
+    if (Ty->isBFloatTy() &&
+        BM->hasCapability(internal::CapabilityBFloat16ArithmeticINTEL))
+      return true;
     if ((!Ty->isFloatTy() && !Ty->isDoubleTy() && !Ty->isHalfTy()) ||
         (!BM->hasCapability(CapabilityVectorAnyINTEL) &&
          ((NumElems > 4) && (NumElems != 8) && (NumElems != 16)))) {
@@ -1863,6 +1871,9 @@ bool checkTypeForSPIRVExtendedInstLowering(IntrinsicInst *II, SPIRVModule *BM) {
       NumElems = VecTy->getNumElements();
       Ty = VecTy->getElementType();
     }
+    if (Ty->isBFloatTy() &&
+        BM->hasCapability(internal::CapabilityBFloat16ArithmeticINTEL))
+      return true;
     if ((!Ty->isIntegerTy()) ||
         (!BM->hasCapability(CapabilityVectorAnyINTEL) &&
          ((NumElems > 4) && (NumElems != 8) && (NumElems != 16)))) {

lib/SPIRV/SPIRVWriter.cpp

@@ -3630,6 +3630,20 @@ SPIRVValue *LLVMToSPIRVBase::transIntrinsicInst(IntrinsicInst *II,
   // -spirv-allow-unknown-intrinsics work correctly.
   auto IID = II->getIntrinsicID();
   switch (IID) {
+  case Intrinsic::fabs:
+  case Intrinsic::fma:
+  case Intrinsic::maxnum:
+  case Intrinsic::minnum:
+  case Intrinsic::fmuladd: {
+    Type *Ty = II->getType();
+    if (Ty->isBFloatTy())
+      BM->addCapability(internal::CapabilityBFloat16ArithmeticINTEL);
+    break;
+  }
+  default:
+    break;
+  }
+  switch (IID) {
   case Intrinsic::assume: {
     // llvm.assume translation is currently supported only within
     // SPV_KHR_expect_assume extension, ignore it otherwise, since it's
@@ -4413,6 +4427,11 @@ SPIRVValue *LLVMToSPIRVBase::transDirectCallInst(CallInst *CI,
   SmallVector<std::string, 2> Dec;
   if (isBuiltinTransToExtInst(CI->getCalledFunction(), &ExtSetKind, &ExtOp,
                               &Dec)) {
+    if (const auto *FirstArg = F->getArg(0)) {
+      const auto *Type = FirstArg->getType();
+      if (Type->isBFloatTy())
+        BM->addCapability(internal::CapabilityBFloat16ArithmeticINTEL);
+    }
     if (DemangledName.find("__spirv_ocl_printf") != StringRef::npos) {
       auto *FormatStrPtr = cast<PointerType>(CI->getArgOperand(0)->getType());
       if (FormatStrPtr->getAddressSpace() !=

lib/SPIRV/libSPIRV/SPIRVEntry.h

@@ -904,6 +904,8 @@ class SPIRVCapability : public SPIRVEntryNoId<OpCapability> {
     case CapabilityVectorComputeINTEL:
     case CapabilityVectorAnyINTEL:
       return ExtensionID::SPV_INTEL_vector_compute;
+    case internal::CapabilityBFloat16ArithmeticINTEL:
+      return ExtensionID::SPV_INTEL_bfloat16_arithmetic;
     default:
       return {};
     }

lib/SPIRV/libSPIRV/SPIRVEnum.h

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

lib/SPIRV/libSPIRV/SPIRVModule.cpp

@@ -1543,6 +1543,8 @@ SPIRVInstruction *SPIRVModuleImpl::addBinaryInst(Op TheOpCode, SPIRVType *Type,
                                                  SPIRVValue *Op1,
                                                  SPIRVValue *Op2,
                                                  SPIRVBasicBlock *BB) {
+  if (Type->isTypeFloat(16, FPEncodingBFloat16KHR) && TheOpCode != OpDot)
+    addCapability(internal::CapabilityBFloat16ArithmeticINTEL);
   return addInstruction(SPIRVInstTemplateBase::create(
                             TheOpCode, Type, getId(),
                             getVec(Op1->getId(), Op2->getId()), BB, this),
@@ -1566,6 +1568,8 @@ SPIRVInstruction *SPIRVModuleImpl::addUnaryInst(Op TheOpCode,
                                                 SPIRVType *TheType,
                                                 SPIRVValue *Op,
                                                 SPIRVBasicBlock *BB) {
+  if (TheType->isTypeFloat(16, FPEncodingBFloat16KHR) && TheOpCode != OpDot)
+    addCapability(internal::CapabilityBFloat16ArithmeticINTEL);
   return addInstruction(
       SPIRVInstTemplateBase::create(TheOpCode, TheType, getId(),
                                     getVec(Op->getId()), BB, this),

lib/SPIRV/libSPIRV/SPIRVNameMapEnum.h

@@ -646,6 +646,7 @@ template <> inline void SPIRVMap<Capability, std::string>::init() {
   add(internal::CapabilityCooperativeMatrixCheckedInstructionsINTEL,
       "CooperativeMatrixCheckedInstructionsINTEL");
   add(internal::CapabilityBindlessImagesINTEL, "BindlessImagesINTEL");
+  add(internal::CapabilityBFloat16ArithmeticINTEL, "BFloat16ArithmeticINTEL");
 }
 SPIRV_DEF_NAMEMAP(Capability, SPIRVCapabilityNameMap)
 

lib/SPIRV/libSPIRV/spirv_internal.hpp

@@ -114,6 +114,7 @@ enum InternalCapability {
   ICapFPArithmeticFenceINTEL = 6144,
   ICapGlobalVariableDecorationsINTEL = 6146,
   ICapabilityCooperativeMatrixCheckedInstructionsINTEL = 6192,
+  ICapabilityBFloat16ArithmeticINTEL = 6226,
   ICapabilityCooperativeMatrixPrefetchINTEL = 6411,
   ICapabilityComplexFloatMulDivINTEL = 6414,
   ICapabilityTensorFloat32RoundingINTEL = 6425,
@@ -304,6 +305,8 @@ constexpr Capability CapabilityGlobalVariableDecorationsINTEL =
     static_cast<Capability>(ICapGlobalVariableDecorationsINTEL);
 constexpr Capability CapabilityRegisterLimitsINTEL =
     static_cast<Capability>(ICapRegisterLimitsINTEL);
+constexpr Capability CapabilityBFloat16ArithmeticINTEL =
+    static_cast<Capability>(ICapabilityBFloat16ArithmeticINTEL);
 
 constexpr FunctionControlMask FunctionControlOptNoneINTELMask =
     static_cast<FunctionControlMask>(IFunctionControlOptNoneINTELMask);