[Mlir-commits] [mlir] [mlir][analysis] Enhance the analysis precision of local alias analysis. (PR #95922)

[donald chen]

https://github.com/cxy-1993 created https://github.com/llvm/llvm-project/pull/95922

Local alias analysis utilizes view interfaces, branch interfaces, and region branch interfaces to identify underlying values. However, it overlooks operations like arith.select.

This patch enhances the underlying value identification process by introducing support for buffer view interfaces, leading to more accurate alias results.

>From da1495849a257016b5310995a1f5f9d642b612c6 Mon Sep 17 00:00:00 2001
From: chenxunyu <chenxunyu at cambricon.com>
Date: Tue, 18 Jun 2024 13:18:27 +0000
Subject: [PATCH] [mlir][analysis] Enhance the analysis precision of local
 alias analysis.

Local alias analysis utilizes view interfaces, branch interfaces, and region
branch interfaces to identify underlying values. However, it overlooks
operations like arith.select.

This patch enhances the underlying value identification process by introducing
support for buffer view interfaces, leading to more accurate alias results.
---
 .../AliasAnalysis/LocalAliasAnalysis.cpp      | 21 +++++++++++++++++++
 mlir/test/Analysis/test-alias-analysis.mlir   | 18 ++++++++++++++++
 2 files changed, 39 insertions(+)

diff --git a/mlir/lib/Analysis/AliasAnalysis/LocalAliasAnalysis.cpp b/mlir/lib/Analysis/AliasAnalysis/LocalAliasAnalysis.cpp
index b22f3fdc9fb7a..bc8990afacaee 100644
--- a/mlir/lib/Analysis/AliasAnalysis/LocalAliasAnalysis.cpp
+++ b/mlir/lib/Analysis/AliasAnalysis/LocalAliasAnalysis.cpp
@@ -9,6 +9,7 @@
 #include "mlir/Analysis/AliasAnalysis/LocalAliasAnalysis.h"
 
 #include "mlir/Analysis/AliasAnalysis.h"
+#include "mlir/Dialect/Bufferization/IR/BufferViewFlowOpInterface.h"
 #include "mlir/IR/Attributes.h"
 #include "mlir/IR/Block.h"
 #include "mlir/IR/Matchers.h"
@@ -131,6 +132,26 @@ static void collectUnderlyingAddressValues(OpResult result, unsigned maxDepth,
   if (ViewLikeOpInterface view = dyn_cast<ViewLikeOpInterface>(op))
     return collectUnderlyingAddressValues(view.getViewSource(), maxDepth,
                                           visited, output);
+
+  // If this is a buffer view, unwarp to the operands viewed from.
+  if (mlir::bufferization::BufferViewFlowOpInterface bufferView =
+          dyn_cast<mlir::bufferization::BufferViewFlowOpInterface>(op)) {
+    SmallVector<Value> viewSources;
+    bufferView.populateDependencies(
+        [&](ValueRange operands, ValueRange results) {
+          if (results.size() == 1 && results[0] == result) {
+            viewSources.append(operands.begin(), operands.end());
+          }
+        });
+
+    if (!viewSources.empty()) {
+      for (Value viewSource : viewSources) {
+        collectUnderlyingAddressValues(viewSource, maxDepth, visited, output);
+      }
+      return;
+    }
+  }
+
   // Check to see if we can reason about the control flow of this op.
   if (auto branch = dyn_cast<RegionBranchOpInterface>(op)) {
     return collectUnderlyingAddressValues(branch, /*region=*/nullptr, result,
diff --git a/mlir/test/Analysis/test-alias-analysis.mlir b/mlir/test/Analysis/test-alias-analysis.mlir
index 8cbee61c78b45..459ede3bdcb53 100644
--- a/mlir/test/Analysis/test-alias-analysis.mlir
+++ b/mlir/test/Analysis/test-alias-analysis.mlir
@@ -256,3 +256,21 @@ func.func @constants(%arg: memref<2xf32>) attributes {test.ptr = "func"} {
 
   return
 }
+
+// -----
+
+// CHECK-LABEL: Testing : "buffer_view"
+// CHECK-DAG: alloc_1#0 <-> select#0: NoAlias
+// CHECK-DAG: alloca_2#0 <-> select#0: MayAlias
+// CHECK-DAG: alloca_3#0 <-> select#0: MayAlias
+// CHECK-DAG: select#0 <-> func.region0#0: NoAlias
+// CHECK-DAG: select#0 <-> func.region0#1: NoAlias
+// CHECK-DAG: select#0 <-> func.region0#2: NoAlias
+func.func @buffer_view(%arg: memref<2xf32>, %size: index, %cond : i1) attributes {test.ptr = "func"} {
+  %1 = memref.alloc() {test.ptr = "alloc_1"} : memref<8x64xf32>
+  %c0 = arith.constant 0 : index
+  %2 = memref.alloca (%size) {test.ptr = "alloca_2"} : memref<?xi8>
+  %3 = memref.alloca (%size) {test.ptr = "alloca_3"} : memref<?xi8>
+  %4 = arith.select %cond, %2, %3 {test.ptr = "select"}: memref<?xi8>
+  return
+}