[InferAS] Support getAssumedAddrSpace for Arguments for NVPTX

Author: AlexMaclean

llvm/lib/Target/NVPTX/NVPTXLowerArgs.cpp

@@ -678,11 +678,8 @@ static bool runOnKernelFunction(const NVPTXTargetMachine &TM, Function &F) {
 
   LLVM_DEBUG(dbgs() << "Lowering kernel args of " << F.getName() << "\n");
   for (Argument &Arg : F.args()) {
-    if (Arg.getType()->isPointerTy()) {
-      if (Arg.hasByValAttr())
-        handleByValParam(TM, &Arg);
-      else if (TM.getDrvInterface() == NVPTX::CUDA)
-        markPointerAsGlobal(&Arg);
+    if (Arg.getType()->isPointerTy() && Arg.hasByValAttr()) {
+      handleByValParam(TM, &Arg);
     } else if (Arg.getType()->isIntegerTy() &&
                TM.getDrvInterface() == NVPTX::CUDA) {
       HandleIntToPtr(Arg);
@@ -699,10 +696,9 @@ static bool runOnDeviceFunction(const NVPTXTargetMachine &TM, Function &F) {
       cast<NVPTXTargetLowering>(TM.getSubtargetImpl()->getTargetLowering());
 
   for (Argument &Arg : F.args())
-    if (Arg.getType()->isPointerTy() && Arg.hasByValAttr()) {
-      markPointerAsAS(&Arg, ADDRESS_SPACE_LOCAL);
+    if (Arg.getType()->isPointerTy() && Arg.hasByValAttr())
       adjustByValArgAlignment(&Arg, &Arg, TLI);
-    }
+
   return true;
 }
 

llvm/lib/Target/NVPTX/NVPTXTargetTransformInfo.cpp

@@ -599,6 +599,21 @@ unsigned NVPTXTTIImpl::getAssumedAddrSpace(const Value *V) const {
   if (isa<AllocaInst>(V))
     return ADDRESS_SPACE_LOCAL;
 
+  if (const Argument *Arg = dyn_cast<Argument>(V)) {
+    if (isKernelFunction(*Arg->getParent())) {
+      const NVPTXTargetMachine &TM =
+          static_cast<const NVPTXTargetMachine &>(getTLI()->getTargetMachine());
+      if (TM.getDrvInterface() == NVPTX::CUDA && !Arg->hasByValAttr())
+        return ADDRESS_SPACE_GLOBAL;
+    } else {
+      // We assume that all device parameters that are passed byval will be
+      // placed in the local AS. Very simple cases will be updated after ISel to
+      // use the device param space where possible.
+      if (Arg->hasByValAttr())
+        return ADDRESS_SPACE_LOCAL;
+    }
+  }
+
   return -1;
 }
 

llvm/lib/Transforms/Scalar/InferAddressSpaces.cpp

@@ -305,10 +305,15 @@ static bool isNoopPtrIntCastPair(const Operator *I2P, const DataLayout &DL,
 }
 
 // Returns true if V is an address expression.
-// TODO: Currently, we consider only phi, bitcast, addrspacecast, and
-// getelementptr operators.
+// TODO: Currently, we consider only arguments and phi, bitcast, addrspacecast,
+// and getelementptr operators.
 static bool isAddressExpression(const Value &V, const DataLayout &DL,
                                 const TargetTransformInfo *TTI) {
+
+  if (const Argument *Arg = dyn_cast<Argument>(&V))
+    return Arg->getType()->isPointerTy() &&
+           TTI->getAssumedAddrSpace(&V) != UninitializedAddressSpace;
+
   const Operator *Op = dyn_cast<Operator>(&V);
   if (!Op)
     return false;
@@ -341,6 +346,9 @@ static bool isAddressExpression(const Value &V, const DataLayout &DL,
 static SmallVector<Value *, 2>
 getPointerOperands(const Value &V, const DataLayout &DL,
                    const TargetTransformInfo *TTI) {
+  if (isa<Argument>(&V))
+    return {};
+
   const Operator &Op = cast<Operator>(V);
   switch (Op.getOpcode()) {
   case Instruction::PHI: {
@@ -505,13 +513,11 @@ void InferAddressSpacesImpl::appendsFlatAddressExpressionToPostorderStack(
     if (Visited.insert(V).second) {
       PostorderStack.emplace_back(V, false);
 
-      Operator *Op = cast<Operator>(V);
-      for (unsigned I = 0, E = Op->getNumOperands(); I != E; ++I) {
-        if (ConstantExpr *CE = dyn_cast<ConstantExpr>(Op->getOperand(I))) {
-          if (isAddressExpression(*CE, *DL, TTI) && Visited.insert(CE).second)
-            PostorderStack.emplace_back(CE, false);
-        }
-      }
+      if (auto *Op = dyn_cast<Operator>(V))
+        for (auto &O : Op->operands())
+          if (ConstantExpr *CE = dyn_cast<ConstantExpr>(O))
+            if (isAddressExpression(*CE, *DL, TTI) && Visited.insert(CE).second)
+              PostorderStack.emplace_back(CE, false);
     }
   }
 }
@@ -828,6 +834,18 @@ Value *InferAddressSpacesImpl::cloneValueWithNewAddressSpace(
   assert(V->getType()->getPointerAddressSpace() == FlatAddrSpace &&
          isAddressExpression(*V, *DL, TTI));
 
+  if (auto *Arg = dyn_cast<Argument>(V)) {
+    // Arguments are address space casted in the function body, as we do not
+    // want to change the function signature.
+    Function *F = Arg->getParent();
+    BasicBlock::iterator Insert = F->getEntryBlock().getFirstNonPHIIt();
+
+    Type *NewPtrTy = PointerType::get(Arg->getContext(), NewAddrSpace);
+    auto *NewI = new AddrSpaceCastInst(Arg, NewPtrTy);
+    NewI->insertBefore(Insert);
+    return NewI;
+  }
+
   if (Instruction *I = dyn_cast<Instruction>(V)) {
     Value *NewV = cloneInstructionWithNewAddressSpace(
         I, NewAddrSpace, ValueWithNewAddrSpace, PredicatedAS, PoisonUsesToFix);
@@ -966,8 +984,9 @@ bool InferAddressSpacesImpl::updateAddressSpace(
   // of all its pointer operands.
   unsigned NewAS = UninitializedAddressSpace;
 
-  const Operator &Op = cast<Operator>(V);
-  if (Op.getOpcode() == Instruction::Select) {
+  if (isa<Operator>(V) &&
+      cast<Operator>(V).getOpcode() == Instruction::Select) {
+    const Operator &Op = cast<Operator>(V);
     Value *Src0 = Op.getOperand(1);
     Value *Src1 = Op.getOperand(2);
 
@@ -1275,7 +1294,7 @@ void InferAddressSpacesImpl::performPointerReplacement(
     // This instruction may contain multiple uses of V, update them all.
     CurUser->replaceUsesOfWith(
         V, new AddrSpaceCastInst(NewV, V->getType(), "", InsertPos));
-  } else {
+  } else if (isa<Constant>(V)) {
     CurUserI->replaceUsesOfWith(
         V, ConstantExpr::getAddrSpaceCast(cast<Constant>(NewV), V->getType()));
   }

llvm/test/Transforms/InferAddressSpaces/NVPTX/arguments.ll

@@ -0,0 +1,35 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 5
+; RUN: opt -S -passes=infer-address-spaces %s | FileCheck %s
+
+target triple = "nvptx64-nvidia-cuda"
+
+
+define ptx_kernel i32 @test_kernel(ptr %a, ptr byval(i32) %b) {
+; CHECK-LABEL: define ptx_kernel i32 @test_kernel(
+; CHECK-SAME: ptr [[A:%.*]], ptr byval(i32) [[B:%.*]]) {
+; CHECK-NEXT:    [[TMP1:%.*]] = addrspacecast ptr [[A]] to ptr addrspace(1)
+; CHECK-NEXT:    [[V1:%.*]] = load i32, ptr addrspace(1) [[TMP1]], align 4
+; CHECK-NEXT:    [[V2:%.*]] = load i32, ptr [[B]], align 4
+; CHECK-NEXT:    [[SUM:%.*]] = add i32 [[V1]], [[V2]]
+; CHECK-NEXT:    ret i32 [[SUM]]
+;
+  %v1 = load i32, ptr %a
+  %v2 = load i32, ptr %b
+  %sum = add i32 %v1, %v2
+  ret i32 %sum
+}
+
+define i32 @test_device(ptr %a, ptr byval(i32) %b) {
+; CHECK-LABEL: define i32 @test_device(
+; CHECK-SAME: ptr [[A:%.*]], ptr byval(i32) [[B:%.*]]) {
+; CHECK-NEXT:    [[TMP1:%.*]] = addrspacecast ptr [[B]] to ptr addrspace(5)
+; CHECK-NEXT:    [[V1:%.*]] = load i32, ptr [[A]], align 4
+; CHECK-NEXT:    [[V2:%.*]] = load i32, ptr addrspace(5) [[TMP1]], align 4
+; CHECK-NEXT:    [[SUM:%.*]] = add i32 [[V1]], [[V2]]
+; CHECK-NEXT:    ret i32 [[SUM]]
+;
+  %v1 = load i32, ptr %a
+  %v2 = load i32, ptr %b
+  %sum = add i32 %v1, %v2
+  ret i32 %sum
+}