[AggressiveInstCombine] Add support for select instruction.

[aymanmus]

Summary

Add support for select instructions in the sequences that may be optimized in the AggressiveInstCombine pass.

[lebedev.ri]

This is certainly missing tests

[aymanmus]

Somehow the test wasn't included with the differential text file.
Adding now.

[spatel]

Verify that the translation is correct, but I don't think this patch is safe:

define i16 @cmp_select_sext_const(i8 %a) {
  %conv = sext i8 %a to i32
  %cmp = icmp slt i32 %conv, 32768
  %cond = select i1 %cmp, i32 32768, i32 %conv
  %conv4 = trunc i32 %cond to i16
  ret i16 %conv4
}

$ opt -aggressive-instcombine agg.ll -S
define i16 @cmp_select_sext_const(i8 %a) {
  %conv = sext i8 %a to i16
  %cmp = icmp slt i16 %conv, -32768
  %cond = select i1 %cmp, i16 -32768, i16 %conv
  ret i16 %cond
}

https://rise4fun.com/Alive/JCo

[aymanmus]

Fix @spatel's test case. Thanks for the review.
Involve constants in cmp operations in the MinBitWidth calculations.

[nikic]

Is it possible to split this into a patch that adds just the select handling, and another that extends the icmp support? It seems like the select part should be testable independently, as long as you don't derive the select condition from the select operands.

[aymanmus]

Following @nikic's comment:
I agree, the patch is missing some select exclusive tests.
Adding such test cases.
IMHO, there is no added value to splitting to 2 patches, the select changed are too obvious, and the compare is strongly related to the select changes.

[nikic]

In D72837#1840886, @aymanmus wrote:

IMHO, there is no added value to splitting to 2 patches, the select changed are too obvious, and the compare is strongly related to the select changes.

That's kind of the point: The select change is obvious and can land right away, the icmp part is trickier.

I'm not really familiar with this code, but coupling the icmp and the select handling seems like a bit of a hack: The transform is also possible if the select condition is unrelated to the trunc chain, and the transform is also possible if there is an icmp dependent on the trunc chain but not used in a select condition. I think the more principled approach would be to extend the code in https://github.com/llvm/llvm-project/blob/0e0c65264aeb6f66b6f711884c55cdbf66d975f6/llvm/lib/Transforms/AggressiveInstCombine/TruncInstCombine.cpp#L241-L256 to collect the extra use icmp.

[spatel]

In D72837#1841046, @nikic wrote:

In D72837#1840886, @aymanmus wrote:

IMHO, there is no added value to splitting to 2 patches, the select changed are too obvious, and the compare is strongly related to the select changes.

That's kind of the point: The select change is obvious and can land right away, the icmp part is trickier.

+1. Minimal patches are easier to review/debug, and we've already shown that this patch has tricky edge cases (and I don't know this code very well either). Also, please push the test file with baseline CHECKs as a preliminary commit.

[aymanmus]

Removed all cmp instruction handling to get a pure SELECT patch support.

[lebedev.ri]

Please upload all patches with full context (-U99999)

[aymanmus]

Adding full context to patch.

[lebedev.ri]

llvm/lib/Transforms/AggressiveInstCombine/TruncInstCombine.cpp

225–243

Why is this being changed?

[aymanmus]

llvm/lib/Transforms/AggressiveInstCombine/TruncInstCombine.cpp

225–243

