Allow InstCombiner to eliminate truncates even if it will require inserting additional instructions

[igor-laevsky]

Summary

Currently InstCombiner can eliminate truncate when entire expression tree can be evaluated in a narrower type without adding additional instructions.

However there are some cases when it is profitable to evaluate expression in a narrower type, even if it will require adding additional instructions. (See test case in the patch)

In this patch I split "CanEvaluateTruncated" function into two: "EstimateCostForTruncatedEvaluation" and "IsProfitableToEvaluateTruncated". First one is basically an old "CanEvaluateTruncated", but it also calculates cost for truncating expression tree. Second function basically just calls "EstimateCostForTruncatedEvaluation" and compares cost of the truncation with some threshold.

Currently cost model is very simple - we want to truncate expression tree only if we will remove more instructions than we will add. Probably it can be somehow adjusted in the future.

[sanjoy]

Overall, I think this is going in the right direction.

Some comments inline.

lib/Transforms/InstCombine/InstCombineCasts.cpp

234

Don't construct a std::string here -- just using I->getName() + ".truncated" will construct an llvm::Twine for you.

237

Use InsertBefore = I->getNextNode()

239

I'm not sure I understand this assert -- InsertBefore is not the instruction instcombine visited, right?

342

Nit: "increase"

345

If we change the bit handing the !isa<Instruction> case below, will we need this to return a bool? Or will it always "succeed"?

352

Why does the This is not a recognized instruction bit apply here?

471

I'd just move this within InstCombiner. Then you won't need IC and will be able to use ShouldChangeType without making it public.

[igor-laevsky]

lib/Transforms/InstCombine/InstCombineCasts.cpp

234

done

237

done

239

Naming is slightly confusing here. I have changed it a bit.

342

done

352

"V" could be Argument, MetadataAsValue, InlineAsm or BasicBlock.

We can handle Argument by adding special case into "EvaluateInDifferentType".

Probably for MetadataAsValue and BasicBlock we can just assert false - should not meet this instructions in arithmetic expression trees.

InlineAsm we can probably handle similarly to the Argument, but I am not absolutely sure how it works.

This are possible extensions, but they feel more like a separate change.

471

Agree. Actually both "IsProfitableToEvaluateTruncated" and "EstimateCostForTruncatedEvaluation" should live inside InstCombiner. I don't want to add too much unrelated refactoring, so I will do this in a separate change and move "ShouldChangeType" back to the private section.

[sanjoy]

lib/Transforms/InstCombine/InstCombineCasts.cpp

228

Spelling: guaranteed

238

I still don't understand this assert: why can't you have IR like:

%i = add i32 ...  ;; I == %i
br label %foo

foo:
use(%i)

340

Nit: "arrive at the same value"

342

Nit: "increase cost"

343

Nit: decrease

352

We can handle Argument by adding special case into "EvaluateInDifferentType"

Why would you need any special casing in EvaluateInDifferentType? Doesn't it have a catch-all that will take any iN to iM by using a trunc instruction?

360–365

We can eliminate the extension only if it has at most one use. Right now, if I'm reading the code correctly, you'll consider this case to be profitable:

X = f()
Y = g()

A = zext(B) from i32 to i64
C = zext(D) from i32 to i64

M = A + C + X + Y

use(A)
use(C)

and translate it to

X = f()
Y = g()

A = zext(B) from i32 to i64
C = zext(D) from i32 to i64

M = B + D + trunc(X) + trunc(Y)

use(A)
use(C)

454–455

I'd either remove the default: break; or move this bit of logic into the default: clause.

465

What is Ty?

471

Use either Instruction *I or Value *V.

[igor-laevsky]

lib/Transforms/InstCombine/InstCombineCasts.cpp

238

I->getParent()->end() is end of the instruction list, not a terminator instruction. We will fail this assertion if "I" was a terminator. In other words "I++ == BB->end()"

352

Yes, but special case is required to determine place where to insert this additional explicit trunc. I am not saying this is anyhow complicated, just seems more like a separate change.

