1 files changed, 22 insertions, 11 deletions

diff --git a/llvm/lib/Target/WebAssembly/WebAssemblyStoreResults.cpp b/llvm/lib/Target/WebAssembly/WebAssemblyStoreResults.cpp
index 3a7f50e3b14..4a8fc09878c 100644
--- a/llvm/lib/Target/WebAssembly/WebAssemblyStoreResults.cpp
+++ b/llvm/lib/Target/WebAssembly/WebAssemblyStoreResults.cpp
@@ -10,9 +10,12 @@
 /// \file
 /// \brief This file implements an optimization pass using store result values.
 ///
-/// WebAssembly's store instructions return the stored value, specifically to
-/// enable the optimization of reducing get_local/set_local traffic, which is
-/// what we're doing here.
+/// WebAssembly's store instructions return the stored value. This is to enable
+/// an optimization wherein uses of the stored value can be replaced by uses of
+/// the store's result value, making the stored value register more likely to
+/// be single-use, thus more likely to be useful to register stackifying, and
+/// potentially also exposing the store to register stackifying. These both can
+/// reduce get_local/set_local traffic.
 ///
 //===----------------------------------------------------------------------===//
 
@@ -89,14 +92,22 @@ bool WebAssemblyStoreResults::runOnMachineFunction(MachineFunction &MF) {
         for (auto I = MRI.use_begin(FromReg), E = MRI.use_end(); I != E;) {
           MachineOperand &O = *I++;
           MachineInstr *Where = O.getParent();
-          if (Where->getOpcode() == TargetOpcode::PHI)
-            Where = Where->getOperand(&O - &Where->getOperand(0) + 1)
-                        .getMBB()
-                        ->getFirstTerminator();
-          if (&MI == Where || !MDT.dominates(&MI, Where))
-            continue;
-          DEBUG(dbgs() << "Setting operand " << O << " in " << *Where <<
-                " from " << MI <<"\n");
+          if (Where->getOpcode() == TargetOpcode::PHI) {
+            // PHIs use their operands on their incoming CFG edges rather than
+            // in their parent blocks. Get the basic block paired with this use
+            // of FromReg and check that MI's block dominates it.
+            MachineBasicBlock *Pred =
+                Where->getOperand(&O - &Where->getOperand(0) + 1).getMBB();
+            if (!MDT.dominates(&MBB, Pred))
+              continue;
+          } else {
+            // For a non-PHI, check that MI dominates the instruction in the
+            // normal way.
+            if (&MI == Where || !MDT.dominates(&MI, Where))
+              continue;
+          }
+          DEBUG(dbgs() << "Setting operand " << O << " in " << *Where
+                       << " from " << MI << "\n");
           O.setReg(ToReg);
         }
       }