[LoopFlatten] Move it from LPM2 to LPM1

[SjoerdMeijer]

Summary

In D110057 we moved LoopFlatten to a LoopPassManager. This caused a regression for our 64-bit targets (the 32-bit were unaffected), the pass is no longer triggering for a motivating example. Long story short, the IR is just very different than expected; we try to match loop statements and particular uses of induction variables. The easiest is to just move LoopFlatten to a place in the pipeline where the IR is as expected, which is just before IndVarSimplify. This means we move it from LPM2 to LPM1, so that it actually runs just a bit earlier from where it was running before. IndVarSimplify is responsible for significant rewrites that are difficult to "look through" in LoopFlatten.

One thing I had to do, and am unsure about, is the necessary change to add MemorySSAAnalysis as preserved. Without this, I am running in this assertion failure:

fatal error: error in backend: Loop pass manager using MemorySSA contains a pass that does not preserve MemorySSA

It looks correct to me to add MemorySSA as preserved, but am still investigating this, and thought about sharing this in the mean time.

[lebedev.ri]

Does LoopFlatten not make use of SCEV?

[SjoerdMeijer]

In D116612#3220448, @lebedev.ri wrote:

Does LoopFlatten not make use of SCEV?

I was expecting this question. :-)

But yes, it does. In a nutshell, we discover the induction phi, loop increment, the loop test statement, and verify with SCEV that we have found the right ones. In addition, we only allow particular uses of the induction variable for the transformation to be valid. The latter is where IndVarSimplify is causing troubles; it introduces extra uses of the IV that are unexpected for LoopFlatten. Moving LoopFlatten earlier solves this problem, and I don't see any disadvantages of doing that.

[nikic]

Could you please add a PhaseOrdering test that demonstrates the issue?

llvm/lib/Transforms/Scalar/LoopFlatten.cpp

828 ↗

