[TypePromotion] Search from ZExt + PHI

Expand TypePromotion pass to try to promote PHI-nodes in loops that are the
operand of a ZExt, using the ZExt's result type to determine the Promote Width.

Differential Revision: https://reviews.llvm.org/D111237

llvm/lib/CodeGen/TypePromotion.cpp

@@ -17,6 +17,7 @@
 
 #include "llvm/ADT/SetVector.h"
 #include "llvm/ADT/StringRef.h"
+#include "llvm/Analysis/LoopInfo.h"
 #include "llvm/Analysis/TargetTransformInfo.h"
 #include "llvm/CodeGen/Passes.h"
 #include "llvm/CodeGen/TargetLowering.h"
@@ -175,9 +176,11 @@ class TypePromotion : public FunctionPass {
   TypePromotion() : FunctionPass(ID) {}
 
   void getAnalysisUsage(AnalysisUsage &AU) const override {
+    AU.addRequired<LoopInfoWrapperPass>();
     AU.addRequired<TargetTransformInfoWrapperPass>();
     AU.addRequired<TargetPassConfig>();
     AU.setPreservesCFG();
+    AU.addPreserved<LoopInfoWrapperPass>();
   }
 
   StringRef getPassName() const override { return PASS_NAME; }
@@ -872,7 +875,8 @@ bool TypePromotion::TryToPromote(Value *V, unsigned PromotedWidth) {
 
   // DAG optimizations should be able to handle these cases better, especially
   // for function arguments.
-  if (ToPromote < 2 || (Blocks.size() == 1 && (NonFreeArgs > SafeWrap.size())))
+  if (!isa<PHINode>(V) && (ToPromote < 2 || (Blocks.size() == 1 &&
+                                             (NonFreeArgs > SafeWrap.size()))))
     return false;
 
   IRPromoter Promoter(*Ctx, PromotedWidth, CurrentVisited, Sources, Sinks,
@@ -901,6 +905,7 @@ bool TypePromotion::runOnFunction(Function &F) {
   const TargetLowering *TLI = SubtargetInfo->getTargetLowering();
   const TargetTransformInfo &TII =
       getAnalysis<TargetTransformInfoWrapperPass>().getTTI(F);
+  const LoopInfo &LI = getAnalysis<LoopInfoWrapperPass>().getLoopInfo();
   RegisterBitWidth =
       TII.getRegisterBitWidth(TargetTransformInfo::RGK_Scalar).getFixedSize();
   Ctx = &F.getParent()->getContext();
@@ -929,28 +934,39 @@ bool TypePromotion::runOnFunction(Function &F) {
     return PromotedVT.getFixedSizeInBits();
   };
 
-  // Search up from icmps to try to promote their operands.
+  auto BBIsInLoop = [&](BasicBlock *BB) -> bool {
+    for (auto *L : LI)
+      if (L->contains(BB))
+        return true;
+    return false;
+  };
+
   for (BasicBlock &BB : F) {
     for (Instruction &I : BB) {
       if (AllVisited.count(&I))
         continue;
 
-      if (!isa<ICmpInst>(&I))
-        continue;
-
-      auto *ICmp = cast<ICmpInst>(&I);
-
-      // Skip signed
-      if (ICmp->isSigned())
-        continue;
-
-      LLVM_DEBUG(dbgs() << "IR Promotion: Searching from: " << *ICmp << "\n");
-
-      for (auto &Op : ICmp->operands()) {
-        if (auto *OpI = dyn_cast<Instruction>(Op)) {
-          if (auto PromotedWidth = GetPromoteWidth(OpI)) {
-            MadeChange |= TryToPromote(OpI, PromotedWidth);
-            break;
+      if (isa<ZExtInst>(&I) && isa<PHINode>(I.getOperand(0)) &&
+          BBIsInLoop(&BB)) {
+        LLVM_DEBUG(dbgs() << "IR Promotion: Searching from: " << I.getOperand(0)
+                          << "\n");
+        EVT ZExtVT = TLI->getValueType(DL, I.getType());
+        Instruction *Phi = static_cast<Instruction *>(I.getOperand(0));
+        MadeChange |= TryToPromote(Phi, ZExtVT.getFixedSizeInBits());
+      } else if (auto *ICmp = dyn_cast<ICmpInst>(&I)) {
+        // Search up from icmps to try to promote their operands.
+        // Skip signed or pointer compares
+        if (ICmp->isSigned())
+          continue;
+
+        LLVM_DEBUG(dbgs() << "IR Promotion: Searching from: " << *ICmp << "\n");
+
+        for (auto &Op : ICmp->operands()) {
+          if (auto *OpI = dyn_cast<Instruction>(Op)) {
+            if (auto PromotedWidth = GetPromoteWidth(OpI)) {
+              MadeChange |= TryToPromote(OpI, PromotedWidth);
+              break;
+            }
           }
         }
       }

llvm/test/CodeGen/AArch64/O3-pipeline.ll

@@ -90,8 +90,8 @@
 ; CHECK-NEXT:       Interleaved Load Combine Pass
 ; CHECK-NEXT:       Dominator Tree Construction
 ; CHECK-NEXT:       Interleaved Access Pass
-; CHECK-NEXT:       Type Promotion
 ; CHECK-NEXT:       Natural Loop Information
+; CHECK-NEXT:       Type Promotion
 ; CHECK-NEXT:       CodeGen Prepare
 ; CHECK-NEXT:       Dominator Tree Construction
 ; CHECK-NEXT:       Exception handling preparation

llvm/test/Transforms/TypePromotion/AArch64/phi-zext-gep2.ll

@@ -0,0 +1,44 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
+; RUN: opt -mtriple=aarch64 -type-promotion -verify -S %s -o - | FileCheck %s
+target datalayout = "e-m:e-i8:8:32-i16:16:32-i64:64-i128:128-n32:64-S128"
+
+; Function Attrs: mustprogress nofree nosync nounwind uwtable
+define dso_local void @foo(ptr noundef %ptr0, ptr nocapture noundef readonly %ptr1, ptr nocapture noundef %dest) local_unnamed_addr {
+; CHECK-LABEL: @foo(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[TMP0:%.*]] = load i8, ptr [[PTR0:%.*]], align 1
+; CHECK-NEXT:    [[TMP1:%.*]] = zext i8 [[TMP0]] to i64
+; CHECK-NEXT:    br label [[DO_BODY:%.*]]
+; CHECK:       do.body:
+; CHECK-NEXT:    [[TO_PROMOTE:%.*]] = phi i64 [ [[TMP1]], [[ENTRY:%.*]] ], [ [[TMP4:%.*]], [[DO_BODY]] ]
+; CHECK-NEXT:    [[ARRAYIDX1:%.*]] = getelementptr inbounds i8, ptr [[PTR1:%.*]], i64 [[TO_PROMOTE]]
+; CHECK-NEXT:    [[TMP2:%.*]] = load i8, ptr [[ARRAYIDX1]], align 2
+; CHECK-NEXT:    [[TMP3:%.*]] = zext i8 [[TMP2]] to i64
+; CHECK-NEXT:    [[COND_IN_I:%.*]] = getelementptr inbounds i8, ptr [[PTR1]], i64 [[TMP3]]
+; CHECK-NEXT:    [[COND_I:%.*]] = load i8, ptr [[COND_IN_I]], align 1
+; CHECK-NEXT:    [[TMP4]] = zext i8 [[COND_I]] to i64
+; CHECK-NEXT:    store i8 [[TMP2]], ptr [[DEST:%.*]], align 1
+; CHECK-NEXT:    [[CMP:%.*]] = icmp ult i64 [[TMP4]], 0
+; CHECK-NEXT:    br i1 [[CMP]], label [[DO_BODY]], label [[DO_END:%.*]]
+; CHECK:       do.end:
+; CHECK-NEXT:    ret void
+;
+entry:
+  %0 = load i8, ptr %ptr0, align 1
+  br label %do.body
+
+do.body:                                          ; preds = %do.body, %entry
+  %to_promote = phi i8 [ %0, %entry ], [ %cond.i, %do.body ]
+  %ext0 = zext i8 %to_promote to i64
+  %arrayidx1 = getelementptr inbounds i8, ptr %ptr1, i64 %ext0
+  %1 = load i8, ptr %arrayidx1, align 2
+  %2 = zext i8 %1 to i64
+  %cond.in.i = getelementptr inbounds i8, ptr %ptr1, i64 %2
+  %cond.i = load i8, ptr %cond.in.i, align 1
+  store i8 %1, ptr %dest, align 1
+  %cmp = icmp ult i8 %cond.i, 0
+  br i1 %cmp, label %do.body, label %do.end
+
+do.end:                                           ; preds = %do.body
+  ret void
+}

llvm/test/Transforms/TypePromotion/ARM/casts.ll

@@ -57,24 +57,26 @@ define i8 @icmp_i32_zext(i8* %ptr) {
 ; CHECK-NEXT:  entry:
 ; CHECK-NEXT:    [[GEP:%.*]] = getelementptr inbounds i8, i8* [[PTR:%.*]], i32 0
 ; CHECK-NEXT:    [[TMP0:%.*]] = load i8, i8* [[GEP]], align 1
-; CHECK-NEXT:    [[TMP1:%.*]] = sub nuw nsw i8 [[TMP0]], 1
+; CHECK-NEXT:    [[TMP1:%.*]] = zext i8 [[TMP0]] to i32
+; CHECK-NEXT:    [[TMP2:%.*]] = sub nuw nsw i32 [[TMP1]], 1
 ; CHECK-NEXT:    [[CONV44:%.*]] = zext i8 [[TMP0]] to i32
 ; CHECK-NEXT:    br label [[PREHEADER:%.*]]
 ; CHECK:       preheader:
 ; CHECK-NEXT:    br label [[BODY:%.*]]
 ; CHECK:       body:
-; CHECK-NEXT:    [[TMP2:%.*]] = phi i8 [ [[TMP1]], [[PREHEADER]] ], [ [[TMP3:%.*]], [[IF_END:%.*]] ]
+; CHECK-NEXT:    [[TMP3:%.*]] = phi i32 [ [[TMP2]], [[PREHEADER]] ], [ [[TMP5:%.*]], [[IF_END:%.*]] ]
 ; CHECK-NEXT:    [[SI_0274:%.*]] = phi i32 [ [[CONV44]], [[PREHEADER]] ], [ [[INC:%.*]], [[IF_END]] ]
-; CHECK-NEXT:    [[CONV51266:%.*]] = zext i8 [[TMP2]] to i32
-; CHECK-NEXT:    [[CMP52267:%.*]] = icmp eq i32 [[SI_0274]], [[CONV51266]]
+; CHECK-NEXT:    [[CMP52267:%.*]] = icmp eq i32 [[SI_0274]], [[TMP3]]
 ; CHECK-NEXT:    br i1 [[CMP52267]], label [[IF_END]], label [[EXIT:%.*]]
 ; CHECK:       if.end:
 ; CHECK-NEXT:    [[INC]] = add i32 [[SI_0274]], 1
 ; CHECK-NEXT:    [[GEP1:%.*]] = getelementptr inbounds i8, i8* [[PTR]], i32 [[INC]]
-; CHECK-NEXT:    [[TMP3]] = load i8, i8* [[GEP1]], align 1
+; CHECK-NEXT:    [[TMP4:%.*]] = load i8, i8* [[GEP1]], align 1
+; CHECK-NEXT:    [[TMP5]] = zext i8 [[TMP4]] to i32
 ; CHECK-NEXT:    br label [[BODY]]
 ; CHECK:       exit:
-; CHECK-NEXT:    ret i8 [[TMP2]]
+; CHECK-NEXT:    [[TMP6:%.*]] = trunc i32 [[TMP3]] to i8
+; CHECK-NEXT:    ret i8 [[TMP6]]
 ;
 entry:
   %gep = getelementptr inbounds i8, i8* %ptr, i32 0