[ARM] Enable mixed types in ARM CGP

[samparker]

Summary

Previously, during the search, all values had to have the same 'TypeSize', which is equal to number of bits of the integer type of the icmp operand. All values in the tree had to match this size; meaning that, if we searched from i16, we wouldn't accept i8s. A change in type size requires zext and truncs to perform the casts so, to allow mixed narrow types, the handling of these instructions is now different:

• zexts are sinks if the result is larger than TypeSize.
• truncs are sinks if their result is less or equal to TypeSize.
• truncs are sources if their incoming value is greater or equal than TypeSize.

[SjoerdMeijer]

lib/Target/ARM/ARMCodeGenPrepare.cpp

239

I don't really understand the 'Equal' part in the GreaterOrEqual comparison. Can these types be equal, and is that what we want?

251

By briefly looking at these new cases here and the tests, I think we are missing some tests.

[samparker]

lib/Target/ARM/ARMCodeGenPrepare.cpp

239

Good question, looks like I've messed up there.

251

Agreed.

[samparker]

Now disallowing icmps that operate on types that are smaller than TypeSize. The handling of truncs being sources and sinks has also been reverted.

• we allow casts if their result or operand is <= TypeSize.
• zexts are sinks if their result > TypeSize.
• truncs are still sinks if their operand == TypeSize.
• truncs are still sources if their result == TypeSize.

[SjoerdMeijer]

LGTM

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

392

nit: I think we should have 2 checks for uxtb here.

Same for the other tests below.

lib/Target/ARM/ARMCodeGenPrepare.cpp

void PrepareConstants(SmallPtrSetImpl<Value*> &Visited,
                          SmallPtrSetImpl<Instruction*> &SafeToPromote);
   void ExtendSources(SmallPtrSetImpl<Value*> &Sources);
   void PromoteTree(SmallPtrSetImpl<Value*> &Visited,
                    SmallPtrSetImpl<Value*> &Sources,
                    SmallPtrSetImpl<Instruction*> &Sinks,
                    SmallPtrSetImpl<Instruction*> &SafeToPromote);
   void TruncateSinks(SmallPtrSetImpl<Value*> &Sources,
                      SmallPtrSetImpl<Instruction*> &Sinks);
