[ARM] Allow zext in ARMCodeGenPrepare

[samparker]

Summary

Treat zext instructions as roots, like we do for truncs.

[SjoerdMeijer]

Looks okay to me.

test/CodeGen/ARM/arm-cgp-zext-truncs.ll

88

typo: Won't don't

lib/Target/ARM/ARMCodeGenPrepare.cpp

   };
 
   if (auto *Store = dyn_cast<StoreInst>(V))
     return UsesNarrowValue(Store->getValueOperand());
   if (auto *Return = dyn_cast<ReturnInst>(V))
     return UsesNarrowValue(Return->getReturnValue());
   if (auto *Trunc = dyn_cast<TruncInst>(V))
     return UsesNarrowValue(Trunc->getOperand(0));
+  if (auto *ZExt = dyn_cast<ZExtInst>(V))
+    return UsesNarrowValue(ZExt->getOperand(0));
   if (auto *ICmp = dyn_cast<ICmpInst>(V))
     return ICmp->isSigned();
 
   return isa<CallInst>(V);
 }
 
 /// Return whether the instruction can be promoted within any modifications to
 /// it's operands or result.

     InsertDSPIntrinsic(cast<Instruction>(V));
   }
 
   auto InsertTrunc = [&](Value *V, Type *TruncTy) -> Instruction* {
     if (TruncTy == ExtTy || !isa<Instruction>(V) ||
         !isa<IntegerType>(V->getType()))
       return nullptr;
 
-    if (!Promoted.count(V) && !NewInsts.count(V))
+    if (Leaves.count(V) || (!Promoted.count(V) && !NewInsts.count(V)))
       return nullptr;
 
     LLVM_DEBUG(dbgs() << "ARM CGP: Creating " << *TruncTy << " Trunc for "
                << *V << "\n");
     Builder.SetInsertPoint(cast<Instruction>(V));
     auto *Trunc = cast<Instruction>(Builder.CreateTrunc(V, TruncTy));
     NewInsts.insert(Trunc);
     return Trunc;

       auto *PtrTy = cast<PointerType>(Store->getPointerOperandType());
       TruncTy = PtrTy->getElementType();
     } else if (isa<ReturnInst>(I)) {
       Function *F = I->getParent()->getParent();
       TruncTy = F->getFunctionType()->getReturnType();
     }
 
     // These will only have one operand to fix.
-    if (isa<StoreInst>(I) || isa<ReturnInst>(I) || isa<TruncInst>(I)) {
+    if (isa<StoreInst>(I) || isa<ReturnInst>(I) || isa<CastInst>(I)) {
       if (Instruction *Trunc = InsertTrunc(I->getOperand(0), TruncTy)) {
         Trunc->moveBefore(I);
         I->setOperand(0, Trunc);
       }
       continue;
     }
 
     // Now handle calls and signed icmps.
     for (unsigned i = 0; i < I->getNumOperands(); ++i) {
       if (auto *Call = dyn_cast<CallInst>(I))
         TruncTy = Call->getFunctionType()->getParamType(i);
       else
         TruncTy = TruncTysMap[I->getOperand(i)];
 
       if (Instruction *Trunc = InsertTrunc(I->getOperand(i), TruncTy)) {
         Trunc->moveBefore(I);
         I->setOperand(i, Trunc);
       }
     }
   }
