[Reland][LSR] Hoist IVInc to loop header if its all uses are in the loop header

[bcl5980]

Summary

Original code will cause crash when the load/store memory type is structure because isIndexedLoadLegal/isIndexedStore doesn't support struct type.
So we limit the load/store memory type to integer.

Origin commit message:
When the latch block is different from header block, IVInc will be expanded in the latch loop. We can't generate the post index load/store this case.
But if the IVInc only used in the loop, actually we still can use the post index load/store because when exit loop we don't care the last IVInc value.
So, trying to hoist IVInc to help backend to generate more post index load/store.

Fix #53625

[bcl5980]

Update tests.

[bcl5980]

Ping.

[eopXD]

Minor comments on style and redundant code, the change looks reasonable to me and resolves the problem.

llvm/lib/Transforms/Scalar/LoopStrengthReduce.cpp

5646

Personally I think isAllUsedInBB (or isUserAllWithinBB) is a more suitable name that follows general naming rules across the codebase.

5646

Also, maybe a comment for the function for clarity.

5658

CanHoistIVInc -> canHoistIVInc.

5668

The condition here makes the previous early-return condition redundant.

5668

Possibly a negative test case for this?

5671

For the sake of my own knowledge, may I ask why does this necessary has to be an Instruction?

5675

This looks familiar with mayUsePostIncMode

5701

No need to change the nested loop here, can do

/* ... */ = canHoistIVInc(TTI, Fixup, Uses[LUIdx], /* ... */);

[eopXD]

llvm/test/Transforms/LoopStrengthReduce/ivincs-hoist.ll

4 ↗

