[WebAssembly] Change WebAssemblyMCLowerPrePass to ModulePass

It was a FunctionPass before, which subverted its purpose to collect ALL symbols before MCLowering, depending on how LLVM schedules function passes.
Fixes https://bugs.llvm.org/show_bug.cgi?id=51555

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

llvm/lib/Target/WebAssembly/WebAssembly.h

@@ -56,7 +56,7 @@ FunctionPass *createWebAssemblyLowerBrUnless();
 FunctionPass *createWebAssemblyRegNumbering();
 FunctionPass *createWebAssemblyDebugFixup();
 FunctionPass *createWebAssemblyPeephole();
-FunctionPass *createWebAssemblyMCLowerPrePass();
+ModulePass *createWebAssemblyMCLowerPrePass();
 
 // PassRegistry initialization declarations.
 void initializeWebAssemblyAddMissingPrototypesPass(PassRegistry &);

llvm/lib/Target/WebAssembly/WebAssemblyMCLowerPrePass.cpp

@@ -33,21 +33,21 @@ using namespace llvm;
 #define DEBUG_TYPE "wasm-mclower-prepass"
 
 namespace {
-class WebAssemblyMCLowerPrePass final : public MachineFunctionPass {
+class WebAssemblyMCLowerPrePass final : public ModulePass {
   StringRef getPassName() const override {
     return "WebAssembly MC Lower Pre Pass";
   }
 
   void getAnalysisUsage(AnalysisUsage &AU) const override {
     AU.setPreservesCFG();
-    MachineFunctionPass::getAnalysisUsage(AU);
+    ModulePass::getAnalysisUsage(AU);
   }
 
-  bool runOnMachineFunction(MachineFunction &MF) override;
+  bool runOnModule(Module &M) override;
 
 public:
   static char ID; // Pass identification, replacement for typeid
-  WebAssemblyMCLowerPrePass() : MachineFunctionPass(ID) {}
+  WebAssemblyMCLowerPrePass() : ModulePass(ID) {}
 };
 } // end anonymous namespace
 
@@ -57,30 +57,43 @@ INITIALIZE_PASS(
     "Collects information ahead of time for MC lowering",
     false, false)
 
-FunctionPass *llvm::createWebAssemblyMCLowerPrePass() {
+ModulePass *llvm::createWebAssemblyMCLowerPrePass() {
   return new WebAssemblyMCLowerPrePass();
 }
 
-bool WebAssemblyMCLowerPrePass::runOnMachineFunction(MachineFunction &MF) {
-  LLVM_DEBUG(dbgs() << "********** MC Lower Pre Pass **********\n"
-                       "********** Function: "
-                    << MF.getName() << '\n');
+// NOTE: this is a ModulePass since we need to enforce that this code has run
+// for all functions before AsmPrinter. If this way of doing things is ever
+// suboptimal, we could opt to make it a MachineFunctionPass and instead use
+// something like createBarrierNoopPass() to enforce ordering.
+bool WebAssemblyMCLowerPrePass::runOnModule(Module &M) {
+  auto *MMIWP = getAnalysisIfAvailable<MachineModuleInfoWrapperPass>();
+  if (!MMIWP)
+    return true;
 
-  MachineModuleInfo &MMI = MF.getMMI();
+  MachineModuleInfo &MMI = MMIWP->getMMI();
   MachineModuleInfoWasm &MMIW = MMI.getObjFileInfo<MachineModuleInfoWasm>();
 
-  for (MachineBasicBlock &MBB : MF) {
-    for (auto &MI : MBB) {
-      // FIXME: what should all be filtered out beyond these?
-      if (MI.isDebugInstr() || MI.isInlineAsm())
-        continue;
-      for (MachineOperand &MO : MI.uses()) {
-        if (MO.isSymbol()) {
-          MMIW.MachineSymbolsUsed.insert(MO.getSymbolName());
+  for (Function &F : M) {
+    MachineFunction *MF = MMI.getMachineFunction(F);
+    if (!MF)
+      continue;
+
+    LLVM_DEBUG(dbgs() << "********** MC Lower Pre Pass **********\n"
+                         "********** Function: "
+                      << MF->getName() << '\n');
+
+    for (MachineBasicBlock &MBB : *MF) {
+      for (auto &MI : MBB) {
+        // FIXME: what should all be filtered out beyond these?
+        if (MI.isDebugInstr() || MI.isInlineAsm())
+          continue;
+        for (MachineOperand &MO : MI.uses()) {
+          if (MO.isSymbol()) {
+            MMIW.MachineSymbolsUsed.insert(MO.getSymbolName());
+          }
         }
       }
     }
   }
-
   return true;
 }

llvm/test/MC/WebAssembly/stack-ptr-mclower.ll

@@ -0,0 +1,43 @@
+; RUN: llc --mtriple=wasm32-unknown-unknown -asm-verbose=false %s -o - | FileCheck --check-prefixes CHECK -DPTR=i32 %s
+; RUN: llc --mtriple=wasm64-unknown-unknown -asm-verbose=false %s -o - | FileCheck --check-prefixes CHECK -DPTR=i64 %s
+
+; empty function that does not use stack, to check WebAssemblyMCLowerPrePass
+; works correctly.
+define hidden void @bar() #0 {
+  ret void
+}
+
+; Function that uses explict stack, and should generate a reference to
+; __stack_pointer, along with the corresponding reloction entry.
+define hidden void @foo() #0 {
+entry:
+  alloca i32, align 4
+  ret void
+}
+
+; CHECK:              .text
+; CHECK-NEXT:         .file   "stack-ptr-mclower.ll"
+; CHECK-NEXT:         .section        .text.bar,"",@
+; CHECK-NEXT:         .hidden bar
+; CHECK-NEXT:         .globl  bar
+; CHECK-NEXT:         .globaltype     __stack_pointer, [[PTR]]
+; CHECK-NEXT:         .type   bar,@function
+; CHECK-NEXT: bar:
+; CHECK-NEXT:         .functype       bar () -> ()
+; CHECK-NEXT:         end_function
+; CHECK-NEXT: .Lfunc_end0:
+; CHECK-NEXT:         .size   bar, .Lfunc_end0-bar
+
+; CHECK:              .section        .text.foo,"",@
+; CHECK-NEXT:         .hidden foo
+; CHECK-NEXT:         .globl  foo
+; CHECK-NEXT:         .type   foo,@function
+; CHECK-NEXT: foo:
+; CHECK-NEXT:         .functype       foo () -> ()
+; CHECK-NEXT:         global.get      __stack_pointer
+; CHECK-NEXT:         [[PTR]].const       16
+; CHECK-NEXT:         [[PTR]].sub
+; CHECK-NEXT:         drop
+; CHECK-NEXT:         end_function
+; CHECK-NEXT: .Lfunc_end1:
+; CHECK-NEXT:         .size   foo, .Lfunc_end1-foo