-  void Cleanup(SmallPtrSetImpl<Instruction*> &Sinks);
+  void Cleanup(SmallPtrSetImpl<Value*> &Visited);
 
 public:
   IRPromoter(Module *M) : M(M), Ctx(M->getContext()),
                           ExtTy(Type::getInt32Ty(Ctx)) { }
 
 
   void Mutate(Type *OrigTy,
               SmallPtrSetImpl<Value*> &Visited,
static bool generateSignBits(Value *V) {
   if (!isa<Instruction>(V))
     return false;
 
   unsigned Opc = cast<Instruction>(V)->getOpcode();
   return Opc == Instruction::AShr || Opc == Instruction::SDiv ||
          Opc == Instruction::SRem;
 }
 
+static bool LessOrEqualTypeSize(Value *V) {
+  return V->getType()->getScalarSizeInBits() <= ARMCodeGenPrepare::TypeSize;
+}
+
+static bool LessThanTypeSize(Value *V) {
+  return V->getType()->getScalarSizeInBits() < ARMCodeGenPrepare::TypeSize;
+}
+
+static bool GreaterOrEqualTypeSize(Value *V) {
+  return V->getType()->getScalarSizeInBits() >= ARMCodeGenPrepare::TypeSize;
+}
+
+static bool GreaterThanTypeSize(Value *V) {
+  return V->getType()->getScalarSizeInBits() > ARMCodeGenPrepare::TypeSize;
+}
+
 /// Some instructions can use 8- and 16-bit operands, and we don't need to
 /// promote anything larger. We disallow booleans to make life easier when
 /// dealing with icmps but allow any other integer that is <= 16 bits. Void
 /// types are accepted so we can handle switches.
 static bool isSupportedType(Value *V) {
   Type *Ty = V->getType();
 
   // Allow voids and pointers, these won't be promoted.
   if (Ty->isVoidTy() || Ty->isPointerTy())
     return true;
 
   if (auto *Ld = dyn_cast<LoadInst>(V))
     Ty = cast<PointerType>(Ld->getPointerOperandType())->getElementType();
 
-  const IntegerType *IntTy = dyn_cast<IntegerType>(Ty);
-  if (!IntTy)
+  if (!isa<IntegerType>(Ty))
     return false;
 
-  return IntTy->getBitWidth() == ARMCodeGenPrepare::TypeSize;
+  return LessOrEqualTypeSize(V);
 }
 
 /// Return true if the given value is a source in the use-def chain, producing
 /// a narrow (i8, i16) value. These values will be zext to start the promotion
 /// of the tree to i32. We guarantee that these won't populate the upper bits
 /// of the register. ZExt on the loads will be free, and the same for call
 /// return values because we only accept ones that guarantee a zeroext ret val.
 /// Many arguments will have the zeroext attribute too, so those would be free
 /// too.
 static bool isSource(Value *V) {
   if (!isa<IntegerType>(V->getType()))
     return false;
   // TODO Allow zext to be sources.
   if (isa<Argument>(V))
     return true;
   else if (isa<LoadInst>(V))
     return true;
   else if (isa<BitCastInst>(V))
     return true;
   else if (auto *Call = dyn_cast<CallInst>(V))
     return Call->hasRetAttr(Attribute::AttrKind::ZExt);
   else if (auto *Trunc = dyn_cast<TruncInst>(V))
-    return isSupportedType(Trunc);
+    return GreaterOrEqualTypeSize(Trunc->getOperand(0));
   return false;
 }
 
 /// Return true if V will require any promoted values to be truncated for the
 /// the IR to remain valid. We can't mutate the value type of these
 /// instructions.
 static bool isSink(Value *V) {
   // TODO The truncate also isn't actually necessary because we would already
   // proved that the data value is kept within the range of the original data
   // type.
-  auto UsesNarrowValue = [](Value *V) {
-    return V->getType()->getScalarSizeInBits() == ARMCodeGenPrepare::TypeSize;
-  };
 
   if (auto *Store = dyn_cast<StoreInst>(V))
-    return UsesNarrowValue(Store->getValueOperand());
+    return LessOrEqualTypeSize(Store->getValueOperand());
   if (auto *Return = dyn_cast<ReturnInst>(V))
-    return UsesNarrowValue(Return->getReturnValue());
+    return LessOrEqualTypeSize(Return->getReturnValue());
   if (auto *Trunc = dyn_cast<TruncInst>(V))
-    return UsesNarrowValue(Trunc->getOperand(0));
+    return LessOrEqualTypeSize(Trunc);
   if (auto *ZExt = dyn_cast<ZExtInst>(V))
-    return UsesNarrowValue(ZExt->getOperand(0));
+    return GreaterThanTypeSize(ZExt);
   if (auto *ICmp = dyn_cast<ICmpInst>(V))
-    return ICmp->isSigned();
+    return ICmp->isSigned() || LessThanTypeSize(ICmp->getOperand(0));
 
   return isa<CallInst>(V);
 }
 
 /// Return whether the instruction can be promoted within any modifications to
 /// its operands or result.
 bool ARMCodeGenPrepare::isSafeOverflow(Instruction *I) {
   // FIXME Do we need NSW too?
for (auto I : Sinks) {
     // Now handle the others.
     for (unsigned i = 0; i < I->getNumOperands(); ++i) {
       if (Instruction *Trunc = InsertTrunc(I->getOperand(i))) {
         Trunc->moveBefore(I);
         I->setOperand(i, Trunc);
       }
     }
   }
-
 }
 
-void IRPromoter::Cleanup(SmallPtrSetImpl<Instruction*> &Sinks) {
-  // Some zext sinks will now have become redundant, along with their trunc
-  // operands, so remove them.
-  for (auto I : Sinks) {
-    if (auto *ZExt = dyn_cast<ZExtInst>(I)) {
-      if (ZExt->getDestTy() != ExtTy)
-        continue;
+void IRPromoter::Cleanup(SmallPtrSetImpl<Value*> &Visited) {
+  // Some zexts will now have become redundant, along with their trunc
+  // operands, so remove them
+  for (auto V : Visited) {
+    if (!isa<ZExtInst>(V))
+      continue;
+
+    auto ZExt = cast<ZExtInst>(V);
 
-      Value *Src = ZExt->getOperand(0);
-      if (ZExt->getSrcTy() == ZExt->getDestTy()) {
-        LLVM_DEBUG(dbgs() << "ARM CGP: Removing unnecessary zext\n");
-        ReplaceAllUsersOfWith(ZExt, Src);
-        InstsToRemove.push_back(ZExt);
-        continue;
-      }
+      LLVM_DEBUG(dbgs() << "ARM CGP: Removing unnecessary cast.\n");
 
-      // 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<TruncInst>(Src)) {
-        auto *Trunc = cast<TruncInst>(Src);
-        assert(Trunc->getOperand(0)->getType() == ExtTy &&
-               "expected inserted trunc to be operating on i32");
-        LLVM_DEBUG(dbgs() << "ARM CGP: Replacing zext with trunc operand: "
-                   << *Trunc->getOperand(0));
-        ReplaceAllUsersOfWith(ZExt, Trunc->getOperand(0));
-        InstsToRemove.push_back(ZExt);
-      }
-    }
-  }
 
   for (auto *I : InstsToRemove) {
     LLVM_DEBUG(dbgs() << "ARM CGP: Removing " << *I << "\n");
     I->dropAllReferences();
     I->eraseFromParent();
   }
 
   InstsToRemove.clear();
void IRPromoter::Mutate(Type *OrigTy,
   // promote.
   PromoteTree(Visited, Sources, Sinks, SafeToPromote);
 
   // Insert trunc instructions for use by calls, stores etc...
   TruncateSinks(Sources, Sinks);
 
   // Finally, remove unecessary zexts and truncs, delete old instructions and
   // clear the data structures.
-  Cleanup(Sinks);
+  Cleanup(Visited);
 
-  LLVM_DEBUG(dbgs() << "ARM CGP: Mutation complete:\n");
-  LLVM_DEBUG(dbgs();
-             for (auto *V : Sources)
-               V->dump();
-             for (auto *I : NewInsts)
-               I->dump();
-             for (auto *V : Visited) {
-               if (!Sources.count(V))
-                 V->dump();
-              });
+  LLVM_DEBUG(dbgs() << "ARM CGP: Mutation complete\n");
 }
 
 /// We accept most instructions, as well as Arguments and ConstantInsts. We
 /// Disallow casts other than zext and truncs and only allow calls if their
 /// return value is zeroext. We don't allow opcodes that can introduce sign
 /// bits.
 bool ARMCodeGenPrepare::isSupportedValue(Value *V) {
   if (isa<ICmpInst>(V))
bool ARMCodeGenPrepare::isSupportedValue(Value *V) {
   // Non-instruction values that we can handle.
   if ((isa<Constant>(V) && !isa<ConstantExpr>(V)) || isa<Argument>(V))
     return isSupportedType(V);
 
   if (isa<PHINode>(V) || isa<SelectInst>(V) || isa<ReturnInst>(V) ||
       isa<LoadInst>(V))
     return isSupportedType(V);
 
-  // Truncs can be either sources or sinks.
-  if (auto *Trunc = dyn_cast<TruncInst>(V))
-    return isSupportedType(Trunc) || isSupportedType(Trunc->getOperand(0));
+  if (isa<SExtInst>(V))
+    return false;
 
-  if (isa<CastInst>(V) && !isa<SExtInst>(V))
-    return isSupportedType(cast<CastInst>(V)->getOperand(0));
+  if (auto *Cast = dyn_cast<CastInst>(V))
+    return isSupportedType(Cast) || isSupportedType(Cast->getOperand(0));
 
   // Special cases for calls as we need to check for zeroext
   // TODO We should accept calls even if they don't have zeroext, as they can
   // still be sinks.
   if (auto *Call = dyn_cast<CallInst>(V))
     return isSupportedType(Call) &&
            Call->hasRetAttr(Attribute::AttrKind::ZExt);
 
for (auto &I : Insts) {
 
       if (isa<ICmpInst>(I)) {
         auto &CI = cast<ICmpInst>(I);
 
         // Skip signed or pointer compares
         if (CI.isSigned() || !isa<IntegerType>(CI.getOperand(0)->getType()))
           continue;
 
+        LLVM_DEBUG(dbgs() << "ARM CGP: Searching from: " << CI << "\n");
+
         for (auto &Op : CI.operands()) {
           if (auto *I = dyn_cast<Instruction>(Op))
             MadeChange |= TryToPromote(I);
         }
       }
     }
     LLVM_DEBUG(if (verifyFunction(F, &dbgs())) {
                 dbgs() << F;

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

 ; CHECK-NEXT: sxtb r1, r0
 define i16 @test_call(i8 zeroext %arg) {
   %call = call i8 @dummy_i8(i8 %arg)
   %cmp = icmp ult i8 %call, 128
   %conv = zext i1 %cmp to i16
   ret i16 %conv
 }
 
-; Test that the transformation bails when it finds that i16 is larger than i8.
-; TODO: We should be able to remove the uxtb in these cases.
 ; CHECK-LABEL: promote_i8_sink_i16_1
 ; CHECK: bl dummy_i8
 ; CHECK: add{{.*}} r0, #1
-; CHECK: uxtb r0, r0
+; CHECK-NOT: uxt
 ; CHECK: cmp r0
 define i16 @promote_i8_sink_i16_1(i8 zeroext %arg0, i16 zeroext %arg1, i16 zeroext %arg2) {
   %call = tail call zeroext i8 @dummy_i8(i8 %arg0)
   %add = add nuw i8 %call, 1
   %conv = zext i8 %add to i16
   %cmp = icmp ne i16 %conv, %arg1
   %sel = select i1 %cmp, i16 %arg1, i16 %arg2
   %res = tail call zeroext i16 @dummy3(i16 %sel)

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

 declare %class.x* @_ZNK2ae2afEv(%class.ae*) local_unnamed_addr
 declare %class.v* @_ZN1x2acEv(%class.x*) local_unnamed_addr
 declare i32 @dummy(i32, i32)
 
 @d_uch = hidden local_unnamed_addr global [16 x i8] zeroinitializer, align 1
 @sh1 = hidden local_unnamed_addr global i16 0, align 2
 @d_sh = hidden local_unnamed_addr global [16 x i16] zeroinitializer, align 2
 
-; CHECK-LABEL: two_stage_zext_trunc_mix
-; CHECK-NOT: uxt
+; CHECK-COMMON-LABEL: two_stage_zext_trunc_mix
+; CHECK-COMMON-NOT: uxt
 define i8* @two_stage_zext_trunc_mix(i32* %this, i32 %__pos1, i32 %__n1, i32** %__str, i32 %__pos2, i32 %__n2) {
 entry:
   %__size_.i.i.i.i = bitcast i32** %__str to i8*
   %0 = load i8, i8* %__size_.i.i.i.i, align 4
   %1 = and i8 %0, 1
   %tobool.i.i.i.i = icmp eq i8 %1, 0
   %__size_.i5.i.i = getelementptr inbounds i32*, i32** %__str, i32 %__n1
   %cast = bitcast i32** %__size_.i5.i.i to i32*
   %2 = load i32, i32* %cast, align 4
   %3 = lshr i8 %0, 1
   %4 = zext i8 %3 to i32
   %cond.i.i = select i1 %tobool.i.i.i.i, i32 %4, i32 %2
   %__size_.i.i.i.i.i = bitcast i32* %this to i8*
   %5 = load i8, i8* %__size_.i.i.i.i.i, align 4
   %6 = and i8 %5, 1
   %tobool.i.i.i.i.i = icmp eq i8 %6, 0
   %7 = getelementptr inbounds i8, i8* %__size_.i.i.i.i, i32 %__pos1
   %8 = getelementptr inbounds i8, i8* %__size_.i.i.i.i, i32 %__pos2
   %res = select i1 %tobool.i.i.i.i.i,  i8* %7, i8* %8
   ret i8* %res
 }
+
+; CHECK-COMMON-LABEL: search_through_zext_1
+; CHECK-COMMON-NOT: uxt
+define i8 @search_through_zext_1(i8 zeroext %a, i8 zeroext %b, i16 zeroext %c) {
+entry:
+  %add = add nuw i8 %a, %b
+  %conv = zext i8 %add to i16
+  %cmp = icmp ult i16 %conv, %c
+  br i1 %cmp, label %if.then, label %if.end
+
+if.then:
+  %sub = sub nuw i8 %b, %a
+  %conv2 = zext i8 %sub to i16
+  %cmp2 = icmp ugt i16 %conv2, %c
+  %res = select i1 %cmp2, i8 %a, i8 %b
+  br label %if.end
+
+if.end:
+  %retval = phi i8 [ 0, %entry ], [ %res, %if.then ]
+  ret i8 %retval
+}
+
+; TODO: We should be able to remove the uxtb here. The transform fails because
+; the icmp ugt uses an i32, which is too large... but this doesn't matter
+; because it won't be writing a large value to a register as a result.
+; CHECK-COMMON-LABEL: search_through_zext_2
+; CHECK-COMMON: uxtb
+define i8 @search_through_zext_2(i8 zeroext %a, i8 zeroext %b, i16 zeroext %c, i32 %d) {
+entry:
+  %add = add nuw i8 %a, %b
+  %conv = zext i8 %add to i16
+  %cmp = icmp ult i16 %conv, %c
+  br i1 %cmp, label %if.then, label %if.end
+
+if.then:
+  %sub = sub nuw i8 %b, %a
+  %conv2 = zext i8 %sub to i32
+  %cmp2 = icmp ugt i32 %conv2, %d
+  %res = select i1 %cmp2, i8 %a, i8 %b
+  br label %if.end
+
+if.end:
+  %retval = phi i8 [ 0, %entry ], [ %res, %if.then ]
+  ret i8 %retval
+}
+
+; TODO: We should be able to remove the uxtb here as all the calculations are
+; performed on i8s. The promotion of i8 to i16 and then the later truncation
+; results in the uxtb.
+; CHECK-COMMON-LABEL: search_through_zext_3
+; CHECK-COMMON: uxtb
+define i8 @search_through_zext_3(i8 zeroext %a, i8 zeroext %b, i16 zeroext %c, i32 %d) {
+entry:
+  %add = add nuw i8 %a, %b
+  %conv = zext i8 %add to i16
+  %cmp = icmp ult i16 %conv, %c
+  br i1 %cmp, label %if.then, label %if.end
+
+if.then:
+  %trunc = trunc i16 %conv to i8
+  %sub = sub nuw i8 %b, %trunc
+  %conv2 = zext i8 %sub to i32
+  %cmp2 = icmp ugt i32 %conv2, %d
+  %res = select i1 %cmp2, i8 %a, i8 %b
+  br label %if.end
+
+if.end:
+  %retval = phi i8 [ 0, %entry ], [ %res, %if.then ]
+  ret i8 %retval
+}
+
+; CHECK-COMMON-LABEL: search_through_zext_cmp
+; CHECK-COMMON-NOT: uxt
+define i8 @search_through_zext_cmp(i8 zeroext %a, i8 zeroext %b, i16 zeroext %c) {
+entry:
+  %cmp = icmp ne i8 %a, %b
+  %conv = zext i1 %cmp to i16
+  %cmp1 = icmp ult i16 %conv, %c
+  br i1 %cmp1, label %if.then, label %if.end
+
+if.then:
+  %sub = sub nuw i8 %b, %a
+  %conv2 = zext i8 %sub to i16
+  %cmp3 = icmp ugt i16 %conv2, %c
+  %res = select i1 %cmp3, i8 %a, i8 %b
+  br label %if.end
+
+if.end:
+  %retval = phi i8 [ 0, %entry ], [ %res, %if.then ]
+  ret i8 %retval
+}
+
+; CHECK-COMMON-LABEL: search_through_zext_load
+; CHECK-COMMON-NOT: uxt
+define i8 @search_through_zext_load(i8* %a, i8 zeroext %b, i16 zeroext %c) {
+entry:
+  %load = load i8, i8* %a
+  %conv = zext i8 %load to i16
+  %cmp1 = icmp ult i16 %conv, %c
+  br i1 %cmp1, label %if.then, label %if.end
+
+if.then:
+  %sub = sub nuw i8 %b, %load
+  %conv2 = zext i8 %sub to i16
+  %cmp3 = icmp ugt i16 %conv2, %c
+  %res = select i1 %cmp3, i8 %load, i8 %b
+  br label %if.end
+
+if.end:
+  %retval = phi i8 [ 0, %entry ], [ %res, %if.then ]
+  ret i8 %retval
+}
+
+; CHECK-COMMON-LABEL: trunc_sink_less_than
+; CHECK-COMMON-NOT: uxth
+; CHECK-COMMON: cmp
+; CHECK-COMMON: uxtb
+define i16 @trunc_sink_less_than_cmp(i16 zeroext %a, i16 zeroext %b, i16 zeroext %c, i8 zeroext %d) {
+entry:
+  %sub = sub nuw i16 %b, %a
+  %cmp = icmp ult i16 %sub, %c
+  br i1 %cmp, label %if.then, label %if.end
+
+if.then:
+  %trunc = trunc i16 %sub to i8
+  %add = add nuw i8 %d, 1
+  %cmp2 = icmp ugt i8 %trunc, %add
+  %res = select i1 %cmp2, i16 %a, i16 %b
+  br label %if.end
+
+if.end:
+  %retval = phi i16 [ 0, %entry ], [ %res, %if.then ]
+  ret i16 %retval
+}
+
+; CHECK-COMMON-LABEL: trunc_sink_less_than_arith
+; CHECK-COMMON-NOT: uxth
+; CHECK-COMMON: cmp
+; CHECK-COMMON: uxtb
+define i16 @trunc_sink_less_than_arith(i16 zeroext %a, i16 zeroext %b, i16 zeroext %c, i8 zeroext %d, i8 zeroext %e) {
+entry:
+  %sub = sub nuw i16 %b, %a
+  %cmp = icmp ult i16 %sub, %c
+  br i1 %cmp, label %if.then, label %if.end
+
+if.then:
+  %trunc = trunc i16 %sub to i8
+  %add = add nuw i8 %d, %trunc
+  %cmp2 = icmp ugt i8 %e, %add
+  %res = select i1 %cmp2, i16 %a, i16 %b
+  br label %if.end
+
+if.end:
+  %retval = phi i16 [ 0, %entry ], [ %res, %if.then ]
+  ret i16 %retval
+}