[InstCombine] Fold bitcast [x,...,x] to iN to mul x, C

llvm/lib/Transforms/InstCombine/InstCombineCasts.cpp

@@ -3083,6 +3083,70 @@ static Value *foldCopySignIdioms(BitCastInst &CI,
   return Builder.CreateCopySign(Builder.CreateBitCast(Y, FTy), X);
 }
 
+/// bitcast (shuf X, Y, splat_mask) to iN --> (zext x) * C
+/// where x is the splatted integer source element with bitwidth W and
+///       C = 1 + 2^W + 2^(2W) + ... = (2^N - 1)/(2^W - 1)
+/// E.g.,
+///   x: i1, y = bitcast [x, x, x, x] --> y = x * 15
+///   x: i8, y = bitcast [x, x, x, x] --> y = x * 16843009
+static Instruction *foldSplatShuffleToMul(const ShuffleVectorInst &Shuf,
+                                          IntegerType *DstTy,
+                                          InstCombiner::BuilderTy &Builder) {
+  // If Shuf has other user besides the bitcast, bail out.
+  if (!Shuf.hasOneUse())
+    return nullptr;
+
+  assert(isa<BitCastInst>(*Shuf.user_begin()) &&
+         "The sole user of shuf must be a bitcast");
+
+  auto *ShufTy = dyn_cast<FixedVectorType>(Shuf.getType());
+  // Cannot support scalable vector.
+  if (!ShufTy)
+    return nullptr;
+  auto *EltTy = dyn_cast<IntegerType>(ShufTy->getElementType());
+  // Restrict this fold to integer splats. Reinterpreting a non-integer splat
+  // element as an integer and then multiplying by C is algebraically sound, but
+  // llvm-mca shows that it can generate worse code than keeping the
+  // splat-vector bitcast form.
+  if (!EltTy)
+    return nullptr;
+
+  const unsigned DstWidth = DstTy->getBitWidth();
+  assert(DstWidth == ShufTy->getPrimitiveSizeInBits().getFixedValue() &&
+         "bitcast width mismatch");
+  // It would be less beneficial when the dest type is so large that it needs to
+  // be legalized in the backend.
+  if (!Shuf.getDataLayout().fitsInLegalInteger(DstWidth))
+    return nullptr;
+
+  ArrayRef<int> Mask = Shuf.getShuffleMask();
+
+  // Check if this is a splat-shuffle with a valid index
+  if (!all_equal(Mask) || Mask[0] == PoisonMaskElem)
+    return nullptr;
+
+  // Get the value to splat via the splat index.
+  unsigned SplatIndex = static_cast<unsigned>(Mask[0]);
+  Value *SplatSource = Shuf.getOperand(0);
+  unsigned NumElts =
+      cast<FixedVectorType>(SplatSource->getType())->getNumElements();
+  if (SplatIndex >= NumElts) {
+    SplatSource = Shuf.getOperand(1);
+    SplatIndex -= NumElts;
+  }
+
+  assert(SplatIndex < NumElts &&
+         "splat index must be within the selected shuffle source");
+
+  // bitcast (splat x) to integer is:
+  //   y = x * C, where C = 1 + 2^W + 2^(2W) + ...
+  // and W is the source element width.
+  Value *Splat = Builder.CreateExtractElement(SplatSource, SplatIndex);
+  APInt MulC = APInt::getSplat(DstWidth, APInt(EltTy->getBitWidth(), 1));
+  Value *WideSplat = Builder.CreateZExt(Splat, DstTy);
+  return BinaryOperator::CreateMul(WideSplat, ConstantInt::get(DstTy, MulC));
+}
+
 Instruction *InstCombinerImpl::visitBitCast(BitCastInst &CI) {
   // If the operands are integer typed then apply the integer transforms,
   // otherwise just apply the common ones.
@@ -3161,8 +3225,14 @@ Instruction *InstCombinerImpl::visitBitCast(BitCastInst &CI) {
   }
 
   if (auto *Shuf = dyn_cast<ShuffleVectorInst>(Src)) {
-    // Okay, we have (bitcast (shuffle ..)).  Check to see if this is
-    // a bitcast to a vector with the same # elts.
+    // Okay, we have (bitcast (shuffle ..)).
+
+    // Try to fold `bitcast [x, ..., x] to iN` into `(zext x) * C`
+    if (auto *DstIntTy = dyn_cast<IntegerType>(DestTy))
+      if (Instruction *I = foldSplatShuffleToMul(*Shuf, DstIntTy, Builder))
+        return I;
+
+    // Check to see if this is a bitcast to a vector with the same # elts.
     Value *ShufOp0 = Shuf->getOperand(0);
     Value *ShufOp1 = Shuf->getOperand(1);
     auto ShufElts = cast<VectorType>(Shuf->getType())->getElementCount();

llvm/test/Transforms/InstCombine/bitcast-splat-mul.ll

@@ -0,0 +1,104 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 6
+; RUN: opt < %s -passes=instcombine -S | FileCheck %s
+
+; splat-vec to mul is related to target arch.
+; E.g., i256 mul is not efficient in x86_64.
+target triple = "x86_64-pc-linux-gnu"
+
+define i32 @issue185694_i1(i64 %arg0, i64 %arg1) {
+; CHECK-LABEL: define i32 @issue185694_i1(
+; CHECK-SAME: i64 [[ARG0:%.*]], i64 [[ARG1:%.*]]) {
+; CHECK-NEXT:    [[CMP:%.*]] = icmp eq i64 [[ARG1]], [[ARG0]]
+; CHECK-NEXT:    [[EXT:%.*]] = select i1 [[CMP]], i32 15, i32 0
+; CHECK-NEXT:    ret i32 [[EXT]]
+;
+  %cmp = icmp eq i64 %arg1, %arg0
+  %ins = insertelement <4 x i1> poison, i1 %cmp, i64 0
+  %splat = shufflevector <4 x i1> %ins, <4 x i1> poison, <4 x i32> zeroinitializer
+  %bc = bitcast <4 x i1> %splat to i4
+  %ext = zext i4 %bc to i32
+  ret i32 %ext
+}
+
+define i32 @splat_i8_nonzero_lane(i8 %x) {
+; CHECK-LABEL: define i32 @splat_i8_nonzero_lane(
+; CHECK-SAME: i8 [[X:%.*]]) {
+; CHECK-NEXT:    [[TMP1:%.*]] = zext i8 [[X]] to i32
+; CHECK-NEXT:    [[MUL:%.*]] = mul nuw i32 [[TMP1]], 16843009
+; CHECK-NEXT:    ret i32 [[MUL]]
+;
+  %ins = insertelement <4 x i8> poison, i8 %x, i64 2
+  %splat = shufflevector <4 x i8> %ins, <4 x i8> poison, <4 x i32> <i32 2, i32 2, i32 2, i32 2>
+  %bc = bitcast <4 x i8> %splat to i32
+  ret i32 %bc
+}
+
+define i64 @zext_splat_i8_to_i64(i8 %x) {
+; CHECK-LABEL: define i64 @zext_splat_i8_to_i64(
+; CHECK-SAME: i8 [[X:%.*]]) {
+; CHECK-NEXT:    [[TMP1:%.*]] = zext i8 [[X]] to i64
+; CHECK-NEXT:    [[EXT:%.*]] = mul nuw nsw i64 [[TMP1]], 257
+; CHECK-NEXT:    ret i64 [[EXT]]
+;
+  %ins = insertelement <2 x i8> poison, i8 %x, i64 0
+  %splat = shufflevector <2 x i8> %ins, <2 x i8> poison, <2 x i32> zeroinitializer
+  %bc = bitcast <2 x i8> %splat to i16
+  %ext = zext i16 %bc to i64
+  ret i64 %ext
+}
+
+define i64 @splat_float(float %x) {
+; CHECK-LABEL: define i64 @splat_float(
+; CHECK-SAME: float [[X:%.*]]) {
+; CHECK-NEXT:    [[TMP1:%.*]] = insertelement <2 x float> poison, float [[X]], i64 0
+; CHECK-NEXT:    [[SPLAT:%.*]] = shufflevector <2 x float> [[TMP1]], <2 x float> poison, <2 x i32> zeroinitializer
+; CHECK-NEXT:    [[BC:%.*]] = bitcast <2 x float> [[SPLAT]] to i64
+; CHECK-NEXT:    ret i64 [[BC]]
+;
+  %ins = insertelement <2 x float> poison, float %x, i64 1
+  %splat = shufflevector <2 x float> %ins, <2 x float> poison, <2 x i32> <i32 1, i32 1>
+  %bc = bitcast <2 x float> %splat to i64
+  ret i64 %bc
+}
+
+define i128 @splat_float_to_i128(float %x) {
+; CHECK-LABEL: define i128 @splat_float_to_i128(
+; CHECK-SAME: float [[X:%.*]]) {
+; CHECK-NEXT:    [[TMP1:%.*]] = insertelement <4 x float> poison, float [[X]], i64 0
+; CHECK-NEXT:    [[SPLAT:%.*]] = shufflevector <4 x float> [[TMP1]], <4 x float> poison, <4 x i32> zeroinitializer
+; CHECK-NEXT:    [[BC:%.*]] = bitcast <4 x float> [[SPLAT]] to i128
+; CHECK-NEXT:    ret i128 [[BC]]
+;
+  %ins = insertelement <4 x float> poison, float %x, i64 3
+  %splat = shufflevector <4 x float> %ins, <4 x float> poison,
+  <4 x i32> <i32 3, i32 3, i32 3, i32 3>
+  %bc = bitcast <4 x float> %splat to i128
+  ret i128 %bc
+}
+
+define i128 @splat_i32_to_i128(i32 %x) {
+; CHECK-LABEL: define i128 @splat_i32_to_i128(
+; CHECK-SAME: i32 [[X:%.*]]) {
+; CHECK-NEXT:    [[TMP1:%.*]] = insertelement <4 x i32> poison, i32 [[X]], i64 0
+; CHECK-NEXT:    [[SPLAT:%.*]] = shufflevector <4 x i32> [[TMP1]], <4 x i32> poison, <4 x i32> zeroinitializer
+; CHECK-NEXT:    [[MUL:%.*]] = bitcast <4 x i32> [[SPLAT]] to i128
+; CHECK-NEXT:    ret i128 [[MUL]]
+;
+  %ins = insertelement <4 x i32> poison, i32 %x, i64 1
+  %splat = shufflevector <4 x i32> %ins, <4 x i32> poison,
+  <4 x i32> <i32 1, i32 1, i32 1, i32 1>
+  %bc = bitcast <4 x i32> %splat to i128
+  ret i128 %bc
+}
+
+define i32 @nonsplat_shuffle(<4 x i8> %x) {
+; CHECK-LABEL: define i32 @nonsplat_shuffle(
+; CHECK-SAME: <4 x i8> [[X:%.*]]) {
+; CHECK-NEXT:    [[SHUF:%.*]] = shufflevector <4 x i8> [[X]], <4 x i8> poison, <4 x i32> <i32 0, i32 1, i32 0, i32 1>
+; CHECK-NEXT:    [[BC:%.*]] = bitcast <4 x i8> [[SHUF]] to i32
+; CHECK-NEXT:    ret i32 [[BC]]
+;
+  %shuf = shufflevector <4 x i8> %x, <4 x i8> poison, <4 x i32> <i32 0, i32 1, i32 0, i32 1>
+  %bc = bitcast <4 x i8> %shuf to i32
+  ret i32 %bc
+}