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

[via llvm-commits]

https://github.com/Chengjunp created https://github.com/llvm/llvm-project/pull/128117

In current `FlattenCFG`, the alias check between two instructions is imprecise. Specifically, `AA->isNoAlias` should not be called with two instructions directly. When passing a store instruction and a load instruction directly into `AA->isNoAlias` always returns `NoAlia`s. 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`.

In this patch, we take the `MemoryLocation`s from the instructions and then do the alias check on these two `MemoryLocation`s. Related tests are also added.

>From 2d987ba1b4f7e556eed4d04ef2138f1a5d6d53c1 Mon Sep 17 00:00:00 2001
From: chengjunp <chengjunp at nvidia.com>
Date: Fri, 21 Feb 2025 03:04:18 +0000
Subject: [PATCH 1/2] Fix a wrong use of AA in FlattenCFG

---
 llvm/lib/Transforms/Utils/FlattenCFG.cpp | 8 ++++++--
 1 file changed, 6 insertions(+), 2 deletions(-)

diff --git a/llvm/lib/Transforms/Utils/FlattenCFG.cpp b/llvm/lib/Transforms/Utils/FlattenCFG.cpp
index 16b4bb1981d8b..151ca06f3498a 100644
--- a/llvm/lib/Transforms/Utils/FlattenCFG.cpp
+++ b/llvm/lib/Transforms/Utils/FlattenCFG.cpp
@@ -357,8 +357,12 @@ bool FlattenCFGOpt::CompareIfRegionBlock(BasicBlock *Block1, BasicBlock *Block2,
       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))
-            return false;
+          if (AA) {
+            MemoryLocation Loc1 = MemoryLocation::get(&*iter1);
+            MemoryLocation Loc2 = MemoryLocation::get(&*BI);
+            if (!AA->isNoAlias(Loc1, Loc2))
+              return false;
+          }
         }
       }
     }

>From 5cd91b8f94d822b1c2af389d006f3e721b1f177a Mon Sep 17 00:00:00 2001
From: chengjunp <chengjunp at nvidia.com>
Date: Fri, 21 Feb 2025 03:22:11 +0000
Subject: [PATCH 2/2] Update test

---
 llvm/test/Transforms/Util/flatten-cfg.ll | 88 +++++++++++++++++++++++-
 1 file changed, 87 insertions(+), 1 deletion(-)

diff --git a/llvm/test/Transforms/Util/flatten-cfg.ll b/llvm/test/Transforms/Util/flatten-cfg.ll
index 038dcaa47419a..bb592a7031031 100644
--- a/llvm/test/Transforms/Util/flatten-cfg.ll
+++ b/llvm/test/Transforms/Util/flatten-cfg.ll
@@ -1,5 +1,5 @@
 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 2
-; RUN: opt -passes=flatten-cfg -S < %s | FileCheck %s
+; RUN: opt -passes='require<aa>,flatten-cfg' -S < %s | FileCheck %s
 
 
 ; This test checks whether the pass completes without a crash.
@@ -309,3 +309,89 @@ if.then.y:
 exit:
   ret i1 %cmp.y
 }
+
+; 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 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:    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
+  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
+  br label %if.end2
+
+if.end2:
+  ret i32 0
+}
\ No newline at end of file