360–365

That's right. I fixed the patch to update cost only for casts with single use.

[sanjoy]

Overall, this LGTM given that the cleanup changes discussed happen promptly. However, I'd wait for @majnemer to take a final look before checking in.

lib/Transforms/InstCombine/InstCombineCasts.cpp

225

It isn't clear what 'truncating expression type' means here. Can you turn this into an assert?

238

Yes, you're right. I confused myself.

371

Given that we're now tracking a generalized cost, does it make sense to remove this and have

case Inst::Add:
  bool Left = EstimateCostFor(LHS);
  bool Right = EstimateCostFor(LHS);
  if (Left && Right) {
    if (!I->hasOneUse())
      Cost++;  // new add instruction  (or may be +=2 because of code duplication)
    return true;
  } else return false;

(Not necessarily in this change, but in a later change?)

[igor-laevsky]

lib/Transforms/InstCombine/InstCombineCasts.cpp

225

There is an assertion for this on the next line. It checks that we can reach destination type by applying truncate instruction.

371

It is possible. However I would be careful about such transform since it increases register pressure, some targets might not respond to this change positively.

[igor-laevsky]

Gentle ping. @majnemer, can you please take a quick look at this?

[igor-laevsky]

ping?

[reames]

Looking through the patch, this looks entirely reasonable to me. However, I'm not comfortable enough with this code to sign off. Can we get someone familiar with InstCombine and our canonicalization rules around truncation to take a quick look?

lib/Transforms/InstCombine/InstCombineCasts.cpp

488

You might consider passing the CostThreshold in to CanEvaluateTruncated and return early if the running cost estimate ever gets too high above CostThreshold. We want to allow analysis of cases close to the threshold, but not burn a lot of compile time. Having said that, I'm just being pedantic. I don't think this is an actual issue.

[igor-laevsky]

lib/Transforms/InstCombine/InstCombineCasts.cpp

488

Thanks for taking a look.

I think it would be easier to just bailout on the number of visited instruction. It would give direct control on the maximum number of iterations which directly affects compile time.

Anyway I am not changing iteration count for this method. I guess if no one ever got a problem with it, I should not address it either. At least not in this change.

[majnemer]

Could this be split into two changes: one which adds the cost model and another which teaches EvaluateInDifferentType how to insert truncates?

[igor-laevsky]

In D10571#222989, @majnemer wrote:

Could this be split into two changes: one which adds the cost model and another which teaches EvaluateInDifferentType how to insert truncates?

Sure. Please take a look at http://reviews.llvm.org/D12012 and http://reviews.llvm.org/D12013

lib/Transforms/InstCombine/InstCombineCasts.cpp

if (!AI.hasOneUse()) {
     // allocation instruction, also pointer typed. Thus, cast to use is BitCast.
     Value *NewCast = AllocaBuilder.CreateBitCast(New, AI.getType(), "tmpcast");
     ReplaceInstUsesWith(AI, NewCast);
   }
   return ReplaceInstUsesWith(CI, New);
 }
 
 /// EvaluateInDifferentType - Given an expression that
