@llvm/pr-subscribers-flang-openmp

Author: Chi-Chun, Chen (chichunchen)

diff --git a/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp b/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp
index 06026582538a2..200e6897fd6ab 100644
--- a/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp
+++ b/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp
@@ -6573,8 +6573,17 @@ static void removeUnusedBlocksFromParent(ArrayRef<BasicBlock *> BBs) {
     return false;
   };
 
-  while (BBsToErase.remove_if(HasRemainingUses)) {
-    // Try again if anything was removed.
+  // Compute the keep set before mutating BBsToErase; otherwise the result can
+  // depend on SmallPtrSet iteration order.
+  bool Changed = true;
+  while (Changed) {
+    Changed = false;
+    SmallVector<BasicBlock *, 4> BBsToKeep;
+    for (BasicBlock *BB : BBsToErase)
+      if (HasRemainingUses(BB))
+        BBsToKeep.push_back(BB);
+    for (BasicBlock *BB : BBsToKeep)
+      Changed |= BBsToErase.erase(BB);
   }
 
   SmallVector<BasicBlock *, 7> BBVec(BBsToErase.begin(), BBsToErase.end());
diff --git a/llvm/unittests/Frontend/OpenMPIRBuilderTest.cpp b/llvm/unittests/Frontend/OpenMPIRBuilderTest.cpp
index acfa993855b2e..a7e572bc427a3 100644
--- a/llvm/unittests/Frontend/OpenMPIRBuilderTest.cpp
+++ b/llvm/unittests/Frontend/OpenMPIRBuilderTest.cpp
@@ -8461,4 +8461,55 @@ TEST_F(OpenMPIRBuilderTest, ScopeDirectiveNowait) {
   }
 }
 
+TEST_F(OpenMPIRBuilderTest, FuseLoopRangeNoDeadBlocks) {
+  using InsertPointTy = OpenMPIRBuilder::InsertPointTy;
+  OpenMPIRBuilder OMPBuilder(*M);
+  OMPBuilder.initialize();
+  F->setName("func");
+
+  IRBuilder<> Builder(BB);
+  Type *LCTy = F->getArg(0)->getType();
+  Value *TripCount = F->getArg(0);
+
+  auto LoopBodyGenCB = [&](InsertPointTy CodeGenIP, Value *LC) {
+    Builder.restoreIP(CodeGenIP);
+    createPrintfCall(Builder, "%d\n", {LC});
+    return Error::success();
+  };
+
+  OpenMPIRBuilder::LocationDescription Loc1({Builder.saveIP(), DL});
+  ASSERT_EXPECTED_INIT(
+      CanonicalLoopInfo *, Loop1,
+      OMPBuilder.createCanonicalLoop(Loc1, LoopBodyGenCB, TripCount, "loop1"));
+  Builder.restoreIP(Loop1->getAfterIP());
+
+  Value *TripCount2 = Builder.CreateAdd(TripCount, ConstantInt::get(LCTy, 1));
+  OpenMPIRBuilder::LocationDescription Loc2({Builder.saveIP(), DL});
+  ASSERT_EXPECTED_INIT(
+      CanonicalLoopInfo *, Loop2,
+      OMPBuilder.createCanonicalLoop(Loc2, LoopBodyGenCB, TripCount2, "loop2"));
+  Builder.restoreIP(Loop2->getAfterIP());
+
+  Value *TripCount3 = Builder.CreateAdd(TripCount, ConstantInt::get(LCTy, 2));
+  OpenMPIRBuilder::LocationDescription Loc3({Builder.saveIP(), DL});
+  ASSERT_EXPECTED_INIT(
+      CanonicalLoopInfo *, Loop3,
+      OMPBuilder.createCanonicalLoop(Loc3, LoopBodyGenCB, TripCount3, "loop3"));
+  Builder.restoreIP(Loop3->getAfterIP());
+  Builder.CreateRetVoid();
+
+  CanonicalLoopInfo *Fused = OMPBuilder.fuseLoops(DL, {Loop1, Loop2});
+
+  OMPBuilder.finalize();
+  ASSERT_FALSE(verifyModule(*M, &errs()));
+
+  Fused->assertOK();
+  for (BasicBlock &Block : *F) {
+    if (&Block == &F->getEntryBlock())
+      continue;
+    EXPECT_TRUE(Block.hasNPredecessorsOrMore(1))
+        << "Dead block found: " << Block.getName().str();
+  }
+}
+
 } // namespace

