[WebAssembly] Optimizes [zext](and 63|32) for shl, srl and sra with selectShiftMask

llvm/lib/Target/WebAssembly/WebAssemblyISelDAGToDAG.cpp

@@ -70,6 +70,15 @@ class WebAssemblyDAGToDAGISel final : public SelectionDAGISel {
   bool SelectAddrOperands32(SDValue Op, SDValue &Offset, SDValue &Addr);
   bool SelectAddrOperands64(SDValue Op, SDValue &Offset, SDValue &Addr);
 
+  bool selectShiftMask(SDValue N, unsigned ShiftWidth, SDValue &ShAmt);
+
+  bool selectShiftMask32(SDValue N, SDValue &ShAmt) {
+    return selectShiftMask(N, 32, ShAmt);
+  }
+
+  bool selectShiftMask64(SDValue N, SDValue &ShAmt) {
+    return selectShiftMask(N, 64, ShAmt);
+  }
 // Include the pieces autogenerated from the target description.
 #include "WebAssemblyGenDAGISel.inc"
 
@@ -548,6 +557,49 @@ bool WebAssemblyDAGToDAGISel::SelectAddrOperands64(SDValue Op, SDValue &Offset,
   return SelectAddrOperands(MVT::i64, WebAssembly::CONST_I64, Op, Offset, Addr);
 }
 
+bool WebAssemblyDAGToDAGISel::selectShiftMask(SDValue N, unsigned ShiftWidth,
+                                              SDValue &ShAmt) {
+
+  ShAmt = N;
+
+  bool HasZext = false;
+  // Peek through zext.
+  if (ShAmt->getOpcode() == ISD::ZERO_EXTEND) {
+    ShAmt = ShAmt.getOperand(0);
+    HasZext = true;
+  }
+
+  if (ShAmt.getOpcode() == ISD::AND &&
+      isa<ConstantSDNode>(ShAmt.getOperand(1))) {
+    const APInt &AndMask = ShAmt.getConstantOperandAPInt(1);
+
+    // Since the max shift amount is a power of 2 we can subtract 1 to make a
+    // mask that covers the bits needed to represent all shift amounts.
+    assert(isPowerOf2_32(ShiftWidth) && "Unexpected max shift amount!");
+    APInt ShMask(AndMask.getBitWidth(), ShiftWidth - 1);
+
+    if (ShMask.isSubsetOf(AndMask)) {
+      ShAmt = ShAmt.getOperand(0);
+    } else {
+      // TODO: port computeKnownBits from riscv in another PR about rotr and
+      // rotl
+      return false;
+    }
+
+    // Only reinstate zext if it's i32 -> i64, WebAssembly would have legalize
+    // i16 to i32 in the dag otherwise.
+    if (HasZext && ShiftWidth == 64) {
+      ShAmt = SDValue(CurDAG->getMachineNode(WebAssembly::I64_EXTEND_U_I32,
+                                             SDLoc(N), MVT::i64, ShAmt),
+                      0);
+    }
+    return true;
+  }
+
+  // TODO: Port rest of riscv if applicable
+  return false;
+}
+
 /// This pass converts a legalized DAG into a WebAssembly-specific DAG, ready
 /// for instruction scheduling.
 FunctionPass *llvm::createWebAssemblyISelDag(WebAssemblyTargetMachine &TM,

llvm/lib/Target/WebAssembly/WebAssemblyInstrInteger.td

@@ -44,6 +44,10 @@ multiclass ComparisonInt<CondCode cond, string name, bits<32> i32Inst, bits<32>
                 !strconcat("i64.", name), i64Inst>;
 }
 
+// ComplexPattern
+def shiftMask32 : ComplexPattern<i32, 1, "selectShiftMask32", [], [], 0>;
+def shiftMask64 : ComplexPattern<i64, 1, "selectShiftMask64", [], [], 0>;
+
 // The spaces after the names are for aesthetic purposes only, to make
 // operands line up vertically after tab expansion.
 let isCommutable = 1 in
@@ -94,26 +98,19 @@ defm EQZ_I64 : I<(outs I32:$dst), (ins I64:$src), (outs), (ins),
                  "i64.eqz \t$dst, $src", "i64.eqz", 0x50>;
 
 // Optimize away an explicit mask on a shift count.
-def : Pat<(shl I32:$lhs, (and I32:$rhs, 31)), (SHL_I32 I32:$lhs, I32:$rhs)>;
-def : Pat<(sra I32:$lhs, (and I32:$rhs, 31)), (SHR_S_I32 I32:$lhs, I32:$rhs)>;
-def : Pat<(srl I32:$lhs, (and I32:$rhs, 31)), (SHR_U_I32 I32:$lhs, I32:$rhs)>;
-def : Pat<(shl I64:$lhs, (and I64:$rhs, 63)), (SHL_I64 I64:$lhs, I64:$rhs)>;
-def : Pat<(sra I64:$lhs, (and I64:$rhs, 63)), (SHR_S_I64 I64:$lhs, I64:$rhs)>;
-def : Pat<(srl I64:$lhs, (and I64:$rhs, 63)), (SHR_U_I64 I64:$lhs, I64:$rhs)>;
+def : Pat<(shl I32:$lhs, (shiftMask32 I32:$rhs)), (SHL_I32 I32:$lhs, I32:$rhs)>;
+def : Pat<(sra I32:$lhs, (shiftMask32 I32:$rhs)), (SHR_S_I32 I32:$lhs, I32:$rhs)>;
+def : Pat<(srl I32:$lhs, (shiftMask32 I32:$rhs)), (SHR_U_I32 I32:$lhs, I32:$rhs)>;
+def : Pat<(shl I64:$lhs, (shiftMask64 I64:$rhs)), (SHL_I64 I64:$lhs, I64:$rhs)>;
+def : Pat<(sra I64:$lhs, (shiftMask64 I64:$rhs)), (SHR_S_I64 I64:$lhs, I64:$rhs)>;
+def : Pat<(srl I64:$lhs, (shiftMask64 I64:$rhs)), (SHR_U_I64 I64:$lhs, I64:$rhs)>;
 
 // Optimize away an explicit mask on a rotate count.
 def : Pat<(rotl I32:$lhs, (and I32:$rhs, 31)), (ROTL_I32 I32:$lhs, I32:$rhs)>;
 def : Pat<(rotr I32:$lhs, (and I32:$rhs, 31)), (ROTR_I32 I32:$lhs, I32:$rhs)>;
 def : Pat<(rotl I64:$lhs, (and I64:$rhs, 63)), (ROTL_I64 I64:$lhs, I64:$rhs)>;
 def : Pat<(rotr I64:$lhs, (and I64:$rhs, 63)), (ROTR_I64 I64:$lhs, I64:$rhs)>;
 
-def : Pat<(shl I64:$lhs, (zext (and I32:$rhs, 63))),
-                               (SHL_I64 I64:$lhs, (I64_EXTEND_U_I32 I32:$rhs))>;
-def : Pat<(sra I64:$lhs, (zext (and I32:$rhs, 63))),
-                               (SHR_S_I64 I64:$lhs, (I64_EXTEND_U_I32 I32:$rhs))>;
-def : Pat<(srl I64:$lhs, (zext (and I32:$rhs, 63))),
-                               (SHR_U_I64 I64:$lhs, (I64_EXTEND_U_I32 I32:$rhs))>;
-
 defm SELECT_I32 : I<(outs I32:$dst), (ins I32:$lhs, I32:$rhs, I32:$cond),
                     (outs), (ins),
                     [(set I32:$dst, (select I32:$cond, I32:$lhs, I32:$rhs))],

llvm/test/CodeGen/WebAssembly/disable-feature.ll

@@ -13,8 +13,6 @@ define i8 @not_use_extend8_s(i8 %v, i8 %x) {
 ; CHECK-NEXT:    i32.const 24
 ; CHECK-NEXT:    i32.shr_s
 ; CHECK-NEXT:    local.get 1
-; CHECK-NEXT:    i32.const 255
-; CHECK-NEXT:    i32.and
 ; CHECK-NEXT:    i32.shr_s
 ; CHECK-NEXT:    # fallthrough-return
   %a = ashr i8 %v, %x

llvm/test/CodeGen/WebAssembly/legalize.ll

@@ -14,9 +14,7 @@ define i3 @shl_i3(i3 %a, i3 %b, ptr %p) {
 }
 
 ; CHECK-LABEL: shl_i53:
-; CHECK: i64.const   $push0=, 9007199254740991{{$}}
-; CHECK: i64.and     $push1=, $1, $pop0{{$}}
-; CHECK: i64.shl     $push2=, $0, $pop1{{$}}
+; CHECK: i64.shl     $push0=, $0, $1
 define i53 @shl_i53(i53 %a, i53 %b, ptr %p) {
   %t = shl i53 %a, %b
   ret i53 %t