(On Diff #397354)

This is very likely not correct without additional changes. This pass changes control flow, and that usually requires updates to MSSA as well.

[SjoerdMeijer]

Thanks for the pointers! I have added a PhaseOrdering test, while I now look into updating MSSA.

[SjoerdMeijer]

I decided to do the updating of MSSA separately in D116660 (and will remove it from here).

[SjoerdMeijer]

Ping

[SjoerdMeijer]

@dmgreen : You reverted "rG86825fc2fb36: [LoopFlatten] Move it to a LoopPassManager" because of a downstream regression. With D116660 now accepted/committed, that unblocks this change. This should address this regression, so are you happy with this?

[dmgreen]

I'm happy so long as the performance remains OK! Does this need to be rebased now?

[SjoerdMeijer]

In D116612#3254116, @dmgreen wrote:

I'm happy so long as the performance remains OK! Does this need to be rebased now?

Ah, yeah, this needs a bit of rebasing. So I will do that and also double check performance before landing this, and will recommit D110057 and apply this change in one go.

llvm/lib/Passes/PassBuilderPipelines.cpp

PassBuilder::buildO1FunctionSimplificationPipeline(OptimizationLevel Level,
   // TODO: Investigate promotion cap for O1.
   LPM1.addPass(LICMPass(PTO.LicmMssaOptCap, PTO.LicmMssaNoAccForPromotionCap));
 
   LPM1.addPass(LoopRotatePass(/* Disable header duplication */ true,
                               isLTOPreLink(Phase)));
   // TODO: Investigate promotion cap for O1.
   LPM1.addPass(LICMPass(PTO.LicmMssaOptCap, PTO.LicmMssaNoAccForPromotionCap));
   LPM1.addPass(SimpleLoopUnswitchPass());
+  if (EnableLoopFlatten)
+    LPM1.addPass(LoopFlattenPass());
 
   LPM2.addPass(LoopIdiomRecognizePass());
   LPM2.addPass(IndVarSimplifyPass());
 
   for (auto &C : LateLoopOptimizationsEPCallbacks)
     C(LPM2, Level);
 
   LPM2.addPass(LoopDeletionPass());
 
   if (EnableLoopInterchange)
     LPM2.addPass(LoopInterchangePass());
-  if (EnableLoopFlatten)
-    LPM2.addPass(LoopFlattenPass());
 
   // Do not enable unrolling in PreLinkThinLTO phase during sample PGO
   // because it changes IR to makes profile annotation in back compile
   // inaccurate. The normal unroller doesn't pay attention to forced full unroll
   // attributes so we need to make sure and allow the full unroll pass to pay
   // attention to it.
   if (Phase != ThinOrFullLTOPhase::ThinLTOPreLink || !PGOOpt ||
       PGOOpt->Action != PGOOptions::SampleUse)
PassBuilder::buildFunctionSimplificationPipeline(OptimizationLevel Level,
   // Disable header duplication in loop rotation at -Oz.
   LPM1.addPass(
       LoopRotatePass(Level != OptimizationLevel::Oz, isLTOPreLink(Phase)));
   // TODO: Investigate promotion cap for O1.
   LPM1.addPass(LICMPass(PTO.LicmMssaOptCap, PTO.LicmMssaNoAccForPromotionCap));
   LPM1.addPass(
       SimpleLoopUnswitchPass(/* NonTrivial */ Level == OptimizationLevel::O3 &&
                              EnableO3NonTrivialUnswitching));
+  if (EnableLoopFlatten)
+    LPM1.addPass(LoopFlattenPass());
+
   LPM2.addPass(LoopIdiomRecognizePass());
   LPM2.addPass(IndVarSimplifyPass());
 
   for (auto &C : LateLoopOptimizationsEPCallbacks)
     C(LPM2, Level);
 
   LPM2.addPass(LoopDeletionPass());
 
   if (EnableLoopInterchange)
     LPM2.addPass(LoopInterchangePass());
-  if (EnableLoopFlatten)
-    LPM2.addPass(LoopFlattenPass());
 
   // Do not enable unrolling in PreLinkThinLTO phase during sample PGO
   // because it changes IR to makes profile annotation in back compile
   // inaccurate. The normal unroller doesn't pay attention to forced full unroll
   // attributes so we need to make sure and allow the full unroll pass to pay
   // attention to it.
   if (Phase != ThinOrFullLTOPhase::ThinLTOPreLink || !PGOOpt ||
       PGOOpt->Action != PGOOptions::SampleUse)
PassBuilder::buildLTODefaultPipeline(OptimizationLevel Level,
   MainFPM.addPass(DSEPass());
   MainFPM.addPass(MergedLoadStoreMotionPass());
 
 
   if (EnableConstraintElimination)
     MainFPM.addPass(ConstraintEliminationPass());
 
   LoopPassManager LPM;
-  LPM.addPass(IndVarSimplifyPass());
-  LPM.addPass(LoopDeletionPass());
   if (EnableLoopFlatten && Level.getSpeedupLevel() > 1)
     LPM.addPass(LoopFlattenPass());
+  LPM.addPass(IndVarSimplifyPass());
+  LPM.addPass(LoopDeletionPass());
   // FIXME: Add loop interchange.
 
   // Unroll small loops and perform peeling.
   LPM.addPass(LoopFullUnrollPass(Level.getSpeedupLevel(),
                                  /* OnlyWhenForced= */ !PTO.LoopUnrolling,
                                  PTO.ForgetAllSCEVInLoopUnroll));
   // The loop passes in LPM (LoopFullUnrollPass) do not preserve MemorySSA.
   // *All* loop passes must preserve it, in order to be able to use it.

llvm/test/Transforms/PhaseOrdering/AArch64/loopflatten.ll

+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
+; RUN: opt -passes='default<O3>' -enable-loop-flatten -loop-flatten-cost-threshold=3 -S %s | FileCheck %s
+
+target datalayout = "e-m:e-i8:8:32-i16:16:32-i64:64-i128:128-n32:64-S128"
+target triple = "aarch64"
+
+define dso_local void @_Z3fooPiii(i32* %A, i32 %N, i32 %M) #0 {
+; CHECK-LABEL: @_Z3fooPiii(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[CMP3:%.*]] = icmp sgt i32 [[N:%.*]], 0
+; CHECK-NEXT:    [[CMP21:%.*]] = icmp sgt i32 [[M:%.*]], 0
+; CHECK-NEXT:    [[OR_COND:%.*]] = select i1 [[CMP3]], i1 [[CMP21]], i1 false
+; CHECK-NEXT:    br i1 [[OR_COND]], label [[FOR_COND1_PREHEADER_LR_PH_SPLIT_US:%.*]], label [[FOR_COND_CLEANUP:%.*]]
+; CHECK:       for.cond1.preheader.lr.ph.split.us:
+; CHECK-NEXT:    [[TMP0:%.*]] = zext i32 [[M]] to i64
+; CHECK-NEXT:    [[TMP1:%.*]] = zext i32 [[N]] to i64
+; CHECK-NEXT:    [[FLATTEN_TRIPCOUNT:%.*]] = mul nuw nsw i64 [[TMP0]], [[TMP1]]
+; CHECK-NEXT:    br label [[FOR_COND1_PREHEADER_US:%.*]]
+; CHECK:       for.cond1.preheader.us:
+; CHECK-NEXT:    [[INDVAR6:%.*]] = phi i64 [ [[INDVAR_NEXT7:%.*]], [[FOR_COND1_PREHEADER_US]] ], [ 0, [[FOR_COND1_PREHEADER_LR_PH_SPLIT_US]] ]
+; CHECK-NEXT:    [[ARRAYIDX_US:%.*]] = getelementptr inbounds i32, i32* [[A:%.*]], i64 [[INDVAR6]]
+; CHECK-NEXT:    [[TMP2:%.*]] = load i32, i32* [[ARRAYIDX_US]], align 4
+; CHECK-NEXT:    tail call void @_Z1fi(i32 [[TMP2]])
+; CHECK-NEXT:    [[INDVAR_NEXT7]] = add nuw nsw i64 [[INDVAR6]], 1
+; CHECK-NEXT:    [[EXITCOND_NOT:%.*]] = icmp eq i64 [[INDVAR_NEXT7]], [[FLATTEN_TRIPCOUNT]]
+; CHECK-NEXT:    br i1 [[EXITCOND_NOT]], label [[FOR_COND_CLEANUP]], label [[FOR_COND1_PREHEADER_US]]
+; CHECK:       for.cond.cleanup:
+; CHECK-NEXT:    ret void
+;
+entry:
+  br label %for.cond
+
+for.cond:
+  %i.0 = phi i32 [ 0, %entry ], [ %inc6, %for.cond.cleanup3 ]
+  %cmp = icmp slt i32 %i.0, %N
+  br i1 %cmp, label %for.body, label %for.cond.cleanup
+
+for.cond.cleanup:
+  ret void
+
+for.body:
+  br label %for.cond1
+
+for.cond1:
+  %j.0 = phi i32 [ 0, %for.body ], [ %inc, %for.body4 ]
+  %cmp2 = icmp slt i32 %j.0, %M
+  br i1 %cmp2, label %for.body4, label %for.cond.cleanup3
+
+for.cond.cleanup3:
+  %inc6 = add nsw i32 %i.0, 1
+  br label %for.cond
+
+for.body4:
+  %mul = mul nsw i32 %i.0, %M
+  %add = add nsw i32 %mul, %j.0
+  %idxprom = sext i32 %add to i64
+  %arrayidx = getelementptr inbounds i32, i32* %A, i64 %idxprom
+  %0 = load i32, i32* %arrayidx, align 4
+  call void @_Z1fi(i32 %0)
+  %inc = add nsw i32 %j.0, 1
+  br label %for.cond1
+}
+
+declare dso_local void @_Z1fi(i32) #2