The max operation is not needed here.
If we reach this point of the code, then ValidBitWidth MUST be bigger than IOpBitwidth (because other wise we'll hit the "continue" at line 202.

[lebedev.ri]

llvm/lib/Transforms/AggressiveInstCombine/TruncInstCombine.cpp

225–243

Can you submit that as a separate change then please?

[nikic]

LGTM

llvm/lib/Transforms/AggressiveInstCombine/TruncInstCombine.cpp

79

Style nit: Omit {} if not necessary.

225–243

This change looks fine, but please commit it independently as NFC.

[aymanmus]

Extract the std::max omission to a separate patch as requested (patch D74476)

llvm/lib/Transforms/AggressiveInstCombine/TruncInstCombine.cpp

 #include "llvm/Analysis/TargetLibraryInfo.h"
 #include "llvm/IR/DataLayout.h"
 #include "llvm/IR/Dominators.h"
 #include "llvm/IR/IRBuilder.h"
 using namespace llvm;
 
 #define DEBUG_TYPE "aggressive-instcombine"
 
+// This function returns true if Value V is a constant or if it's a type
+// extension node.
+static bool isConstOrExt(Value *V) {
+  if (isa<Constant>(V))
+    return true;
+
+  if (Instruction *I = dyn_cast<Instruction>(V)) {
+    switch(I->getOpcode()) {
+    case Instruction::ZExt:
+    case Instruction::SExt:
+      return true;
+    default:
+      return false;
+    }
+  }
+  return false;
+}
+
+
 /// Given an instruction and a container, it fills all the relevant operands of
 /// that instruction, with respect to the Trunc expression dag optimizaton.
 static void getRelevantOperands(Instruction *I, SmallVectorImpl<Value *> &Ops) {
   unsigned Opc = I->getOpcode();
   switch (Opc) {
   case Instruction::Trunc:
   case Instruction::ZExt:
   case Instruction::SExt:
     // These CastInst are considered leaves of the evaluated expression, thus,
     // 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::ICmp:
     Ops.push_back(I->getOperand(0));
     Ops.push_back(I->getOperand(1));
     break;
+  case Instruction::Select: {
+    Value *Op0 = I->getOperand(0);
+    Ops.push_back(I->getOperand(1));
+    Ops.push_back(I->getOperand(2));
+    // In case the condition is a compare instruction, that both of its operands
+    // are a type extension/truncate or a constant, that can be shrinked without
+    // loosing information in the compare instruction, add them as well.
+    if (CmpInst *C = dyn_cast<CmpInst>(Op0))
+      if (isConstOrExt(C->getOperand(0)) && isConstOrExt(C->getOperand(1)))
+        Ops.push_back(Op0);
+    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::ICmp: {
       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
+      // 1. shufflevector, extractelement, insertelement
       // 2. udiv, urem
       // 3. shl, lshr, ashr
       // 4. phi node(and loop handling)
       // ...
       return false;
     }
   }
   return true;
while (!Worklist.empty()) {
     unsigned ValidBitWidth = Info.ValidBitWidth;
 
     // Update minimum bit-width before handling its operands. This is required
     // when the instruction is part of a loop.
     Info.MinBitWidth = std::max(Info.MinBitWidth, Info.ValidBitWidth);
 
     for (auto *Operand : Operands)
       if (auto *IOp = dyn_cast<Instruction>(Operand)) {
+        if (isa<CmpInst>(I)) {
+          // Cmp instructions kind of resets the valid bits analysis for its
+          // operands, as it does not continue with the same calculation chain
+          // but rather creates a new chain of its own.
+          switch (IOp->getOpcode()) {
+          case Instruction::SExt:
+          case Instruction::ZExt:
+            InstInfoMap[IOp].ValidBitWidth =
+                cast<CastInst>(IOp)->getSrcTy()->getScalarSizeInBits();
+            break;
+          }
+        } else {
           // If we already calculated the minimum bit-width for this valid
           // bit-width, or for a smaller valid bit-width, then just keep the
           // answer we already calculated.
           unsigned IOpBitwidth = InstInfoMap.lookup(IOp).ValidBitWidth;
           if (IOpBitwidth >= ValidBitWidth)
             continue;
-        InstInfoMap[IOp].ValidBitWidth = std::max(ValidBitWidth, IOpBitwidth);
+          InstInfoMap[IOp].ValidBitWidth = ValidBitWidth;
         Worklist.push_back(IOp);
         }
       }
+  }
   unsigned MinBitWidth = InstInfoMap.lookup(cast<Instruction>(Src)).MinBitWidth;
   assert(MinBitWidth >= TruncBitWidth);
 
   if (MinBitWidth > TruncBitWidth) {
     // In this case reducing expression with vector type might generate a new
     // vector type, which is not preferable as it might result in generating
     // sub-optimal code.
     if (DstTy->isVectorTy())
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: {
+      Value *Op0 = I->getOperand(0);
+      if (ICmpInst *C = dyn_cast<ICmpInst>(Op0))
+        if (isConstOrExt(C->getOperand(0)) && isConstOrExt(C->getOperand(1)))
+          Op0 = getReducedOperand(Op0, SclTy);
+      Value *LHS = getReducedOperand(I->getOperand(1), SclTy);
+      Value *RHS = getReducedOperand(I->getOperand(2), SclTy);
+      Res = Builder.CreateSelect(Op0, LHS, RHS);
+      break;
+    }
+    case Instruction::ICmp: {
+      auto ICmp = cast<ICmpInst>(I);
+      Value *LHS = getReducedOperand(ICmp->getOperand(0), SclTy);
+      Value *RHS = getReducedOperand(ICmp->getOperand(1), SclTy);
+      Res = Builder.CreateICmp(ICmp->getPredicate(), 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_select_cmp.ll

+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
+; RUN: opt < %s -aggressive-instcombine -S | FileCheck %s
+; RUN: opt < %s -passes=aggressive-instcombine -S | FileCheck %s
+
+target datalayout = "e-m:m-p1:64:64:64-p:32:32:32-n16:32"
+
+define dso_local i16 @cmp_select_sext_const(i8 %a) {
+; CHECK-LABEL: @cmp_select_sext_const(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[CONV:%.*]] = sext i8 [[A:%.*]] to i16
+; CHECK-NEXT:    [[CMP:%.*]] = icmp slt i16 [[CONV]], 109
+; CHECK-NEXT:    [[COND:%.*]] = select i1 [[CMP]], i16 109, i16 [[CONV]]
+; CHECK-NEXT:    ret i16 [[COND]]
+;
+entry:
+  %conv = sext i8 %a to i32
+  %cmp = icmp slt i32 %conv, 109
+  %cond = select i1 %cmp, i32 109, i32 %conv
+  %conv4 = trunc i32 %cond to i16
+  ret i16 %conv4
+}
+
+define dso_local i16 @cmp_select_sext(i8 %a, i8 %b) {
+; CHECK-LABEL: @cmp_select_sext(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[CONV:%.*]] = sext i8 [[A:%.*]] to i16
+; CHECK-NEXT:    [[CONV2:%.*]] = sext i8 [[B:%.*]] to i16
+; CHECK-NEXT:    [[CMP:%.*]] = icmp slt i16 [[CONV]], [[CONV2]]
+; CHECK-NEXT:    [[COND:%.*]] = select i1 [[CMP]], i16 [[CONV2]], i16 [[CONV]]
+; CHECK-NEXT:    ret i16 [[COND]]
+;
+entry:
+  %conv = sext i8 %a to i32
+  %conv2 = sext i8 %b to i32
+  %cmp = icmp slt i32 %conv, %conv2
+  %cond = select i1 %cmp, i32 %conv2, i32 %conv
+  %conv4 = trunc i32 %cond to i16
+  ret i16 %conv4
+}
+
+define dso_local i16 @cmp_select_zext(i8 %a, i8 %b) {
+; CHECK-LABEL: @cmp_select_zext(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[CONV:%.*]] = zext i8 [[A:%.*]] to i16
+; CHECK-NEXT:    [[CONV2:%.*]] = zext i8 [[B:%.*]] to i16
+; CHECK-NEXT:    [[CMP:%.*]] = icmp slt i16 [[CONV]], [[CONV2]]
+; CHECK-NEXT:    [[COND:%.*]] = select i1 [[CMP]], i16 [[CONV2]], i16 [[CONV]]
+; CHECK-NEXT:    ret i16 [[COND]]
+;
+entry:
+  %conv = zext i8 %a to i32
+  %conv2 = zext i8 %b to i32
+  %cmp = icmp slt i32 %conv, %conv2
+  %cond = select i1 %cmp, i32 %conv2, i32 %conv
+  %conv4 = trunc i32 %cond to i16
+  ret i16 %conv4
+}
+
+define dso_local i16 @cmp_select_zext_sext(i8 %a, i8 %b) {
+; CHECK-LABEL: @cmp_select_zext_sext(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[CONV:%.*]] = zext i8 [[A:%.*]] to i16
+; CHECK-NEXT:    [[CONV2:%.*]] = sext i8 [[B:%.*]] to i16
+; CHECK-NEXT:    [[CMP:%.*]] = icmp slt i16 [[CONV]], [[CONV2]]
+; CHECK-NEXT:    [[COND:%.*]] = select i1 [[CMP]], i16 [[CONV2]], i16 [[CONV]]
+; CHECK-NEXT:    ret i16 [[COND]]
+;
+entry:
+  %conv = zext i8 %a to i32
+  %conv2 = sext i8 %b to i32
+  %cmp = icmp slt i32 %conv, %conv2
+  %cond = select i1 %cmp, i32 %conv2, i32 %conv
+  %conv4 = trunc i32 %cond to i16
+  ret i16 %conv4
+}
+
+define dso_local i16 @cmp_select_zext_sext_diffOrigTy(i8 %a, i16 %b) {
+; CHECK-LABEL: @cmp_select_zext_sext_diffOrigTy(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[CONV:%.*]] = zext i8 [[A:%.*]] to i16
+; CHECK-NEXT:    [[CMP:%.*]] = icmp slt i16 [[CONV]], [[B:%.*]]
+; CHECK-NEXT:    [[COND:%.*]] = select i1 [[CMP]], i16 [[B]], i16 [[CONV]]
+; CHECK-NEXT:    ret i16 [[COND]]
+;
+entry:
+  %conv = zext i8 %a to i32
+  %conv2 = sext i16 %b to i32
+  %cmp = icmp slt i32 %conv, %conv2
+  %cond = select i1 %cmp, i32 %conv2, i32 %conv
+  %conv4 = trunc i32 %cond to i16
+  ret i16 %conv4
+}
+
+define dso_local i16 @my_abs_sext(i8 %a) {
+; CHECK-LABEL: @my_abs_sext(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[CONV:%.*]] = sext i8 [[A:%.*]] to i16
+; CHECK-NEXT:    [[CMP:%.*]] = icmp slt i16 [[CONV]], 0
+; CHECK-NEXT:    [[SUB:%.*]] = sub i16 0, [[CONV]]
+; CHECK-NEXT:    [[COND:%.*]] = select i1 [[CMP]], i16 [[SUB]], i16 [[CONV]]
+; CHECK-NEXT:    ret i16 [[COND]]
+;
+entry:
+  %conv = sext i8 %a to i32
+  %cmp = icmp slt i32 %conv, 0
+  %sub = sub nsw i32 0, %conv
+  %cond = select i1 %cmp, i32 %sub, i32 %conv
+  %conv4 = trunc i32 %cond to i16
+  ret i16 %conv4
+}
+
+define dso_local i16 @my_abs_zext(i8 %a) {
+; CHECK-LABEL: @my_abs_zext(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[CONV:%.*]] = zext i8 [[A:%.*]] to i16
+; CHECK-NEXT:    [[CMP:%.*]] = icmp slt i16 [[CONV]], 0
+; CHECK-NEXT:    [[SUB:%.*]] = sub i16 0, [[CONV]]
+; CHECK-NEXT:    [[COND:%.*]] = select i1 [[CMP]], i16 [[SUB]], i16 [[CONV]]
+; CHECK-NEXT:    ret i16 [[COND]]
+;
+entry:
+  %conv = zext i8 %a to i32
+  %cmp = icmp slt i32 %conv, 0
+  %sub = sub nsw i32 0, %conv
+  %cond = select i1 %cmp, i32 %sub, i32 %conv
+  %conv4 = trunc i32 %cond to i16
+  ret i16 %conv4
+}
+
+define dso_local i16 @select_sext(i8 %a, i1 %cond) {
+; CHECK-LABEL: @select_sext(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[CONV:%.*]] = sext i8 [[A:%.*]] to i16
+; CHECK-NEXT:    [[SUB:%.*]] = sub i16 0, [[CONV]]
+; CHECK-NEXT:    [[SEL:%.*]] = select i1 [[COND:%.*]], i16 [[SUB]], i16 [[CONV]]
+; CHECK-NEXT:    ret i16 [[SEL]]
+;
+entry:
+  %conv = sext i8 %a to i32
+  %sub = sub nsw i32 0, %conv
+  %sel = select i1 %cond, i32 %sub, i32 %conv
+  %conv4 = trunc i32 %sel to i16
+  ret i16 %conv4
+}
+
+define dso_local i16 @select_zext(i8 %a, i1 %cond) {
+; CHECK-LABEL: @select_zext(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[CONV:%.*]] = zext i8 [[A:%.*]] to i16
+; CHECK-NEXT:    [[SUB:%.*]] = sub i16 0, [[CONV]]
+; CHECK-NEXT:    [[SEL:%.*]] = select i1 [[COND:%.*]], i16 [[SUB]], i16 [[CONV]]
+; CHECK-NEXT:    ret i16 [[SEL]]
+;
+entry:
+  %conv = zext i8 %a to i32
+  %sub = sub nsw i32 0, %conv
+  %sel = select i1 %cond, i32 %sub, i32 %conv
+  %conv4 = trunc i32 %sel to i16
+  ret i16 %conv4
+}
+
+