[AggresiveInstCombine] Added support of select and ShuffleVector to TruncInstCombine expression pattern

[aaboud]

Summary

Added to TruncInstCombine class the support of SelectInst and ShuffleVectorInst instructions.

[spatel]

I apologize for not noticing this before, but do all of the tests here fail in regular instcombine because there's an instruction that uses the same value more than once? Or are there more complicated patterns?

If it's just the simpler case (or even if it's not), can we enhance instcombine to catch those cases like this:

Index: lib/Transforms/InstCombine/InstCombineCasts.cpp
===================================================================
--- lib/Transforms/InstCombine/InstCombineCasts.cpp	(revision 323427)
+++ lib/Transforms/InstCombine/InstCombineCasts.cpp	(working copy)
@@ -184,8 +184,14 @@
   case Instruction::Shl:
   case Instruction::UDiv:
   case Instruction::URem: {
-    Value *LHS = EvaluateInDifferentType(I->getOperand(0), Ty, isSigned);
-    Value *RHS = EvaluateInDifferentType(I->getOperand(1), Ty, isSigned);
+    Value *LHS, *RHS;
+    if (I->getOperand(0) == I->getOperand(1)) {
+      // Don't create an unnecessary value if the operands are repeated.
+      LHS = RHS = EvaluateInDifferentType(I->getOperand(0), Ty, isSigned);
+    } else {
+      LHS = EvaluateInDifferentType(I->getOperand(0), Ty, isSigned);
+      RHS = EvaluateInDifferentType(I->getOperand(1), Ty, isSigned);
+    }
     Res = BinaryOperator::Create((Instruction::BinaryOps)Opc, LHS, RHS);
     break;
   }
@@ -315,9 +321,12 @@
       I->getOperand(0)->getType() == Ty)
     return true;
 
-  // We can't extend or shrink something that has multiple uses: doing so would
-  // require duplicating the instruction in general, which isn't profitable.
-  if (!I->hasOneUse()) return false;
+  // We can't extend or shrink something that has multiple uses -- unless those
+  // multiple uses are all in the same instruction -- doing so would require
+  // duplicating the instruction which isn't profitable.
+  if (!I->hasOneUse())
+    if (any_of(I->users(), [&](User *U) { return U != I->user_back(); }))
+      return false;
 
   unsigned Opc = I->getOpcode();
   switch (Opc) {

[aaboud]

In D42536#988261, @spatel wrote:

I apologize for not noticing this before, but do all of the tests here fail in regular instcombine because there's an instruction that uses the same value more than once? Or are there more complicated patterns?

Maybe, the test I added are for the simple case, but there are more complicated cases.
Consider this example:

define void @multi_uses_add(i32 %X, i32 %Y) {
  %Ax = zext i32 %X to i64
  %Ay = zext i32 %Y to i64
  %B = add i64 %Ax, %Ay
  %C = mul i64 %B, %Ax
  %T = trunc i64 %C to i32
  call i32 @use32(i32 %T)
  ret void
}

Also, remember that I will be adding the PHINode, in the coming patches, then we need to handle this case:

    A
 /     \
B       C
 \     /
 PHINode

[spatel]

In D42536#988277, @aaboud wrote:

define void @multi_uses_add(i32 %X, i32 %Y) {
  %Ax = zext i32 %X to i64
  %Ay = zext i32 %Y to i64
  %B = add i64 %Ax, %Ay
  %C = mul i64 %B, %Ax
  %T = trunc i64 %C to i32
  call i32 @use32(i32 %T)
  ret void
}

That's too easy - instcombine gets that. I think you need at least one more intermediate binop to show a case that instcombine can't handle.

We need to make sure aggressive-instcombine is showing its unique capability in each test. If you agree that the case with a repeated op is just an oversight for instcombine, lets:

1. Move the existing tests over to instcombine
2. Fix instcombine
3. Add more complex tests for aggressive-instcombine's current functionality
4. Enhance aggressive-instcombine for more opcodes/patterns.

[aaboud]

In D42536#988422, @spatel wrote:

That's too easy - instcombine gets that. I think you need at least one more intermediate binop to show a case that instcombine can't handle.

You are right, I forgot that instcombine algorithm can handle zext/sext with multi use in case they have same type as trunc instruction destination type.
Here is a better test suggestion:

define void @multi_uses_add(i32 %X) {
  %A = zext i32 %X to i64
  %B = add i64 %A, 5
  %M1 = mul i64 %B, 66
  %M2 = mul i64 %B, %A
  %C = and i64 %M1, %M2
  %T = trunc i64 %C to i32
  call i32 @use32(i32 %T)
  ret void
}

We need to make sure aggressive-instcombine is showing its unique capability in each test. If you agree that the case with a repeated op is just an oversight for instcombine, lets:

1. Move the existing tests over to instcombine
2. Fix instcombine

I agree that we should improve instcombine where ever it make sense, such as this case, but I am afraid that I have no time to do it myself.
So, if you would like to take this task, I will be happy to review the code.

3. Add more complex tests for aggressive-instcombine's current functionality

Do you agree to the above example? I can duplicate it to handle all other instructions.

4. Enhance aggressive-instcombine for more opcodes/patterns.

I will be adding more opcodes that trunc expression should support, and with each new opcode I will be adding more tests.
Also, I already added few tests that were driven from implementation bugs due to missing handle of constant expressions.
see https://reviews.llvm.org/D42622.

[spatel]

Yes, I'll take the task of fixing the gap in instcombine. (Hoping it's not too much more than what I posted already.)