Can you help me understand the problem?

The patch is based on a commit before #197637 where SmallPtrSet::remove_if has nothing that would make the data structure invalid between calls of the predicate. I could not make your test fail without your change in OpenMPIRBuilder.
The algorithm is supposed to be a fixpoint iteration: It should not matter which order BBs are iterated over; if a BB is still marked for deletion in one round due to different iteration order, it will be unmarked in the next.

Did the contract actually change,

Yes. I think this was an oversight, as the line is explicitly mentioned in one of the comments..

The explanation in the PR is wrong. The problem is that during the execution of SmallPtrSet::remove_if(), lookup in the set doesn't work correctly until the rehash at the end of remove_if.

Did the contract actually change,

Yes. I think this was an oversight, as the line is explicitly mentioned in one of the comments..

The explanation in the PR is wrong. The problem is that during the execution of SmallPtrSet::remove_if(), lookup in the set doesn't work correctly until the rehash at the end of remove_if.

That makes sense, thank you!

It seems like copilot may have helped find this (now) broken use of remove_if, but the stochastically reverse-engineered "analysis" of it is junk (go figure)

If the remove_if change is purely for performance can we just revert for now?

Edit: I'm just now reading your comment on the other change. I will comment there instead

Since the input is a small array of BasicBlocks (otherwise it probably makes more sense to just do a DFS/simple DCE?), we could also work in terms of indices into BBs and use SmallBitVectors:

/// Determine which blocks in \p BBs are reachable from outside and remove the
/// ones that are not reachable from the function.
static void removeUnusedBlocksFromParent(ArrayRef<BasicBlock *> BBs) {
  DenseMap<BasicBlock *, unsigned> BB2Index(BBs.size());
  for (const auto [I, V] : enumerate(BBs))
    BB2Index.insert({V, I});

  // Optimistically assume all BBs are dead. We iteratively remove
  // immediately-used blocks from this set until we reach a fixed point.
  SmallBitVector BBsToErase(BBs.size(), true);
  // Accumulates immediately-used blocks for one iteration.
  SmallBitVector BBsToKeep(BBs.size());
  do {
    BBsToKeep.reset();
    for (unsigned BBToEraseIndex : BBsToErase.set_bits()) {
      for (Use &U : BBs[BBToEraseIndex]->uses()) {
        auto *UseInst = dyn_cast<Instruction>(U.getUser());
        if (!UseInst)
          continue;
        BasicBlock *UseBB = UseInst->getParent();
        auto UseBBIndexI = BB2Index.find(UseBB);
        // If this use is from an external block, or from a block in BBs we have
        // already proved is transitively used by an external block, then
        // this block is also used.
        if (UseBBIndexI == BB2Index.end() ||
            !BBsToErase.test(UseBBIndexI->second)) {
          BBsToKeep.set(BBToEraseIndex);
          break;
        }
      }
    }
    BBsToErase &= ~BBsToKeep;
  } while (BBsToKeep.any());

  SmallVector<BasicBlock *> FinalBBsToErase;
  for (unsigned BBToEraseIndex : BBsToErase.set_bits())
    FinalBBsToErase.push_back(BBs[BBToEraseIndex]);
  DeleteDeadBlocks(FinalBBsToErase);
}

This looks like AI-overcomplexity. BBsToKeep is not even needed because it's finding a fixpoint, just reset that bit in BBsToErase directly. In constrast to SmallPtrSet you can actually do that while iterating over it. Also, building the DenseMap over a small list of an estimated 7 elements max is probably overkill.

Here someting simpler by just inverting BBsToErase to BBsToKeep:

