[ConstantFolding] add simplifications for funnel shift intrinsics

[spatel]

Summary

This is another step towards being able to canonicalize to the funnel shift intrinsics in IR (see D49242 for the initial patch). We should not have any loss of simplification power in IR between these and the equivalent IR constructs.

[fabiang]

LGTM.

[craig.topper]

lib/Analysis/ConstantFolding.cpp

2096 ↗

(On Diff #161039)

Can we use APInt::urem(uint64_t) version here since bitwidth is guaranteed 32-bits? Then your ShAmt would be unsigned. Then LshrAmt and ShlAmt can be unsigned. And we get rid of a lot of temporary APInts which could save heap allocations if bit width is more than 64.

[spatel]

Patch updated:
Use native integers where possible to make the math more efficient.

[craig.topper]

LGTM

llvm/trunk/lib/Analysis/ConstantFolding.cpp

bool llvm::canConstantFoldCallTo(ImmutableCallSite CS, const Function *F) {
   case Intrinsic::rint:
   case Intrinsic::nearbyint:
   case Intrinsic::pow:
   case Intrinsic::powi:
   case Intrinsic::bswap:
   case Intrinsic::ctpop:
   case Intrinsic::ctlz:
   case Intrinsic::cttz:
+  case Intrinsic::fshl:
+  case Intrinsic::fshr:
   case Intrinsic::fma:
   case Intrinsic::fmuladd:
   case Intrinsic::copysign:
   case Intrinsic::launder_invariant_group:
   case Intrinsic::strip_invariant_group:
   case Intrinsic::round:
   case Intrinsic::masked_load:
   case Intrinsic::sadd_with_overflow:
if (const auto *Op2 = dyn_cast<ConstantFP>(Operands[1])) {
 
           return nullptr;
         }
         }
       }
     }
   }
 
