[NVPTX] Fix poorly designed assertion introduced in D120129

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.

Differential Revision: https://reviews.llvm.org/D122562

Author: kovdan01

llvm/lib/Target/NVPTX/NVPTXISelLowering.cpp

@@ -4260,40 +4260,11 @@ Align NVPTXTargetLowering::getFunctionParamOptimizedAlign(
 
   // 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;
-
-        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:
+  if (!F->hasLocalLinkage())
     return Align(ABITypeAlign);
-  }
+
+  assert(!isKernelFunction(*F) && "Expect kernels to have non-local linkage");
+  return Align(std::max(uint64_t(16), ABITypeAlign));
 }
 
 /// isLegalAddressingMode - Return true if the addressing mode represented

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

@@ -0,0 +1,34 @@
+; 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}