[AggressiveInstCombine] Add `{extract/insert}element` to `TruncInstCombine` DAG

[anton-afanasyev]

Summary

Alive2 for {extract/insert}element: https://alive2.llvm.org/ce/z/hwy_E-

Actually, no one file of test suite is touched by this change,
which means that is rare pattern not generated by frontend. But
it's worth being in place.

[lebedev.ri]

Can you add a test with insert into poison / unary shuffle?

[RKSimon]

llvm/lib/Transforms/AggressiveInstCombine/TruncInstCombine.cpp

452

Are we allowed to create new shuffles inside aggressiveinstcombine? A different type, even with the same mask, means that this might have a very different performance.

[anton-afanasyev]

In D109236#2982345, @lebedev.ri wrote:

Can you add a test with insert into poison / unary shuffle?

Ok, added

llvm/lib/Transforms/AggressiveInstCombine/TruncInstCombine.cpp

452

Hmm, if that's the case, the only stuff we could do is to remove shuffles processing back. We can check perfomance difference for specific target, but that can't be done before NewBitWidth choosing, when it is too late (the DAG is started reducing already). One could change whole algorithm, but that's not the case to do it specially for shufflevector support.

[anton-afanasyev]

Add tests

[anton-afanasyev]

Remove shufflevector

[xbolva00]

llvm/lib/Transforms/AggressiveInstCombine/TruncInstCombine.cpp

149

Leave shufflevector

[anton-afanasyev]

llvm/lib/Transforms/AggressiveInstCombine/TruncInstCombine.cpp

149

Ok, done

[anton-afanasyev]

Leave shufflevector in TODO

[anton-afanasyev]

Add new test demonstrating AIC power vs IC

[spatel]

LGTM

llvm/lib/Transforms/AggressiveInstCombine/TruncInstCombine.cpp

static void getRelevantOperands(Instruction *I, SmallVectorImpl<Value *> &Ops) {
   case Instruction::And:
   case Instruction::Or:
   case Instruction::Xor:
   case Instruction::Shl:
   case Instruction::LShr:
   case Instruction::AShr:
   case Instruction::UDiv:
   case Instruction::URem:
+  case Instruction::InsertElement:
     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));
     Ops.push_back(I->getOperand(2));
     break;
   default:
     llvm_unreachable("Unreachable!");
   }
 }
while (!Worklist.empty()) {
     case Instruction::And:
     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::ExtractElement:
     case Instruction::Select: {
       SmallVector<Value *, 2> Operands;
       getRelevantOperands(I, Operands);
       append_range(Worklist, Operands);
       break;
     }
     default:
       // TODO: Can handle more cases here:
-      // 1. shufflevector, extractelement, insertelement
+      // 1. shufflevector
       // 2. sdiv, srem
       // 3. phi node(and loop handling)
       // ...
       return false;
     }
   }
   return true;
 }
case Instruction::URem: {
       Value *RHS = getReducedOperand(I->getOperand(1), SclTy);
       Res = Builder.CreateBinOp((Instruction::BinaryOps)Opc, LHS, RHS);
       // Preserve `exact` flag since truncation doesn't change exactness
       if (auto *PEO = dyn_cast<PossiblyExactOperator>(I))
         if (auto *ResI = dyn_cast<Instruction>(Res))
           ResI->setIsExact(PEO->isExact());
       break;
     }
+    case Instruction::ExtractElement: {
+      Value *Vec = getReducedOperand(I->getOperand(0), SclTy);
+      Value *Idx = I->getOperand(1);
+      Res = Builder.CreateExtractElement(Vec, Idx);
+      break;
+    }
+    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::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;
     }
     default:
       llvm_unreachable("Unhandled instruction");
     }
 
     NodeInfo.NewValue = Res;
     if (auto *ResI = dyn_cast<Instruction>(Res))
       ResI->takeName(I);

llvm/test/Transforms/AggressiveInstCombine/trunc_vector_instrs.ll

