[WebAssembly] Optimize away mask of 63 for shl ( zext (and i32 63)))

llvm/lib/Target/WebAssembly/WebAssemblyISelDAGToDAG.cpp

@@ -70,6 +70,11 @@ 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 selectShiftMask64FromI32(SDValue N, SDValue &ShAmt) {
+    return selectShiftMask(N, 64, ShAmt);
+  }
 // Include the pieces autogenerated from the target description.
 #include "WebAssemblyGenDAGISel.inc"
 
@@ -539,6 +544,36 @@ 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;
+
+  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 {
+      // SimplifyDemandedBits may have optimized the mask so try restoring any
+      // bits that are known zero.
+      KnownBits Known = CurDAG->computeKnownBits(ShAmt.getOperand(0));
+      if (ShMask.isSubsetOf(AndMask | Known.Zero))
+        ShAmt = ShAmt.getOperand(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/WebAssemblyISelLowering.cpp

@@ -3343,7 +3343,6 @@ static SDValue performAnyAllCombine(SDNode *N, SelectionDAG &DAG) {
       Ret = DAG.getNOT(DL, Ret, MVT::i1);
     return DAG.getZExtOrTrunc(Ret, DL, N->getValueType(0));
   };
-
   if (SDValue AnyTrueEQ = CombineSetCC(Intrinsic::wasm_anytrue, ISD::SETEQ,
                                        Intrinsic::wasm_alltrue))
     return AnyTrueEQ;

llvm/lib/Target/WebAssembly/WebAssemblyInstrInfo.td

@@ -460,8 +460,8 @@ def : Pat<(i64 (WebAssemblyWrapperREL texternalsym:$addr)),
 include "WebAssemblyInstrMemory.td"
 include "WebAssemblyInstrCall.td"
 include "WebAssemblyInstrControl.td"
-include "WebAssemblyInstrInteger.td"
 include "WebAssemblyInstrConv.td"
+include "WebAssemblyInstrInteger.td"
 include "WebAssemblyInstrFloat.td"
 include "WebAssemblyInstrAtomics.td"
 include "WebAssemblyInstrSIMD.td"

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 shiftMask64FromI32
+    : ComplexPattern<i32, 1, "selectShiftMask64FromI32", [], [], 0>;
+
 // The spaces after the names are for aesthetic purposes only, to make
 // operands line up vertically after tab expansion.
 let isCommutable = 1 in
@@ -101,6 +105,9 @@ 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 I64:$lhs, (zext(shiftMask64FromI32 I32:$rhs))),
+          (SHL_I64 I64:$lhs, (I64_EXTEND_U_I32 I32:$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)>;

llvm/test/CodeGen/WebAssembly/masked-shifts.ll

@@ -18,6 +18,21 @@ define i32 @shl_i32(i32 %v, i32 %x) {
   ret i32 %a
 }
 
+define i64 @shl_i64_i32(i64 %v, i32 %x) {
+; CHECK-LABEL: shl_i64_i32:
+; CHECK:         .functype shl_i64_i32 (i64, i32) -> (i64)
+; CHECK-NEXT:  # %bb.0:
+; CHECK-NEXT:    local.get 0
+; CHECK-NEXT:    local.get 1
+; CHECK-NEXT:    i64.extend_i32_u
+; CHECK-NEXT:    i64.shl
+; CHECK-NEXT:    # fallthrough-return
+  %m = and i32 %x, 63
+  %z = zext i32 %m to i64
+  %a = shl i64 %v, %z
+  ret i64 %a
+}
+
 define i32 @sra_i32(i32 %v, i32 %x) {
 ; CHECK-LABEL: sra_i32:
 ; CHECK:         .functype sra_i32 (i32, i32) -> (i32)