[NVPTX] Fix poorly designed assertion introduced in D120129

[kovdan01]

Summary

NVPTXTargetLowering::getFunctionParamOptimizedAlign, which was introduces in
D120129, contained a poorly designed assertion checking that a function with
internal or private linkage is not a kernel. It relied on invariants that
were not actually guaranteed, and that resulted in compiler crash with some
CUDA versions (see discussion with @jdoerfert in D120129). This patch changes
that assertion and makes it use isKernelFunction which is designed exactly for
such checks. This patch also includes a test with IR that caused compiler crash
before.

[jdoerfert]

LG, one nit.

(Also much cleaner, thanks)

llvm/lib/Target/NVPTX/NVPTXISelLowering.cpp

4266

Add some text, && "Expected kernels to have external linkage"), or similar.

[kovdan01]

@jdoerfert Merged, hope your CI will be fixed. Could you please change "Changes requested" mark in D120129 when you verify that everything is OK? Just to make that revision closed again.

llvm/lib/Target/NVPTX/NVPTXISelLowering.cpp

4266

Ah, forgot that, thanks!

llvm/lib/Target/NVPTX/NVPTXISelLowering.cpp

 /// alignment. To allow using 128-bit vectorized loads/stores, this function
 /// ensures that alignment is 16 or greater.
 Align NVPTXTargetLowering::getFunctionParamOptimizedAlign(
     const Function *F, Type *ArgTy, const DataLayout &DL) const {
   const uint64_t ABITypeAlign = DL.getABITypeAlign(ArgTy).value();
 
   // If a function has linkage different from internal or private, we
   // must use default ABI alignment as external users rely on it.
-  switch (F->getLinkage()) {
-  case GlobalValue::InternalLinkage:
-  case GlobalValue::PrivateLinkage: {
-    // Check that if a function has internal or private linkage
-    // it is not a kernel.
-#ifndef NDEBUG
-    const NamedMDNode *NMDN =
-        F->getParent()->getNamedMetadata("nvvm.annotations");
-    if (NMDN) {
-      for (const MDNode *MDN : NMDN->operands()) {
-        assert(MDN->getNumOperands() == 3);
-
-        const Metadata *MD0 = MDN->getOperand(0).get();
-        const auto *MDV0 = cast<ConstantAsMetadata>(MD0)->getValue();
-        const auto *MDFn = cast<Function>(MDV0);
-        if (MDFn != F)
-          continue;
-
-        const Metadata *MD1 = MDN->getOperand(1).get();
-        const MDString *MDStr = cast<MDString>(MD1);
-        if (MDStr->getString() != "kernel")
-          continue;
+  if (!F->hasLocalLinkage())
+    return Align(ABITypeAlign);
 
-        const Metadata *MD2 = MDN->getOperand(2).get();
-        const auto *MDV2 = cast<ConstantAsMetadata>(MD2)->getValue();
-        assert(!cast<ConstantInt>(MDV2)->isZero());
-      }
-    }
-#endif
-    return Align(std::max(uint64_t(16), ABITypeAlign));
-  }
-  default:
-    return Align(ABITypeAlign);
-  }
-}
+  assert(!isKernelFunction(*F) && "Expect kernels to have non-local linkage");
 
 /// isLegalAddressingMode - Return true if the addressing mode represented
 /// by AM is legal for this target, for a load/store of the specified type.
 /// Used to guide target specific optimizations, like loop strength reduction
 /// (LoopStrengthReduce.cpp) and memory optimization for address mode
 /// (CodeGenPrepare.cpp)
 bool NVPTXTargetLowering::isLegalAddressingMode(const DataLayout &DL,
                                                 const AddrMode &AM, Type *Ty,

llvm/test/CodeGen/NVPTX/nvvm-annotations-D120129.ll

+; RUN: llc < %s -mtriple=nvptx-unknown-unknown | FileCheck %s
+;
+; NVPTXTargetLowering::getFunctionParamOptimizedAlign, which was introduces in
+; D120129, contained a poorly designed assertion checking that a function with
+; internal or private linkage is not a kernel. It relied on invariants that
+; were not actually guaranteed, and that resulted in compiler crash with some
+; CUDA versions (see discussion with @jdoerfert in D120129). This test contains
+; metadata that caused compiler crash and a function with internal linkage
+; which purpose is to let compiler run on path where the crash happened.
+; Metadata was obtained from libdevice.10.bc shipped with cuda-11-0.
+
+
+define internal i32 @foo() {
+  ; CHECK-LABEL: .func (.param .b32 func_retval0) foo()
+  ret i32 42
+}
+
+define i32 @bar() {
+  ; CHECK-LABEL: .visible .func (.param .b32 func_retval0) bar()
+  %x = call i32 @foo()
+  ret i32 %x
+}
+
+!nvvmir.version = !{!0}
+!nvvm.annotations = !{!1, !2, !1, !3, !3, !3, !3, !4, !4, !3}
+
+!0 = !{i32 1, i32 4}
+!1 = !{null, !"align", i32 8}
+!2 = !{null, !"align", i32 8, !"align", i32 65544, !"align", i32 131080}
+!3 = !{null, !"align", i32 16}
+!4 = !{null, !"align", i32 16, !"align", i32 65552, !"align", i32 131088}
+!5 = distinct !{!5, !6}
+!6 = !{!"llvm.loop.unroll.count", i32 1}
+!7 = distinct !{!7, !6}