-/// CanEvaluateTruncated or CanEvaluateSExtd returns true for, actually
-/// insert the code to evaluate the expression.
+/// IsProfitableToEvaluateTruncated or CanEvaluateSExtd returns true for,
+/// actually insert the code to evaluate the expression.
 Value *InstCombiner::EvaluateInDifferentType(Value *V, Type *Ty,
                                              bool isSigned) {
   if (Constant *C = dyn_cast<Constant>(V)) {
     C = ConstantExpr::getIntegerCast(C, Ty, isSigned /*Sext or ZExt*/);
     // If we got a constantexpr back, try to simplify it with DL info.
     if (ConstantExpr *CE = dyn_cast<ConstantExpr>(C))
       C = ConstantFoldConstantExpression(CE, DL, TLI);
     return C;
for (unsigned i = 0, e = OPN->getNumIncomingValues(); i != e; ++i) {
           EvaluateInDifferentType(OPN->getIncomingValue(i), Ty, isSigned);
       NPN->addIncoming(V, OPN->getIncomingBlock(i));
     }
     Res = NPN;
     break;
   }
   default:
     // TODO: Can handle more cases here.
-    llvm_unreachable("Unreachable!");
+    // This is not a recognized instruction. We need an explicit cast here.
+    // Currently we can reach this only when truncating expression type.
+    assert(CastInst::castIsValid(Instruction::Trunc, I, Ty) &&
+           "can not insert explicit cast");
+    // hasOneUse property is guaranteed by "EstimateCostForTruncatedEvaluation"
+    assert(I->hasOneUse() &&
+           "can not truncate branching expression trees");
+
+    // Create truncate from "I" type to the Ty
+    Instruction *Trunc = CastInst::CreateTruncOrBitCast(I, Ty);
+    Trunc->setName(I->getName() + ".truncated");
+
+    // Insert new truncate after "I"
+    Instruction *NextI = I->getNextNode();
+    assert(NextI != I->getParent()->end() &&
+           "can not insert explicit truncate");
+    return InsertNewInstBefore(Trunc, *NextI);
   }
   
+  assert(Res && "failed to evaluate instruction in different type");
   Res->takeName(I);
   return InsertNewInstWith(Res, *I);
 }
 
 
 /// This function is a wrapper around CastInst::isEliminableCastPair. It
 /// simply extracts arguments and returns what that function returns.
 static Instruction::CastOps
if (!Src->getType()->isIntegerTy() || !CI.getType()->isIntegerTy() ||
         ShouldChangeType(CI.getType(), Src->getType()))
       if (Instruction *NV = FoldOpIntoPhi(CI))
         return NV;
   }
 
   return nullptr;
 }
 