(On Diff #477856)

As a generic test case here (not under a specific target), is there other ways to show the effect of this patch without adding this target triple?

[bcl5980]

llvm/lib/Transforms/Scalar/LoopStrengthReduce.cpp

5668

I'm not sure how to get the assertion. Why not in latch means it is in the header?

5671

I can change to Value here, but I feel strange if it is a constant.

5675

It looks mayUsePostIncMode is not what I want. It only checks the instruction is legal but haven't check the instruction opcode.

llvm/test/Transforms/LoopStrengthReduce/ivincs-hoist.ll

4 ↗

(On Diff #477856)

The headache thing is we must have a target to enable the post-index load/store to enable the optimization. I don't know how to show the code change without target triple info.

[bcl5980]

addresss comments.

[bcl5980]

Ping

[bcl5980]

Ping

[eopXD]

Sorry for the delay. Minor comments left and I think we are good to land this.

llvm/lib/Transforms/Scalar/LoopStrengthReduce.cpp

5668

Use LLVM_DEBUG and have the following at the first line of the test case (.ll file)

; REQUIRES: asserts

5675

if ((isa<LoadInst>(I) || isa<StoreInst>(I)) && mayUsePostIncMode(...)

Just trying my best to reuse existing code :)

llvm/test/Transforms/LoopStrengthReduce/ivincs-hoist.ll

4 ↗

(On Diff #477856)

Then I guess the target triple here should not be added. You already have the test case under AaArch64 for this.

[bcl5980]

llvm/lib/Transforms/Scalar/LoopStrengthReduce.cpp

5675

It looks I can't find an easy way to reuse mayUsePostIncMode. It need the paramter SCEV that I can't find here.

[eopXD]

LGTM, thank you.

[fhahn]

llvm/test/Transforms/LoopStrengthReduce/AArch64/pr53625.ll

3

This should not be needed, the tests don't require asserts.

[fhahn]

llvm/lib/Transforms/Scalar/LoopStrengthReduce.cpp

5646

nit: Check if all users of \p V are in \p BB?

name should use are instead of is and Users, as this checks all users.

5648

nit: use any_of?

5656

...all IVInc users are ...

5677

can just return (isa<LoadInst>(I) &&.....

llvm/test/Transforms/LoopStrengthReduce/AArch64/pr53625.ll

100

nit: drop redundant indvars. prefix to make IR more concise.

[bcl5980]

address comments

[bcl5980]

@fhahn , do you have any other comments for the patch?

[tobiasjs]

This commit seems to cause a regression in numba tests. The call to TTI.isIndexedLoadLegal(...) calls with type that has TypeID StructTyID, which causes an abort in getValueType because AllowUnknown is false by default.

[bcl5980]

In D138636#4039899, @tobiasjs wrote:

This commit seems to cause a regression in numba tests. The call to TTI.isIndexedLoadLegal(...) calls with type that has TypeID StructTyID, which causes an abort in getValueType because AllowUnknown is false by default.

@tobiasjs , can you help to upload the regression case?

[bcl5980]

Fix crash when the load/store memory type is struct.

[bcl5980]

rebase code.

llvm/lib/Transforms/Scalar/LoopStrengthReduce.cpp

if (PHINode *PN = dyn_cast<PHINode>(LF.UserInst)) {
     else
       LF.UserInst->replaceUsesOfWith(LF.OperandValToReplace, FullV);
   }
 
   if (auto *OperandIsInstr = dyn_cast<Instruction>(LF.OperandValToReplace))
     DeadInsts.emplace_back(OperandIsInstr);
 }
 
+// Trying to hoist the IVInc to loop header if all IVInc users are in
+// the loop header. It will help backend to generate post index load/store
+// when the latch block is different from loop header block.
+static bool canHoistIVInc(const TargetTransformInfo &TTI, const LSRFixup &Fixup,
+                          const LSRUse &LU, Instruction *IVIncInsertPos,
+                          Loop *L) {
+  if (LU.Kind != LSRUse::Address)
+    return false;
+
+  // For now this code do the conservative optimization, only work for
+  // the header block. Later we can hoist the IVInc to the block post
+  // dominate all users.
+  BasicBlock *LHeader = L->getHeader();
+  if (IVIncInsertPos->getParent() == LHeader)
+    return false;
+
+  if (!Fixup.OperandValToReplace ||
+      any_of(Fixup.OperandValToReplace->users(), [&LHeader](User *U) {
+        Instruction *UI = cast<Instruction>(U);
+        return UI->getParent() != LHeader;
+      }))
+    return false;
+
+  Instruction *I = Fixup.UserInst;
+  Type *Ty = I->getType();
+  return Ty->isIntegerTy() &&
+         ((isa<LoadInst>(I) && TTI.isIndexedLoadLegal(TTI.MIM_PostInc, Ty)) ||
+          (isa<StoreInst>(I) && TTI.isIndexedStoreLegal(TTI.MIM_PostInc, Ty)));
+}
+
 /// Rewrite all the fixup locations with new values, following the chosen
 /// solution.
 void LSRInstance::ImplementSolution(
     const SmallVectorImpl<const Formula *> &Solution) {
   // Keep track of instructions we may have made dead, so that
   // we can remove them after we are done working.
   SmallVector<WeakTrackingVH, 16> DeadInsts;
 
-  Rewriter.setIVIncInsertPos(L, IVIncInsertPos);
-
   // Mark phi nodes that terminate chains so the expander tries to reuse them.
   for (const IVChain &Chain : IVChainVec) {
     if (PHINode *PN = dyn_cast<PHINode>(Chain.tailUserInst()))
       Rewriter.setChainedPhi(PN);
   }
 
   // Expand the new value definitions and update the users.
   for (size_t LUIdx = 0, NumUses = Uses.size(); LUIdx != NumUses; ++LUIdx)
     for (const LSRFixup &Fixup : Uses[LUIdx].Fixups) {
+      Instruction *InsertPos =
+          canHoistIVInc(TTI, Fixup, Uses[LUIdx], IVIncInsertPos, L)
+              ? L->getHeader()->getTerminator()
+              : IVIncInsertPos;
+      Rewriter.setIVIncInsertPos(L, InsertPos);
       Rewrite(Uses[LUIdx], Fixup, *Solution[LUIdx], DeadInsts);
       Changed = true;
     }
 
   for (const IVChain &Chain : IVChainVec) {
     GenerateIVChain(Chain, DeadInsts);
     Changed = true;
   }
 
   for (const WeakVH &IV : Rewriter.getInsertedIVs())
     if (IV && dyn_cast<Instruction>(&*IV)->getParent())
       ScalarEvolutionIVs.push_back(IV);

llvm/test/Transforms/LoopStrengthReduce/AArch64/pr53625.ll

 ; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
 ; RUN: llc < %s -mtriple=aarch64-unknown-unknown | FileCheck %s
 
 target datalayout = "e-m:w-p:64:64-i32:32-i64:64-i128:128-n32:64-S128"
 
 ; we need to generate post index load for this case
 define i32 @test(i32 %c, ptr %a, ptr %b) {
 ; CHECK-LABEL: test:
 ; CHECK:       // %bb.0: // %entry
 ; CHECK-NEXT:    cmp w0, #1
 ; CHECK-NEXT:    b.lt .LBB0_4
 ; CHECK-NEXT:  // %bb.1: // %for.body.preheader
 ; CHECK-NEXT:    mov w8, w0
 ; CHECK-NEXT:  .LBB0_2: // %for.body
 ; CHECK-NEXT:    // =>This Inner Loop Header: Depth=1
-; CHECK-NEXT:    ldr w9, [x1]
+; CHECK-NEXT:    ldr w9, [x1], #4
 ; CHECK-NEXT:    cbnz w9, .LBB0_5
 ; CHECK-NEXT:  // %bb.3: // %for.cond
 ; CHECK-NEXT:    // in Loop: Header=BB0_2 Depth=1
-; CHECK-NEXT:    add x1, x1, #4
 ; CHECK-NEXT:    subs x8, x8, #1
 ; CHECK-NEXT:    b.ne .LBB0_2
 ; CHECK-NEXT:  .LBB0_4:
 ; CHECK-NEXT:    mov w0, wzr
 ; CHECK-NEXT:    ret
 ; CHECK-NEXT:  .LBB0_5:
 ; CHECK-NEXT:    mov w0, #1
 ; CHECK-NEXT:    ret
 entry:
   %cmp13 = icmp sgt i32 %c, 0
   br i1 %cmp13, label %for.body.preheader, label %return
 
 for.body.preheader:                               ; preds = %entry
   %wide.trip.count = zext i32 %c to i64
   br label %for.body
 
 for.cond:                                         ; preds = %for.body
-  %indvars.iv.next = add nuw nsw i64 %indvars.iv, 1
-  %exitcond.not = icmp eq i64 %indvars.iv.next, %wide.trip.count
+  %iv.next = add nuw nsw i64 %iv, 1
+  %exitcond.not = icmp eq i64 %iv.next, %wide.trip.count
   br i1 %exitcond.not, label %return, label %for.body
 
 for.body:                                         ; preds = %for.body.preheader, %for.cond
-  %indvars.iv = phi i64 [ 0, %for.body.preheader ], [ %indvars.iv.next, %for.cond ]
-  %arrayidx = getelementptr inbounds i32, ptr %a, i64 %indvars.iv
+  %iv = phi i64 [ 0, %for.body.preheader ], [ %iv.next, %for.cond ]
+  %arrayidx = getelementptr inbounds i32, ptr %a, i64 %iv
   %val = load i32, ptr %arrayidx, align 4
   %tobool3.not = icmp eq i32 %val, 0
   br i1 %tobool3.not, label %for.cond, label %return
 
 return:                                           ; preds = %for.cond, %for.body, %entry
   %retval.1 = phi i32 [ 0, %entry ], [ 0, %for.cond ], [ 1, %for.body ]
   ret i32 %retval.1
 }
+
+; negative case: %arrayidx.b is not in header
+
+define i64 @IVIncHoist_not_all_user_in_header(i32 %c, ptr %a, ptr %b) {
+; CHECK-LABEL: IVIncHoist_not_all_user_in_header:
+; CHECK:       // %bb.0: // %entry
+; CHECK-NEXT:    cmp w0, #1
+; CHECK-NEXT:    b.lt .LBB1_5
+; CHECK-NEXT:  // %bb.1: // %for.body.preheader
+; CHECK-NEXT:    mov x8, xzr
+; CHECK-NEXT:    mov w9, w0
+; CHECK-NEXT:    add x10, x1, #4
+; CHECK-NEXT:    add x11, x2, #8
+; CHECK-NEXT:    mov w0, #1
+; CHECK-NEXT:  .LBB1_2: // %for.body
+; CHECK-NEXT:    // =>This Inner Loop Header: Depth=1
+; CHECK-NEXT:    ldr w12, [x10, x8, lsl #2]
+; CHECK-NEXT:    cbnz w12, .LBB1_7
+; CHECK-NEXT:  // %bb.3: // %if.then
+; CHECK-NEXT:    // in Loop: Header=BB1_2 Depth=1
+; CHECK-NEXT:    ldr w12, [x11, x8, lsl #2]
+; CHECK-NEXT:    cbnz w12, .LBB1_6
+; CHECK-NEXT:  // %bb.4: // %for.cond
+; CHECK-NEXT:    // in Loop: Header=BB1_2 Depth=1
+; CHECK-NEXT:    add x8, x8, #1
+; CHECK-NEXT:    cmp x9, x8
+; CHECK-NEXT:    b.ne .LBB1_2
+; CHECK-NEXT:  .LBB1_5:
+; CHECK-NEXT:    mov x0, xzr
+; CHECK-NEXT:    ret
+; CHECK-NEXT:  .LBB1_6: // %if.then.return.loopexit_crit_edge
+; CHECK-NEXT:    add x0, x8, #3
+; CHECK-NEXT:  .LBB1_7: // %return
+; CHECK-NEXT:    ret
+entry:
+  %cmp13 = icmp sgt i32 %c, 0
+  br i1 %cmp13, label %for.body.preheader, label %return
+
+for.body.preheader:                               ; preds = %entry
+  %wide.trip.count = zext i32 %c to i64
+  br label %for.body
+
+for.cond:                                         ; preds = %for.body
+  %iv.next = add nuw nsw i64 %iv, 1
+  %exitcond.not = icmp eq i64 %iv.next, %wide.trip.count
+  br i1 %exitcond.not, label %return, label %for.body
+
+for.body:                                         ; preds = %for.body.preheader, %for.cond
+  %iv = phi i64 [ 0, %for.body.preheader ], [ %iv.next, %for.cond ]
+  %iv.a = phi i64 [ 1, %for.body.preheader ], [ %iv.next.a, %for.cond ]
+  %iv.b = phi i64 [ 2, %for.body.preheader ], [ %iv.next.b, %for.cond ]
+  %arrayidx.a = getelementptr inbounds i32, ptr %a, i64 %iv.a
+  %iv.next.a = add nuw nsw i64 %iv.a, 1
+  %val.a = load i32, ptr %arrayidx.a, align 4
+  %tobool3.not = icmp eq i32 %val.a, 0
+  br i1 %tobool3.not, label %if.then, label %return
+
+if.then:
+  %arrayidx.b = getelementptr inbounds i32, ptr %b, i64 %iv.b
+  %iv.next.b = add nuw nsw i64 %iv.b, 1
+  %val.b = load i32, ptr %arrayidx.b, align 4
+  %tobool4.not = icmp eq i32 %val.b, 0
+  br i1 %tobool4.not, label %for.cond, label %return
+
+return:                                           ; preds = %for.cond, %for.body, %entry
+  %retval.1 = phi i64 [ 0, %entry ], [ 0, %for.cond ], [ 1, %for.body ], [ %iv.next.b, %if.then ]
+  ret i64 %retval.1
+}
+
+%struct.A = type { i32 }
+
+; this used to crash because the load type is struct
+
+define i32 @negative_test_type_is_struct(i32 %c, ptr %a, ptr %b) {
+; CHECK-LABEL: negative_test_type_is_struct:
+; CHECK:       // %bb.0: // %entry
+; CHECK-NEXT:    cmp w0, #1
+; CHECK-NEXT:    b.lt .LBB2_4
+; CHECK-NEXT:  // %bb.1: // %for.body.preheader
+; CHECK-NEXT:    mov w8, w0
+; CHECK-NEXT:  .LBB2_2: // %for.body
+; CHECK-NEXT:    // =>This Inner Loop Header: Depth=1
+; CHECK-NEXT:    ldr w9, [x1]
+; CHECK-NEXT:    cbnz w9, .LBB2_5
+; CHECK-NEXT:  // %bb.3: // %for.cond
+; CHECK-NEXT:    // in Loop: Header=BB2_2 Depth=1
+; CHECK-NEXT:    add x1, x1, #4
+; CHECK-NEXT:    subs x8, x8, #1
+; CHECK-NEXT:    b.ne .LBB2_2
+; CHECK-NEXT:  .LBB2_4:
+; CHECK-NEXT:    mov w0, wzr
+; CHECK-NEXT:    ret
+; CHECK-NEXT:  .LBB2_5:
+; CHECK-NEXT:    mov w0, #1
+; CHECK-NEXT:    ret
+entry:
+  %cmp13 = icmp sgt i32 %c, 0
+  br i1 %cmp13, label %for.body.preheader, label %return
+
+for.body.preheader:                               ; preds = %entry
+  %wide.trip.count = zext i32 %c to i64
+  br label %for.body
+
+for.cond:                                         ; preds = %for.body
+  %iv.next = add nuw nsw i64 %iv, 1
+  %exitcond.not = icmp eq i64 %iv.next, %wide.trip.count
+  br i1 %exitcond.not, label %return, label %for.body
+
+for.body:                                         ; preds = %for.body.preheader, %for.cond
+  %iv = phi i64 [ 0, %for.body.preheader ], [ %iv.next, %for.cond ]
+  %arrayidx = getelementptr inbounds i32, ptr %a, i64 %iv
+  %str.val = load %struct.A, ptr %arrayidx
+  %val = extractvalue %struct.A %str.val, 0
+  %tobool3.not = icmp eq i32 %val, 0
+  br i1 %tobool3.not, label %for.cond, label %return
+
+return:                                           ; preds = %for.cond, %for.body, %entry
+  %retval.1 = phi i32 [ 0, %entry ], [ 0, %for.cond ], [ 1, %for.body ]
+  ret i32 %retval.1
+}
+