Add SPV_INTEL_masked_gather_scatter extension (#1580)

This extension allows TypeVector to have a Physical Pointer Type
Component Type and introduces gather/scatter instructions.
It will be useful for explicitly vectorized kernels.

Spec: intel/llvm#6613

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

Author: MrSidims

include/LLVMSPIRVExtensions.inc

@@ -54,3 +54,4 @@ EXT(SPV_INTEL_global_variable_decorations)
 EXT(SPV_INTEL_non_constant_addrspace_printf)
 EXT(SPV_INTEL_complex_float_mul_div)
 EXT(SPV_INTEL_split_barrier)
+EXT(SPV_INTEL_masked_gather_scatter)

lib/SPIRV/SPIRVReader.cpp

@@ -2130,7 +2130,12 @@ Value *SPIRVToLLVM::transValueWithoutDecoration(SPIRVValue *BV, Function *F,
   case OpInBoundsPtrAccessChain: {
     auto AC = static_cast<SPIRVAccessChainBase *>(BV);
     auto Base = transValue(AC->getBase(), F, BB);
-    Type *BaseTy = transType(AC->getBase()->getType()->getPointerElementType());
+    SPIRVType *BaseSPVTy = AC->getBase()->getType();
+    Type *BaseTy =
+        BaseSPVTy->isTypeVector()
+            ? transType(
+                  BaseSPVTy->getVectorComponentType()->getPointerElementType())
+            : transType(BaseSPVTy->getPointerElementType());
     auto Index = transValue(AC->getIndices(), F, BB);
     if (!AC->hasPtrIndex())
       Index.insert(Index.begin(), getInt32(M, 0));
@@ -2543,6 +2548,29 @@ Value *SPIRVToLLVM::transValueWithoutDecoration(SPIRVValue *BV, Function *F,
     return mapValue(
         BV, Builder.CreateIntrinsic(Intrinsic::arithmetic_fence, RetTy, Val));
   }
+  case internal::OpMaskedGatherINTEL: {
+    IRBuilder<> Builder(BB);
+    auto *Inst = static_cast<SPIRVMaskedGatherINTELInst *>(BV);
+    Type *RetTy = transType(Inst->getType());
+    Value *PtrVector = transValue(Inst->getOperand(0), F, BB);
+    uint32_t Alignment = Inst->getOpWord(1);
+    Value *Mask = transValue(Inst->getOperand(2), F, BB);
+    Value *FillEmpty = transValue(Inst->getOperand(3), F, BB);
+    return mapValue(BV, Builder.CreateMaskedGather(RetTy, PtrVector,
+                                                   Align(Alignment), Mask,
+                                                   FillEmpty));
+  }
+
+  case internal::OpMaskedScatterINTEL: {
+    IRBuilder<> Builder(BB);
+    auto *Inst = static_cast<SPIRVMaskedScatterINTELInst *>(BV);
+    Value *InputVector = transValue(Inst->getOperand(0), F, BB);
+    Value *PtrVector = transValue(Inst->getOperand(1), F, BB);
+    uint32_t Alignment = Inst->getOpWord(2);
+    Value *Mask = transValue(Inst->getOperand(3), F, BB);
+    return mapValue(BV, Builder.CreateMaskedScatter(InputVector, PtrVector,
+                                                    Align(Alignment), Mask));
+  }
 
   default: {
     auto OC = BV->getOpCode();
@@ -3150,7 +3178,11 @@ std::string getSPIRVFuncSuffix(SPIRVInstruction *BI) {
   }
   if (BI->getOpCode() == OpGenericCastToPtrExplicit) {
     Suffix += kSPIRVPostfix::Divider;
-    auto GenericCastToPtrInst = BI->getType()->getPointerStorageClass();
+    auto *Ty = BI->getType();
+    auto GenericCastToPtrInst =
+        Ty->isTypeVectorPointer()
+            ? Ty->getVectorComponentType()->getPointerStorageClass()
+            : Ty->getPointerStorageClass();
     switch (GenericCastToPtrInst) {
     case StorageClassCrossWorkgroup:
       Suffix += std::string(kSPIRVPostfix::ToGlobal);

lib/SPIRV/SPIRVRegularizeLLVM.cpp

@@ -543,11 +543,19 @@ bool SPIRVRegularizeLLVMBase::regularize() {
         // Add an additional bitcast in case address space cast also changes
         // pointer element type.
         if (auto *ASCast = dyn_cast<AddrSpaceCastInst>(&II)) {
-          PointerType *DestTy = cast<PointerType>(ASCast->getDestTy());
-          PointerType *SrcTy = cast<PointerType>(ASCast->getSrcTy());
-          if (!DestTy->hasSameElementTypeAs(SrcTy)) {
-            PointerType *InterTy = PointerType::getWithSamePointeeType(
-                DestTy, SrcTy->getPointerAddressSpace());
+          Type *DestTy = ASCast->getDestTy();
+          Type *SrcTy = ASCast->getSrcTy();
+          if (!II.getContext().supportsTypedPointers())
+            continue;
+          if (DestTy->getScalarType()->getNonOpaquePointerElementType() !=
+              SrcTy->getScalarType()->getNonOpaquePointerElementType()) {
+            Type *InterTy = PointerType::getWithSamePointeeType(
+                cast<PointerType>(DestTy->getScalarType()),
+                cast<PointerType>(SrcTy->getScalarType())
+                    ->getPointerAddressSpace());
+            if (DestTy->isVectorTy())
+              InterTy = VectorType::get(
+                  InterTy, cast<VectorType>(DestTy)->getElementCount());
             BitCastInst *NewBCast = new BitCastInst(
                 ASCast->getPointerOperand(), InterTy, /*NameStr=*/"", ASCast);
             AddrSpaceCastInst *NewASCast =

lib/SPIRV/SPIRVWriter.cpp

@@ -342,9 +342,28 @@ SPIRVType *LLVMToSPIRVBase::transType(Type *T) {
     return transPointerType(ET, AddrSpc);
   }
 
-  if (auto *VecTy = dyn_cast<FixedVectorType>(T))
+  if (auto *VecTy = dyn_cast<FixedVectorType>(T)) {
+    if (VecTy->getElementType()->isPointerTy()) {
+      // SPV_INTEL_masked_gather_scatter extension changes 2.16.1. Universal
+      // Validation Rules:
+      // Vector types must be parameterized only with numerical types,
+      // [Physical Pointer Type] types or the [OpTypeBool] type.
+      // Without it vector of pointers is not allowed in SPIR-V.
+      if (!BM->isAllowedToUseExtension(
+              ExtensionID::SPV_INTEL_masked_gather_scatter)) {
+        BM->getErrorLog().checkError(
+            false, SPIRVEC_RequiresExtension,
+            "SPV_INTEL_masked_gather_scatter\n"
+            "NOTE: LLVM module contains vector of pointers, translation "
+            "of which requires this extension");
+        return nullptr;
+      }
+      BM->addExtension(ExtensionID::SPV_INTEL_masked_gather_scatter);
+      BM->addCapability(internal::CapabilityMaskedGatherScatterINTEL);
+    }
     return mapType(T, BM->addVectorType(transType(VecTy->getElementType()),
                                         VecTy->getNumElements()));
+  }
 
   if (T->isArrayTy()) {
     // SPIR-V 1.3 s3.32.6: Length is the number of elements in the array.
@@ -3835,6 +3854,47 @@ SPIRVValue *LLVMToSPIRVBase::transIntrinsicInst(IntrinsicInst *II,
     }
     return Op;
   }
+  case Intrinsic::masked_gather: {
+    if (!BM->isAllowedToUseExtension(
+            ExtensionID::SPV_INTEL_masked_gather_scatter)) {
+      BM->getErrorLog().checkError(
+          BM->isUnknownIntrinsicAllowed(II), SPIRVEC_InvalidFunctionCall, II,
+          "Translation of llvm.masked.gather intrinsic requires "
+          "SPV_INTEL_masked_gather_scatter extension or "
+          "-spirv-allow-unknown-intrinsics option.");
+      return nullptr;
+    }
+    SPIRVType *Ty = transType(II->getType());
+    auto *PtrVector = transValue(II->getArgOperand(0), BB);
+    uint32_t Alignment =
+        cast<ConstantInt>(II->getArgOperand(1))->getZExtValue();
+    auto *Mask = transValue(II->getArgOperand(2), BB);
+    auto *FillEmpty = transValue(II->getArgOperand(3), BB);
+    std::vector<SPIRVWord> Ops = {PtrVector->getId(), Alignment, Mask->getId(),
+                                  FillEmpty->getId()};
+    return BM->addInstTemplate(internal::OpMaskedGatherINTEL, Ops, BB, Ty);
+  }
+  case Intrinsic::masked_scatter: {
+    if (!BM->isAllowedToUseExtension(
+            ExtensionID::SPV_INTEL_masked_gather_scatter)) {
+      BM->getErrorLog().checkError(
+          BM->isUnknownIntrinsicAllowed(II), SPIRVEC_InvalidFunctionCall, II,
+          "Translation of llvm.masked.scatter intrinsic requires "
+          "SPV_INTEL_masked_gather_scatter extension or "
+          "-spirv-allow-unknown-intrinsics option.");
+      return nullptr;
+    }
+    auto *InputVector = transValue(II->getArgOperand(0), BB);
+    auto *PtrVector = transValue(II->getArgOperand(1), BB);
+    uint32_t Alignment =
+        cast<ConstantInt>(II->getArgOperand(2))->getZExtValue();
+    auto *Mask = transValue(II->getArgOperand(3), BB);
+    std::vector<SPIRVWord> Ops = {InputVector->getId(), PtrVector->getId(),
+                                  Alignment, Mask->getId()};
+    return BM->addInstTemplate(internal::OpMaskedScatterINTEL, Ops, BB,
+                               nullptr);
+  }
+
   default:
     if (BM->isUnknownIntrinsicAllowed(II))
       return BM->addCallInst(

lib/SPIRV/libSPIRV/SPIRVInstruction.h

@@ -3413,6 +3413,154 @@ class SPIRVComplexFloatInst
 _SPIRV_OP(ComplexFMulINTEL)
 _SPIRV_OP(ComplexFDivINTEL)
 #undef _SPIRV_OP
+
+class SPIRVMaskedGatherScatterINTELInstBase : public SPIRVInstTemplateBase {
+protected:
+  SPIRVCapVec getRequiredCapability() const override {
+    return getVec(internal::CapabilityMaskedGatherScatterINTEL);
+  }
+  llvm::Optional<ExtensionID> getRequiredExtension() const override {
+    return ExtensionID::SPV_INTEL_masked_gather_scatter;
+  }
+};
+
+class SPIRVMaskedGatherINTELInst
+    : public SPIRVMaskedGatherScatterINTELInstBase {
+  void validate() const override {
+    SPIRVInstruction::validate();
+    SPIRVErrorLog &SPVErrLog = this->getModule()->getErrorLog();
+    std::string InstName = "MaskedGatherINTEL";
+
+    SPIRVType *ResTy = this->getType();
+    SPVErrLog.checkError(ResTy->isTypeVector(), SPIRVEC_InvalidInstruction,
+                         InstName + "\nResult must be a vector type\n");
+    SPIRVWord ResCompCount = ResTy->getVectorComponentCount();
+    SPIRVType *ResCompTy = ResTy->getVectorComponentType();
+
+    SPIRVValue *PtrVec =
+        const_cast<SPIRVMaskedGatherINTELInst *>(this)->getOperand(0);
+    SPIRVType *PtrVecTy = PtrVec->getType();
+    SPVErrLog.checkError(
+        PtrVecTy->isTypeVectorPointer(), SPIRVEC_InvalidInstruction,
+        InstName + "\nPtrVector must be a vector of pointers type\n");
+    SPIRVWord PtrVecCompCount = PtrVecTy->getVectorComponentCount();
+    SPIRVType *PtrVecCompTy = PtrVecTy->getVectorComponentType();
+    SPIRVType *PtrElemTy = PtrVecCompTy->getPointerElementType();
+
+    SPVErrLog.checkError(
+        this->isOperandLiteral(1), SPIRVEC_InvalidInstruction,
+        InstName + "\nAlignment must be a constant expression integer\n");
+    const uint32_t Align =
+        static_cast<SPIRVConstant *>(
+            const_cast<SPIRVMaskedGatherINTELInst *>(this)->getOperand(2))
+            ->getZExtIntValue();
+    SPVErrLog.checkError(
+        ((Align & (Align - 1)) == 0), SPIRVEC_InvalidInstruction,
+        InstName + "\nAlignment must be 0 or power-of-two integer\n");
+
+    SPIRVValue *Mask =
+        const_cast<SPIRVMaskedGatherINTELInst *>(this)->getOperand(2);
+    SPIRVType *MaskTy = Mask->getType();
+    SPVErrLog.checkError(MaskTy->isTypeVector(), SPIRVEC_InvalidInstruction,
+                         InstName + "\nMask must be a vector type\n");
+    SPIRVType *MaskCompTy = MaskTy->getVectorComponentType();
+    SPVErrLog.checkError(MaskCompTy->isTypeBool(), SPIRVEC_InvalidInstruction,
+                         InstName + "\nMask must be a boolean vector type\n");
+    SPIRVWord MaskCompCount = MaskTy->getVectorComponentCount();
+
+    SPIRVValue *FillEmpty =
+        const_cast<SPIRVMaskedGatherINTELInst *>(this)->getOperand(3);
+    SPIRVType *FillEmptyTy = FillEmpty->getType();
+    SPVErrLog.checkError(FillEmptyTy->isTypeVector(),
+                         SPIRVEC_InvalidInstruction,
+                         InstName + "\nFillEmpty must be a vector type\n");
+    SPIRVWord FillEmptyCompCount = FillEmptyTy->getVectorComponentCount();
+    SPIRVType *FillEmptyCompTy = FillEmptyTy->getVectorComponentType();
+
+    SPVErrLog.checkError(
+        ResCompCount == PtrVecCompCount &&
+            PtrVecCompCount == FillEmptyCompCount &&
+            FillEmptyCompCount == MaskCompCount,
+        SPIRVEC_InvalidInstruction,
+        InstName + "\nResult, PtrVector, Mask and FillEmpty vectors must have "
+                   "the same size\n");
+
+    SPVErrLog.checkError(
+        ResCompTy == PtrElemTy && PtrElemTy == FillEmptyCompTy,
+        SPIRVEC_InvalidInstruction,
+        InstName + "\nComponent Type of Result and FillEmpty vector must be "
+                   "same as base type of PtrVector the same base type\n");
+  }
+};
+
+class SPIRVMaskedScatterINTELInst
+    : public SPIRVMaskedGatherScatterINTELInstBase {
+  void validate() const override {
+    SPIRVInstruction::validate();
+    SPIRVErrorLog &SPVErrLog = this->getModule()->getErrorLog();
+    std::string InstName = "MaskedScatterINTEL";
+
+    SPIRVValue *InputVec =
+        const_cast<SPIRVMaskedScatterINTELInst *>(this)->getOperand(0);
+    SPIRVType *InputVecTy = InputVec->getType();
+    SPVErrLog.checkError(
+        InputVecTy->isTypeVector(), SPIRVEC_InvalidInstruction,
+        InstName + "\nInputVector must be a vector of pointers type\n");
+    SPIRVWord InputVecCompCount = InputVecTy->getVectorComponentCount();
+    SPIRVType *InputVecCompTy = InputVecTy->getVectorComponentType();
+
+    SPIRVValue *PtrVec =
+        const_cast<SPIRVMaskedScatterINTELInst *>(this)->getOperand(1);
+    SPIRVType *PtrVecTy = PtrVec->getType();
+    SPVErrLog.checkError(
+        PtrVecTy->isTypeVectorPointer(), SPIRVEC_InvalidInstruction,
+        InstName + "\nPtrVector must be a vector of pointers type\n");
+    SPIRVWord PtrVecCompCount = PtrVecTy->getVectorComponentCount();
+    SPIRVType *PtrVecCompTy = PtrVecTy->getVectorComponentType();
+    SPIRVType *PtrElemTy = PtrVecCompTy->getPointerElementType();
+
+    SPVErrLog.checkError(
+        this->isOperandLiteral(2), SPIRVEC_InvalidInstruction,
+        InstName + "\nAlignment must be a constant expression integer\n");
+    const uint32_t Align =
+        static_cast<SPIRVConstant *>(
+            const_cast<SPIRVMaskedScatterINTELInst *>(this)->getOperand(2))
+            ->getZExtIntValue();
+    SPVErrLog.checkError(
+        ((Align & (Align - 1)) == 0), SPIRVEC_InvalidInstruction,
+        InstName + "\nAlignment must be 0 or power-of-two integer\n");
+
+    SPIRVValue *Mask =
+        const_cast<SPIRVMaskedScatterINTELInst *>(this)->getOperand(2);
+    SPIRVType *MaskTy = Mask->getType();
+    SPVErrLog.checkError(MaskTy->isTypeVector(), SPIRVEC_InvalidInstruction,
+                         InstName + "\nMask must be a vector type\n");
+    SPIRVType *MaskCompTy = MaskTy->getVectorComponentType();
+    SPVErrLog.checkError(MaskCompTy->isTypeBool(), SPIRVEC_InvalidInstruction,
+                         InstName + "\nMask must be a boolean vector type\n");
+    SPIRVWord MaskCompCount = MaskTy->getVectorComponentCount();
+
+    SPVErrLog.checkError(
+        InputVecCompCount == PtrVecCompCount &&
+            PtrVecCompCount == MaskCompCount,
+        SPIRVEC_InvalidInstruction,
+        InstName + "\nInputVector, PtrVector and Mask vectors must have "
+                   "the same size\n");
+
+    SPVErrLog.checkError(
+        InputVecCompTy == PtrElemTy, SPIRVEC_InvalidInstruction,
+        InstName + "\nComponent Type of InputVector must be "
+                   "same as base type of PtrVector the same base type\n");
+  }
+};
+
+#define _SPIRV_OP(x, ...)                                                      \
+  typedef SPIRVInstTemplate<SPIRVMaskedGatherScatterINTELInstBase,             \
+                            internal::Op##x##INTEL, __VA_ARGS__>               \
+      SPIRV##x##INTEL;
+_SPIRV_OP(MaskedGather, true, 7)
+_SPIRV_OP(MaskedScatter, false, 5)
+#undef _SPIRV_OP
 } // namespace SPIRV
 
 #endif // SPIRV_LIBSPIRV_SPIRVINSTRUCTION_H

lib/SPIRV/libSPIRV/SPIRVNameMapEnum.h

@@ -617,6 +617,7 @@ template <> inline void SPIRVMap<Capability, std::string>::init() {
   add(internal::CapabilityNonConstantAddrspacePrintfINTEL,
       "NonConstantAddrspacePrintfINTEL");
   add(internal::CapabilityComplexFloatMulDivINTEL, "ComplexFloatMulDivINTEL");
+  add(internal::CapabilityMaskedGatherScatterINTEL, "MaskedGatherScatterINTEL");
 }
 SPIRV_DEF_NAMEMAP(Capability, SPIRVCapabilityNameMap)
 

lib/SPIRV/libSPIRV/SPIRVOpCodeEnumInternal.h

@@ -13,3 +13,5 @@ _SPIRV_OP_INTERNAL(JointMatrixWorkItemLengthINTEL,
                    internal::OpJointMatrixWorkItemLengthINTEL)
 _SPIRV_OP_INTERNAL(ComplexFMulINTEL, internal::ComplexFMulINTEL)
 _SPIRV_OP_INTERNAL(ComplexFDivINTEL, internal::ComplexFDivINTEL)
+_SPIRV_OP_INTERNAL(MaskedGatherINTEL, internal::OpMaskedGatherINTEL)
+_SPIRV_OP_INTERNAL(MaskedScatterINTEL, internal::OpMaskedScatterINTEL)

lib/SPIRV/libSPIRV/SPIRVType.cpp

@@ -218,6 +218,10 @@ bool SPIRVType::isTypeVectorOrScalarBool() const {
   return isTypeBool() || isTypeVectorBool();
 }
 
+bool SPIRVType::isTypeVectorPointer() const {
+  return isTypeVector() && getVectorComponentType()->isTypePointer();
+}
+
 bool SPIRVType::isTypeSubgroupAvcINTEL() const {
   return isSubgroupAvcINTELTypeOpCode(OpCode);
 }

lib/SPIRV/libSPIRV/SPIRVType.h

@@ -102,6 +102,7 @@ class SPIRVType : public SPIRVEntry {
   bool isTypeVectorOrScalarInt() const;
   bool isTypeVectorOrScalarFloat() const;
   bool isTypeVectorOrScalarBool() const;
+  bool isTypeVectorPointer() const;
   bool isTypeSubgroupAvcINTEL() const;
   bool isTypeSubgroupAvcMceINTEL() const;
 };

lib/SPIRV/libSPIRV/spirv_internal.hpp

@@ -46,6 +46,8 @@ enum InternalOp {
   IOpJointMatrixWorkItemLengthINTEL = 6410,
   IOpComplexFMulINTEL = 6415,
   IOpComplexFDivINTEL = 6416,
+  IOpMaskedGatherINTEL = 6428,
+  IOpMaskedScatterINTEL = 6429,
   IOpPrev = OpMax - 2,
   IOpForward
 };
@@ -76,7 +78,8 @@ enum InternalCapability {
   ICapFPArithmeticFenceINTEL = 6144,
   ICapGlobalVariableDecorationsINTEL = 6146,
   ICapabilityNonConstantAddrspacePrintfINTEL = 6411,
-  ICapabilityComplexFloatMulDivINTEL = 6414
+  ICapabilityComplexFloatMulDivINTEL = 6414,
+  ICapabilityMaskedGatherScatterINTEL = 6427
 };
 
 enum InternalFunctionControlMask { IFunctionControlOptNoneINTELMask = 0x10000 };
@@ -124,6 +127,10 @@ _SPIRV_OP(Capability, NonConstantAddrspacePrintfINTEL)
 _SPIRV_OP(Capability, ComplexFloatMulDivINTEL)
 _SPIRV_OP(Op, ComplexFMulINTEL)
 _SPIRV_OP(Op, ComplexFDivINTEL)
+
+_SPIRV_OP(Capability, MaskedGatherScatterINTEL)
+_SPIRV_OP(Op, MaskedGatherINTEL)
+_SPIRV_OP(Op, MaskedScatterINTEL)
 #undef _SPIRV_OP
 
 constexpr Op OpForward = static_cast<Op>(IOpForward);