[InstCombine] Try to recognize bswap pattern when calling funnel shifts

[junaire]

Summary

Alive2: https://alive2.llvm.org/ce/z/dxxD7B
Fixes: https://github.com/llvm/llvm-project/issues/60690

Signed-off-by: Jun Zhang <[e-mail]>

[RKSimon]

llvm/lib/Transforms/InstCombine/InstCombineCalls.cpp

1779

Why not use test matchBSwapOrBitReverse here?

[nikic]

As @RKSimon said, this transform is in the wrong place. It has nothing to do with shuffled intrinsic operands...

[junaire]

Put the transform in the right place.

[junaire]

Update existing tests

[nikic]

llvm/lib/Transforms/InstCombine/InstCombineCalls.cpp

1796

This code block can probably be dropped now, if it's already handled by the preceding?

[junaire]

llvm/test/Transforms/InstCombine/fsh.ll

677 ↗

(On Diff #507635)

Looks like this transform somehow breaks the existing optimization? Any thoughts about how to fix it? My idea is to add another default parameter called bool ForceTrunc = true and ask matchBSwapOrBitReverse not to truncate in this case. However, I don't know if this is appropriate and please let me know if you have better solutions...

[nikic]

llvm/lib/Transforms/InstCombine/InstCombineCalls.cpp

1796

Or you can move your code after the existing one, which will avoid the regression in fshl_mask_args_same2. It shouldn't be necessary to handle this specially, but clearly we're missing a zext(shl(trunc)) fold.

[junaire]

llvm/lib/Transforms/InstCombine/InstCombineCalls.cpp

1796

No that doesn't work since the existing transform is highly depend on the bit-width (only valid if it's 16 but in the test that's i32)

Or maybe we can leave it as is and add a zext(shl(trunc)) fold in the future?

[nikic]

LGTM

llvm/test/Transforms/InstCombine/fsh.ll

677 ↗

(On Diff #507635)

We should just add support for folding this trunc+shl+zext pattern. I've filed https://github.com/llvm/llvm-project/issues/61650 to track.

llvm/lib/Transforms/InstCombine/InstCombineCalls.cpp

static Instruction *factorizeMinMaxTree(IntrinsicInst *II) {
 
   Module *Mod = II->getModule();
   Function *MinMax = Intrinsic::getDeclaration(Mod, MinMaxID, II->getType());
   return CallInst::Create(MinMax, { MinMaxOp, ThirdOp });
 }
 
 /// If all arguments of the intrinsic are unary shuffles with the same mask,
 /// try to shuffle after the intrinsic.
-static Instruction *
-foldShuffledIntrinsicOperands(IntrinsicInst *II,
-                              InstCombiner::BuilderTy &Builder) {
+static Instruction *foldShuffledIntrinsicOperands(
+    IntrinsicInst *II, InstCombiner::BuilderTy &Builder, InstCombinerImpl &IC) {
   // TODO: This should be extended to handle other intrinsics like fshl, ctpop,
   //       etc. Use llvm::isTriviallyVectorizable() and related to determine
   //       which intrinsics are safe to shuffle?
   switch (II->getIntrinsicID()) {
   case Intrinsic::smax:
   case Intrinsic::smin:
   case Intrinsic::umax:
   case Intrinsic::umin:
   case Intrinsic::fma:
-  case Intrinsic::fshl:
   case Intrinsic::fshr:
     break;
+  case Intrinsic::fshl: {
+    if (Instruction *BitOp =
+            IC.matchBSwapOrBitReverse(*II, /*MatchBSwaps*/ true,
+                                      /*MatchBitReversals*/ true))
+      return BitOp;
+    break;
+  }
   default:
     return nullptr;
   }
 
   Value *X;
   ArrayRef<int> Mask;
   if (!match(II->getArgOperand(0),
              m_Shuffle(m_Value(X), m_Undef(), m_Mask(Mask))))
if (match(II->getArgOperand(2), m_ImmConstant(ShAmtC))) {
         // fshr X, Y, C --> fshl X, Y, (BitWidth - C)
         Constant *LeftShiftC = ConstantExpr::getSub(WidthC, ShAmtC);
         Module *Mod = II->getModule();
         Function *Fshl = Intrinsic::getDeclaration(Mod, Intrinsic::fshl, Ty);
         return CallInst::Create(Fshl, { Op0, Op1, LeftShiftC });
       }
       assert(IID == Intrinsic::fshl &&
              "All funnel shifts by simple constants should go left");
 
       // fshl(X, 0, C) --> shl X, C
       // fshl(X, undef, C) --> shl X, C
       if (match(Op1, m_ZeroInt()) || match(Op1, m_Undef()))
         return BinaryOperator::CreateShl(Op0, ShAmtC);
 
       // fshl(0, X, C) --> lshr X, (BW-C)
       // fshl(undef, X, C) --> lshr X, (BW-C)
       if (match(Op0, m_ZeroInt()) || match(Op0, m_Undef()))
         return BinaryOperator::CreateLShr(Op1,
                                           ConstantExpr::getSub(WidthC, ShAmtC));
 
       // fshl i16 X, X, 8 --> bswap i16 X (reduce to more-specific form)
       if (Op0 == Op1 && BitWidth == 16 && match(ShAmtC, m_SpecificInt(8))) {
         Module *Mod = II->getModule();
         Function *Bswap = Intrinsic::getDeclaration(Mod, Intrinsic::bswap, Ty);
         return CallInst::Create(Bswap, { Op0 });
       }
     }
 
     // Left or right might be masked.
     if (SimplifyDemandedInstructionBits(*II))
       return &CI;
 
     // The shift amount (operand 2) of a funnel shift is modulo the bitwidth,
     // so only the low bits of the shift amount are demanded if the bitwidth is
default: {
     // Handle target specific intrinsics
     std::optional<Instruction *> V = targetInstCombineIntrinsic(*II);
     if (V)
       return *V;
     break;
   }
   }
 
-  if (Instruction *Shuf = foldShuffledIntrinsicOperands(II, Builder))
+  if (Instruction *Shuf = foldShuffledIntrinsicOperands(II, Builder, *this))
     return Shuf;
 
   // Some intrinsics (like experimental_gc_statepoint) can be used in invoke
   // context, so it is handled in visitCallBase and we should trigger it.
   return visitCallBase(*II);
 }
 
 // Fence instruction simplification

llvm/test/Transforms/InstCombine/bswap.ll

;
   %t4 = lshr i64 %t0, 56
   %t5 = or i64 %t3, %t4
   %t6 = trunc i64 %t5 to i32
   ret i32 %t6
 }
 
 define i64 @call_fshl(i64 %result) {
 ; CHECK-LABEL: @call_fshl(
-; CHECK-NEXT:    [[AND_I:%.*]] = lshr i64 [[RESULT:%.*]], 8
-; CHECK-NEXT:    [[SHR_I:%.*]] = and i64 [[AND_I]], 71777214294589695
-; CHECK-NEXT:    [[AND1_I:%.*]] = shl i64 [[RESULT]], 8
-; CHECK-NEXT:    [[SHL_I:%.*]] = and i64 [[AND1_I]], -71777214294589696
-; CHECK-NEXT:    [[OR_I:%.*]] = or i64 [[SHR_I]], [[SHL_I]]
-; CHECK-NEXT:    [[AND_I7:%.*]] = shl i64 [[OR_I]], 16
-; CHECK-NEXT:    [[SHL_I8:%.*]] = and i64 [[AND_I7]], -281470681808896
-; CHECK-NEXT:    [[AND1_I9:%.*]] = lshr i64 [[OR_I]], 16
-; CHECK-NEXT:    [[SHR_I10:%.*]] = and i64 [[AND1_I9]], 281470681808895
-; CHECK-NEXT:    [[OR_I11:%.*]] = or i64 [[SHL_I8]], [[SHR_I10]]
-; CHECK-NEXT:    [[OR_I12:%.*]] = tail call i64 @llvm.fshl.i64(i64 [[OR_I11]], i64 [[OR_I11]], i64 32)
+; CHECK-NEXT:    [[OR_I12:%.*]] = call i64 @llvm.bswap.i64(i64 [[RESULT:%.*]])
 ; CHECK-NEXT:    ret i64 [[OR_I12]]
 ;
   %and.i = lshr i64 %result, 8
   %shr.i = and i64 %and.i, 71777214294589695
   %and1.i = shl i64 %result, 8
   %shl.i = and i64 %and1.i, -71777214294589696
   %or.i = or i64 %shr.i, %shl.i
   %and.i7 = shl i64 %or.i, 16
   %shl.i8 = and i64 %and.i7, -281470681808896
   %and1.i9 = lshr i64 %or.i, 16
   %shr.i10 = and i64 %and1.i9, 281470681808895
   %or.i11 = or i64 %shl.i8, %shr.i10
   %or.i12 = tail call i64 @llvm.fshl.i64(i64 %or.i11, i64 %or.i11, i64 32)
   ret i64 %or.i12
 }
 declare i64 @llvm.fshl.i64(i64, i64, i64)