+  if (IntrinsicID == Intrinsic::fshl || IntrinsicID == Intrinsic::fshr) {
+    auto *C0 = dyn_cast<ConstantInt>(Operands[0]);
+    auto *C1 = dyn_cast<ConstantInt>(Operands[1]);
+    auto *C2 = dyn_cast<ConstantInt>(Operands[2]);
+    if (!(C0 && C1 && C2))
+      return nullptr;
+
+    // The shift amount is interpreted as modulo the bitwidth. If the shift
+    // amount is effectively 0, avoid UB due to oversized inverse shift below.
+    unsigned BitWidth = C0->getBitWidth();
+    unsigned ShAmt = C2->getValue().urem(BitWidth);
+    bool IsRight = IntrinsicID == Intrinsic::fshr;
+    if (!ShAmt)
+      return IsRight ? C1 : C0;
+
+    // (X << ShlAmt) | (Y >> LshrAmt)
+    const APInt &X = C0->getValue();
+    const APInt &Y = C1->getValue();
+    unsigned LshrAmt = IsRight ? ShAmt : BitWidth - ShAmt;
+    unsigned ShlAmt = !IsRight ? ShAmt : BitWidth - ShAmt;
+    return ConstantInt::get(Ty->getContext(), X.shl(ShlAmt) | Y.lshr(LshrAmt));
+  }
+
   return nullptr;
 }
 
 Constant *ConstantFoldVectorCall(StringRef Name, unsigned IntrinsicID,
                                  VectorType *VTy, ArrayRef<Constant *> Operands,
                                  const DataLayout &DL,
                                  const TargetLibraryInfo *TLI,
                                  ImmutableCallSite CS) {

llvm/trunk/test/Analysis/ConstantFolding/funnel-shift.ll

 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
 ; RUN: opt < %s -constprop -S -o - | FileCheck %s
 
 declare i32 @llvm.fshl.i32(i32, i32, i32)
 declare i32 @llvm.fshr.i32(i32, i32, i32)
 declare i7 @llvm.fshl.i7(i7, i7, i7)
 declare i7 @llvm.fshr.i7(i7, i7, i7)
 declare <4 x i8> @llvm.fshl.v4i8(<4 x i8>, <4 x i8>, <4 x i8>)
 declare <4 x i8> @llvm.fshr.v4i8(<4 x i8>, <4 x i8>, <4 x i8>)
 
 ; extract(concat(0x12345678, 0xABCDEF01) << 5) = 0x468ACF15
 
 define i32 @fshl_i32() {
 ; CHECK-LABEL: @fshl_i32(
-; CHECK-NEXT:    [[F:%.*]] = call i32 @llvm.fshl.i32(i32 305419896, i32 -1412567295, i32 5)
-; CHECK-NEXT:    ret i32 [[F]]
+; CHECK-NEXT:    ret i32 1183502101
 ;
   %f = call i32 @llvm.fshl.i32(i32 305419896, i32 2882400001, i32 5)
   ret i32 %f
 }
 
 ; extract(concat(0x12345678, 0xABCDEF01) >> 5) = 0xC55E6F78
 ; Try an oversized shift to test modulo functionality.
 
 define i32 @fshr_i32() {
 ; CHECK-LABEL: @fshr_i32(
-; CHECK-NEXT:    [[F:%.*]] = call i32 @llvm.fshr.i32(i32 305419896, i32 -1412567295, i32 37)
-; CHECK-NEXT:    ret i32 [[F]]
+; CHECK-NEXT:    ret i32 -983666824
 ;
   %f = call i32 @llvm.fshr.i32(i32 305419896, i32 2882400001, i32 37)
   ret i32 %f
 }
 
 ; Use a weird type.
 ; Try an oversized shift to test modulo functionality.
 
 ; extract(concat(0b1110000, 0b1111111) << 2) = 0b1000011
 
 define i7 @fshl_i7() {
 ; CHECK-LABEL: @fshl_i7(
-; CHECK-NEXT:    [[F:%.*]] = call i7 @llvm.fshl.i7(i7 -16, i7 -1, i7 9)
-; CHECK-NEXT:    ret i7 [[F]]
+; CHECK-NEXT:    ret i7 -61
 ;
   %f = call i7 @llvm.fshl.i7(i7 112, i7 127, i7 9)
   ret i7 %f
 }
 
 ; extract(concat(0b1110000, 0b1111111) >> 2) = 0b0011111
 ; Try an oversized shift to test modulo functionality.
 
 define i7 @fshr_i7() {
 ; CHECK-LABEL: @fshr_i7(
-; CHECK-NEXT:    [[F:%.*]] = call i7 @llvm.fshr.i7(i7 -16, i7 -1, i7 16)
-; CHECK-NEXT:    ret i7 [[F]]
+; CHECK-NEXT:    ret i7 31
 ;
   %f = call i7 @llvm.fshr.i7(i7 112, i7 127, i7 16)
   ret i7 %f
 }
 
 ; Vectors are folded by handling each scalar element individually, so this is the equivalent of 4 scalar tests:
 ; extract(concat(0x00, 0xFF) << 0) = 0x00
 ; extract(concat(0xFF, 0x00) << 0) = 0xFF
 ; extract(concat(0x10, 0x55) << 1) = 0x20
 ; extract(concat(0x11, 0xAA) << 2) = 0x46
 
 define <4 x i8> @fshl_v4i8() {
 ; CHECK-LABEL: @fshl_v4i8(
-; CHECK-NEXT:    [[F:%.*]] = call <4 x i8> @llvm.fshl.v4i8(<4 x i8> <i8 0, i8 -1, i8 16, i8 17>, <4 x i8> <i8 -1, i8 0, i8 85, i8 -86>, <4 x i8> <i8 0, i8 8, i8 9, i8 10>)
-; CHECK-NEXT:    ret <4 x i8> [[F]]
+; CHECK-NEXT:    ret <4 x i8> <i8 0, i8 -1, i8 32, i8 70>
 ;
   %f = call <4 x i8> @llvm.fshl.v4i8(<4 x i8> <i8 0, i8 -1, i8 16, i8 17>, <4 x i8> <i8 -1, i8 0, i8 85, i8 170>, <4 x i8> <i8 0, i8 8, i8 9, i8 10>)
   ret <4 x i8> %f
 }
 
 ; Vectors are folded by handling each scalar element individually, so this is the equivalent of 4 scalar tests:
 ; extract(concat(0x00, 0xFF) >> 0) = 0xFF
 ; extract(concat(0xFF, 0x00) >> 0) = 0x00
 ; extract(concat(0x10, 0x55) >> 1) = 0x2A
 ; extract(concat(0x11, 0xAA) >> 2) = 0x6A
 
 define <4 x i8> @fshr_v4i8() {
 ; CHECK-LABEL: @fshr_v4i8(
-; CHECK-NEXT:    [[F:%.*]] = call <4 x i8> @llvm.fshr.v4i8(<4 x i8> <i8 0, i8 -1, i8 16, i8 17>, <4 x i8> <i8 -1, i8 0, i8 85, i8 -86>, <4 x i8> <i8 0, i8 8, i8 9, i8 10>)
-; CHECK-NEXT:    ret <4 x i8> [[F]]
+; CHECK-NEXT:    ret <4 x i8> <i8 -1, i8 0, i8 42, i8 106>
 ;
   %f = call <4 x i8> @llvm.fshr.v4i8(<4 x i8> <i8 0, i8 -1, i8 16, i8 17>, <4 x i8> <i8 -1, i8 0, i8 85, i8 170>, <4 x i8> <i8 0, i8 8, i8 9, i8 10>)
   ret <4 x i8> %f
 }