[Bacport to 21] Implement SPV_INTEL_bfloat16_arithmetic (#3290)

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

@@ -77,6 +77,7 @@ 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)
 EXT(SPV_INTEL_ternary_bitwise_function)
 EXT(SPV_INTEL_int4)
 EXT(SPV_INTEL_function_variants)

lib/SPIRV/SPIRVUtil.cpp

@@ -525,6 +525,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)) {
@@ -1913,6 +1918,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)))) {
@@ -1929,6 +1937,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

@@ -4201,6 +4201,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
@@ -5485,6 +5499,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

@@ -916,6 +916,8 @@ class SPIRVCapability : public SPIRVEntryNoId<OpCapability> {
     case CapabilityFunctionVariantsINTEL:
     case CapabilitySpecConditionalINTEL:
       return ExtensionID::SPV_INTEL_function_variants;
+    case internal::CapabilityBFloat16ArithmeticINTEL:
+      return ExtensionID::SPV_INTEL_bfloat16_arithmetic;
     default:
       return {};
     }

lib/SPIRV/libSPIRV/SPIRVEnum.h

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

lib/SPIRV/libSPIRV/SPIRVModule.cpp

@@ -1732,6 +1732,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),
@@ -1755,6 +1757,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

@@ -690,6 +690,7 @@ template <> inline void SPIRVMap<Capability, std::string>::init() {
   add(CapabilityInt4CooperativeMatrixINTEL, "Int4CooperativeMatrixINTEL");
   add(CapabilityFunctionVariantsINTEL, "FunctionVariantsINTEL");
   add(CapabilitySpecConditionalINTEL, "SpecConditionalINTEL");
+  add(internal::CapabilityBFloat16ArithmeticINTEL, "BFloat16ArithmeticINTEL");
 }
 SPIRV_DEF_NAMEMAP(Capability, SPIRVCapabilityNameMap)
 

lib/SPIRV/libSPIRV/spirv_internal.hpp

@@ -106,6 +106,7 @@ enum InternalCapability {
   ICapGlobalVariableDecorationsINTEL = 6146,
   ICapabilityTaskSequenceINTEL = 6162,
   ICapabilityCooperativeMatrixCheckedInstructionsINTEL = 6192,
+  ICapabilityBFloat16ArithmeticINTEL = 6226,
   ICapabilityCooperativeMatrixOffsetInstructionsINTEL = 6238,
   ICapabilityAtomicBFloat16AddINTEL = 6255,
   ICapabilityAtomicBFloat16MinMaxINTEL = 6256,
@@ -268,6 +269,8 @@ constexpr Capability CapabilityTokenTypeINTEL =
     static_cast<Capability>(ICapTokenTypeINTEL);
 constexpr Capability CapabilityGlobalVariableDecorationsINTEL =
     static_cast<Capability>(ICapGlobalVariableDecorationsINTEL);
+constexpr Capability CapabilityBFloat16ArithmeticINTEL =
+    static_cast<Capability>(ICapabilityBFloat16ArithmeticINTEL);
 
 constexpr ExecutionMode ExecutionModeNamedSubgroupSizeINTEL =
     static_cast<ExecutionMode>(IExecModeNamedSubgroupSizeINTEL);