[SPIR-V] Preserve pointer address space for load/gep instructions

[iliya-diyachkov]

The patch looks good to me, thanks Michal.

llvm/include/llvm/IR/IntrinsicsSPIRV.td

 //===- IntrinsicsSPIRV.td - Defines SPIRV intrinsics -------*- tablegen -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // This file defines all of the SPIRV-specific intrinsics.
 //
 //===----------------------------------------------------------------------===//
 
 let TargetPrefix = "spv" in {
   def int_spv_assign_type : Intrinsic<[], [llvm_any_ty, llvm_metadata_ty]>;
-  def int_spv_assign_ptr_type : Intrinsic<[], [llvm_any_ty, llvm_metadata_ty]>;
+  def int_spv_assign_ptr_type : Intrinsic<[], [llvm_any_ty, llvm_metadata_ty, llvm_i32_ty], [ImmArg<ArgIndex<2>>]>;
   def int_spv_assign_name : Intrinsic<[], [llvm_any_ty, llvm_vararg_ty]>;
 
   def int_spv_track_constant : Intrinsic<[llvm_any_ty], [llvm_any_ty, llvm_metadata_ty]>;
   def int_spv_init_global : Intrinsic<[], [llvm_any_ty, llvm_any_ty]>;
   def int_spv_unref_global : Intrinsic<[], [llvm_any_ty]>;
 
   def int_spv_gep : Intrinsic<[llvm_anyptr_ty], [llvm_i1_ty, llvm_any_ty, llvm_vararg_ty], [ImmArg<ArgIndex<0>>]>;
   def int_spv_load : Intrinsic<[llvm_i32_ty], [llvm_anyptr_ty, llvm_i16_ty, llvm_i8_ty], [ImmArg<ArgIndex<1>>, ImmArg<ArgIndex<2>>]>;

llvm/lib/Target/SPIRV/SPIRVEmitIntrinsics.cpp

class SPIRVEmitIntrinsics
   IRBuilder<> *IRB = nullptr;
   Function *F = nullptr;
   bool TrackConstants = true;
   DenseMap<Instruction *, Constant *> AggrConsts;
   DenseSet<Instruction *> AggrStores;
   void preprocessCompositeConstants();
   void preprocessUndefs();
   CallInst *buildIntrWithMD(Intrinsic::ID IntrID, ArrayRef<Type *> Types,
-                            Value *Arg, Value *Arg2) {
+                            Value *Arg, Value *Arg2,
+                            ArrayRef<Constant *> Imms) {
     ConstantAsMetadata *CM = ValueAsMetadata::getConstant(Arg);
     MDTuple *TyMD = MDNode::get(F->getContext(), CM);
     MetadataAsValue *VMD = MetadataAsValue::get(F->getContext(), TyMD);
-    return IRB->CreateIntrinsic(IntrID, {Types}, {Arg2, VMD});
+    SmallVector<Value *, 4> Args;
+    Args.push_back(Arg2);
+    Args.push_back(VMD);
+    for (auto *Imm : Imms)
+      Args.push_back(Imm);
+    return IRB->CreateIntrinsic(IntrID, {Types}, Args);
   }
   void replaceMemInstrUses(Instruction *Old, Instruction *New);
   void processInstrAfterVisit(Instruction *I);
   void insertAssignPtrTypeIntrs(Instruction *I);
   void insertAssignTypeIntrs(Instruction *I);
   void processGlobalValue(GlobalVariable &GV);
 
 public:

 
 static void setInsertPointSkippingPhis(IRBuilder<> &B, Instruction *I) {
   if (isa<PHINode>(I))
     B.SetInsertPoint(I->getParent(), I->getParent()->getFirstInsertionPt());
   else
     B.SetInsertPoint(I);
 }
 
+static bool requireAssignPtrType(Instruction *I) {
+  if (isa<AllocaInst>(I) || isa<GetElementPtrInst>(I))
+    return true;
+
+  return false;
+}
+
 static bool requireAssignType(Instruction *I) {
   IntrinsicInst *Intr = dyn_cast<IntrinsicInst>(I);
   if (Intr) {
     switch (Intr->getIntrinsicID()) {
     case Intrinsic::invariant_start:
     case Intrinsic::invariant_end:
       return false;
     }
if (GV.hasInitializer() && !isa<UndefValue>(GV.getInitializer())) {
     InitInst->setArgOperand(1, Init);
   }
   if ((!GV.hasInitializer() || isa<UndefValue>(GV.getInitializer())) &&
       GV.getNumUses() == 0)
     IRB->CreateIntrinsic(Intrinsic::spv_unref_global, GV.getType(), &GV);
 }
 
 void SPIRVEmitIntrinsics::insertAssignPtrTypeIntrs(Instruction *I) {
-  if (I->getType()->isVoidTy() || !requireAssignType(I))
+  if (I->getType()->isVoidTy() || !requireAssignPtrType(I))
     return;
 
   setInsertPointSkippingPhis(*IRB, I->getNextNode());
+
+  Constant *EltTyConst;
+  unsigned AddressSpace = 0;
   if (auto *AI = dyn_cast<AllocaInst>(I)) {
-    Constant *Const = Constant::getNullValue(AI->getAllocatedType());
-    buildIntrWithMD(Intrinsic::spv_assign_ptr_type, {I->getType()}, Const, I);
+    EltTyConst = Constant::getNullValue(AI->getAllocatedType());
+    AddressSpace = AI->getAddressSpace();
+  } else if (auto *GEP = dyn_cast<GetElementPtrInst>(I)) {
+    EltTyConst = Constant::getNullValue(GEP->getSourceElementType());
+    AddressSpace = GEP->getPointerAddressSpace();
+  } else {
+    llvm_unreachable("Unexpected instruction!");
   }
+
+  buildIntrWithMD(Intrinsic::spv_assign_ptr_type, {I->getType()}, EltTyConst, I,
+                  {IRB->getInt32(AddressSpace)});
 }
 
 void SPIRVEmitIntrinsics::insertAssignTypeIntrs(Instruction *I) {
   Type *Ty = I->getType();
-  if (!Ty->isVoidTy() && requireAssignType(I) &&
-      I->getOpcode() != Instruction::Alloca) {
+  if (!Ty->isVoidTy() && requireAssignType(I) && !requireAssignPtrType(I)) {
     setInsertPointSkippingPhis(*IRB, I->getNextNode());
     Type *TypeToAssign = Ty;
     if (auto *II = dyn_cast<IntrinsicInst>(I)) {
       if (II->getIntrinsicID() == Intrinsic::spv_const_composite ||
           II->getIntrinsicID() == Intrinsic::spv_undef) {
         auto t = AggrConsts.find(II);
         assert(t != AggrConsts.end());
         TypeToAssign = t->second->getType();
       }
     }
     Constant *Const = Constant::getNullValue(TypeToAssign);
-    buildIntrWithMD(Intrinsic::spv_assign_type, {Ty}, Const, I);
+    buildIntrWithMD(Intrinsic::spv_assign_type, {Ty}, Const, I, {});
   }
   for (const auto &Op : I->operands()) {
     if (isa<ConstantPointerNull>(Op) || isa<UndefValue>(Op) ||
         // Check GetElementPtrConstantExpr case.
         (isa<ConstantExpr>(Op) && isa<GEPOperator>(Op))) {
       setInsertPointSkippingPhis(*IRB, I);
       if (isa<UndefValue>(Op) && Op->getType()->isAggregateType())
         buildIntrWithMD(Intrinsic::spv_assign_type, {IRB->getInt32Ty()}, Op,
-                        UndefValue::get(IRB->getInt32Ty()));
+                        UndefValue::get(IRB->getInt32Ty()), {});
       else
-        buildIntrWithMD(Intrinsic::spv_assign_type, {Op->getType()}, Op, Op);
+        buildIntrWithMD(Intrinsic::spv_assign_type, {Op->getType()}, Op, Op,
+                        {});
     }
   }
 }
 
 void SPIRVEmitIntrinsics::processInstrAfterVisit(Instruction *I) {
   auto *II = dyn_cast<IntrinsicInst>(I);
   if (II && II->getIntrinsicID() == Intrinsic::spv_const_composite &&
       TrackConstants) {
     IRB->SetInsertPoint(I->getNextNode());
     Type *Ty = IRB->getInt32Ty();
     auto t = AggrConsts.find(I);
     assert(t != AggrConsts.end());
-    auto *NewOp =
-        buildIntrWithMD(Intrinsic::spv_track_constant, {Ty, Ty}, t->second, I);
+    auto *NewOp = buildIntrWithMD(Intrinsic::spv_track_constant, {Ty, Ty},
+                                  t->second, I, {});
     I->replaceAllUsesWith(NewOp);
     NewOp->setArgOperand(0, I);
   }
   for (const auto &Op : I->operands()) {
     if ((isa<ConstantAggregateZero>(Op) && Op->getType()->isVectorTy()) ||
         isa<PHINode>(I) || isa<SwitchInst>(I))
       TrackConstants = false;
     if ((isa<ConstantData>(Op) || isa<ConstantExpr>(Op)) && TrackConstants) {
       unsigned OpNo = Op.getOperandNo();
       if (II && ((II->getIntrinsicID() == Intrinsic::spv_gep && OpNo == 0) ||
                  (II->paramHasAttr(OpNo, Attribute::ImmArg))))
         continue;
       IRB->SetInsertPoint(I);
       auto *NewOp = buildIntrWithMD(Intrinsic::spv_track_constant,
-                                    {Op->getType(), Op->getType()}, Op, Op);
+                                    {Op->getType(), Op->getType()}, Op, Op, {});
       I->setOperand(OpNo, NewOp);
     }
   }
   if (I->hasName()) {
     setInsertPointSkippingPhis(*IRB, I->getNextNode());
     std::vector<Value *> Args = {I};
     addStringImm(I->getName(), *IRB, Args);
     IRB->CreateIntrinsic(Intrinsic::spv_assign_name, {I->getType()}, Args);

llvm/lib/Target/SPIRV/SPIRVPreLegalizer.cpp

for (auto MII = std::prev(MBB->end()), Begin = MBB->begin();
       MachineInstr &MI = *MII;
 
       if (isSpvIntrinsic(MI, Intrinsic::spv_assign_ptr_type)) {
         Register Reg = MI.getOperand(1).getReg();
         MIB.setInsertPt(*MI.getParent(), MI.getIterator());
         SPIRVType *BaseTy = GR->getOrCreateSPIRVType(
             getMDOperandAsType(MI.getOperand(2).getMetadata(), 0), MIB);
         SPIRVType *AssignedPtrType = GR->getOrCreateSPIRVPointerType(
-            BaseTy, MI, *MF.getSubtarget<SPIRVSubtarget>().getInstrInfo());
+            BaseTy, MI, *MF.getSubtarget<SPIRVSubtarget>().getInstrInfo(),
+            addressSpaceToStorageClass(MI.getOperand(3).getImm()));
         MachineInstr *Def = MRI.getVRegDef(Reg);
         assert(Def && "Expecting an instruction that defines the register");
         insertAssignInstr(Reg, nullptr, AssignedPtrType, GR, MIB,
                           MF.getRegInfo());
         ToErase.push_back(&MI);
       } else if (isSpvIntrinsic(MI, Intrinsic::spv_assign_type)) {
         Register Reg = MI.getOperand(1).getReg();
         Type *Ty = getMDOperandAsType(MI.getOperand(2).getMetadata(), 0);

llvm/test/CodeGen/SPIRV/pointers/getelementptr-addressspace.ll

+; RUN: llc -O0 -mtriple=spirv64-unknown-unknown %s -o - | FileCheck %s
+
+; CHECK:  %[[#INT8:]] = OpTypeInt 8 0
+; CHECK:  %[[#PTR1:]] = OpTypePointer CrossWorkgroup %[[#INT8]]
+; CHECK:  %[[#PTR2:]] = OpTypePointer UniformConstant %[[#INT8]]
+; CHECK:  %[[#]] = OpInBoundsPtrAccessChain %[[#PTR1]] %[[#]] %[[#]]
+; CHECK:  %[[#]] = OpInBoundsPtrAccessChain %[[#PTR2]] %[[#]] %[[#]]
+
+define spir_kernel void @foo(ptr addrspace(1) %a, ptr addrspace(2) %b) {
+entry:
+  %c = getelementptr inbounds i8, ptr addrspace(1) %a, i32 1
+  %d = getelementptr inbounds i8, ptr addrspace(2) %b, i32 2
+  ret void
+}

llvm/test/CodeGen/SPIRV/pointers/load-addressspace.ll

+; RUN: llc -O0 -mtriple=spirv64-unknown-unknown %s -o - | FileCheck %s
+
+; CHECK:  %[[#INT8:]] = OpTypeInt 8 0
+; CHECK:  %[[#PTR1:]] = OpTypePointer CrossWorkgroup %[[#INT8]]
+; CHECK:  %[[#PTR2:]] = OpTypePointer UniformConstant %[[#INT8]]
+; CHECK:  %[[#FNP1:]] = OpFunctionParameter %[[#PTR1]]
+; CHECK:  %[[#FNP2:]] = OpFunctionParameter %[[#PTR2]]
+; CHECK:  %[[#]] = OpLoad %[[#INT8]] %[[#FNP1]] Aligned 1
+; CHECK:  %[[#]] = OpLoad %[[#INT8]] %[[#FNP2]] Aligned 1
+
+define spir_kernel void @foo(ptr addrspace(1) %a, ptr addrspace(2) %b) {
+entry:
+  %c = load i8, ptr addrspace(1) %a
+  %d = load i8, ptr addrspace(2) %b
+  ret void
+}