Address reviewer feedback

Author: dakersnar

llvm/lib/Target/NVPTX/NVPTXISelDAGToDAG.cpp

@@ -1028,7 +1028,7 @@ static inline bool isAddLike(const SDValue V) {
 }
 
 static SDValue stripAssertAlign(SDValue N) {
-  while (N.getOpcode() == ISD::AssertAlign)
+  if (N.getOpcode() == ISD::AssertAlign)
     N = N.getOperand(0);
   return N;
 }
@@ -1051,7 +1051,6 @@ static SDValue selectBaseADDR(SDValue N, SelectionDAG *DAG) {
 }
 
 static SDValue accumulateOffset(SDValue &Addr, SDLoc DL, SelectionDAG *DAG) {
-  Addr = stripAssertAlign(Addr);
   APInt AccumulatedOffset(64u, 0);
   while (isAddLike(Addr)) {
     const auto *CN = dyn_cast<ConstantSDNode>(Addr.getOperand(1));
@@ -1063,7 +1062,7 @@ static SDValue accumulateOffset(SDValue &Addr, SDLoc DL, SelectionDAG *DAG) {
       break;
 
     AccumulatedOffset += CI;
-    Addr = Addr->getOperand(0);
+    Addr = stripAssertAlign(Addr->getOperand(0));
   }
   return DAG->getSignedTargetConstant(AccumulatedOffset.getSExtValue(), DL,
                                       MVT::i32);

llvm/lib/Target/NVPTX/NVPTXLowerArgs.cpp

@@ -412,6 +412,23 @@ static void adjustByValArgAlignment(Argument *Arg, Value *ArgInParamAS,
   }
 }
 
+// Create a call to the nvvm_internal_addrspace_wrap intrinsic and set the
+// alignment of the return value based on the alignment of the argument.
+static CallInst *createNVVMInternalAddrspaceWrap(IRBuilder<> &IRB,
+                                                 Argument &Arg,
+                                                 const Twine &Name = "") {
+  CallInst *ArgInParam =
+      IRB.CreateIntrinsic(Intrinsic::nvvm_internal_addrspace_wrap,
+                          {IRB.getPtrTy(ADDRESS_SPACE_PARAM), Arg.getType()},
+                          &Arg, {}, Name);
+
+  if (MaybeAlign ParamAlign = Arg.getParamAlign())
+    ArgInParam->addRetAttr(
+        Attribute::getWithAlignment(ArgInParam->getContext(), *ParamAlign));
+
+  return ArgInParam;
+}
+
 namespace {
 struct ArgUseChecker : PtrUseVisitor<ArgUseChecker> {
   using Base = PtrUseVisitor<ArgUseChecker>;
@@ -515,14 +532,7 @@ void copyByValParam(Function &F, Argument &Arg) {
       Arg.getParamAlign().value_or(DL.getPrefTypeAlign(StructType)));
   Arg.replaceAllUsesWith(AllocA);
 
-  CallInst *ArgInParam =
-      IRB.CreateIntrinsic(Intrinsic::nvvm_internal_addrspace_wrap,
-                          {IRB.getPtrTy(ADDRESS_SPACE_PARAM), Arg.getType()},
-                          &Arg, {}, Arg.getName());
-
-  if (MaybeAlign ParamAlign = Arg.getParamAlign())
-    ArgInParam->addRetAttr(
-        Attribute::getWithAlignment(ArgInParam->getContext(), *ParamAlign));
+  CallInst *ArgInParam = createNVVMInternalAddrspaceWrap(IRB, Arg, Arg.getName());
 
   // Be sure to propagate alignment to this load; LLVM doesn't know that NVPTX
   // addrspacecast preserves alignment.  Since params are constant, this load
@@ -553,13 +563,7 @@ static void handleByValParam(const NVPTXTargetMachine &TM, Argument *Arg) {
     SmallVector<Use *, 16> UsesToUpdate(llvm::make_pointer_range(Arg->uses()));
 
     IRBuilder<> IRB(&*FirstInst);
-    CallInst *ArgInParamAS = IRB.CreateIntrinsic(
-        Intrinsic::nvvm_internal_addrspace_wrap,
-        {IRB.getPtrTy(ADDRESS_SPACE_PARAM), Arg->getType()}, {Arg});
-
-    if (MaybeAlign ParamAlign = Arg->getParamAlign())
-      ArgInParamAS->addRetAttr(
-          Attribute::getWithAlignment(ArgInParamAS->getContext(), *ParamAlign));
+    CallInst *ArgInParamAS = createNVVMInternalAddrspaceWrap(IRB, *Arg);
 
     for (Use *U : UsesToUpdate)
       convertToParamAS(U, ArgInParamAS, HasCvtaParam, IsGridConstant);
@@ -589,14 +593,7 @@ static void handleByValParam(const NVPTXTargetMachine &TM, Argument *Arg) {
     // argument already in the param address space, we need to use the noop
     // intrinsic, this had the added benefit of preventing other optimizations
     // from folding away this pair of addrspacecasts.
-    auto *ParamSpaceArg =
-        IRB.CreateIntrinsic(Intrinsic::nvvm_internal_addrspace_wrap,
-                            {IRB.getPtrTy(ADDRESS_SPACE_PARAM), Arg->getType()},
-                            Arg, {}, Arg->getName() + ".param");
-
-    if (MaybeAlign ParamAlign = Arg->getParamAlign())
-      ParamSpaceArg->addRetAttr(Attribute::getWithAlignment(
-          ParamSpaceArg->getContext(), *ParamAlign));
+    auto *ParamSpaceArg = createNVVMInternalAddrspaceWrap(IRB, *Arg, Arg->getName() + ".param");
 
     // Cast param address to generic address space.
     Value *GenericArg = IRB.CreateAddrSpaceCast(

llvm/test/CodeGen/NVPTX/lower-args-alignment.ll

@@ -1,7 +1,6 @@
 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 5
 ; RUN: opt < %s -passes=nvptx-lower-args,infer-alignment -S | FileCheck %s
 
-target datalayout = "e-p:64:64:64-p3:32:32:32-i1:8:8-i8:8:8-i16:16:16-i32:32:32-i64:64:64-i128:128:128-f32:32:32-f64:64:64-v16:16:16-v32:32:32-v64:64:64-v128:128:128-n16:32:64"
 target triple = "nvptx64-nvidia-cuda"
 
 ; ------------------------------------------------------------------------------