For the tests, we don't need the extra-use helper function? The thing that regular instcombine can't see through is when an intermediate value has multiple uses:

define i32 @multi_uses_add(i32 %X) {
    %A = zext i32 %X to i64
    %B = add i64 %A, 5
    %M = mul i64 %B, %A
    %C = and i64 %B, %M
    %T = trunc i64 %C to i32
    ret i32 %T
  }

lib/Transforms/AggressiveInstCombine/TruncInstCombine.cpp

case Instruction::SExt:
     // their operands are not relevent.
     break;
   case Instruction::Add:
   case Instruction::Sub:
   case Instruction::Mul:
   case Instruction::And:
   case Instruction::Or:
   case Instruction::Xor:
+  case Instruction::ShuffleVector:
     Ops.push_back(I->getOperand(0));
     Ops.push_back(I->getOperand(1));
     break;
+  case Instruction::Select: {
+    SelectInst *SI = cast<SelectInst>(I);
+    Ops.push_back(SI->getTrueValue());
+    Ops.push_back(SI->getFalseValue());
+    break;
+  }
   default:
     llvm_unreachable("Unreachable!");
   }
 }
 
 bool TruncInstCombine::buildTruncExpressionDag() {
   SmallVector<Value *, 8> Worklist;
   SmallVector<Instruction *, 8> Stack;
case Instruction::SExt:
       // trunc(ext(x)) -> trunc(x) if the source type is larger than the new
       // dest
       break;
     case Instruction::Add:
     case Instruction::Sub:
     case Instruction::Mul:
     case Instruction::And:
     case Instruction::Or:
-    case Instruction::Xor: {
+    case Instruction::Xor:
+    case Instruction::Select:
+    case Instruction::ShuffleVector: {
       SmallVector<Value *, 2> Operands;
       getRelevantOperands(I, Operands);
       for (Value *Operand : Operands)
         Worklist.push_back(Operand);
       break;
     }
     default:
       // TODO: Can handle more cases here:
-      // 1. select, shufflevector, extractelement, insertelement
-      // 2. udiv, urem
-      // 3. shl, lshr, ashr
-      // 4. phi node(and loop handling)
+      // 1. udiv, urem
+      // 2. shl, lshr, ashr
+      // 3. phi node(and loop handling)
+      // 4. extractelement, insertelement
       // ...
       return false;
     }
   }
   return true;
 }
 
 unsigned TruncInstCombine::getMinBitWidth() {
for (auto &Itr : InstInfoMap) { // Forward
     case Instruction::And:
     case Instruction::Or:
     case Instruction::Xor: {
       Value *LHS = getReducedOperand(I->getOperand(0), SclTy);
       Value *RHS = getReducedOperand(I->getOperand(1), SclTy);
       Res = Builder.CreateBinOp((Instruction::BinaryOps)Opc, LHS, RHS);
       break;
     }
+    case Instruction::Select: {
+      SelectInst *SI = cast<SelectInst>(I);
+      Value *TrueValue = getReducedOperand(SI->getTrueValue(), SclTy);
+      Value *FalseValue = getReducedOperand(SI->getFalseValue(), SclTy);
+      Res = Builder.CreateSelect(SI->getCondition(), TrueValue, FalseValue);
+      break;
+    }
+    case Instruction::ShuffleVector: {
+      ShuffleVectorInst *SI = cast<ShuffleVectorInst>(I);
+      Value *LHS = getReducedOperand(I->getOperand(0), SclTy);
+      Value *RHS = getReducedOperand(I->getOperand(1), SclTy);
+      Res = Builder.CreateShuffleVector(LHS, RHS, SI->getMask());
+      break;
+    }
     default:
       llvm_unreachable("Unhandled instruction");
     }
 
     NodeInfo.NewValue = Res;
     if (auto *ResI = dyn_cast<Instruction>(Res))
       ResI->takeName(I);
   }

