diff --git a/llvm/lib/Transforms/InstCombine/InstructionCombining.cpp b/llvm/lib/Transforms/InstCombine/InstructionCombining.cpp index 206d41e30db2c..ed94866482517 100644 --- a/llvm/lib/Transforms/InstCombine/InstructionCombining.cpp +++ b/llvm/lib/Transforms/InstCombine/InstructionCombining.cpp @@ -2271,6 +2271,27 @@ Instruction *InstCombinerImpl::foldVectorBinop(BinaryOperator &Inst) { } } + // Similar to the combine above, but handles the case for scalable vectors + // where both shuffle(V1, 0) and C are splats. + // + // Op(shuffle(V1, 0), (splat C)) -> shuffle(Op(V1, (splat C)), 0) + if (isa(Inst.getType()) && + match(&Inst, m_c_BinOp(m_OneUse(m_Shuffle(m_Value(V1), m_Poison(), + m_ZeroMask())), + m_ImmConstant(C)))) { + if (Constant *Splat = C->getSplatValue()) { + bool ConstOp1 = isa(RHS); + VectorType *V1Ty = cast(V1->getType()); + Constant *NewC = ConstantVector::getSplat(V1Ty->getElementCount(), Splat); + + Value *NewLHS = ConstOp1 ? V1 : NewC; + Value *NewRHS = ConstOp1 ? NewC : V1; + VectorType *VTy = cast(Inst.getType()); + SmallVector Mask(VTy->getElementCount().getKnownMinValue(), 0); + return createBinOpShuffle(NewLHS, NewRHS, Mask); + } + } + // Try to reassociate to sink a splat shuffle after a binary operation. if (Inst.isAssociative() && Inst.isCommutative()) { // Canonicalize shuffle operand as LHS. diff --git a/llvm/test/Transforms/InstCombine/getelementptr.ll b/llvm/test/Transforms/InstCombine/getelementptr.ll index feba952919b9a..61236df80bfa6 100644 --- a/llvm/test/Transforms/InstCombine/getelementptr.ll +++ b/llvm/test/Transforms/InstCombine/getelementptr.ll @@ -282,8 +282,8 @@ define <2 x i1> @test13_fixed_scalable(i64 %X, ptr %P, <2 x i64> %y) nounwind { define @test13_scalable_scalable(i64 %X, ptr %P, %y) nounwind { ; CHECK-LABEL: @test13_scalable_scalable( ; CHECK-NEXT: [[DOTSPLATINSERT:%.*]] = insertelement poison, i64 [[X:%.*]], i64 0 -; CHECK-NEXT: [[DOTSPLAT:%.*]] = shufflevector [[DOTSPLATINSERT]], poison, zeroinitializer -; CHECK-NEXT: [[A_IDX:%.*]] = shl nsw [[DOTSPLAT]], splat (i64 3) +; CHECK-NEXT: [[TMP3:%.*]] = shl nsw [[DOTSPLATINSERT]], splat (i64 3) +; CHECK-NEXT: [[A_IDX:%.*]] = shufflevector [[TMP3]], poison, zeroinitializer ; CHECK-NEXT: [[TMP1:%.*]] = call i64 @llvm.vscale.i64() ; CHECK-NEXT: [[TMP2:%.*]] = shl i64 [[TMP1]], 4 ; CHECK-NEXT: [[DOTSPLATINSERT1:%.*]] = insertelement poison, i64 [[TMP2]], i64 0 diff --git a/llvm/test/Transforms/InstCombine/vec_shuffle-inseltpoison.ll b/llvm/test/Transforms/InstCombine/vec_shuffle-inseltpoison.ll index c6329af164623..74d07fea9793a 100644 --- a/llvm/test/Transforms/InstCombine/vec_shuffle-inseltpoison.ll +++ b/llvm/test/Transforms/InstCombine/vec_shuffle-inseltpoison.ll @@ -1789,3 +1789,39 @@ define <4 x i32> @PR46872(<4 x i32> %x) { ret <4 x i32> %a } +define @scalable_splat_binop_constant_rhs( %x) { +; CHECK-LABEL: @scalable_splat_binop_constant_rhs( +; CHECK-NEXT: [[R1:%.*]] = add [[R:%.*]], splat (i32 42) +; CHECK-NEXT: [[R2:%.*]] = shufflevector [[R1]], poison, zeroinitializer +; CHECK-NEXT: ret [[R2]] +; + + %splatx = shufflevector %x, poison, zeroinitializer + %r = add %splatx, splat (i32 42) + ret %r +} + +define @scalable_splat_binop_constant_lhs( %x) { +; CHECK-LABEL: @scalable_splat_binop_constant_lhs( +; CHECK-NEXT: [[R1:%.*]] = fadd [[R:%.*]], splat (float 4.200000e+01) +; CHECK-NEXT: [[R2:%.*]] = shufflevector [[R1]], poison, zeroinitializer +; CHECK-NEXT: ret [[R2]] +; + + %splatx = shufflevector %x, poison, zeroinitializer + %r = fadd splat (float 42.0), %splatx + ret %r +} + +; Negative test - shouldn't pull shuffle out as it udiv isn't safe to speculate. +define @scalable_splat_binop_constant_ub( %x) { +; CHECK-LABEL: @scalable_splat_binop_constant_ub( +; CHECK-NEXT: [[SPLATX:%.*]] = shufflevector [[X:%.*]], poison, zeroinitializer +; CHECK-NEXT: [[R:%.*]] = udiv splat (i32 42), [[SPLATX]] +; CHECK-NEXT: ret [[R]] +; + + %splatx = shufflevector %x, poison, zeroinitializer + %r = udiv splat (i32 42), %splatx + ret %r +}