Bugfix for application of trivial loop optimization in LoopIdiomRecognize(Github Issue #50308)

[razetime]

Summary

Github issue: Trivial memset optimization not applied to loops under -Oz (LoopIdiomRecognize)

When opt -loop-idiom is run on an LLVM IR file containing a countable unrotated loop containing stores of a simple constant to a single array, a memset optimization is not applied under -Oz or -loop-idiom. The optimization is not applied since runOnLoopBlock stops before the memset optimization is applied on the loop.

runOnLoopBlock stops before the optimization because it checks whether each basic block dominates all exit blocks in the loop. For unrotated loops, this check causes an early return. To work around this, a boolean representing the rotation form of the loop is passed to runOnLoopBlock. Based on the given boolean, the check for domination of the exit blocks is done. Otherwise, the check is skipped and the optimization is applied directly.

[lebedev.ri]

llvm/lib/Transforms/Scalar/LoopIdiomRecognize.cpp

628–636

I do not understand why we can just not do this check for unrotated loops?

[razetime]

llvm/lib/Transforms/Scalar/LoopIdiomRecognize.cpp

628–636

The check unconditionally failed on the unrotated test case I was using from https://llvm.godbolt.org/z/YrGqGhnE9.

The place where it failed was when BB =loop.latch and ExitBlock = exit. pre-rotating the loop would work, but it seemed unintuitive and costlier than skipping the check.

[ChuanqiXu]

llvm/lib/Transforms/Scalar/LoopIdiomRecognize.cpp

628–636

It might not be a good reason to remove a check simply if it fails on a false negative condition. Since the condition might do many other true negative checks.

llvm/test/Transforms/LoopIdiom/memset-unrotated-loop.ll

7–11

This transformation looks not good. Given %start is not no alias and %start == %end, then the original behavior would be no-op. However, it would set value 1 at %start in the transformed code. See https://llvm.godbolt.org/z/a9czn4z4b, which contains the original semantics.

[razetime]

llvm/lib/Transforms/Scalar/LoopIdiomRecognize.cpp

628–636

This is the dominator tree for the given IR file:

The problem here is that loop.latch cannot dominate exit since they're on the same level i.e. this check returns false:

From llvm/include/llvm/Support/GenericDomTree.h

// A can only dominate B if it is higher in the tree.
if (A->getLevel() >= B->getLevel()) return false;

I'm not sure how this check can be modified to accept the IR for optimization. Maybe a completely different check should be added?

[razetime]

llvm/lib/Transforms/Scalar/LoopIdiomRecognize.cpp

628–636

Since the IR optimization is incorrect, is pre-rotating the loop a valid strategy here?

[ChuanqiXu]

llvm/lib/Transforms/Scalar/LoopIdiomRecognize.cpp

628–636

Yeah, when we made a change, not only we need to consider how could we make the false negative case to pass, we need to consider if the change would affect other cases, too. You are talking about the first part and we are talking about the second part. For the concrete suggestion, maybe @fhahn could help?

[fhahn]

llvm/lib/Transforms/Scalar/LoopIdiomRecognize.cpp

628–636

So to support the loop-not-rotated-case here, we could check if it can trivially rotated (header is the only exiting block, all instructions in header can moved/duplicated).

Then proceed with the analysis. If we can replace the loop, we need to rotate the loop before we perform a transform. To start with, we should only rotate if the loop will go away entirely, i.e. it only contains memory operations that will be removed.

[arsenm]

I'm not the right person for loop transforms

[lebedev.ri]

This review may be stuck/dead, consider abandoning if no longer relevant.
Removing myself as reviewer in attempt to clean dashboard.

llvm/lib/Transforms/Scalar/LoopIdiomRecognize.cpp

enum LegalStoreKind {
     DontUse // Dummy retval never to be used. Allows catching errors in retval
             // handling.
   };
 
   /// \name Countable Loop Idiom Handling
   /// @{
 
   bool runOnCountableLoop();
-  bool runOnLoopBlock(BasicBlock *BB, const SCEV *BECount,
+  bool runOnLoopBlock(bool isFromRotatedLoop, BasicBlock *BB,
+                      const SCEV *BECount,
                       SmallVectorImpl<BasicBlock *> &ExitBlocks);
 
   void collectStores(BasicBlock *BB);
   LegalStoreKind isLegalStore(StoreInst *SI);
   enum class ForMemset { No, Yes };
   bool processLoopStores(SmallVectorImpl<StoreInst *> &SL, const SCEV *BECount,
                          ForMemset For);
 
bool LoopIdiomRecognize::runOnCountableLoop() {
   bool MadeChange = false;
 
   // Scan all the blocks in the loop that are not in subloops.
   for (auto *BB : CurLoop->getBlocks()) {
     // Ignore blocks in subloops.
     if (LI->getLoopFor(BB) != CurLoop)
       continue;
 
-    MadeChange |= runOnLoopBlock(BB, BECount, ExitBlocks);
+    MadeChange |=
+        runOnLoopBlock(CurLoop->isRotatedForm(), BB, BECount, ExitBlocks);
   }
   return MadeChange;
 }
 
 static APInt getStoreStride(const SCEVAddRecExpr *StoreEv) {
   const SCEVConstant *ConstStride = cast<SCEVConstant>(StoreEv->getOperand(1));
   return ConstStride->getAPInt();
 }
for (Instruction &I : *BB) {
     }
   }
 }
 
 /// runOnLoopBlock - Process the specified block, which lives in a counted loop
 /// with the specified backedge count.  This block is known to be in the current
 /// loop and not in any subloops.
 bool LoopIdiomRecognize::runOnLoopBlock(
-    BasicBlock *BB, const SCEV *BECount,
+    bool isFromRotatedLoop, BasicBlock *BB, const SCEV *BECount,
     SmallVectorImpl<BasicBlock *> &ExitBlocks) {
+
   // We can only promote stores in this block if they are unconditionally
-  // executed in the loop.  For a block to be unconditionally executed, it has
-  // to dominate all the exit blocks of the loop.  Verify this now.
+  // executed in the loop. For a block to be unconditionally executed, it has
+  // to dominate all the exit blocks of the loop. Verify this now.
+  // On countable unrotated loops, checking the exit blocks fails.
+  // Hence, this is only run on rotated loops.
+  if (isFromRotatedLoop)
     for (BasicBlock *ExitBlock : ExitBlocks)
       if (!DT->dominates(BB, ExitBlock))
         return false;
 
   bool MadeChange = false;
   // Look for store instructions, which may be optimized to memset/memcpy.
   collectStores(BB);
 
   // Look for a single store or sets of stores with a common base, which can be
   // optimized into a memset (memset_pattern).  The latter most commonly happens
   // with structs and handunrolled loops.

llvm/test/Transforms/LoopIdiom/memset-unrotated-loop.ll

+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
+; RUN: opt -loop-idiom -S %s | FileCheck %s
+
+define void @test(i8* noalias nonnull align 1 %start, i8* %end) unnamed_addr {
+; CHECK-LABEL: @test(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[START2:%.*]] = ptrtoint i8* [[START:%.*]] to i64
+; CHECK-NEXT:    [[END1:%.*]] = ptrtoint i8* [[END:%.*]] to i64
+; CHECK-NEXT:    [[TMP0:%.*]] = add i64 [[END1]], 1
+; CHECK-NEXT:    [[TMP1:%.*]] = sub i64 [[TMP0]], [[START2]]
+; CHECK-NEXT:    call void @llvm.memset.p0i8.i64(i8* align 1 [[START]], i8 1, i64 [[TMP1]], i1 false)
+; CHECK-NEXT:    br label [[LOOP_HEADER:%.*]]
+; CHECK:       loop.header:
+; CHECK-NEXT:    [[PTR_IV:%.*]] = phi i8* [ [[START]], [[ENTRY:%.*]] ], [ [[PTR_IV_NEXT:%.*]], [[LOOP_LATCH:%.*]] ]
+; CHECK-NEXT:    [[_12_I:%.*]] = icmp eq i8* [[PTR_IV]], [[END]]
+; CHECK-NEXT:    br i1 [[_12_I]], label [[EXIT:%.*]], label [[LOOP_LATCH]]
+; CHECK:       loop.latch:
+; CHECK-NEXT:    [[PTR_IV_NEXT]] = getelementptr inbounds i8, i8* [[PTR_IV]], i64 1
+; CHECK-NEXT:    br label [[LOOP_HEADER]]
+; CHECK:       exit:
+; CHECK-NEXT:    ret void
+;
+entry:
+  br label %loop.header
+
+loop.header:
+  %ptr.iv = phi i8* [ %start, %entry ], [ %ptr.iv.next, %loop.latch ]
+  %_12.i = icmp eq i8* %ptr.iv, %end
+  br i1 %_12.i, label %exit, label %loop.latch
+
+loop.latch:
+  %ptr.iv.next = getelementptr inbounds i8, i8* %ptr.iv, i64 1
+  store i8 1, i8* %ptr.iv, align 1
+  br label %loop.header
+
+exit:
+  ret void
+}