test/Transforms/AggressiveInstCombine/trunc_multi_uses.ll

;
   %T1 = trunc i64 %C1 to i32
   call i32 @use32(i32 %T1)
   ; make sure zext have another use that is not post-dominated by the TruncInst.
   call i32 @use64(i64 %A1)
   call i32 @use64(i64 %A2)
   ret void
 }
 
+define void @multi_uses_select(i32 %X, i32 %Y, i1 %COND) {
+; CHECK-LABEL: @multi_uses_select(
+; CHECK-NEXT:    [[A1:%.*]] = zext i32 [[X:%.*]] to i64
+; CHECK-NEXT:    [[A2:%.*]] = zext i32 [[Y:%.*]] to i64
+; CHECK-NEXT:    [[B1:%.*]] = select i1 [[COND:%.*]], i32 [[X]], i32 [[Y]]
+; CHECK-NEXT:    [[C1:%.*]] = mul i32 [[B1]], [[B1]]
+; CHECK-NEXT:    [[TMP1:%.*]] = call i32 @use32(i32 [[C1]])
+; CHECK-NEXT:    [[TMP2:%.*]] = call i32 @use64(i64 [[A1]])
+; CHECK-NEXT:    [[TMP3:%.*]] = call i32 @use64(i64 [[A2]])
+; CHECK-NEXT:    ret void
+;
+  %A1 = zext i32 %X to i64
+  %A2 = zext i32 %Y to i64
+  %B1 = select i1 %COND, i64 %A1, i64 %A2
+  %C1 = mul i64 %B1, %B1
+  %T1 = trunc i64 %C1 to i32
+  call i32 @use32(i32 %T1)
+  ; make sure zext have another use that is not post-dominated by the TruncInst.
+  call i32 @use64(i64 %A1)
+  call i32 @use64(i64 %A2)
+  ret void
+}
+
 define void @multi_use_vec_add(<2 x i32> %X) {
 ; CHECK-LABEL: @multi_use_vec_add(
 ; CHECK-NEXT:    [[A1:%.*]] = zext <2 x i32> [[X:%.*]] to <2 x i64>
 ; CHECK-NEXT:    [[B1:%.*]] = add <2 x i32> [[X]], <i32 15, i32 15>
 ; CHECK-NEXT:    [[C1:%.*]] = mul <2 x i32> [[B1]], [[B1]]
 ; CHECK-NEXT:    [[TMP1:%.*]] = call <2 x i32> @use32_vec(<2 x i32> [[C1]])
 ; CHECK-NEXT:    [[TMP2:%.*]] = call <2 x i32> @use64_vec(<2 x i64> [[A1]])
 ; CHECK-NEXT:    ret void
;
   %C1 = mul <2 x i64> %B1, %B1
   %T1 = trunc <2 x i64> %C1 to <2 x i32>
   call <2 x i32> @use32_vec(<2 x i32> %T1)
   ; make sure zext have another use that is not post-dominated by the TruncInst.
   call <2 x i32> @use64_vec(<2 x i64> %A1)
   call <2 x i32> @use64_vec(<2 x i64> %A2)
   ret void
 }
+
+define void @multi_use_vec_select(<2 x i32> %X, <2 x i32> %Y, <2 x i1> %COND) {
+; CHECK-LABEL: @multi_use_vec_select(
+; CHECK-NEXT:    [[A1:%.*]] = zext <2 x i32> [[X:%.*]] to <2 x i64>
+; CHECK-NEXT:    [[A2:%.*]] = zext <2 x i32> [[Y:%.*]] to <2 x i64>
+; CHECK-NEXT:    [[B1:%.*]] = select <2 x i1> [[COND:%.*]], <2 x i32> [[X]], <2 x i32> [[Y]]
+; CHECK-NEXT:    [[C1:%.*]] = mul <2 x i32> [[B1]], [[B1]]
+; CHECK-NEXT:    [[TMP1:%.*]] = call <2 x i32> @use32_vec(<2 x i32> [[C1]])
+; CHECK-NEXT:    [[TMP2:%.*]] = call <2 x i32> @use64_vec(<2 x i64> [[A1]])
+; CHECK-NEXT:    [[TMP3:%.*]] = call <2 x i32> @use64_vec(<2 x i64> [[A2]])
+; CHECK-NEXT:    ret void
+;
+  %A1 = zext <2 x i32> %X to <2 x i64>
+  %A2 = zext <2 x i32> %Y to <2 x i64>
+  %B1 = select <2 x i1> %COND, <2 x i64> %A1, <2 x i64> %A2
+  %C1 = mul <2 x i64> %B1, %B1
+  %T1 = trunc <2 x i64> %C1 to <2 x i32>
+  call <2 x i32> @use32_vec(<2 x i32> %T1)
+  ; make sure zext have another use that is not post-dominated by the TruncInst.
+  call <2 x i32> @use64_vec(<2 x i64> %A1)
+  call <2 x i32> @use64_vec(<2 x i64> %A2)
+  ret void
+}
+
+define void @multi_use_vec_shufflevector(<2 x i32> %X, <2 x i32> %Y) {
+; CHECK-LABEL: @multi_use_vec_shufflevector(
+; CHECK-NEXT:    [[A1:%.*]] = zext <2 x i32> [[X:%.*]] to <2 x i64>
+; CHECK-NEXT:    [[A2:%.*]] = zext <2 x i32> [[Y:%.*]] to <2 x i64>
+; CHECK-NEXT:    [[B1:%.*]] = shufflevector <2 x i32> [[X]], <2 x i32> [[Y]], <2 x i32> <i32 0, i32 2>
+; CHECK-NEXT:    [[C1:%.*]] = mul <2 x i32> [[B1]], [[B1]]
+; CHECK-NEXT:    [[TMP1:%.*]] = call <2 x i32> @use32_vec(<2 x i32> [[C1]])
+; CHECK-NEXT:    [[TMP2:%.*]] = call <2 x i32> @use64_vec(<2 x i64> [[A1]])
+; CHECK-NEXT:    [[TMP3:%.*]] = call <2 x i32> @use64_vec(<2 x i64> [[A2]])
+; CHECK-NEXT:    ret void
+;
+  %A1 = zext <2 x i32> %X to <2 x i64>
+  %A2 = zext <2 x i32> %Y to <2 x i64>
+  %B1 = shufflevector <2 x i64> %A1, <2 x i64> %A2, <2 x i32> <i32 0, i32 2>
+  %C1 = mul <2 x i64> %B1, %B1
+  %T1 = trunc <2 x i64> %C1 to <2 x i32>
+  call <2 x i32> @use32_vec(<2 x i32> %T1)
+  ; make sure zext have another use that is not post-dominated by the TruncInst.
+  call <2 x i32> @use64_vec(<2 x i64> %A1)
+  call <2 x i32> @use64_vec(<2 x i64> %A2)
+  ret void
+}
+