Support arbitrary address space for intrinsics

[delena]

Summary

This is an alternative proposal for supporting address space in intrinsics. It's applicable for any intrinsic, not only masked-load-store.
Related to http://reviews.llvm.org/D17270

I this proposal I add non-zero address space suffix to intrinsic name. The original name looks like
@llvm.xxx.xxx
The same name with address space 1 :
@llvm.xxx.xxx.a_1

(I did not update documentation. I'll do this if the proposed change looks reasonable for reviewers ).

[apilipenko]

Given the fact that you introduce a generic mangling scheme other intrinsics might benefit from I suggest you bringing this up on llvm-dev for a wider discussion. I didn't review the code thoroughly since we haven't yet agreed on the scheme.

[reames]

Marking as request changes to get it off my queue. I believe we've agreed to abandon this approach and possibly consider forward references instead.

../include/llvm/IR/IRBuilder.h

CallInst *CreateGCRelocate(Instruction *Statepoint,
                              int DerivedOffset,
                              Type *ResultType,
                              const Twine &Name = "");
 
 private:
   /// \brief Create a call to a masked intrinsic with given Id.
   /// Masked intrinsic has only one overloaded type - data type.
   CallInst *CreateMaskedIntrinsic(Intrinsic::ID Id, ArrayRef<Value *> Ops,
-                                  Type *DataTy, const Twine &Name = "");
+                                  Type *DataTy, unsigned AS,
+                                  const Twine &Name = "");
 
   Value *getCastedInt8PtrValue(Value *Ptr);
 };
 
 /// \brief This provides a uniform API for creating instructions and inserting
 /// them into a basic block: either at the end of a BasicBlock, or at a specific
 /// iterator location in a block.
 ///

../include/llvm/IR/Intrinsics.h

static IITDescriptor get(IITDescriptorKind K, unsigned Field) {
       return Result;
     }
   };
 
   /// Return the IIT table descriptor for the specified intrinsic into an array
   /// of IITDescriptors.
   void getIntrinsicInfoTableEntries(ID id, SmallVectorImpl<IITDescriptor> &T);
 
+  /// \brief Attach Address Space suffix to intrinsic \p Name
+  void attachAddressSpace(std::string& Name, unsigned AS);
+  
+  /// \brief Detach Address Space suffixes from intrinsic \p Name
+  void detachAddressSpace(std::string& Name);
 } // End Intrinsic namespace
 
 } // End llvm namespace
 
 #endif

../lib/IR/Function.cpp

if (llvm::GlobalValue::isLocalLinkage(Linkage)) {
     // version control in different locations.
     if (FileName.empty())
       NewFuncName = NewFuncName.insert(0, "<unknown>:");
     else
       NewFuncName = NewFuncName.insert(0, FileName.str() + ":");
   }
   return NewFuncName;
 }
+
+void Intrinsic::attachAddressSpace(std::string& Name, unsigned AddrSpace) {
+  Name += ".a_" + std::to_string(AddrSpace);
+}
+
+void Intrinsic::detachAddressSpace(std::string& Name) {
+  std::string::size_type Pos = Name.rfind(".a_");
+  while (Pos != std::string::npos) {
+    Name.erase(Pos, Name.size());
+    Pos = Name.rfind(".a_");
+  }
+}

../lib/IR/IRBuilder.cpp

 //===----------------------------------------------------------------------===//
 
 #include "llvm/IR/Function.h"
 #include "llvm/IR/GlobalVariable.h"
 #include "llvm/IR/IRBuilder.h"
 #include "llvm/IR/Intrinsics.h"
 #include "llvm/IR/LLVMContext.h"
 #include "llvm/IR/Statepoint.h"