-  LLVM_DEBUG(dbgs() << "ARM CGP: Mutation complete.\n");
+  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) {
   LLVM_DEBUG(dbgs() << "ARM CGP: Is " << *V << " supported?\n");

   // 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);
 
-  // Currently, Trunc is the only cast we support.
   if (auto *Trunc = dyn_cast<TruncInst>(V))
     return isSupportedType(Trunc->getOperand(0));
 
+  if (auto *ZExt = dyn_cast<ZExtInst>(V))
+    return isSupportedType(ZExt->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 roots.
   if (auto *Call = dyn_cast<CallInst>(V))
     return isSupportedType(Call) &&
            Call->hasRetAttr(Attribute::AttrKind::ZExt);
 
   if (!isa<BinaryOperator>(V)) {

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

define i16 @promote_i8_sink_i16_1(i8 zeroext %arg0, i16 zeroext %arg1, i16 zeroext %arg2) {
   %sel = select i1 %cmp, i16 %arg1, i16 %arg2
   %res = tail call zeroext i16 @dummy3(i16 %sel)
   ret i16 %res
 }
 
 ; CHECK-COMMON-LABEL: promote_i8_sink_i16_2
 ; CHECK-COMMON: bl dummy_i8
 ; CHECK-COMMON: adds r0, #1
-; CHECK-COMMON: uxtb r0, r0
+; CHECK-COMMON-NOT: uxt
 ; CHECK-COMMON: cmp r0
 define i16 @promote_i8_sink_i16_2(i8 zeroext %arg0, i8 zeroext %arg1, i16 zeroext %arg2) {
   %call = tail call zeroext i8 @dummy_i8(i8 %arg0)
   %add = add nuw i8 %call, 1
   %cmp = icmp ne i8 %add, %arg1
   %conv = zext i8 %arg1 to i16
   %sel = select i1 %cmp, i16 %conv, i16 %arg2
   %res = tail call zeroext i16 @dummy3(i16 %sel)

test/CodeGen/ARM/arm-cgp-zext-truncs.ll

entry:
   %0 = load i16, i16* %ptr
   %1 = add i16 %0, %arg0
   %2 = trunc i16 %1 to i8
   %3 = icmp ugt i8 %2, %arg1
   %4 = select i1 %3, i8 %2, i8 %arg1
   ret i8 %4
 }
 
-; The pass will bail because of the zext, otherwise we'd want something like:
-; ldrb [[LD:r[^ ]+]], [r0]
-; subs [[SUB:r[^ ]+]], [[LD]], #1
-; cmp [[LD]], [[SUB]]
+; The pass perform the transform, but a uxtb will still be inserted to handle
+; the zext to the icmp.
 ; CHECK-COMMON-LABEL: icmp_i32_zext:
+; CHECK-COMMON: sub
 ; CHECK-COMMON: uxtb
+; CHECK-COMMON: cmp
 define i8 @icmp_i32_zext(i8* %ptr) {
 entry:
   %gep = getelementptr inbounds i8, i8* %ptr, i32 0
   %0 = load i8, i8* %gep, align 1
   %1 = sub nuw nsw i8 %0, 1
   %conv44 = zext i8 %0 to i32
   br label %preheader
 
if.end:
   %gep1 = getelementptr inbounds i8, i8* %ptr, i32 %inc
   %3 = load i8, i8* %gep1, align 1
   br label %body
 
 exit:
   ret i8 %2
 }
 
-; Won't handle zext or sext
+; Won't don't handle sext
 ; CHECK-COMMON-LABEL: icmp_sext_zext_store_i8_i16
+; CHECK-COMMON: ldrb
+; CHECK-COMMON: ldrsh
 define i32 @icmp_sext_zext_store_i8_i16() {
 entry:
   %0 = load i8, i8* getelementptr inbounds ([16 x i8], [16 x i8]* @d_uch, i32 0, i32 2), align 1
   %conv = zext i8 %0 to i16
   store i16 %conv, i16* @sh1, align 2
   %conv1 = zext i8 %0 to i32
   %1 = load i16, i16* getelementptr inbounds ([16 x i16], [16 x i16]* @d_sh, i32 0, i32 2), align 2
   %conv2 = sext i16 %1 to i32
   %cmp = icmp eq i32 %conv1, %conv2
   %conv3 = zext i1 %cmp to i32
   ret i32 %conv3
 }
 
-; Pass will bail because of the zext, otherwise:
-; ldrb [[LD:r[^ ]+]], [r1]
-; subs [[SUB:r[^ ]+]], #1
-; cmp [[SUB]], #3
 ; CHECK-COMMON-LABEL: or_icmp_ugt:
-; CHECK-COMMON: uxt
+; CHECK-COMMON:     ldrb
+; CHECK-COMMON:     sub.w
+; CHECK-COMMON-NOT: uxt
+; CHECK-COMMON:     cmp.w
+; CHECK-COMMON-NOT: uxt
+; CHECK-COMMON:     cmp
 define i1 @or_icmp_ugt(i32 %arg, i8* %ptr) {
 entry:
   %0 = load i8, i8* %ptr
   %1 = zext i8 %0 to i32
   %mul = shl nuw nsw i32 %1, 1
   %add0 = add nuw nsw i32 %mul, 6
   %cmp0 = icmp ne i32 %arg, %add0
   %add1 = add i8 %0, -1

 default:
   br label %exit
 
 exit:
   %res = phi i16 [ %select, %sw.bb ], [ %select.i, %sw.bb.i ], [ %mul, %default ]
   ret i16 %res
 }
 
-; Pass will bail because of the zext
+; We currently only handle truncs as sinks, so a uxt will still be needed for
+; the icmp ugt instruction.
 ; CHECK-COMMON-LABEL: urem_trunc_icmps
+; CHECK-COMMON: cmp
 ; CHECK-COMMON: uxt
+; CHECK-COMMON: cmp
 define void @urem_trunc_icmps(i16** %in, i32* %g, i32* %k) {
 entry:
   %ptr = load i16*, i16** %in, align 4
   %ld = load i16, i16* %ptr, align 2
   %cmp.i = icmp eq i16 %ld, 0
   br i1 %cmp.i, label %exit, label %cond.false.i
 
 cond.false.i:
latch:
   %count = add nuw i8 %phi.2, 1
   store i8 %count, i8* %store, align 1
   br label %header
 
 exit:
   ret void
 }
 
-; Again, zexts will prevent the transform.
 ; Check that %exp requires uxth in all cases, and will also be required to
 ; promote %1 for the call - unless we can generate a uadd16.
 ; CHECK-COMMON-LABEL: zext_load_sink_call:
 ; CHECK-COMMON: uxt
 ; uadd16
 ; cmp
 ; CHECK-COMMON: uxt
 define i32 @zext_load_sink_call(i16* %ptr, i16 %exp) {