AlignmentLowering: Handle all possible cases for i64, f32, f64 (#3008)

Previously we only handled i32. That was enough for all real-world code
people have run through wasm2js apparently (which is the only place the
pass is needed - it lowers unaligned loads to individual loads etc., as
unaligned operations fail in JS). Apparently it's pretty rare to have
unaligned f32 loads for example.

This will be useful in fuzzing wasm2js, as without this we can't compare
results to the interpreter (which does alignment properly).

1 files changed, 160 insertions, 18 deletions

diff --git a/src/passes/AlignmentLowering.cpp b/src/passes/AlignmentLowering.cpp
index 818189056..26815b389 100644
--- a/src/passes/AlignmentLowering.cpp
+++ b/src/passes/AlignmentLowering.cpp
@@ -27,16 +27,15 @@
 namespace wasm {
 
 struct AlignmentLowering : public WalkerPass<PostWalker<AlignmentLowering>> {
-  void visitLoad(Load* curr) {
+  // Core lowering of a 32-bit load: ensures it is done using aligned
+  // operations, which means we can leave it alone if it's already aligned, or
+  // else we break it up into smaller loads that are.
+  Expression* lowerLoadI32(Load* curr) {
     if (curr->align == 0 || curr->align == curr->bytes) {
-      return;
+      return curr;
     }
     Builder builder(*getModule());
-    if (curr->type == Type::unreachable) {
-      replaceCurrent(curr->ptr);
-      return;
-    }
-    assert(curr->type == Type::i32); // TODO: i64, f32, f64
+    assert(curr->type == Type::i32);
     auto temp = builder.addVar(getFunction(), Type::i32);
     Expression* ret;
     if (curr->bytes == 2) {
@@ -125,21 +124,16 @@ struct AlignmentLowering : public WalkerPass<PostWalker<AlignmentLowering>> {
     } else {
       WASM_UNREACHABLE("invalid size");
     }
-    replaceCurrent(
-      builder.makeBlock({builder.makeLocalSet(temp, curr->ptr), ret}));
+    return builder.makeBlock({builder.makeLocalSet(temp, curr->ptr), ret});
   }
 
-  void visitStore(Store* curr) {
+  // Core lowering of a 32-bit store.
+  Expression* lowerStoreI32(Store* curr) {
     if (curr->align == 0 || curr->align == curr->bytes) {
-      return;
+      return curr;
     }
     Builder builder(*getModule());
-    if (curr->type == Type::unreachable) {
-      replaceCurrent(builder.makeBlock(
-        {builder.makeDrop(curr->ptr), builder.makeDrop(curr->value)}));
-      return;
-    }
-    assert(curr->value->type == Type::i32); // TODO: i64, f32, f64
+    assert(curr->value->type == Type::i32);
     auto tempPtr = builder.addVar(getFunction(), Type::i32);
     auto tempValue = builder.addVar(getFunction(), Type::i32);
     auto* block =
@@ -222,7 +216,155 @@ struct AlignmentLowering : public WalkerPass<PostWalker<AlignmentLowering>> {
       WASM_UNREACHABLE("invalid size");
     }
     block->finalize();
-    replaceCurrent(block);
+    return block;
+  }
+
+  void visitLoad(Load* curr) {
+    // If unreachable, just remove the load, which removes the unaligned
+    // operation in a trivial way.
+    if (curr->type == Type::unreachable) {
+      replaceCurrent(curr->ptr);
+      return;
+    }
+    if (curr->align == 0 || curr->align == curr->bytes) {
+      // Nothing to do: leave the node unchanged. All code lower down assumes
+      // the operation is unaligned.
+      return;
+    }
+    Builder builder(*getModule());
+    auto type = curr->type.getSingle();
+    Expression* replacement;
+    switch (type) {
+      default:
+        WASM_UNREACHABLE("unhandled unaligned load");
+      case Type::i32:
+        replacement = lowerLoadI32(curr);
+        break;
+      case Type::f32:
+        curr->type = Type::i32;
+        replacement = builder.makeUnary(ReinterpretInt32, lowerLoadI32(curr));
+        break;
+      case Type::i64:
+      case Type::f64:
+        if (type == Type::i64 && curr->bytes != 8) {
+          // A load of <64 bits.
+          curr->type = Type::i32;
+          replacement = builder.makeUnary(
+            curr->signed_ ? ExtendSInt32 : ExtendUInt32, lowerLoadI32(curr));
+          break;
+        }
+        // Load two 32-bit pieces, and combine them.
+        auto temp = builder.addVar(getFunction(), Type::i32);
+        auto* set = builder.makeLocalSet(temp, curr->ptr);
+        Expression* low =
+          lowerLoadI32(builder.makeLoad(4,
+                                        false,
+                                        curr->offset,
+                                        curr->align,
+                                        builder.makeLocalGet(temp, Type::i32),
+                                        Type::i32));
+        low = builder.makeUnary(ExtendUInt32, low);
+        // Note that the alignment is assumed to be the same here, even though
+        // we add an offset of 4. That is because this is an unaligned load, so
+        // the alignment is 1, 2, or 4, which means it stays the same after
+        // adding 4.
+        Expression* high =
+          lowerLoadI32(builder.makeLoad(4,
+                                        false,
+                                        curr->offset + 4,
+                                        curr->align,
+                                        builder.makeLocalGet(temp, Type::i32),
+                                        Type::i32));
+        high = builder.makeUnary(ExtendUInt32, high);
+        high =
+          builder.makeBinary(ShlInt64, high, builder.makeConst(int64_t(32)));
+        auto* combined = builder.makeBinary(OrInt64, low, high);
+        replacement = builder.makeSequence(set, combined);
+        // Ensure the proper output type.
+        if (type == Type::f64) {
+          replacement = builder.makeUnary(ReinterpretInt64, replacement);
+        }
+        break;
+    }
+    replaceCurrent(replacement);
+  }
+
+  void visitStore(Store* curr) {
+    Builder builder(*getModule());
+    // If unreachable, just remove the store, which removes the unaligned
+    // operation in a trivial way.
+    if (curr->type == Type::unreachable) {
+      replaceCurrent(builder.makeBlock(
+        {builder.makeDrop(curr->ptr), builder.makeDrop(curr->value)}));
+      return;
+    }
+    if (curr->align == 0 || curr->align == curr->bytes) {
+      // Nothing to do: leave the node unchanged. All code lower down assumes
+      // the operation is unaligned.
+      return;
+    }
+    auto type = curr->value->type.getSingle();
+    Expression* replacement;
+    switch (type) {
+      default:
+        WASM_UNREACHABLE("unhandled unaligned store");
+      case Type::i32:
+        replacement = lowerStoreI32(curr);
+        break;
+      case Type::f32:
+        curr->type = Type::i32;
+        curr->value = builder.makeUnary(ReinterpretFloat32, curr->value);
+        replacement = lowerStoreI32(curr);
+        break;
+      case Type::i64:
+      case Type::f64:
+        if (type == Type::i64 && curr->bytes != 8) {
+          // A store of <64 bits.
+          curr->type = Type::i32;
+          curr->value = builder.makeUnary(WrapInt64, curr->value);
+          replacement = lowerStoreI32(curr);
+          break;
+        }
+        // Otherwise, fall through to f64 case for a 64-bit load.
+        // Ensure an integer input value.
+        auto* value = curr->value;
+        if (type == Type::f64) {
+          value = builder.makeUnary(ReinterpretFloat64, value);
+        }
+        // Store as two 32-bit pieces.
+        auto tempPtr = builder.addVar(getFunction(), Type::i32);
+        auto* setPtr = builder.makeLocalSet(tempPtr, curr->ptr);
+        auto tempValue = builder.addVar(getFunction(), Type::i64);
+        auto* setValue = builder.makeLocalSet(tempValue, value);
+        Expression* low = builder.makeUnary(
+          WrapInt64, builder.makeLocalGet(tempValue, Type::i64));
+        low = lowerStoreI32(
+          builder.makeStore(4,
+                            curr->offset,
+                            curr->align,
+                            builder.makeLocalGet(tempPtr, Type::i32),
+                            low,
+                            Type::i32));
+        Expression* high =
+          builder.makeBinary(ShrUInt64,
+                             builder.makeLocalGet(tempValue, Type::i64),
+                             builder.makeConst(int64_t(32)));
+        high = builder.makeUnary(WrapInt64, high);
+        // Note that the alignment is assumed to be the same here, even though
+        // we add an offset of 4. That is because this is an unaligned store, so
+        // the alignment is 1, 2, or 4, which means it stays the same after
+        // adding 4.
+        high = lowerStoreI32(
+          builder.makeStore(4,
+                            curr->offset + 4,
+                            curr->align,
+                            builder.makeLocalGet(tempPtr, Type::i32),
+                            high,
+                            Type::i32));
+        replacement = builder.makeBlock({setPtr, setValue, low, high});
+        break;
+    }
+    replaceCurrent(replacement);
   }
 };