Fix PickLoadSigns on SignExt feature instructions (#7069)

I believe the history here is that

1.  We added a PickLoadSigns pass. It checks if a load from memory is stored in
    a local that is only every used in a signed or an unsigned manner. If it is, we can
    adjust the sign of the load (load8_u/s) to do the sign/unsign during the load.
2.  The pass finds each LocalGet and looks either 2 or 3 parents above it. For
    a sign operation, we need to look up 3, since the operation is x << K >> K. For
    an unsigned, we need only 2, since we have x & M. We hardcoded those
    numbers 2 and 3.
3.  We added the SignExt feature, which adds i32.extend8_s. This does a sign
    extend with a single instruction, not two nested ones, so now we can sign-
    extend at depth 2, unlike before. Properties::getSignExtValue was updated
    for this, but not the pass PickLoadSigns.

The bug that is fixed here is that we looked at depth 3 for a sign-extend, and
we blindly accepted it if we found one. So we ended up accepting
(i32.extend8_s (ANYTHING (x))), which is a sign-extend of something, but
not of x, which is bad.

We were also missing an optimization opportunity, as we didn't look for
depth 2 sign extends.

This bug is quite old, from when Properties got SignExt support, in #3910.
But the blame isn't there - to notice this then, we'd have had to check each
caller of getSignExtValue throughout the codebase, which isn't reasonable.
The fault is mine, from the first write-up of PickLoadSigns in 2017: the code
should have been fully general, handling 2/3 and checking the output when
it does so (adding == curr, that the sign/zero-extended value is the one we
expect). That is what this PR does.

1 files changed, 22 insertions, 17 deletions

diff --git a/src/passes/PickLoadSigns.cpp b/src/passes/PickLoadSigns.cpp
index 587314209..693fd4a68 100644
--- a/src/passes/PickLoadSigns.cpp
+++ b/src/passes/PickLoadSigns.cpp
@@ -20,10 +20,11 @@
 
 namespace wasm {
 
+//
 // Adjust load signedness based on usage. If a load only has uses that sign or
 // unsign it anyhow, then it could be either, and picking the popular one can
-// help remove the most sign/unsign operations
-// unsigned, then it could be either
+// help remove the most sign/unsign operations.
+//
 
 struct PickLoadSigns : public WalkerPass<ExpressionStackWalker<PickLoadSigns>> {
   bool isFunctionParallel() override { return true; }
@@ -59,13 +60,20 @@ struct PickLoadSigns : public WalkerPass<ExpressionStackWalker<PickLoadSigns>> {
   }
 
   void visitLocalGet(LocalGet* curr) {
-    // this is a use. check from the context what it is, signed or unsigned,
-    // etc.
+    // This is a use. check from the context what it is, signed or unsigned,
+    // etc. Sign- and zero-extends may have two levels of nesting (x << K >> K)
+    // or one (x & K), so check both.
     auto& usage = usages[curr->index];
     usage.totalUsages++;
-    if (expressionStack.size() >= 2) {
-      auto* parent = expressionStack[expressionStack.size() - 2];
-      if (Properties::getZeroExtValue(parent)) {
+    for (Index i = 2; i <= 3; i++) {
+      if (expressionStack.size() < i) {
+        // We do not have that many expressions above us. (And if we don't have
+        // 2 then we definitely don't have 3.)
+        break;
+      }
+
+      auto* parent = expressionStack[expressionStack.size() - i];
+      if (Properties::getZeroExtValue(parent) == curr) {
         auto bits = Properties::getZeroExtBits(parent);
         if (usage.unsignedUsages == 0) {
           usage.unsignedBits = bits;
@@ -73,17 +81,14 @@ struct PickLoadSigns : public WalkerPass<ExpressionStackWalker<PickLoadSigns>> {
           usage.unsignedBits = 0;
         }
         usage.unsignedUsages++;
-      } else if (expressionStack.size() >= 3) {
-        auto* grandparent = expressionStack[expressionStack.size() - 3];
-        if (Properties::getSignExtValue(grandparent)) {
-          auto bits = Properties::getSignExtBits(grandparent);
-          if (usage.signedUsages == 0) {
-            usage.signedBits = bits;
-          } else if (usage.signedBits != bits) {
-            usage.signedBits = 0;
-          }
-          usage.signedUsages++;
+      } else if (Properties::getSignExtValue(parent) == curr) {
+        auto bits = Properties::getSignExtBits(parent);
+        if (usage.signedUsages == 0) {
+          usage.signedBits = bits;
+        } else if (usage.signedBits != bits) {
+          usage.signedBits = 0;
         }
+        usage.signedUsages++;
       }
     }
   }