[ARM][CGP] Check trunc value size before replacing

[samparker]

Summary

In the last stage of type promotion, we replace any zext that uses a new trunc with the operand of the trunc. This is okay when we only allowed one type to be optimised, but now its the case that the trunc maybe needed to produce a more narrow type than the one we were optimising for. So we need to check this before doing the replacement.

[SjoerdMeijer]

Was just getting up to speed with this again; just a first round of nits before I start looking at the actual functional change.

lib/Target/ARM/ARMCodeGenPrepare.cpp

188

There's not enough context to show the class definition:

namespace {
class IRPromoter {
  SmallPtrSet<Value*, 8> NewInsts;
  SmallPtrSet<Instruction*, 4> InstsToRemove;
  DenseMap<Value*, SmallVector<Type*, 4>> TruncTysMap;
  SmallPtrSet<Value*, 8> Promoted;
  Module *M = nullptr;
  LLVMContext &Ctx;
  IntegerType *ExtTy = nullptr;
  IntegerType *OrigTy = nullptr;
  SmallPtrSetImpl<Value*> *Visited;
  SmallPtrSetImpl<Value*> *Sources;
  SmallPtrSetImpl<Instruction*> *Sinks;
  SmallPtrSetImpl<Instruction*> *SafeToPromote;
  ...

But I started looking at where ExtTy was defined again, because it was mentioned in the comments of this change. It is quite a crucial aspect of the IRPromoter, so perhaps you can spend a few comments on what it is and does.

329

nit: trailing whitespace

607

nit: trailing whitespace here

703

typo: "we can we can"

703

nit: they = zexts?

705

typo: remove "the"?

[samparker]

Cheers, added some comments and removed whitespace.

[SjoerdMeijer]

Thanks, LGTM

lib/Target/ARM/ARMCodeGenPrepare.cpp

 namespace {
 class IRPromoter {
   SmallPtrSet<Value*, 8> NewInsts;
   SmallPtrSet<Instruction*, 4> InstsToRemove;
   DenseMap<Value*, SmallVector<Type*, 4>> TruncTysMap;
   SmallPtrSet<Value*, 8> Promoted;
   Module *M = nullptr;
   LLVMContext &Ctx;
+  // The type we promote to: always i32
   IntegerType *ExtTy = nullptr;
+  // The type of the value that the search began from, either i8 or i16.
+  // This defines the max range of the values that we allow in the promoted
+  // tree.
   IntegerType *OrigTy = nullptr;
   SmallPtrSetImpl<Value*> *Visited;
   SmallPtrSetImpl<Value*> *Sources;
   SmallPtrSetImpl<Instruction*> *Sinks;
   SmallPtrSetImpl<Instruction*> *SafeToPromote;
 
   void ReplaceAllUsersOfWith(Value *From, Value *To);
   void PrepareConstants(void);
if (!isa<Instruction>(V))
     return false;
 
   unsigned Opc = cast<Instruction>(V)->getOpcode();
   return Opc == Instruction::AShr || Opc == Instruction::SDiv ||
          Opc == Instruction::SRem;
 }
 
 static bool EqualTypeSize(Value *V) {
   return V->getType()->getScalarSizeInBits() == ARMCodeGenPrepare::TypeSize;
 }
 
 static bool LessOrEqualTypeSize(Value *V) {
   return V->getType()->getScalarSizeInBits() <= ARMCodeGenPrepare::TypeSize;
 }
 
