[SelectionDAG] Constant fold ISD:FSHL/FSHR nodes

Author: XChy

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

@@ -11281,6 +11281,11 @@ SDValue DAGCombiner::visitFunnelShift(SDNode *N) {
   unsigned BitWidth = VT.getScalarSizeInBits();
   SDLoc DL(N);
 
+  // fold (fshl C0, C1, C2) -> C3
+  if (SDValue C =
+          DAG.FoldConstantArithmetic(N->getOpcode(), DL, VT, {N0, N1, N2}))
+    return C;
+
   // fold (fshl N0, N1, 0) -> N0
   // fold (fshr N0, N1, 0) -> N1
   if (isPowerOf2_32(BitWidth))

llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp

@@ -7175,6 +7175,28 @@ SDValue SelectionDAG::FoldConstantArithmetic(unsigned Opcode, const SDLoc &DL,
     }
   }
 
+  // Handle fshl/fshr special cases.
+  if (Opcode == ISD::FSHL || Opcode == ISD::FSHR) {
+    auto *C1 = dyn_cast<ConstantSDNode>(Ops[0]);
+    auto *C2 = dyn_cast<ConstantSDNode>(Ops[1]);
+    auto *C3 = dyn_cast<ConstantSDNode>(Ops[2]);
+
+    if (C1 && C2 && C3) {
+      if (C1->isOpaque() || C2->isOpaque() || C3->isOpaque())
+        return SDValue();
+      const APInt V1 = C1->getAPIntValue(), V2 = C2->getAPIntValue(),
+                  V3 = C3->getAPIntValue();
+
+      APInt FoldedVal = Opcode == ISD::FSHL ? APIntOps::fshl(V1, V2, V3)
+                                            : APIntOps::fshr(V1, V2, V3);
+
+      SDValue Folded = getConstant(FoldedVal, DL, VT);
+      assert((!Folded || !VT.isVector()) &&
+             "Can't fold vectors ops with scalar operands");
+      return Folded;
+    }
+  }
+
   // This is for vector folding only from here on.
   if (!VT.isVector())
     return SDValue();
@@ -8158,6 +8180,12 @@ SDValue SelectionDAG::getNode(unsigned Opcode, const SDLoc &DL, EVT VT,
     }
     break;
   }
+  case ISD::FSHL:
+  case ISD::FSHR:
+    // Constant folding.
+    if (SDValue V = FoldConstantArithmetic(Opcode, DL, VT, {N1, N2, N3}))
+      return V;
+    break;
   case ISD::BUILD_VECTOR: {
     // Attempt to simplify BUILD_VECTOR.
     SDValue Ops[] = {N1, N2, N3};

llvm/test/CodeGen/X86/fshl-fshr-constant.ll

@@ -10,9 +10,7 @@ define <4 x i32> @test_fshl_constants() {
 ;
 ; CHECK-UNEXPAND-LABEL: test_fshl_constants:
 ; CHECK-UNEXPAND:       # %bb.0:
-; CHECK-UNEXPAND-NEXT:    vpmovsxbd {{.*#+}} xmm1 = [4,5,6,7]
-; CHECK-UNEXPAND-NEXT:    vpmovsxbd {{.*#+}} xmm0 = [0,1,2,3]
-; CHECK-UNEXPAND-NEXT:    vpshldvd {{\.?LCPI[0-9]+_[0-9]+}}(%rip), %xmm1, %xmm0
+; CHECK-UNEXPAND-NEXT:    vpmovsxwd {{.*#+}} xmm0 = [0,512,2048,6144]
 ; CHECK-UNEXPAND-NEXT:    retq
   %res = call <4 x i32> @llvm.fshl.v4i32(<4 x i32> <i32 0, i32 1, i32 2, i32 3>, <4 x i32> <i32 4, i32 5, i32 6, i32 7>, <4 x i32> <i32 8, i32 9, i32 10, i32 11>)
   ret <4 x i32> %res
@@ -79,17 +77,10 @@ define <4 x i32> @test_fshl_none_constant(<4 x i32> %a, <4 x i32> %b, <4 x i32>
 }
 
 define <4 x i32> @test_fshr_constants() {
-; CHECK-EXPAND-LABEL: test_fshr_constants:
-; CHECK-EXPAND:       # %bb.0:
-; CHECK-EXPAND-NEXT:    vmovaps {{.*#+}} xmm0 = [0,8388608,8388608,6291456]
-; CHECK-EXPAND-NEXT:    retq
-;
-; CHECK-UNEXPAND-LABEL: test_fshr_constants:
-; CHECK-UNEXPAND:       # %bb.0:
-; CHECK-UNEXPAND-NEXT:    vpmovsxbd {{.*#+}} xmm1 = [0,1,2,3]
-; CHECK-UNEXPAND-NEXT:    vpmovsxbd {{.*#+}} xmm0 = [4,5,6,7]
-; CHECK-UNEXPAND-NEXT:    vpshrdvd {{\.?LCPI[0-9]+_[0-9]+}}(%rip), %xmm1, %xmm0
-; CHECK-UNEXPAND-NEXT:    retq
+; CHECK-LABEL: test_fshr_constants:
+; CHECK:       # %bb.0:
+; CHECK-NEXT:    vmovaps {{.*#+}} xmm0 = [0,8388608,8388608,6291456]
+; CHECK-NEXT:    retq
   %res = call <4 x i32> @llvm.fshr.v4i32(<4 x i32> <i32 0, i32 1, i32 2, i32 3>, <4 x i32> <i32 4, i32 5, i32 6, i32 7>, <4 x i32> <i32 8, i32 9, i32 10, i32 11>)
   ret <4 x i32> %res
 }