[AggressiveInstCombine] Add `shufflevector` instr support to `TruncInstCombine`

[anton-afanasyev]

Summary

Add shufflevector instruction to the expression graph post-dominated by trunc,
allowing TruncInstCombine to reduce bitwidth of expressions containing these
instructions.

Fixes #54149

[lebedev.ri]

Do we expect that unused inputs of the shuffle has been already replaced with undef?

[anton-afanasyev]

In D122233#3399664, @lebedev.ri wrote:

Do we expect that unused inputs of the shuffle has been already replaced with undef?

I don't see this could be an issue. For instance, @unary_shuffle() test-case contains shuffle with undef (line 25). Do you mean this case?

[lebedev.ri]

In D122233#3399767, @anton-afanasyev wrote:

In D122233#3399664, @lebedev.ri wrote:

Do we expect that unused inputs of the shuffle has been already replaced with undef?

I don't see this could be an issue. For instance, @unary_shuffle() test-case contains shuffle with undef (line 25). Do you mean this case?

I mean, what if we have a two-input shuffle, and one of the operands is unused as per the shuffle mask.
Then, said operand can be replaced with undef, which doesn't affect the narrowing, while the original operand might?
I guess it's a theoretical question, mainly.

[RKSimon]

This might introduce regressions as the shuffle costs for the same mask but different element types can vary considerably (SSE v4i32/v4i16 unary shuffles are really cheap but v4i8 or v4i64 can be a lot more expensive).

[RKSimon]

FWIW we might be able to perform something similar inside VectorCombine

[anton-afanasyev]

In D122233#3513581, @RKSimon wrote:

FWIW we might be able to perform something similar inside VectorCombine

Do you mean using TTI.getShuffleCost()? There's an issue here: we don't know the exact shuffle type at the moment we need to get its cost. We infer this type (given by MinBitWidth) after the expression graph has been built already. Need to refactor whole pass for this case, which looks redundant.

llvm/lib/Transforms/AggressiveInstCombine/TruncInstCombine.cpp