+#include "llvm/IR/Module.h"
 using namespace llvm;
 
 /// CreateGlobalString - Make a new global variable with an initializer that
 /// has array of i8 type filled in with the nul terminated string value
 /// specified.  If Name is specified, it is the name of the global variable
 /// created.
 GlobalVariable *IRBuilderBase::CreateGlobalString(StringRef Str,
                                                   const Twine &Name,
CallInst *IRBuilderBase::CreateMaskedLoad(Value *Ptr, unsigned Align,
                                           Value *Mask, Value *PassThru,
                                           const Twine &Name) {
   // DataTy is the overloaded type
   Type *DataTy = cast<PointerType>(Ptr->getType())->getElementType();
   assert(DataTy->isVectorTy() && "Ptr should point to a vector");
   if (!PassThru)
     PassThru = UndefValue::get(DataTy);
   Value *Ops[] = { Ptr, getInt32(Align), Mask,  PassThru};
-  return CreateMaskedIntrinsic(Intrinsic::masked_load, Ops, DataTy, Name);
+  return CreateMaskedIntrinsic(Intrinsic::masked_load, Ops, DataTy, 
+                               Ptr->getType()->getPointerAddressSpace(),
+                               Name);
 }
 
 /// \brief Create a call to a Masked Store intrinsic.
 /// \p Val   - data to be stored,
 /// \p Ptr   - base pointer for the store
 /// \p Align - alignment of the destination location
 /// \p Mask  - vector of booleans which indicates what vector lanes should
 ///            be accessed in memory
 CallInst *IRBuilderBase::CreateMaskedStore(Value *Val, Value *Ptr,
                                            unsigned Align, Value *Mask) {
   Value *Ops[] = { Val, Ptr, getInt32(Align), Mask };
   // Type of the data to be stored - the only one overloaded type
-  return CreateMaskedIntrinsic(Intrinsic::masked_store, Ops, Val->getType());
+  return CreateMaskedIntrinsic(Intrinsic::masked_store, Ops, Val->getType(),
+                               Ptr->getType()->getPointerAddressSpace());
+}
+
+static void applyAddressSpace(std::string& IntrName, FunctionType *& FTy,
+                              unsigned AS) {
+  if (AS == 0)
+    return;
+  Intrinsic::attachAddressSpace(IntrName, AS);
+
+  auto changePtrTy = [](Type *T, unsigned AS) {
+    if (T->isPointerTy())
+      return cast<Type>(PointerType::get(T->getPointerElementType(), AS));
+    if (T->isVectorTy() && T->getVectorElementType()->isPointerTy()) {
+      Type *OldPtrTy = T->getVectorElementType()->getPointerElementType();
+      Type *NewPtrTy = PointerType::get(OldPtrTy, AS);
+      return cast<Type>(VectorType::get(NewPtrTy, T->getVectorNumElements()));
+    }
+    return T;
+  };
+
+  Type *RetTy = changePtrTy(FTy->getReturnType(), AS);
+  SmallVector<Type *, 8> ParamTypes;
+  for (unsigned i = 0; i != FTy->getNumParams(); ++i) {
+    Type *T = changePtrTy(FTy->getParamType(i), AS);
+    ParamTypes.push_back(T);
+  }
+  FTy = FunctionType::get(RetTy, ParamTypes, false);
 }
 
 /// Create a call to a Masked intrinsic, with given intrinsic Id,
 /// an array of operands - Ops, and one overloaded type - DataTy
 CallInst *IRBuilderBase::CreateMaskedIntrinsic(Intrinsic::ID Id,
                                                ArrayRef<Value *> Ops,
-                                               Type *DataTy,
+                                               Type *DataTy, unsigned AS,
                                                const Twine &Name) {
   Module *M = BB->getParent()->getParent();
   Type *OverloadedTypes[] = { DataTy };
-  Value *TheFn = Intrinsic::getDeclaration(M, Id, OverloadedTypes);
+  std::string IntrName = Intrinsic::getName(Id, OverloadedTypes);
+  FunctionType *FTy = Intrinsic::getType(M->getContext(), Id, OverloadedTypes);
+
+  applyAddressSpace(IntrName, FTy, AS);
+
+  Value *TheFn = M->getOrInsertFunction(IntrName, FTy);
   return createCallHelper(TheFn, Ops, this, Name);
 }
 
 /// \brief Create a call to a Masked Gather intrinsic.
 /// \p Ptrs     - vector of pointers for loading
 /// \p Align    - alignment for one element
 /// \p Mask     - vector of booleans which indicates what vector lanes should
 ///               be accessed in memory
CallInst *IRBuilderBase::CreateMaskedGather(Value *Ptrs, unsigned Align,
   if (!Mask)
     Mask = Constant::getAllOnesValue(VectorType::get(Type::getInt1Ty(Context),
                                      NumElts));
 
   Value * Ops[] = {Ptrs, getInt32(Align), Mask, UndefValue::get(DataTy)};
 
   // We specify only one type when we create this intrinsic. Types of other
   // arguments are derived from this type.
-  return CreateMaskedIntrinsic(Intrinsic::masked_gather, Ops, DataTy, Name);
+  return CreateMaskedIntrinsic(Intrinsic::masked_gather, Ops, DataTy, 
+                               PtrTy->getPointerAddressSpace(), Name);
 }
 
 /// \brief Create a call to a Masked Scatter intrinsic.
 /// \p Data  - data to be stored,
 /// \p Ptrs  - the vector of pointers, where the \p Data elements should be
 ///            stored
 /// \p Align - alignment for one element
 /// \p Mask  - vector of booleans which indicates what vector lanes should
#endif
 
   if (!Mask)
     Mask = Constant::getAllOnesValue(VectorType::get(Type::getInt1Ty(Context),
                                      NumElts));
   Value * Ops[] = {Data, Ptrs, getInt32(Align), Mask};
 
   // We specify only one type when we create this intrinsic. Types of other
   // arguments are derived from this type.
-  return CreateMaskedIntrinsic(Intrinsic::masked_scatter, Ops, DataTy);
+  unsigned AS = PtrsTy->getElementType()->getPointerAddressSpace();
+  return CreateMaskedIntrinsic(Intrinsic::masked_scatter, Ops, DataTy, AS);
 }
 
 template <typename T0, typename T1, typename T2, typename T3>
 static std::vector<Value *>
 getStatepointArgs(IRBuilderBase &B, uint64_t ID, uint32_t NumPatchBytes,
                   Value *ActualCallee, uint32_t Flags, ArrayRef<T0> CallArgs,
                   ArrayRef<T1> TransitionArgs, ArrayRef<T2> DeoptArgs,
                   ArrayRef<T3> GCArgs) {

../lib/IR/Verifier.cpp

#include "llvm/IR/Metadata.def"
   void visitCatchSwitchInst(CatchSwitchInst &CatchSwitch);
   void visitCleanupReturnInst(CleanupReturnInst &CRI);
 
   void verifyCallSite(CallSite CS);
   void verifyMustTailCall(CallInst &CI);
   bool performTypeCheck(Intrinsic::ID ID, Function *F, Type *Ty, int VT,
                         unsigned ArgNo, std::string &Suffix);
   bool verifyIntrinsicType(Type *Ty, ArrayRef<Intrinsic::IITDescriptor> &Infos,
-                           SmallVectorImpl<Type *> &ArgTys);
+                           SmallVectorImpl<Type *> &ArgTys,
+                           bool *AnyAddressSpace);
   bool verifyIntrinsicIsVarArg(bool isVarArg,
                                ArrayRef<Intrinsic::IITDescriptor> &Infos);
   bool verifyAttributeCount(AttributeSet Attrs, unsigned Params);
   void verifyAttributeTypes(AttributeSet Attrs, unsigned Idx, bool isFunction,
                             const Value *V);
   void verifyParameterAttrs(AttributeSet Attrs, unsigned Idx, Type *Ty,
                             bool isReturnValue, const Value *V);
   void verifyFunctionAttrs(FunctionType *FT, AttributeSet Attrs,

 
 /// Verify that the specified type (which comes from an intrinsic argument or
 /// return value) matches the type constraints specified by the .td file (e.g.
 /// an "any integer" argument really is an integer).
 ///
 /// This returns true on error but does not print a message.
 bool Verifier::verifyIntrinsicType(Type *Ty,
                                    ArrayRef<Intrinsic::IITDescriptor> &Infos,
-                                   SmallVectorImpl<Type*> &ArgTys) {
+                                   SmallVectorImpl<Type*> &ArgTys,
+                                   bool *AnyAddressSpace) {
   using namespace Intrinsic;
 
   // If we ran out of descriptors, there are too many arguments.
   if (Infos.empty()) return true;
   IITDescriptor D = Infos.front();
   Infos = Infos.slice(1);
 
   switch (D.Kind) {
   case IITDescriptor::Void: return !Ty->isVoidTy();
   case IITDescriptor::VarArg: return true;
   case IITDescriptor::MMX:  return !Ty->isX86_MMXTy();
   case IITDescriptor::Token: return !Ty->isTokenTy();
   case IITDescriptor::Metadata: return !Ty->isMetadataTy();
   case IITDescriptor::Half: return !Ty->isHalfTy();
   case IITDescriptor::Float: return !Ty->isFloatTy();
   case IITDescriptor::Double: return !Ty->isDoubleTy();
   case IITDescriptor::Integer: return !Ty->isIntegerTy(D.Integer_Width);
   case IITDescriptor::Vector: {
     VectorType *VT = dyn_cast<VectorType>(Ty);
     return !VT || VT->getNumElements() != D.Vector_Width ||
-           verifyIntrinsicType(VT->getElementType(), Infos, ArgTys);
+           verifyIntrinsicType(VT->getElementType(), Infos, ArgTys,
+           AnyAddressSpace);
   }
   case IITDescriptor::Pointer: {
     PointerType *PT = dyn_cast<PointerType>(Ty);
     return !PT || PT->getAddressSpace() != D.Pointer_AddressSpace ||
-           verifyIntrinsicType(PT->getElementType(), Infos, ArgTys);
+           verifyIntrinsicType(PT->getElementType(), Infos, ArgTys,
+                               AnyAddressSpace);
   }
 
   case IITDescriptor::Struct: {
     StructType *ST = dyn_cast<StructType>(Ty);
     if (!ST || ST->getNumElements() != D.Struct_NumElements)
       return true;
 
     for (unsigned i = 0, e = D.Struct_NumElements; i != e; ++i)
-      if (verifyIntrinsicType(ST->getElementType(i), Infos, ArgTys))
+      if (verifyIntrinsicType(ST->getElementType(i), Infos, ArgTys,
+                              AnyAddressSpace))
         return true;
     return false;
   }
 
   case IITDescriptor::Argument:
     // Two cases here - If this is the second occurrence of an argument, verify
     // that the later instance matches the previous instance.
     if (D.getArgumentNumber() < ArgTys.size())
case IITDescriptor::SameVecWidthArgument: {
     VectorType * ReferenceType =
       dyn_cast<VectorType>(ArgTys[D.getArgumentNumber()]);
     VectorType *ThisArgType = dyn_cast<VectorType>(Ty);
     if (!ThisArgType || !ReferenceType || 
         (ReferenceType->getVectorNumElements() !=
          ThisArgType->getVectorNumElements()))
       return true;
     return verifyIntrinsicType(ThisArgType->getVectorElementType(),
-                               Infos, ArgTys);
+                               Infos, ArgTys, AnyAddressSpace);
   }
   case IITDescriptor::PtrToArgument: {
     if (D.getArgumentNumber() >= ArgTys.size())
       return true;
+    *AnyAddressSpace = true;
     Type * ReferenceType = ArgTys[D.getArgumentNumber()];
     PointerType *ThisArgType = dyn_cast<PointerType>(Ty);
     return (!ThisArgType || ThisArgType->getElementType() != ReferenceType);
   }
   case IITDescriptor::VecOfPtrsToElt: {
     if (D.getArgumentNumber() >= ArgTys.size())
       return true;
+    *AnyAddressSpace = true;
     VectorType * ReferenceType =
       dyn_cast<VectorType> (ArgTys[D.getArgumentNumber()]);
     VectorType *ThisArgVecTy = dyn_cast<VectorType>(Ty);
     if (!ThisArgVecTy || !ReferenceType || 
         (ReferenceType->getVectorNumElements() !=
          ThisArgVecTy->getVectorNumElements()))
       return true;
     PointerType *ThisArgEltTy =
void Verifier::visitIntrinsicCallSite(Intrinsic::ID ID, CallSite CS) {
   FunctionType *IFTy = IF->getFunctionType();
   bool IsVarArg = IFTy->isVarArg();
 
   SmallVector<Intrinsic::IITDescriptor, 8> Table;
   getIntrinsicInfoTableEntries(ID, Table);
   ArrayRef<Intrinsic::IITDescriptor> TableRef = Table;
 
   SmallVector<Type *, 4> ArgTys;
-  Assert(!verifyIntrinsicType(IFTy->getReturnType(), TableRef, ArgTys),
+  bool AnyAddressSpace = false;
+  Assert(!verifyIntrinsicType(IFTy->getReturnType(), TableRef, ArgTys,
+                              &AnyAddressSpace),
          "Intrinsic has incorrect return type!", IF);
   for (unsigned i = 0, e = IFTy->getNumParams(); i != e; ++i)
-    Assert(!verifyIntrinsicType(IFTy->getParamType(i), TableRef, ArgTys),
+    Assert(!verifyIntrinsicType(IFTy->getParamType(i), TableRef, ArgTys,
+                                &AnyAddressSpace),
            "Intrinsic has incorrect argument type!", IF);
 
   // Verify if the intrinsic call matches the vararg property.
   if (IsVarArg)
     Assert(!verifyIntrinsicIsVarArg(IsVarArg, TableRef),
            "Intrinsic was not defined with variable arguments!", IF);
   else
     Assert(!verifyIntrinsicIsVarArg(IsVarArg, TableRef),
            "Callsite was not defined with variable arguments!", IF);
 
   // All descriptors should be absorbed by now.
   Assert(TableRef.empty(), "Intrinsic has too few arguments!", IF);
 
   // Now that we have the intrinsic ID and the actual argument types (and we
   // know they are legal for the intrinsic!) get the intrinsic name through the
   // usual means.  This allows us to verify the mangling of argument types into
   // the name.
   const std::string ExpectedName = Intrinsic::getName(ID, ArgTys);
-  Assert(ExpectedName == IF->getName(),
+
+  std::string NameToCheck = IF->getName();
+  if (AnyAddressSpace)
+    Intrinsic::detachAddressSpace(NameToCheck);
+  Assert(ExpectedName == NameToCheck,
          "Intrinsic name not mangled correctly for type arguments! "
          "Should be: " +
              ExpectedName,
          IF);
 
   // If the intrinsic takes MDNode arguments, verify that they are either global
   // or are local to *this* function.
   for (Value *V : CS.args()) 

../test/Transforms/LoopVectorize/X86/gather_scatter.ll

for.inc:                                          ; preds = %if.end
   %inc = add nsw i32 %9, 1
   store i32 %inc, i32* %i, align 4
   br label %for.cond
 
 for.end:                                          ; preds = %for.cond
   ret void
 }
 
+;AVX512-LABEL: @foo2_a_1
+;AVX512: getelementptr %struct.In, %struct.In addrspace(1)* %in, <16 x i64> %induction, i32 1
+;AVX512: llvm.masked.gather.v16f32.a_1
+;AVX512: llvm.masked.store.v16f32
+;AVX512: ret void
+define void @foo2_a_1(%struct.In addrspace(1)* noalias %in, float* noalias %out, i32* noalias %trigger, i32* noalias %index) #0 {
+entry:
+  %in.addr = alloca %struct.In addrspace(1)*, align 8
+  %out.addr = alloca float*, align 8
+  %trigger.addr = alloca i32*, align 8
+  %index.addr = alloca i32*, align 8
+  %i = alloca i32, align 4
+  store %struct.In addrspace(1)* %in, %struct.In addrspace(1)** %in.addr, align 8
+  store float* %out, float** %out.addr, align 8
+  store i32* %trigger, i32** %trigger.addr, align 8
+  store i32* %index, i32** %index.addr, align 8
+  store i32 0, i32* %i, align 4
+  br label %for.cond
+
+for.cond:                                         ; preds = %for.inc, %entry
+  %0 = load i32, i32* %i, align 4
+  %cmp = icmp slt i32 %0, 4096
+  br i1 %cmp, label %for.body, label %for.end
+
+for.body:                                         ; preds = %for.cond
+  %1 = load i32, i32* %i, align 4
+  %idxprom = sext i32 %1 to i64
+  %2 = load i32*, i32** %trigger.addr, align 8
+  %arrayidx = getelementptr inbounds i32, i32* %2, i64 %idxprom
+  %3 = load i32, i32* %arrayidx, align 4
+  %cmp1 = icmp sgt i32 %3, 0
+  br i1 %cmp1, label %if.then, label %if.end
+
+if.then:                                          ; preds = %for.body
+  %4 = load i32, i32* %i, align 4
+  %idxprom2 = sext i32 %4 to i64
+  %5 = load %struct.In addrspace(1)*, %struct.In addrspace(1)** %in.addr, align 8
+  %arrayidx3 = getelementptr inbounds %struct.In, %struct.In addrspace(1)* %5, i64 %idxprom2
+  %b = getelementptr inbounds %struct.In, %struct.In addrspace(1)* %arrayidx3, i32 0, i32 1
+  %6 = load float, float addrspace(1)* %b, align 4
+  %add = fadd float %6, 5.000000e-01
+  %7 = load i32, i32* %i, align 4
+  %idxprom4 = sext i32 %7 to i64
+  %8 = load float*, float** %out.addr, align 8
+  %arrayidx5 = getelementptr inbounds float, float* %8, i64 %idxprom4
+  store float %add, float* %arrayidx5, align 4
+  br label %if.end
+
+if.end:                                           ; preds = %if.then, %for.body
+  br label %for.inc
+
+for.inc:                                          ; preds = %if.end
+  %9 = load i32, i32* %i, align 4
+  %inc = add nsw i32 %9, 1
+  store i32 %inc, i32* %i, align 4
+  br label %for.cond
+
+for.end:                                          ; preds = %for.cond
+  ret void
+}
+
 ; The source code
 ;struct Out {
 ;  float a;
 ;  float b;
 ;};
 ;void foo3 (In * __restrict__ in, Out * __restrict__ out, int * __restrict__ trigger) {
 ;
 ;  for (int i=0; i<SIZE; ++i) {