;
   %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) {
 ; CHECK-LABEL: @extract_insert(
-; 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:    [[EXTR:%.*]] = extractelement <2 x i32> [[ZEXTA]], i32 0
-; CHECK-NEXT:    [[INSR:%.*]] = insertelement <2 x i32> [[ZEXTB]], i32 [[EXTR]], i32 1
-; CHECK-NEXT:    [[TRUNC:%.*]] = trunc <2 x i32> [[INSR]] 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:    [[ZEXTB:%.*]] = zext <2 x i8> [[B:%.*]] to <2 x i16>
+; CHECK-NEXT:    [[EXTR:%.*]] = extractelement <2 x i16> [[ZEXTA]], i32 0
+; CHECK-NEXT:    [[INSR:%.*]] = insertelement <2 x i16> [[ZEXTB]], i16 [[EXTR]], i32 1
+; CHECK-NEXT:    ret <2 x i16> [[INSR]]
 ;
   %zexta = zext <2 x i8> %a to <2 x i32>
   %zextb = zext <2 x i8> %b to <2 x i32>
   %extr = extractelement <2 x i32> %zexta, i32 0
   %insr = insertelement <2 x i32> %zextb, i32 %extr, i32 1
   %trunc = trunc <2 x i32> %insr to <2 x i16>
   ret <2 x i16> %trunc
 }
 
 define <2 x i16> @insert_poison(i8 %a) {
 ; CHECK-LABEL: @insert_poison(
-; CHECK-NEXT:    [[ZEXTA:%.*]] = zext i8 [[A:%.*]] to i32
-; CHECK-NEXT:    [[INSR:%.*]] = insertelement <2 x i32> poison, i32 [[ZEXTA]], i32 0
-; CHECK-NEXT:    [[TRUNC:%.*]] = trunc <2 x i32> [[INSR]] to <2 x i16>
-; CHECK-NEXT:    ret <2 x i16> [[TRUNC]]
+; CHECK-NEXT:    [[ZEXTA:%.*]] = zext i8 [[A:%.*]] to i16
+; CHECK-NEXT:    [[INSR:%.*]] = insertelement <2 x i16> poison, i16 [[ZEXTA]], i32 0
+; CHECK-NEXT:    ret <2 x i16> [[INSR]]
 ;
   %zexta = zext i8 %a to i32
   %insr = insertelement <2 x i32> poison, i32 %zexta, i32 0
   %trunc = trunc <2 x i32> %insr to <2 x i16>
   ret <2 x i16> %trunc
 }
 
 ; This demonstrates test not folded by 'opt -instcombine'
 define <2 x i16> @extract_mul_insert(<2 x i8> %x) {
 ; CHECK-LABEL: @extract_mul_insert(
-; CHECK-NEXT:    [[ZEXT:%.*]] = zext <2 x i8> [[X:%.*]] to <2 x i32>
-; CHECK-NEXT:    [[LSHR:%.*]] = lshr <2 x i32> [[ZEXT]], <i32 4, i32 5>
-; CHECK-NEXT:    [[EXTR:%.*]] = extractelement <2 x i32> [[LSHR]], i32 1
-; CHECK-NEXT:    [[MUL:%.*]] = mul i32 [[EXTR]], 5
-; CHECK-NEXT:    [[INSR:%.*]] = insertelement <2 x i32> [[LSHR]], i32 [[MUL]], i32 1
-; CHECK-NEXT:    [[TRUNC:%.*]] = trunc <2 x i32> [[INSR]] to <2 x i16>
-; CHECK-NEXT:    ret <2 x i16> [[TRUNC]]
+; CHECK-NEXT:    [[ZEXT:%.*]] = zext <2 x i8> [[X:%.*]] to <2 x i16>
+; CHECK-NEXT:    [[LSHR:%.*]] = lshr <2 x i16> [[ZEXT]], <i16 4, i16 5>
+; CHECK-NEXT:    [[EXTR:%.*]] = extractelement <2 x i16> [[LSHR]], i32 1
+; CHECK-NEXT:    [[MUL:%.*]] = mul i16 [[EXTR]], 5
+; CHECK-NEXT:    [[INSR:%.*]] = insertelement <2 x i16> [[LSHR]], i16 [[MUL]], i32 1
+; CHECK-NEXT:    ret <2 x i16> [[INSR]]
 ;
   %zext = zext <2 x i8> %x to <2 x i32>
   %lshr = lshr <2 x i32> %zext, <i32 4, i32 5>
   %extr = extractelement <2 x i32> %lshr, i32 1
   %mul = mul i32 %extr, 5
   %insr = insertelement <2 x i32> %lshr, i32 %mul, i32 1
   %trunc = trunc <2 x i32> %insr to <2 x i16>
   ret <2 x i16> %trunc
 }