static void getRelevantOperands(Instruction *I, SmallVectorImpl<Value *> &Ops) {
   case Instruction::Or:
   case Instruction::Xor:
   case Instruction::Shl:
   case Instruction::LShr:
   case Instruction::AShr:
   case Instruction::UDiv:
   case Instruction::URem:
   case Instruction::InsertElement:
+  case Instruction::ShuffleVector:
     Ops.push_back(I->getOperand(0));
     Ops.push_back(I->getOperand(1));
     break;
   case Instruction::ExtractElement:
     Ops.push_back(I->getOperand(0));
     break;
   case Instruction::Select:
     Ops.push_back(I->getOperand(1));
while (!Worklist.empty()) {
     case Instruction::Or:
     case Instruction::Xor:
     case Instruction::Shl:
     case Instruction::LShr:
     case Instruction::AShr:
     case Instruction::UDiv:
     case Instruction::URem:
     case Instruction::InsertElement:
+    case Instruction::ShuffleVector:
     case Instruction::ExtractElement:
     case Instruction::Select: {
       SmallVector<Value *, 2> Operands;
       getRelevantOperands(I, Operands);
       append_range(Worklist, Operands);
       break;
     }
     case Instruction::PHI: {
       SmallVector<Value *, 2> Operands;
       getRelevantOperands(I, Operands);
       // Add only operands not in Stack to prevent cycle
       for (auto *Op : Operands)
         if (all_of(Stack, [Op](Value *V) { return Op != V; }))
           Worklist.push_back(Op);
       break;
     }
     default:
-      // TODO: Can handle more cases here:
-      // 1. shufflevector
-      // 2. sdiv, srem
-      // ...
+      // TODO: Can handle more cases here: sdiv, srem, ...
       return false;
     }
   }
   return true;
 }
 
 unsigned TruncInstCombine::getMinBitWidth() {
   SmallVector<Value *, 8> Worklist;
for (auto &Itr : InstInfoMap) { // Forward
     }
     case Instruction::InsertElement: {
       Value *Vec = getReducedOperand(I->getOperand(0), SclTy);
       Value *NewElt = getReducedOperand(I->getOperand(1), SclTy);
       Value *Idx = I->getOperand(2);
       Res = Builder.CreateInsertElement(Vec, NewElt, Idx);
       break;
     }
+    case Instruction::ShuffleVector: {
+      Value *Vec1 = getReducedOperand(I->getOperand(0), SclTy);
+      Value *Vec2 = getReducedOperand(I->getOperand(1), SclTy);
+      ArrayRef<int> Mask = cast<ShuffleVectorInst>(I)->getShuffleMask();
+      Res = Builder.CreateShuffleVector(Vec1, Vec2, Mask);
+      break;
+    }
     case Instruction::Select: {
       Value *Op0 = I->getOperand(0);
       Value *LHS = getReducedOperand(I->getOperand(1), SclTy);
       Value *RHS = getReducedOperand(I->getOperand(2), SclTy);
       Res = Builder.CreateSelect(Op0, LHS, RHS);
       break;
     }
     case Instruction::PHI: {

llvm/test/Transforms/AggressiveInstCombine/trunc_vector_instrs.ll

 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
 ; RUN: opt < %s -aggressive-instcombine -S | FileCheck %s
 
 define <4 x i16> @shuffle(<2 x i8> %a, <2 x i8> %b) {
 ; CHECK-LABEL: @shuffle(
-; CHECK-NEXT:    [[ZEXTA:%.*]] = zext <2 x i8> [[A:%.*]] to <2 x i32>
-; CHECK-NEXT:    [[ZEXTB:%.*]] = zext <2 x i8> [[B:%.*]] to <2 x i32>
-; CHECK-NEXT:    [[SHUF:%.*]] = shufflevector <2 x i32> [[ZEXTA]], <2 x i32> [[ZEXTB]], <4 x i32> <i32 3, i32 2, i32 1, i32 0>
-; CHECK-NEXT:    [[TRUNC:%.*]] = trunc <4 x i32> [[SHUF]] to <4 x i16>
-; CHECK-NEXT:    ret <4 x i16> [[TRUNC]]
+; CHECK-NEXT:    [[ZEXTA:%.*]] = zext <2 x i8> [[A:%.*]] to <2 x i16>
+; CHECK-NEXT:    [[ZEXTB:%.*]] = zext <2 x i8> [[B:%.*]] to <2 x i16>
+; CHECK-NEXT:    [[SHUF:%.*]] = shufflevector <2 x i16> [[ZEXTA]], <2 x i16> [[ZEXTB]], <4 x i32> <i32 3, i32 2, i32 1, i32 0>
+; CHECK-NEXT:    ret <4 x i16> [[SHUF]]
 ;
   %zexta = zext <2 x i8> %a to <2 x i32>
   %zextb = zext <2 x i8> %b to <2 x i32>
   %shuf = shufflevector <2 x i32> %zexta, <2 x i32> %zextb, <4 x i32> <i32 3, i32 2, i32 1, i32 0>
   %trunc = trunc <4 x i32> %shuf to <4 x i16>
   ret <4 x i16> %trunc
 }
 
 define <2 x i16> @unary_shuffle(<2 x i8> %a) {
 ; CHECK-LABEL: @unary_shuffle(
-; CHECK-NEXT:    [[ZEXTA:%.*]] = zext <2 x i8> [[A:%.*]] to <2 x i32>
-; CHECK-NEXT:    [[SHUF:%.*]] = shufflevector <2 x i32> [[ZEXTA]], <2 x i32> undef, <2 x i32> <i32 1, i32 0>
-; CHECK-NEXT:    [[TRUNC:%.*]] = trunc <2 x i32> [[SHUF]] to <2 x i16>
-; CHECK-NEXT:    ret <2 x i16> [[TRUNC]]
+; CHECK-NEXT:    [[ZEXTA:%.*]] = zext <2 x i8> [[A:%.*]] to <2 x i16>
+; CHECK-NEXT:    [[SHUF:%.*]] = shufflevector <2 x i16> [[ZEXTA]], <2 x i16> undef, <2 x i32> <i32 1, i32 0>
+; CHECK-NEXT:    ret <2 x i16> [[SHUF]]
 ;
   %zexta = zext <2 x i8> %a to <2 x i32>
   %shuf = shufflevector <2 x i32> %zexta, <2 x i32> undef, <2 x i32> <i32 1, i32 0>
   %trunc = trunc <2 x i32> %shuf to <2 x i16>
   ret <2 x i16> %trunc
 }
 
 define <4 x i16> @const_shuffle() {
 ; CHECK-LABEL: @const_shuffle(
-; CHECK-NEXT:    [[SHUF:%.*]] = shufflevector <2 x i32> <i32 1, i32 2>, <2 x i32> <i32 3, i32 7>, <4 x i32> <i32 3, i32 2, i32 1, i32 0>
-; CHECK-NEXT:    [[TRUNC:%.*]] = trunc <4 x i32> [[SHUF]] to <4 x i16>
-; CHECK-NEXT:    ret <4 x i16> [[TRUNC]]
+; CHECK-NEXT:    ret <4 x i16> <i16 7, i16 3, i16 2, i16 1>
 ;
   %shuf = shufflevector <2 x i32> <i32 1, i32 2>, <2 x i32> <i32 3, i32 7>, <4 x i32> <i32 3, i32 2, i32 1, i32 0>
   %trunc = trunc <4 x i32> %shuf to <4 x i16>
   ret <4 x i16> %trunc
 }
 
 
 define <2 x i16> @extract_insert(<2 x i8> %a, <2 x i8> %b) {