[NVPTX] Use align attribute for kernel pointer arg alignment

Instead of determining the alignment based on the pointer element type (which is incompatible with opaque pointers), make use of alignment annotations added by the frontend. In particular, clang will add alignment attributes to OpenCL kernels since D118894. Other frontends might need to be adjusted to add the attribute as well. Differential Revision: https://reviews.llvm.org/D119247
alexbatashev · Feb 10, 2022 · 1c729d7 · 1c729d7
1 parent ac0f329
commit 1c729d7
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 36 deletions.
diff --git a/llvm/lib/Target/NVPTX/NVPTXAsmPrinter.cpp b/llvm/lib/Target/NVPTX/NVPTXAsmPrinter.cpp
@@ -1335,34 +1335,6 @@ void NVPTXAsmPrinter::emitPTXGlobalVariable(const GlobalVariable *GVar,
   }
 }
 
-static unsigned int getOpenCLAlignment(const DataLayout &DL, Type *Ty) {
-  if (Ty->isSingleValueType())
-    return DL.getPrefTypeAlignment(Ty);
-
-  auto *ATy = dyn_cast<ArrayType>(Ty);
-  if (ATy)
-    return getOpenCLAlignment(DL, ATy->getElementType());
-
-  auto *STy = dyn_cast<StructType>(Ty);
-  if (STy) {
-    unsigned int alignStruct = 1;
-    // Go through each element of the struct and find the
-    // largest alignment.
-    for (unsigned i = 0, e = STy->getNumElements(); i != e; i++) {
-      Type *ETy = STy->getElementType(i);
-      unsigned int align = getOpenCLAlignment(DL, ETy);
-      if (align > alignStruct)
-        alignStruct = align;
-    }
-    return alignStruct;
-  }
-
-  auto *FTy = dyn_cast<FunctionType>(Ty);
-  if (FTy)
-    return DL.getPointerPrefAlignment().value();
-  return DL.getPrefTypeAlignment(Ty);
-}
-
 void NVPTXAsmPrinter::printParamName(Function::const_arg_iterator I,
                                      int paramIndex, raw_ostream &O) {
   getSymbol(I->getParent())->print(O, MAI);
@@ -1454,7 +1426,6 @@ void NVPTXAsmPrinter::emitFunctionParamList(const Function *F, raw_ostream &O) {
 
           if (static_cast<NVPTXTargetMachine &>(TM).getDrvInterface() !=
               NVPTX::CUDA) {
-            Type *ETy = PTy->getPointerElementType();
             int addrSpace = PTy->getAddressSpace();
             switch (addrSpace) {
             default:
@@ -1470,7 +1441,8 @@ void NVPTXAsmPrinter::emitFunctionParamList(const Function *F, raw_ostream &O) {
               O << ".ptr .global ";
               break;
             }
-            O << ".align " << (int)getOpenCLAlignment(DL, ETy) << " ";
+            Align ParamAlign = I->getParamAlign().valueOrOne();
+            O << ".align " << ParamAlign.value() << " ";
           }
           printParamName(I, paramIndex, O);
           continue;

diff --git a/llvm/test/CodeGen/NVPTX/nvcl-param-align.ll b/llvm/test/CodeGen/NVPTX/nvcl-param-align.ll
@@ -2,15 +2,15 @@
 
 target triple = "nvptx-unknown-nvcl"
 
+define void @foo(i64 %img, i64 %sampler, <5 x float>* align 32 %v1, i32* %v2) {
+; The parameter alignment is determined by the align attribute (default 1).
 ; CHECK-LABEL: .entry foo(
-define void @foo(i64 %img, i64 %sampler, <5 x float>* %v) {
-; The parameter alignment should be the next power of 2 of 5xsizeof(float),
-; which is 32.
 ; CHECK: .param .u32 .ptr .align 32 foo_param_2
+; CHECK: .param .u32 .ptr .align 1 foo_param_3
   ret void
 }
 
 !nvvm.annotations = !{!1, !2, !3}
-!1 = !{void (i64, i64, <5 x float>*)* @foo, !"kernel", i32 1}
-!2 = !{void (i64, i64, <5 x float>*)* @foo, !"rdoimage", i32 0}
-!3 = !{void (i64, i64, <5 x float>*)* @foo, !"sampler", i32 1}
+!1 = !{void (i64, i64, <5 x float>*, i32*)* @foo, !"kernel", i32 1}
+!2 = !{void (i64, i64, <5 x float>*, i32*)* @foo, !"rdoimage", i32 0}
+!3 = !{void (i64, i64, <5 x float>*, i32*)* @foo, !"sampler", i32 1}