[NVPTX] Add support for calling aliases (#81170)

The current implementation of aliases tries to remove all the aliases in
the module to prevent the generic version of `AsmPrinter` from emitting
them incorrectly. Unfortunately, if the aliases are used this will fail.
Instead let's override the function to print aliases directly.

In addition, the declarations of the alias functions must occur before
the uses. To fix this we emit alias declarations as part of
`emitDeclarations` and only emit the `.alias` directives at the end
(where we can assume the aliasee has also already been declared).

Author: AlexMaclean

llvm/include/llvm/CodeGen/AsmPrinter.h

@@ -897,7 +897,7 @@ class AsmPrinter : public MachineFunctionPass {
   virtual void emitModuleCommandLines(Module &M);
 
   GCMetadataPrinter *getOrCreateGCPrinter(GCStrategy &S);
-  void emitGlobalAlias(Module &M, const GlobalAlias &GA);
+  virtual void emitGlobalAlias(const Module &M, const GlobalAlias &GA);
   void emitGlobalIFunc(Module &M, const GlobalIFunc &GI);
 
 private:

llvm/lib/CodeGen/AsmPrinter/AsmPrinter.cpp

@@ -2127,7 +2127,7 @@ void AsmPrinter::emitGlobalGOTEquivs() {
     emitGlobalVariable(GV);
 }
 
-void AsmPrinter::emitGlobalAlias(Module &M, const GlobalAlias &GA) {
+void AsmPrinter::emitGlobalAlias(const Module &M, const GlobalAlias &GA) {
   MCSymbol *Name = getSymbol(&GA);
   bool IsFunction = GA.getValueType()->isFunctionTy();
   // Treat bitcasts of functions as functions also. This is important at least

llvm/lib/Target/NVPTX/NVPTXAsmPrinter.cpp

@@ -57,6 +57,7 @@
 #include "llvm/IR/DebugLoc.h"
 #include "llvm/IR/DerivedTypes.h"
 #include "llvm/IR/Function.h"
+#include "llvm/IR/GlobalAlias.h"
 #include "llvm/IR/GlobalValue.h"
 #include "llvm/IR/GlobalVariable.h"
 #include "llvm/IR/Instruction.h"
@@ -605,14 +606,33 @@ void NVPTXAsmPrinter::emitVirtualRegister(unsigned int vr,
   O << getVirtualRegisterName(vr);
 }
 
+void NVPTXAsmPrinter::emitAliasDeclaration(const GlobalAlias *GA,
+                                           raw_ostream &O) {
+  const Function *F = dyn_cast_or_null<Function>(GA->getAliaseeObject());
+  if (!F || isKernelFunction(*F) || F->isDeclaration())
+    report_fatal_error(
+        "NVPTX aliasee must be a non-kernel function definition");
+
+  if (GA->hasLinkOnceLinkage() || GA->hasWeakLinkage() ||
+      GA->hasAvailableExternallyLinkage() || GA->hasCommonLinkage())
+    report_fatal_error("NVPTX aliasee must not be '.weak'");
+
+  emitDeclarationWithName(F, getSymbol(GA), O);
+}
+
 void NVPTXAsmPrinter::emitDeclaration(const Function *F, raw_ostream &O) {
+  emitDeclarationWithName(F, getSymbol(F), O);
+}
+
+void NVPTXAsmPrinter::emitDeclarationWithName(const Function *F, MCSymbol *S,
+                                              raw_ostream &O) {
   emitLinkageDirective(F, O);
   if (isKernelFunction(*F))
     O << ".entry ";
   else
     O << ".func ";
   printReturnValStr(F, O);
-  getSymbol(F)->print(O, MAI);
+  S->print(O, MAI);
   O << "\n";
   emitFunctionParamList(F, O);
   O << "\n";
@@ -759,6 +779,8 @@ void NVPTXAsmPrinter::emitDeclarations(const Module &M, raw_ostream &O) {
     }
     seenMap[&F] = true;
   }
+  for (const GlobalAlias &GA : M.aliases())
+    emitAliasDeclaration(&GA, O);
 }
 
 static bool isEmptyXXStructor(GlobalVariable *GV) {
@@ -853,25 +875,9 @@ void NVPTXAsmPrinter::emitGlobalAlias(const Module &M, const GlobalAlias &GA) {
   raw_svector_ostream OS(Str);
 
   MCSymbol *Name = getSymbol(&GA);
-  const Function *F = dyn_cast<Function>(GA.getAliasee());
-  if (!F || isKernelFunction(*F))
-    report_fatal_error("NVPTX aliasee must be a non-kernel function");
-
-  if (GA.hasLinkOnceLinkage() || GA.hasWeakLinkage() ||
-      GA.hasAvailableExternallyLinkage() || GA.hasCommonLinkage())
-    report_fatal_error("NVPTX aliasee must not be '.weak'");
-
-  OS << "\n";
-  emitLinkageDirective(F, OS);
-  OS << ".func ";
-  printReturnValStr(F, OS);
-  OS << Name->getName();
-  emitFunctionParamList(F, OS);
-  if (shouldEmitPTXNoReturn(F, TM))
-    OS << "\n.noreturn";
-  OS << ";\n";
 
-  OS << ".alias " << Name->getName() << ", " << F->getName() << ";\n";
+  OS << ".alias " << Name->getName() << ", " << GA.getAliaseeObject()->getName()
+     << ";\n";
 
   OutStreamer->emitRawText(OS.str());
 }
@@ -932,16 +938,6 @@ bool NVPTXAsmPrinter::doFinalization(Module &M) {
     GlobalsEmitted = true;
   }
 
-  // If we have any aliases we emit them at the end.
-  SmallVector<GlobalAlias *> AliasesToRemove;
-  for (GlobalAlias &Alias : M.aliases()) {
-    emitGlobalAlias(M, Alias);
-    AliasesToRemove.push_back(&Alias);
-  }
-
-  for (GlobalAlias *A : AliasesToRemove)
-    A->eraseFromParent();
-
   // call doFinalization
   bool ret = AsmPrinter::doFinalization(M);
 

llvm/lib/Target/NVPTX/NVPTXAsmPrinter.h

@@ -27,6 +27,7 @@
 #include "llvm/IR/DebugLoc.h"
 #include "llvm/IR/DerivedTypes.h"
 #include "llvm/IR/Function.h"
+#include "llvm/IR/GlobalAlias.h"
 #include "llvm/IR/GlobalValue.h"
 #include "llvm/IR/Value.h"
 #include "llvm/MC/MCExpr.h"
@@ -174,7 +175,7 @@ class LLVM_LIBRARY_VISIBILITY NVPTXAsmPrinter : public AsmPrinter {
   void printModuleLevelGV(const GlobalVariable *GVar, raw_ostream &O,
                           bool processDemoted, const NVPTXSubtarget &STI);
   void emitGlobals(const Module &M);
-  void emitGlobalAlias(const Module &M, const GlobalAlias &GA);
+  void emitGlobalAlias(const Module &M, const GlobalAlias &GA) override;
   void emitHeader(Module &M, raw_ostream &O, const NVPTXSubtarget &STI);
   void emitKernelFunctionDirectives(const Function &F, raw_ostream &O) const;
   void emitVirtualRegister(unsigned int vr, raw_ostream &);
@@ -222,6 +223,8 @@ class LLVM_LIBRARY_VISIBILITY NVPTXAsmPrinter : public AsmPrinter {
   void emitLinkageDirective(const GlobalValue *V, raw_ostream &O);
   void emitDeclarations(const Module &, raw_ostream &O);
   void emitDeclaration(const Function *, raw_ostream &O);
+  void emitAliasDeclaration(const GlobalAlias *, raw_ostream &O);
+  void emitDeclarationWithName(const Function *, MCSymbol *, raw_ostream &O);
   void emitDemotedVars(const Function *, raw_ostream &);
 
   bool lowerImageHandleOperand(const MachineInstr *MI, unsigned OpNo,

llvm/test/CodeGen/NVPTX/alias.ll

@@ -1,8 +1,10 @@
 ; RUN: llc < %s -march=nvptx64 -mcpu=sm_30 -mattr=+ptx64 | FileCheck %s
+; RUN: %if ptxas %{ llc < %s -march=nvptx64 -mcpu=sm_30 -mattr=+ptx64 | %ptxas-verify %}
 
 define i32 @a() { ret i32 0 }
 @b = internal alias i32 (), ptr @a
 @c = internal alias i32 (), ptr @a
+@d = internal alias i32 (), ptr @c
 
 define void @foo(i32 %0, ptr %1) { ret void }
 @bar = alias i32 (), ptr @foo
@@ -12,8 +14,37 @@ define void @noreturn() #0 {
 }
 @noreturn_alias = alias i32 (), ptr @noreturn
 
+define i32 @z() {
+  %val = call i32 @b()
+  ret i32 %val
+}
+
+
 attributes #0 = { noreturn }
 
+;      CHECK: .visible .func  (.param .b32 func_retval0) b
+; CHECK-NEXT: ()
+; CHECK-NEXT: ;
+
+;      CHECK: .visible .func  (.param .b32 func_retval0) c
+; CHECK-NEXT: ()
+; CHECK-NEXT: ;
+
+;      CHECK: .visible .func  (.param .b32 func_retval0) d
+; CHECK-NEXT: ()
+; CHECK-NEXT: ;
+
+;      CHECK: .visible .func bar
+; CHECK-NEXT: (
+; CHECK-NEXT:         .param .b32 foo_param_0,
+; CHECK-NEXT:         .param .b64 foo_param_1
+; CHECK-NEXT: )
+; CHECK-NEXT: ;
+
+;      CHECK: .visible .func noreturn_alias
+; CHECK-NEXT: ()
+; CHECK-NEXT: .noreturn;
+
 ; CHECK: .visible .func  (.param .b32 func_retval0) a()
 
 ;      CHECK: .visible .func foo(
@@ -24,18 +55,13 @@ attributes #0 = { noreturn }
 ;      CHECK: .visible .func noreturn()
 ; CHECK-NEXT: .noreturn
 
-;      CHECK: .visible .func  (.param .b32 func_retval0) b();
-; CHECK-NEXT: .alias b, a;
+;      CHECK: .visible .func  (.param .b32 func_retval0) z()
+;      CHECK:      call.uni (retval0), 
+; CHECK-NEXT:      b,
 
-;      CHECK: .visible .func  (.param .b32 func_retval0) c();
-; CHECK-NEXT: .alias c, a;
 
-;      CHECK: .visible .func bar(
-; CHECK-NEXT:         .param .b32 foo_param_0,
-; CHECK-NEXT:         .param .b64 foo_param_1
-; CHECK-NEXT: );
-; CHECK-NEXT: .alias bar, foo;
-
-;      CHECK: .visible .func noreturn_alias()
-; CHECK-NEXT: .noreturn;
-; CHECK-NEXT: .alias noreturn_alias, noreturn;
+; CHECK: .alias b, a;
+; CHECK: .alias c, a;
+; CHECK: .alias d, a;
+; CHECK: .alias bar, foo;
+; CHECK: .alias noreturn_alias, noreturn;