[FlattenCFG] Fix an Imprecise Usage of AA (#128117)

In current `FlattenCFG`, using `isNoAlias` for two instructions is
imprecise. For example, when passing a store instruction and a load
instruction directly into `AA->isNoAlias`, it will always return
`NoAlias`. This happens because when checking the types of the two
Values, the store instruction (which has a `void` type) causes the
analysis to return `NoAlias`.

For instructions, we should use `getModRefInfo` instead of `isNoAlias`,
as aliasing is a concept of memory locations.

In this patch, `AAResults::getModRefInfo` is supported to take in two
instructions. It will check whether two instructions may access the same
memory location or not. And in `FlattenCFG`, we use this new helper
function to do the check instead of `isNoAlias`.

Unit tests and lit tests are also included to this patch.

Author: Chengjunp

llvm/include/llvm/Analysis/AliasAnalysis.h

@@ -521,6 +521,10 @@ class AAResults {
   /// the same memory locations.
   ModRefInfo getModRefInfo(const Instruction *I, const CallBase *Call);
 
+  /// Return information about whether two instructions may refer to the same
+  /// memory locations.
+  ModRefInfo getModRefInfo(const Instruction *I1, const Instruction *I2);
+
   /// Return information about whether a particular call site modifies
   /// or reads the specified memory location \p MemLoc before instruction \p I
   /// in a BasicBlock.
@@ -600,6 +604,8 @@ class AAResults {
   ModRefInfo getModRefInfo(const Instruction *I,
                            const std::optional<MemoryLocation> &OptLoc,
                            AAQueryInfo &AAQIP);
+  ModRefInfo getModRefInfo(const Instruction *I1, const Instruction *I2,
+                           AAQueryInfo &AAQI);
   ModRefInfo callCapturesBefore(const Instruction *I,
                                 const MemoryLocation &MemLoc, DominatorTree *DT,
                                 AAQueryInfo &AAQIP);

llvm/lib/Analysis/AliasAnalysis.cpp

@@ -336,6 +336,26 @@ ModRefInfo AAResults::getModRefInfo(const CallBase *Call1,
   return Result;
 }
 
+ModRefInfo AAResults::getModRefInfo(const Instruction *I1,
+                                    const Instruction *I2) {
+  SimpleAAQueryInfo AAQIP(*this);
+  return getModRefInfo(I1, I2, AAQIP);
+}
+
+ModRefInfo AAResults::getModRefInfo(const Instruction *I1,
+                                    const Instruction *I2, AAQueryInfo &AAQI) {
+  // Early-exit if either instruction does not read or write memory.
+  if (!I1->mayReadOrWriteMemory() || !I2->mayReadOrWriteMemory())
+    return ModRefInfo::NoModRef;
+
+  if (const auto *Call2 = dyn_cast<CallBase>(I2))
+    return getModRefInfo(I1, Call2, AAQI);
+
+  // FIXME: We can have a more precise result.
+  ModRefInfo MR = getModRefInfo(I1, MemoryLocation::getOrNone(I2), AAQI);
+  return isModOrRefSet(MR) ? ModRefInfo::ModRef : ModRefInfo::NoModRef;
+}
+
 MemoryEffects AAResults::getMemoryEffects(const CallBase *Call,
                                           AAQueryInfo &AAQI) {
   MemoryEffects Result = MemoryEffects::unknown();

llvm/lib/Transforms/Utils/FlattenCFG.cpp

@@ -356,8 +356,8 @@ bool FlattenCFGOpt::CompareIfRegionBlock(BasicBlock *Block1, BasicBlock *Block2,
     if (iter1->mayWriteToMemory()) {
       for (BasicBlock::iterator BI(PBI2), BE(PTI2); BI != BE; ++BI) {
         if (BI->mayReadFromMemory() || BI->mayWriteToMemory()) {
-          // Check alias with Head2.
-          if (!AA || !AA->isNoAlias(&*iter1, &*BI))
+          // Check whether iter1 and BI may access the same memory location.
+          if (!AA || AA->getModRefInfo(&*iter1, &*BI) != ModRefInfo::NoModRef)
             return false;
         }
       }

llvm/test/Transforms/Util/flatten-cfg.ll

@@ -309,3 +309,143 @@ if.then.y:
 exit:
   ret i1 %cmp.y
 }
+
+declare void @foo()
+declare void @bar() #1
+
+; Test that two if-regions are not merged when there's potential aliasing
+; between a store in the first if-region and a load in the second if-region's header
+define i32 @test_alias(i32 %a, i32 %b, ptr %p1, ptr %p2) {
+; CHECK-LABEL: define i32 @test_alias
+; CHECK-SAME: (i32 [[A:%.*]], i32 [[B:%.*]], ptr [[P1:%.*]], ptr [[P2:%.*]]) {
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    store i32 42, ptr [[P1]], align 4
+; CHECK-NEXT:    [[COND1:%.*]] = icmp eq i32 [[A]], 0
+; CHECK-NEXT:    br i1 [[COND1]], label [[IF_THEN1:%.*]], label [[IF_END1:%.*]]
+; CHECK:       if.then1:
+; CHECK-NEXT:    store i32 100, ptr [[P2]], align 4
+; CHECK-NEXT:    br label [[IF_END1]]
+; CHECK:       if.end1:
+; CHECK-NEXT:    [[VAL:%.*]] = load i32, ptr [[P1]], align 4
+; CHECK-NEXT:    [[COND2:%.*]] = icmp eq i32 [[B]], 0
+; CHECK-NEXT:    br i1 [[COND2]], label [[IF_THEN2:%.*]], label [[IF_END2:%.*]]
+; CHECK:       if.then2:
+; CHECK-NEXT:    store i32 100, ptr [[P2]], align 4
+; CHECK-NEXT:    br label [[IF_END2]]
+; CHECK:       if.end2:
+; CHECK-NEXT:    ret i32 0
+;
+entry:
+  store i32 42, ptr %p1
+  %cond1 = icmp eq i32 %a, 0
+  br i1 %cond1, label %if.then1, label %if.end1
+
+if.then1:
+  store i32 100, ptr %p2  ; May alias with the load below
+  br label %if.end1
+
+if.end1:
+  %val = load i32, ptr %p1  ; This load prevents merging due to potential alias
+  %cond2 = icmp eq i32 %b, 0
+  br i1 %cond2, label %if.then2, label %if.end2
+
+if.then2:
+  store i32 100, ptr %p2
+  br label %if.end2
+
+if.end2:
+  ret i32 0
+}
+
+; Test that two if-regions are not merged when there's potential aliasing
+; between a store in the first if-region and a function call in the second if-region's header
+define i32 @test_alias_2(i32 %a, i32 %b) {
+; CHECK-LABEL: define i32 @test_alias_2
+; CHECK-SAME: (i32 [[A:%.*]], i32 [[B:%.*]]) {
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[P:%.*]] = alloca i32, align 4
+; CHECK-NEXT:    store i32 42, ptr [[P]], align 4
+; CHECK-NEXT:    [[COND1:%.*]] = icmp eq i32 [[A]], 0
+; CHECK-NEXT:    br i1 [[COND1]], label [[IF_THEN1:%.*]], label [[IF_END1:%.*]]
+; CHECK:       if.then1:
+; CHECK-NEXT:    store i32 100, ptr @g, align 4
+; CHECK-NEXT:    br label [[IF_END1]]
+; CHECK:       if.end1:
+; CHECK-NEXT:    call void @foo()
+; CHECK-NEXT:    [[COND2:%.*]] = icmp eq i32 [[B]], 0
+; CHECK-NEXT:    br i1 [[COND2]], label [[IF_THEN2:%.*]], label [[IF_END2:%.*]]
+; CHECK:       if.then2:
+; CHECK-NEXT:    store i32 100, ptr @g, align 4
+; CHECK-NEXT:    br label [[IF_END2]]
+; CHECK:       if.end2:
+; CHECK-NEXT:    ret i32 0
+;
+entry:
+  %p = alloca i32
+  store i32 42, ptr %p
+  %cond1 = icmp eq i32 %a, 0
+  br i1 %cond1, label %if.then1, label %if.end1
+
+if.then1:
+  store i32 100, ptr @g
+  br label %if.end1
+
+if.end1:
+  call void @foo()
+  %cond2 = icmp eq i32 %b, 0
+  br i1 %cond2, label %if.then2, label %if.end2
+
+if.then2:
+  store i32 100, ptr @g
+  br label %if.end2
+
+if.end2:
+  ret i32 0
+}
+
+; Test that two if-regions are merged when there's no potential aliasing
+; between a store in the first if-region and a load in the second if-region's header
+define i32 @test_no_alias(i32 %a, i32 %b) {
+; CHECK-LABEL: define i32 @test_no_alias
+; CHECK-SAME: (i32 [[A:%.*]], i32 [[B:%.*]]) {
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[P:%.*]] = alloca i32, align 4
+; CHECK-NEXT:    store i32 42, ptr [[P]], align 4
+; CHECK-NEXT:    [[COND1:%.*]] = icmp eq i32 [[A]], 0
+; CHECK-NEXT:    [[VAL:%.*]] = load i32, ptr @g, align 4
+; CHECK-NEXT:    [[COND2:%.*]] = icmp eq i32 [[B]], 0
+; CHECK-NEXT:    [[TMP0:%.*]] = or i1 [[COND1]], [[COND2]]
+; CHECK-NEXT:    br i1 [[TMP0]], label [[IF_THEN2:%.*]], label [[IF_END2:%.*]]
+; CHECK:       if.then2:
+; CHECK-NEXT:    store i32 100, ptr [[P]], align 4
+; CHECK-NEXT:    call void @bar()
+; CHECK-NEXT:    br label [[IF_END2]]
+; CHECK:       if.end2:
+; CHECK-NEXT:    ret i32 0
+;
+entry:
+  %p = alloca i32
+  store i32 42, ptr %p
+  %cond1 = icmp eq i32 %a, 0
+  br i1 %cond1, label %if.then1, label %if.end1
+
+if.then1:
+  store i32 100, ptr %p  ; No alias with the load below
+  call void @bar() ; No alias with the load below since it's a pure function
+  br label %if.end1
+
+if.end1:
+  %val = load i32, ptr @g
+  %cond2 = icmp eq i32 %b, 0
+  br i1 %cond2, label %if.then2, label %if.end2
+
+if.then2:
+  store i32 100, ptr %p
+  call void @bar()
+  br label %if.end2
+
+if.end2:
+  ret i32 0
+}
+
+attributes #1 = { readnone willreturn nounwind }

llvm/unittests/Analysis/AliasAnalysisTest.cpp

@@ -186,6 +186,14 @@ TEST_F(AliasAnalysisTest, getModRefInfo) {
       AtomicRMWInst::Xchg, Addr, ConstantInt::get(IntType, 1), Alignment,
       AtomicOrdering::Monotonic, SyncScope::System, BB);
 
+  FunctionType *FooBarTy = FunctionType::get(Type::getVoidTy(C), {}, false);
+  Function::Create(FooBarTy, Function::ExternalLinkage, "foo", &M);
+  auto *BarF = Function::Create(FooBarTy, Function::ExternalLinkage, "bar", &M);
+  BarF->setDoesNotAccessMemory();
+
+  const Instruction *Foo = CallInst::Create(M.getFunction("foo"), {}, BB);
+  const Instruction *Bar = CallInst::Create(M.getFunction("bar"), {}, BB);
+
   ReturnInst::Create(C, nullptr, BB);
 
   auto &AA = getAAResults(*F);
@@ -203,6 +211,13 @@ TEST_F(AliasAnalysisTest, getModRefInfo) {
   EXPECT_EQ(AA.getModRefInfo(CmpXChg1, std::nullopt), ModRefInfo::ModRef);
   EXPECT_EQ(AA.getModRefInfo(AtomicRMW, MemoryLocation()), ModRefInfo::ModRef);
   EXPECT_EQ(AA.getModRefInfo(AtomicRMW, std::nullopt), ModRefInfo::ModRef);
+  EXPECT_EQ(AA.getModRefInfo(Store1, Load1), ModRefInfo::ModRef);
+  EXPECT_EQ(AA.getModRefInfo(Store1, Store1), ModRefInfo::ModRef);
+  EXPECT_EQ(AA.getModRefInfo(Store1, Add1), ModRefInfo::NoModRef);
+  EXPECT_EQ(AA.getModRefInfo(Store1, Foo), ModRefInfo::ModRef);
+  EXPECT_EQ(AA.getModRefInfo(Store1, Bar), ModRefInfo::NoModRef);
+  EXPECT_EQ(AA.getModRefInfo(Foo, Bar), ModRefInfo::NoModRef);
+  EXPECT_EQ(AA.getModRefInfo(Foo, Foo), ModRefInfo::ModRef);
 }
 
 static Instruction *getInstructionByName(Function &F, StringRef Name) {