 static bool GreaterThanTypeSize(Value *V) {
   return V->getType()->getScalarSizeInBits() > ARMCodeGenPrepare::TypeSize;
bool ARMCodeGenPrepare::isSafeOverflow(Instruction *I) {
   // If %a = 0, %sub = -1 == FF == 255
   // As i32:
   // %sub = -1 == FF FF FF FF == 4294967295
   //
   // In this case, the unsigned compare results would be the same and this
   // would also be true for ult, uge and ugt:
   // - (255 < 254) == (0xFFFFFFFF < 254) == false
   // - (255 <= 254) == (0xFFFFFFFF <= 254) == false
   // - (255 > 254) == (0xFFFFFFFF > 254) == true
   // - (255 >= 254) == (0xFFFFFFFF >= 254) == true
   //
   // To demonstrate why we can't handle increasing values:
-  // 
+  //
   // %add = add i8 %a, 2
   // %cmp = icmp ult i8 %add, 127
   //
   // If %a = 254, %add = 256 == (i8 1)
   // As i32:
   // %add = 256
   //
   // (1 < 127) != (256 < 127)
void IRPromoter::PromoteTree() {
   // Finally, any instructions that should be promoted but haven't yet been,
   // need to be handled using intrinsics.
   for (auto *V : *Visited) {
     auto *I = dyn_cast<Instruction>(V);
     if (!I)
       continue;
 
     if (Sources->count(I) || Sinks->count(I))
       continue;
 
     if (!shouldPromote(I) || SafeToPromote->count(I) || NewInsts.count(I))
       continue;
-  
     assert(EnableDSP && "DSP intrinisc insertion not enabled!");
 
     // Replace unsafe instructions with appropriate intrinsic calls.
     LLVM_DEBUG(dbgs() << "ARM CGP: Inserting DSP intrinsic for "
                << *I << "\n");
     Function *DSPInst =
       Intrinsic::getDeclaration(M, getNarrowIntrinsic(I));
     Builder.SetInsertPoint(I);
for (auto I : *Sinks) {
     }
   }
 }
 
 void IRPromoter::Cleanup() {
   // Some zexts will now have become redundant, along with their trunc
   // operands, so remove them
   for (auto V : *Visited) {
-    if (!isa<CastInst>(V))
+    if (!isa<ZExtInst>(V))
       continue;
 
-    auto ZExt = cast<CastInst>(V);
+    auto ZExt = cast<ZExtInst>(V);
     if (ZExt->getDestTy() != ExtTy)
       continue;
 
     Value *Src = ZExt->getOperand(0);
     if (ZExt->getSrcTy() == ZExt->getDestTy()) {
       LLVM_DEBUG(dbgs() << "ARM CGP: Removing unnecessary cast: " << *ZExt
                  << "\n");
       ReplaceAllUsersOfWith(ZExt, Src);
       continue;
     }
 
-    // For any truncs that we insert to handle zexts, we can replace the
-    // result of the zext with the input to the trunc.
-    if (NewInsts.count(Src) && isa<ZExtInst>(V) && isa<TruncInst>(Src)) {
+    // Unless they produce a value that is narrower than ExtTy, we can
+    // replace the result of the zext with the input of a newly inserted
+    // trunc.
+    if (NewInsts.count(Src) && isa<TruncInst>(Src) &&
+        Src->getType() == OrigTy) {
       auto *Trunc = cast<TruncInst>(Src);
       assert(Trunc->getOperand(0)->getType() == ExtTy &&
              "expected inserted trunc to be operating on i32");
       ReplaceAllUsersOfWith(ZExt, Trunc->getOperand(0));
     }
   }
 
   for (auto *I : InstsToRemove) {

test/CodeGen/ARM/CGP/arm-cgp-casts.ll

cond.false:
   %rem.zext = trunc i16 %rem7 to i8
   br label %cond.end
 
 cond.end:
   %cond = phi i8 [ %rem.zext, %cond.false ], [ %1, %entry ]
   store i8 %cond, i8* @a, align 1
   ret void
 }
+
+; CHECK-LABEL: dont_replace_trunc_1
+; CHECK: cmp
+; CHECK: uxtb
+define void @dont_replace_trunc_1(i8* %a, i16* %b, i16* %c, i32* %d, i8* %e, i32* %f) {
+entry:
+  %0 = load i16, i16* %c, align 2
+  %1 = load i16, i16* %b, align 2
+  %conv = sext i16 %1 to i32
+  store i32 %conv, i32* %f, align 4
+  %2 = trunc i16 %1 to i8
+  %conv1 = and i8 %2, 1
+  store i8 %conv1, i8* %e, align 1
+  %3 = load i8, i8* %a, align 1
+  %narrow = mul nuw i8 %3, %conv1
+  %mul = zext i8 %narrow to i32
+  store i32 %mul, i32* %d, align 4
+  %4 = zext i8 %narrow to i16
+  %conv5 = or i16 %0, %4
+  %tobool = icmp eq i16 %conv5, 0
+  br i1 %tobool, label %if.end, label %for.cond
+
+for.cond:                                         ; preds = %entry, %for.cond
+  br label %for.cond
+
+if.end:                                           ; preds = %entry
+  ret void
+}
+
+; CHECK-LABEL: dont_replace_trunc_2
+; CHECK: cmp
+; CHECK: uxtb
+define i32 @dont_replace_trunc_2(i16* %a, i8* %b) {
+entry:
+  %0 = load i16, i16* %a, align 2
+  %cmp = icmp ugt i16 %0, 8
+  %narrow = select i1 %cmp, i16 %0, i16 0
+  %cond = trunc i16 %narrow to i8
+  %1 = load i8, i8* %b, align 1
+  %or = or i8 %1, %cond
+  store i8 %or, i8* %b, align 1
+  %conv5 = zext i8 %or to i32
+  ret i32 %conv5
+}