From 3bf79d0e44f00a1214f2e9784adef565a656d8b7 Mon Sep 17 00:00:00 2001
From: "tatz.j@northeastern.edu" <tatz.j@husky.neu.edu>
Date: Wed, 24 Jun 2020 09:55:09 -0700
Subject: [PATCH] [NVPTX] Fix for NVPTX module asm regression

Currently module asm ends up emitted twice and at the wrong place in the PTX.
This patch moves module asm generation into emitStartOfAsmFile() which puts at
the correct location in the generated PTX.

Differential Revision: https://reviews.llvm.org/D82280
---
 lib/Target/NVPTX/NVPTXAsmPrinter.cpp    | 27 +++++++++----------------
 lib/Target/NVPTX/NVPTXAsmPrinter.h      |  1 +
 test/CodeGen/NVPTX/module-inline-asm.ll | 10 +++++++++
 3 files changed, 20 insertions(+), 18 deletions(-)

diff --git a/lib/Target/NVPTX/NVPTXAsmPrinter.cpp b/lib/Target/NVPTX/NVPTXAsmPrinter.cpp
index c68ac5cd024..da1a398a68f 100644
--- a/lib/Target/NVPTX/NVPTXAsmPrinter.cpp
+++ b/lib/Target/NVPTX/NVPTXAsmPrinter.cpp
@@ -762,13 +762,21 @@ static bool isEmptyXXStructor(GlobalVariable *GV) {
   return InitList->getNumOperands() == 0;
 }
 
-bool NVPTXAsmPrinter::doInitialization(Module &M) {
+void NVPTXAsmPrinter::emitStartOfAsmFile(Module &M) {
   // Construct a default subtarget off of the TargetMachine defaults. The
   // rest of NVPTX isn't friendly to change subtargets per function and
   // so the default TargetMachine will have all of the options.
   const NVPTXTargetMachine &NTM = static_cast<const NVPTXTargetMachine &>(TM);
   const auto* STI = static_cast<const NVPTXSubtarget*>(NTM.getSubtargetImpl());
+  SmallString<128> Str1;
+  raw_svector_ostream OS1(Str1);
 
+  // Emit header before any dwarf directives are emitted below.
+  emitHeader(M, OS1, *STI);
+  OutStreamer->emitRawText(OS1.str());
+}
+
+bool NVPTXAsmPrinter::doInitialization(Module &M) {
   if (M.alias_size()) {
     report_fatal_error("Module has aliases, which NVPTX does not support.");
     return true; // error
@@ -784,26 +792,9 @@ bool NVPTXAsmPrinter::doInitialization(Module &M) {
     return true;  // error
   }
 
-  SmallString<128> Str1;
-  raw_svector_ostream OS1(Str1);
-
   // We need to call the parent's one explicitly.
   bool Result = AsmPrinter::doInitialization(M);
 
-  // Emit header before any dwarf directives are emitted below.
-  emitHeader(M, OS1, *STI);
-  OutStreamer->emitRawText(OS1.str());
-
-  // Emit module-level inline asm if it exists.
-  if (!M.getModuleInlineAsm().empty()) {
-    OutStreamer->AddComment("Start of file scope inline assembly");
-    OutStreamer->AddBlankLine();
-    OutStreamer->emitRawText(StringRef(M.getModuleInlineAsm()));
-    OutStreamer->AddBlankLine();
-    OutStreamer->AddComment("End of file scope inline assembly");
-    OutStreamer->AddBlankLine();
-  }
-
   GlobalsEmitted = false;
 
   return Result;
diff --git a/lib/Target/NVPTX/NVPTXAsmPrinter.h b/lib/Target/NVPTX/NVPTXAsmPrinter.h
index ae2731b025c..5c3a4eb470c 100644
--- a/lib/Target/NVPTX/NVPTXAsmPrinter.h
+++ b/lib/Target/NVPTX/NVPTXAsmPrinter.h
@@ -200,6 +200,7 @@ private:
   const Function *F;
   std::string CurrentFnName;
 
+  void emitStartOfAsmFile(Module &M) override;
   void emitBasicBlockStart(const MachineBasicBlock &MBB) override;
   void emitFunctionEntryLabel() override;
   void emitFunctionBodyStart() override;
diff --git a/test/CodeGen/NVPTX/module-inline-asm.ll b/test/CodeGen/NVPTX/module-inline-asm.ll
index cdbcf2013c0..a2ce6a7f87b 100644
--- a/test/CodeGen/NVPTX/module-inline-asm.ll
+++ b/test/CodeGen/NVPTX/module-inline-asm.ll
@@ -2,9 +2,19 @@
 
 target datalayout = "e-p:32:32:32-i1:8:8-i8:8:8-i16:16:16-i32:32:32-i64:64:64-f32:32:32-f64:64:64-v16:16:16-v32:32:32-v64:64:64-v128:128:128-n16:32:64"
 
+; module asm must come after PTX version/target directives.
+; CHECK-NOT: .global .b32 val;
+
+; CHECK-DAG: .version
+; CHECK-DAG: .target
+
 ; CHECK: .global .b32 val;
 module asm ".global .b32 val;"
 
+; module asm must happen before we emit other things.
+; CHECK-LABEL: .visible .func foo
 define void @foo() {
   ret void
 }
+; Make sure it does not show up anywhere else in the output.
+; CHECK-NOT: .global .b32 val;