From db5f5e09dfd25639256d10d1785b98c2590f98e6 Mon Sep 17 00:00:00 2001 From: David Green Date: Thu, 23 May 2024 18:45:42 +0100 Subject: [PATCH] [VectorCombine] Add Cmp and Select for shuffleToIdentity Other than some additional checks needed for compare predicates and selects with scalar condition operands, these are relatively simple additions to what already exists. --- .../Transforms/Vectorize/VectorCombine.cpp | 19 ++++++++++-- .../AArch64/shuffletoidentity.ll | 30 +++---------------- 2 files changed, 21 insertions(+), 28 deletions(-) diff --git a/llvm/lib/Transforms/Vectorize/VectorCombine.cpp b/llvm/lib/Transforms/Vectorize/VectorCombine.cpp index b5a292841172b3..53fd4bacbe55c3 100644 --- a/llvm/lib/Transforms/Vectorize/VectorCombine.cpp +++ b/llvm/lib/Transforms/Vectorize/VectorCombine.cpp @@ -1742,6 +1742,10 @@ static Value *generateNewInstTree(ArrayRef Item, FixedVectorType *Ty, if (auto *BI = dyn_cast(I)) return Builder.CreateBinOp((Instruction::BinaryOps)BI->getOpcode(), Ops[0], Ops[1]); + if (auto CI = dyn_cast(I)) + return Builder.CreateCmp(CI->getPredicate(), Ops[0], Ops[1]); + if (auto SI = dyn_cast(I)) + return Builder.CreateSelect(Ops[0], Ops[1], Ops[2], "", SI); if (II) return Builder.CreateIntrinsic(DstTy, II->getIntrinsicID(), Ops); assert(isa(I) && "Unexpected instruction type in Generate"); @@ -1810,6 +1814,12 @@ bool VectorCombine::foldShuffleToIdentity(Instruction &I) { return false; if (V->getValueID() != FrontV->getValueID()) return false; + if (auto *CI = dyn_cast(V)) + if (CI->getPredicate() != cast(FrontV)->getPredicate()) + return false; + if (auto *SI = dyn_cast(V)) + if (!isa(SI->getOperand(0)->getType())) + return false; if (isa(V) && !isa(V)) return false; auto *II = dyn_cast(V); @@ -1821,12 +1831,17 @@ bool VectorCombine::foldShuffleToIdentity(Instruction &I) { // Check the operator is one that we support. We exclude div/rem in case // they hit UB from poison lanes. - if (isa(FrontV) && - !cast(FrontV)->isIntDivRem()) { + if ((isa(FrontV) && + !cast(FrontV)->isIntDivRem()) || + isa(FrontV)) { Worklist.push_back(generateInstLaneVectorFromOperand(Item, 0)); Worklist.push_back(generateInstLaneVectorFromOperand(Item, 1)); } else if (isa(FrontV)) { Worklist.push_back(generateInstLaneVectorFromOperand(Item, 0)); + } else if (isa(FrontV)) { + Worklist.push_back(generateInstLaneVectorFromOperand(Item, 0)); + Worklist.push_back(generateInstLaneVectorFromOperand(Item, 1)); + Worklist.push_back(generateInstLaneVectorFromOperand(Item, 2)); } else if (auto *II = dyn_cast(FrontV); II && isTriviallyVectorizable(II->getIntrinsicID())) { for (unsigned Op = 0, E = II->getNumOperands() - 1; Op < E; Op++) { diff --git a/llvm/test/Transforms/VectorCombine/AArch64/shuffletoidentity.ll b/llvm/test/Transforms/VectorCombine/AArch64/shuffletoidentity.ll index 7dadeb5d72de09..46faecacfee05e 100644 --- a/llvm/test/Transforms/VectorCombine/AArch64/shuffletoidentity.ll +++ b/llvm/test/Transforms/VectorCombine/AArch64/shuffletoidentity.ll @@ -427,19 +427,8 @@ define <8 x i8> @extrause_shuffle(<8 x i8> %a, <8 x i8> %b) { define <8 x i8> @icmpsel(<8 x i8> %a, <8 x i8> %b, <8 x i8> %c, <8 x i8> %d) { ; CHECK-LABEL: @icmpsel( -; CHECK-NEXT: [[AB:%.*]] = shufflevector <8 x i8> [[A:%.*]], <8 x i8> poison, <4 x i32> -; CHECK-NEXT: [[AT:%.*]] = shufflevector <8 x i8> [[A]], <8 x i8> poison, <4 x i32> -; CHECK-NEXT: [[BB:%.*]] = shufflevector <8 x i8> [[B:%.*]], <8 x i8> poison, <4 x i32> -; CHECK-NEXT: [[BT:%.*]] = shufflevector <8 x i8> [[B]], <8 x i8> poison, <4 x i32> -; CHECK-NEXT: [[CB:%.*]] = shufflevector <8 x i8> [[C:%.*]], <8 x i8> poison, <4 x i32> -; CHECK-NEXT: [[CT:%.*]] = shufflevector <8 x i8> [[C]], <8 x i8> poison, <4 x i32> -; CHECK-NEXT: [[DB:%.*]] = shufflevector <8 x i8> [[D:%.*]], <8 x i8> poison, <4 x i32> -; CHECK-NEXT: [[DT:%.*]] = shufflevector <8 x i8> [[D]], <8 x i8> poison, <4 x i32> -; CHECK-NEXT: [[ABT1:%.*]] = icmp slt <4 x i8> [[AT]], [[BT]] -; CHECK-NEXT: [[ABB1:%.*]] = icmp slt <4 x i8> [[AB]], [[BB]] -; CHECK-NEXT: [[ABT:%.*]] = select <4 x i1> [[ABT1]], <4 x i8> [[CT]], <4 x i8> [[DT]] -; CHECK-NEXT: [[ABB:%.*]] = select <4 x i1> [[ABB1]], <4 x i8> [[CB]], <4 x i8> [[DB]] -; CHECK-NEXT: [[R:%.*]] = shufflevector <4 x i8> [[ABT]], <4 x i8> [[ABB]], <8 x i32> +; CHECK-NEXT: [[TMP1:%.*]] = icmp slt <8 x i8> [[A:%.*]], [[B:%.*]] +; CHECK-NEXT: [[R:%.*]] = select <8 x i1> [[TMP1]], <8 x i8> [[C:%.*]], <8 x i8> [[D:%.*]] ; CHECK-NEXT: ret <8 x i8> [[R]] ; %ab = shufflevector <8 x i8> %a, <8 x i8> poison, <4 x i32> @@ -493,19 +482,8 @@ define <8 x i8> @icmpsel_diffentcond(<8 x i8> %a, <8 x i8> %b, <8 x i8> %c, <8 x define <8 x i8> @fcmpsel(<8 x half> %a, <8 x half> %b, <8 x i8> %c, <8 x i8> %d) { ; CHECK-LABEL: @fcmpsel( -; CHECK-NEXT: [[AB:%.*]] = shufflevector <8 x half> [[A:%.*]], <8 x half> poison, <4 x i32> -; CHECK-NEXT: [[AT:%.*]] = shufflevector <8 x half> [[A]], <8 x half> poison, <4 x i32> -; CHECK-NEXT: [[BB:%.*]] = shufflevector <8 x half> [[B:%.*]], <8 x half> poison, <4 x i32> -; CHECK-NEXT: [[BT:%.*]] = shufflevector <8 x half> [[B]], <8 x half> poison, <4 x i32> -; CHECK-NEXT: [[CB:%.*]] = shufflevector <8 x i8> [[C:%.*]], <8 x i8> poison, <4 x i32> -; CHECK-NEXT: [[CT:%.*]] = shufflevector <8 x i8> [[C]], <8 x i8> poison, <4 x i32> -; CHECK-NEXT: [[DB:%.*]] = shufflevector <8 x i8> [[D:%.*]], <8 x i8> poison, <4 x i32> -; CHECK-NEXT: [[DT:%.*]] = shufflevector <8 x i8> [[D]], <8 x i8> poison, <4 x i32> -; CHECK-NEXT: [[ABT1:%.*]] = fcmp olt <4 x half> [[AT]], [[BT]] -; CHECK-NEXT: [[ABB1:%.*]] = fcmp olt <4 x half> [[AB]], [[BB]] -; CHECK-NEXT: [[ABT:%.*]] = select <4 x i1> [[ABT1]], <4 x i8> [[CT]], <4 x i8> [[DT]] -; CHECK-NEXT: [[ABB:%.*]] = select <4 x i1> [[ABB1]], <4 x i8> [[CB]], <4 x i8> [[DB]] -; CHECK-NEXT: [[R:%.*]] = shufflevector <4 x i8> [[ABT]], <4 x i8> [[ABB]], <8 x i32> +; CHECK-NEXT: [[TMP1:%.*]] = fcmp olt <8 x half> [[A:%.*]], [[B:%.*]] +; CHECK-NEXT: [[R:%.*]] = select <8 x i1> [[TMP1]], <8 x i8> [[C:%.*]], <8 x i8> [[D:%.*]] ; CHECK-NEXT: ret <8 x i8> [[R]] ; %ab = shufflevector <8 x half> %a, <8 x half> poison, <4 x i32>