-/// CanEvaluateTruncated - Return true if we can evaluate the specified
-/// expression tree as type Ty instead of its larger type, and arrive with the
-/// same value.  This is used by code that tries to eliminate truncates.
-///
-/// Ty will always be a type smaller than V.  We should return true if trunc(V)
-/// can be computed by computing V in the smaller type.  If V is an instruction,
-/// then trunc(inst(x,y)) can be computed as inst(trunc(x),trunc(y)), which only
-/// makes sense if x and y can be efficiently truncated.
-///
-/// This function works on both vectors and scalars.
-///
-static bool CanEvaluateTruncated(Value *V, Type *Ty, InstCombiner &IC,
+/// EstimateCostForTruncatedEvaluation - Return true if we can evaluate the
+/// specified expression tree as type Ty instead of its larger type, and
+/// arrive at the same value. Also calculate cost of this transformation
+/// and store it into the "Cost" argument.
+/// Currently cost model is simple - increase cost by one for each additional
+/// instruction required, decrease cost by one for each instruction that
+/// will be removed after truncating expression tree.
+static bool EstimateCostForTruncatedEvaluation(int &Cost, Value *V,
+                                               Type *Ty, InstCombiner &IC,
                                                Instruction *CxtI) {
   // We can always evaluate constants in another type.
   if (isa<Constant>(V))
     return true;
 
   Instruction *I = dyn_cast<Instruction>(V);
   if (!I) return false;
 
   Type *OrigTy = V->getType();
 
-  // If this is an extension from the dest type, we can eliminate it, even if it
-  // has multiple uses.
+  // If this is an extension from the dest type, we will not need any additional
+  // instructions to truncate it, even if it has multiple uses.
   if ((isa<ZExtInst>(I) || isa<SExtInst>(I)) &&
-      I->getOperand(0)->getType() == Ty)
+      I->getOperand(0)->getType() == Ty) {
+
+    // Add a cost bonus if we can eliminate this cast.
+    if (I->hasOneUse())
+      Cost -= 1;
+
     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;
 
   unsigned Opc = I->getOpcode();
   switch (Opc) {
   case Instruction::Add:
   case Instruction::Sub:
   case Instruction::Mul:
   case Instruction::And:
   case Instruction::Or:
   case Instruction::Xor:
     // These operators can all arbitrarily be extended or truncated.
-    return CanEvaluateTruncated(I->getOperand(0), Ty, IC, CxtI) &&
-           CanEvaluateTruncated(I->getOperand(1), Ty, IC, CxtI);
+    return EstimateCostForTruncatedEvaluation(Cost, I->getOperand(0), Ty, IC,
+                                              CxtI) &&
+           EstimateCostForTruncatedEvaluation(Cost, I->getOperand(1), Ty, IC,
+                                              CxtI);
 
   case Instruction::UDiv:
   case Instruction::URem: {
     // UDiv and URem can be truncated if all the truncated bits are zero.
     uint32_t OrigBitWidth = OrigTy->getScalarSizeInBits();
     uint32_t BitWidth = Ty->getScalarSizeInBits();
     if (BitWidth < OrigBitWidth) {
       APInt Mask = APInt::getHighBitsSet(OrigBitWidth, OrigBitWidth-BitWidth);
       if (IC.MaskedValueIsZero(I->getOperand(0), Mask, 0, CxtI) &&
           IC.MaskedValueIsZero(I->getOperand(1), Mask, 0, CxtI)) {
-        return CanEvaluateTruncated(I->getOperand(0), Ty, IC, CxtI) &&
-               CanEvaluateTruncated(I->getOperand(1), Ty, IC, CxtI);
+        return EstimateCostForTruncatedEvaluation(Cost, I->getOperand(0), Ty,
+                                                  IC, CxtI) &&
+               EstimateCostForTruncatedEvaluation(Cost, I->getOperand(1), Ty,
+                                                  IC, CxtI);
       }
     }
     break;
   }
   case Instruction::Shl:
     // If we are truncating the result of this SHL, and if it's a shift of a
     // constant amount, we can always perform a SHL in a smaller type.
     if (ConstantInt *CI = dyn_cast<ConstantInt>(I->getOperand(1))) {
       uint32_t BitWidth = Ty->getScalarSizeInBits();
       if (CI->getLimitedValue(BitWidth) < BitWidth)
-        return CanEvaluateTruncated(I->getOperand(0), Ty, IC, CxtI);
+        return EstimateCostForTruncatedEvaluation(Cost, I->getOperand(0), Ty,
+                                                  IC, CxtI);
     }
     break;
   case Instruction::LShr:
     // If this is a truncate of a logical shr, we can truncate it to a smaller
     // lshr iff we know that the bits we would otherwise be shifting in are
     // already zeros.
     if (ConstantInt *CI = dyn_cast<ConstantInt>(I->getOperand(1))) {
       uint32_t OrigBitWidth = OrigTy->getScalarSizeInBits();
       uint32_t BitWidth = Ty->getScalarSizeInBits();
       if (IC.MaskedValueIsZero(I->getOperand(0),
             APInt::getHighBitsSet(OrigBitWidth, OrigBitWidth-BitWidth), 0, CxtI) &&
           CI->getLimitedValue(BitWidth) < BitWidth) {
-        return CanEvaluateTruncated(I->getOperand(0), Ty, IC, CxtI);
+        return EstimateCostForTruncatedEvaluation(Cost, I->getOperand(0), Ty,
+                                                  IC, CxtI);
       }
     }
     break;
   case Instruction::Trunc:
     // trunc(trunc(x)) -> trunc(x)
     return true;
   case Instruction::ZExt:
   case Instruction::SExt:
     // trunc(ext(x)) -> ext(x) if the source type is smaller than the new dest
     // trunc(ext(x)) -> trunc(x) if the source type is larger than the new dest
     return true;
   case Instruction::Select: {
     SelectInst *SI = cast<SelectInst>(I);
-    return CanEvaluateTruncated(SI->getTrueValue(), Ty, IC, CxtI) &&
-           CanEvaluateTruncated(SI->getFalseValue(), Ty, IC, CxtI);
+    return EstimateCostForTruncatedEvaluation(Cost, SI->getTrueValue(), Ty, IC,
+                                              CxtI) &&
+           EstimateCostForTruncatedEvaluation(Cost, SI->getFalseValue(), Ty, IC,
+                                              CxtI);
   }
   case Instruction::PHI: {
     // We can change a phi if we can change all operands.  Note that we never
     // get into trouble with cyclic PHIs here because we only consider
     // instructions with a single use.
     PHINode *PN = cast<PHINode>(I);
     for (Value *IncValue : PN->incoming_values())
-      if (!CanEvaluateTruncated(IncValue, Ty, IC, CxtI))
+      if (!EstimateCostForTruncatedEvaluation(Cost, IncValue, Ty, IC, CxtI))
         return false;
     return true;
   }
-  default:
-    // TODO: Can handle more cases here.
-    break;
+
   }
 
+  // This is not a recognized instruction, but we can always truncate it by
+  // adding additional explicit trunc.
+  Cost += 1;
+  return true;
+}
+
+/// IsProfitableToEvaluateTruncated - Given the expression tree starting with
+/// truncate instruction return true if it would be profitable to remove
+/// this truncate and evaluate expression in a narrower type. This is used by
+/// code that tries to eliminate truncates.
+///
+/// This function works on both vectors and scalars.
+///
+static bool IsProfitableToEvaluateTruncated(TruncInst *Trunc,
+                                            InstCombiner &IC) {
+  Type *DestTy = Trunc->getType(), *SrcTy = Trunc->getOperand(0)->getType();
+
+  // Only truncate if the dest type is a simple type, don't convert the
+  // expression tree to something weird like i93 unless the source is also
+  // strange.
+  if (!DestTy->isVectorTy() && !IC.ShouldChangeType(SrcTy, DestTy))
     return false;
+
+  int Cost = 0;
+  bool CanTruncate = EstimateCostForTruncatedEvaluation(
+      Cost, Trunc->getOperand(0), DestTy, IC, Trunc);
+
+  // Threshold picked in a way that it will allow only cases when number of
+  // removed instructions is strictly greater than number of the additional
+  // instructions required.
+  const int CostThreshold = 1;
+  return CanTruncate && Cost < CostThreshold;
 }
 
 Instruction *InstCombiner::visitTrunc(TruncInst &CI) {
   if (Instruction *Result = commonCastTransforms(CI))
     return Result;
 
   // Test if the trunc is the user of a select which is part of a
   // minimum or maximum operation. If so, don't do any more simplification.
   // Even simplifying demanded bits can break the canonical form of a 
   // min/max.
   Value *LHS, *RHS;
   if (SelectInst *SI = dyn_cast<SelectInst>(CI.getOperand(0)))
     if (matchSelectPattern(SI, LHS, RHS).Flavor != SPF_UNKNOWN)
       return nullptr;
   
   // See if we can simplify any instructions used by the input whose sole
   // purpose is to compute bits we don't care about.
   if (SimplifyDemandedInstructionBits(CI))
     return &CI;
 
   Value *Src = CI.getOperand(0);
-  Type *DestTy = CI.getType(), *SrcTy = Src->getType();
+  Type *DestTy = CI.getType();
 
   // Attempt to truncate the entire input expression tree to the destination
-  // type.   Only do this if the dest type is a simple type, don't convert the
-  // expression tree to something weird like i93 unless the source is also
-  // strange.
-  if ((DestTy->isVectorTy() || ShouldChangeType(SrcTy, DestTy)) &&
-      CanEvaluateTruncated(Src, DestTy, *this, &CI)) {
-
+  // type.
+  if (IsProfitableToEvaluateTruncated(&CI, *this)) {
     // If this cast is a truncate, evaluting in a different type always
     // eliminates the cast, so it is always a win.
     DEBUG(dbgs() << "ICE: EvaluateInDifferentType converting expression type"
           " to avoid cast: " << CI << '\n');
     Value *Res = EvaluateInDifferentType(Src, DestTy, false);
     assert(Res->getType() == DestTy);
     return ReplaceInstUsesWith(CI, Res);
   }

lib/Transforms/InstCombine/InstCombineInternal.h

public:
   // The routine checks that the DI parent and DB are different.
   bool dominatesAllUses(const Instruction *DI, const Instruction *UI,
                         const BasicBlock *DB) const;
 
   // Replace select with select operand SIOpd in SI-ICmp sequence when possible
   bool replacedSelectWithOperand(SelectInst *SI, const ICmpInst *Icmp,
                                  const unsigned SIOpd);
 
-private:
   bool ShouldChangeType(Type *From, Type *To) const;
+private:
   Value *dyn_castNegVal(Value *V) const;
   Value *dyn_castFNegVal(Value *V, bool NoSignedZero = false) const;
   Type *FindElementAtOffset(PointerType *PtrTy, int64_t Offset,
                             SmallVectorImpl<Value *> &NewIndices);
   Instruction *FoldOpIntoSelect(Instruction &Op, SelectInst *SI);
 
   /// \brief Classify whether a cast is worth optimizing.
   ///

test/Transforms/InstCombine/trunc.ll

define i8 @test10(i32 %X) {
   %Y = trunc i32 %X to i8
   %Z = and i8 %Y, 42
   ret i8 %Z
 ; CHECK-LABEL: @test10(
 ; CHECK: trunc
 ; CHECK: and
 ; CHECK: ret
 }
+
+; Here by eliminating trunc we can get rid of a redundant zext. However
+; additional trunc should also be inserted. Check that we consider this
+; transform to be profitable.
+define void @test11(i8 *%res, i32* %t, i8* %s) {
+ entry:
+  %m = load i32, i32* %t
+  %n = load i8, i8* %s
+
+  %n.i32 = zext i8 %n to i32
+  %v = sub i32 %n.i32, %m
+  %v.i8 = trunc i32 %v to i8
+
+  store i8 %v.i8, i8* %res
+  ret void
+; CHECK-LABEL: @test11(
+; CHECK: %m = load i32, i32* %t
+; CHECK: %m.truncated = trunc i32 %m to i8
+; CHECK: %n = load i8, i8* %s
+; CHECK: %v = sub i8 %n, %m.truncated
+; CHECK: store i8 %v, i8* %res
+}
+
+; Here by eliminating trunc we will not remove any additional instructions.
+; Check that this case considered to be not profitable.
+define void @test12(i8 *%res, i32* %t, i32* %s) {
+ entry:
+  %m = load i32, i32* %t
+  %n = load i32, i32* %s
+
+  %v = sub i32 %n, %m
+  %v.i8 = trunc i32 %v to i8
+
+  store i8 %v.i8, i8* %res
+  ret void
+; CHECK-LABEL: @test12(
+; CHECK: %m = load i32, i32* %t
+; CHECK: %n = load i32, i32* %s
+; CHECK: %v = sub i32 %n, %m
+; CHECK: %v.i8 = trunc i32 %v to i8
+; CHECK: store i8 %v.i8, i8* %res
+}
+
+; Check that casts with multiple uses are considered not possible to remove.
+declare void @use.test13(i32)
+define void @test13(i8 *%res, i32* %t, i8* %s) {
+ entry:
+  %m = load i32, i32* %t
+  %n = load i8, i8* %s
+
+  %n.i32 = zext i8 %n to i32
+  %v = sub i32 %n.i32, %m
+  %v.i8 = trunc i32 %v to i8
+
+  store i8 %v.i8, i8* %res
+  call void @use.test13(i32 %n.i32)
+  ret void
+; CHECK-LABEL: @test13(
+; CHECK: %m = load i32, i32* %t
+; CHECK: %n = load i8, i8* %s
+; CHECK: %n.i32 = zext i8 %n to i32
+; CHECK: %v = sub i32 %n.i32, %m
+; CHECK: %v.i8 = trunc i32 %v to i8
+; CHECK: store i8 %v.i8, i8* %res
+; CHECK: call void @use.test13(i32 %n.i32)
+; CHECK: ret void
+}