static void removeUnusedBlocksFromParent(ArrayRef<BasicBlock *> BBs) {
  SmallPtrSet<BasicBlock *, 6> BBsToKeep(llvm::from_range, BBs);
  auto HasRemainingUses = [&BBsToKeep](BasicBlock *BB) {
    for (Use &U : BB->uses()) {
      auto *UseInst = dyn_cast<Instruction>(U.getUser());
      if (!UseInst)
        continue;
      if (!BBsToKeep.count(UseInst->getParent()))
        continue;
      return true;
    }
    return false;
  };

  while (true) {
    bool Changed = false;

    for (BasicBlock *BB : BBs) {
      if (BBsToKeep.count(BB) || !HasRemainingUses(BB)) 
        continue;

      BBsToKeep.insert(BB);
      Changed = true;
    }

    if (!Changed)
      break;
  }

  SmallVector<BasicBlock *, 7> BBVec;
  for (BasicBlock *BB : BBs) {
    if (!BBsToKeep.count(BB))
      BBVec.push_back(BB);
  }
  DeleteDeadBlocks(BBVec);
}

No AI used in this case, I'm just overcomplicating things. I was biased towards the likely LLM generated original version, though

I think these are just artifacts of you giving the gist of the approach quickly, but I don't think your code works as-is:

static void removeUnusedBlocksFromParent(ArrayRef<BasicBlock *> BBs) {
  SmallPtrSet<BasicBlock *, 6> BBsToKeep(llvm::from_range, BBs);

I assume this should start empty?

  auto HasRemainingUses = [&BBsToKeep](BasicBlock *BB) {
    for (Use &U : BB->uses()) {
      auto *UseInst = dyn_cast<Instruction>(U.getUser());
      if (!UseInst)
        continue;
      if (!BBsToKeep.count(UseInst->getParent()))
        continue;
      return true;

Assuming a block is in BBsToKeep iff we know it has an external use, this seems slightly wrong. Isn't it if (<UseBB is not in BBs> || BBsToKeep.count(UseBB)) return true;?

    }
    return false;
  };

  while (true) {
    bool Changed = false;

    for (BasicBlock *BB : BBs) {
      if (BBsToKeep.count(BB) || !HasRemainingUses(BB)) 
        continue;

This is the only use of HasRemainingUses anyway, and it isn't passed as a callback, so we probably don't need the lambda, it could just be inline. I guess the return from the lambda is nicer than assign-and-break? If the lambda remains then the Remaining in the name no longer fits.

      BBsToKeep.insert(BB);
      Changed = true;
    }

    if (!Changed)
      break;
  }

  SmallVector<BasicBlock *, 7> BBVec;
  for (BasicBlock *BB : BBs) {
    if (!BBsToKeep.count(BB))
      BBVec.push_back(BB);
  }
  DeleteDeadBlocks(BBVec);
}

I think this works, and is what you were suggesting?

Some other incidental changes I made are because the SmallPtrSet count is rounded up to a power of two anyway, and SmallVector generally doesn't benefit from a count. I also removed an extra set of braces.

static void removeUnusedBlocksFromParent(ArrayRef<BasicBlock *> BBs) {
  SmallPtrSet<BasicBlock *, 8> InternalBBs(from_range, BBs);
  SmallPtrSet<BasicBlock *, 8> BBsToKeep;

  while (true) {
    bool Changed = false;

    for (BasicBlock *BB : BBs) {
      if (BBsToKeep.contains(BB))
        continue;

      for (Use &U : BB->uses()) {
        auto *UseInst = dyn_cast<Instruction>(U.getUser());
        if (!UseInst)
          continue;
        BasicBlock *UseBB = UseInst->getParent();
        if (!InternalBBs.contains(UseBB) || BBsToKeep.contains(UseBB)) {
          BBsToKeep.insert(BB);
          Changed = true;
          break;
        }
      }
    }

    if (!Changed)
      break;
  }

  SmallVector<BasicBlock *> BBVec;
  for (BasicBlock *BB : BBs)
    if (!BBsToKeep.count(BB))
      BBVec.push_back(BB);

  DeleteDeadBlocks(BBVec);
}