diff --git a/llvm/lib/Transforms/InstCombine/InstCombineCasts.cpp b/llvm/lib/Transforms/InstCombine/InstCombineCasts.cpp index 0cd2c09726a2d..b472e412a73bd 100644 --- a/llvm/lib/Transforms/InstCombine/InstCombineCasts.cpp +++ b/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(*Shuf.user_begin()) && + "The sole user of shuf must be a bitcast"); + + auto *ShufTy = dyn_cast(Shuf.getType()); + // Cannot support scalable vector. + if (!ShufTy) + return nullptr; + auto *EltTy = dyn_cast(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 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(Mask[0]); + Value *SplatSource = Shuf.getOperand(0); + unsigned NumElts = + cast(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(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(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(Shuf->getType())->getElementCount(); diff --git a/llvm/test/Transforms/InstCombine/bitcast-splat-mul.ll b/llvm/test/Transforms/InstCombine/bitcast-splat-mul.ll new file mode 100644 index 0000000000000..ecd3b84897e7b --- /dev/null +++ b/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> + %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> + %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> + %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> + %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> +; 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> + %bc = bitcast <4 x i8> %shuf to i32 + ret i32 %bc +}