[mlir][bufferization] Module bufferization: Delete obsolete code

[matthias-springer]

Delete equivalenceAnalysis, which has been incorporated into the getAliasingValues API. Also add an additional test case to ensure that equivalence is properly propagated across function boundaries.

[llvmbot]

@llvm/pr-subscribers-mlir-bufferization

@llvm/pr-subscribers-mlir

Author: Matthias Springer (matthias-springer)

Changes

Delete equivalenceAnalysis, which has been incorporated into the getAliasingValues API. Also add an additional test case to ensure that equivalence is properly propagated across function boundaries.

---

Full diff: https://github.com/llvm/llvm-project/pull/127455.diff

2 Files Affected:

• (modified) mlir/lib/Dialect/Bufferization/Transforms/OneShotModuleBufferize.cpp (-35)
• (modified) mlir/test/Dialect/Bufferization/Transforms/one-shot-module-bufferize-analysis.mlir (+19)

diff --git a/mlir/lib/Dialect/Bufferization/Transforms/OneShotModuleBufferize.cpp b/mlir/lib/Dialect/Bufferization/Transforms/OneShotModuleBufferize.cpp
index 71ea0fd9d43cd..77840690e6a26 100644
--- a/mlir/lib/Dialect/Bufferization/Transforms/OneShotModuleBufferize.cpp
+++ b/mlir/lib/Dialect/Bufferization/Transforms/OneShotModuleBufferize.cpp
@@ -289,35 +289,6 @@ static func::FuncOp getCalledFunction(func::CallOp callOp) {
       SymbolTable::lookupNearestSymbolFrom(callOp, sym));
 }
 
-/// Gather equivalence info of CallOps.
-/// Note: This only adds new equivalence info if the called function was already
-/// analyzed.
-// TODO: This does not handle cyclic function call graphs etc.
-static void equivalenceAnalysis(func::FuncOp funcOp,
-                                OneShotAnalysisState &state,
-                                FuncAnalysisState &funcState) {
-  funcOp->walk([&](func::CallOp callOp) {
-    func::FuncOp calledFunction = getCalledFunction(callOp);
-    assert(calledFunction && "could not retrieved called func::FuncOp");
-
-    // No equivalence info available for the called function.
-    if (!funcState.equivalentFuncArgs.count(calledFunction))
-      return WalkResult::skip();
-
-    for (auto it : funcState.equivalentFuncArgs[calledFunction]) {
-      int64_t returnIdx = it.first;
-      int64_t bbargIdx = it.second;
-      if (!state.isInPlace(callOp->getOpOperand(bbargIdx)))
-        continue;
-      Value returnVal = callOp.getResult(returnIdx);
-      Value argVal = callOp->getOperand(bbargIdx);
-      state.unionEquivalenceClasses(returnVal, argVal);
-    }
-
-    return WalkResult::advance();
-  });
-}
-
 /// Return "true" if the given function signature has tensor semantics.
 static bool hasTensorSignature(func::FuncOp funcOp) {
   return llvm::any_of(funcOp.getFunctionType().getInputs(),
@@ -493,9 +464,6 @@ mlir::bufferization::analyzeModuleOp(ModuleOp moduleOp,
     // Now analyzing function.
     funcState.startFunctionAnalysis(funcOp);
 
-    // Gather equivalence info for CallOps.
-    equivalenceAnalysis(funcOp, state, funcState);
-
     // Analyze funcOp.
     if (failed(analyzeOp(funcOp, state, statistics)))
       return failure();
@@ -514,9 +482,6 @@ mlir::bufferization::analyzeModuleOp(ModuleOp moduleOp,
     if (!state.getOptions().isOpAllowed(funcOp))
       continue;
 
-    // Gather equivalence info for CallOps.
-    equivalenceAnalysis(funcOp, state, funcState);
-
     // Analyze funcOp.
     if (failed(analyzeOp(funcOp, state, statistics)))
       return failure();
diff --git a/mlir/test/Dialect/Bufferization/Transforms/one-shot-module-bufferize-analysis.mlir b/mlir/test/Dialect/Bufferization/Transforms/one-shot-module-bufferize-analysis.mlir
index 2ca7f7109005c..c947407c63e74 100644
--- a/mlir/test/Dialect/Bufferization/Transforms/one-shot-module-bufferize-analysis.mlir
+++ b/mlir/test/Dialect/Bufferization/Transforms/one-shot-module-bufferize-analysis.mlir
@@ -1,4 +1,5 @@
 // RUN: mlir-opt %s -one-shot-bufferize="bufferize-function-boundaries test-analysis-only" -split-input-file | FileCheck %s
+// RUN: mlir-opt %s -one-shot-bufferize="bufferize-function-boundaries test-analysis-only dump-alias-sets" -split-input-file | FileCheck %s --check-prefix=CHECK-ALIAS
 
 // Run fuzzer with different seeds.
 // RUN: mlir-opt %s -one-shot-bufferize="bufferize-function-boundaries test-analysis-only analysis-heuristic=fuzzer analysis-fuzzer-seed=23" -split-input-file -o /dev/null
@@ -1406,3 +1407,21 @@ func.func @caller(%c: i1, %t0: tensor<5xf32>, %t1: tensor<5xf32>, %t2: tensor<5x
   return %r : tensor<5xf32>
 }
 
+// -----
+
+// CHECK-ALIAS-LABEL: func @foo
+func.func @foo(%arg0: tensor<?xf32>) -> tensor<?xf32> {
+  // CHECK-ALIAS: return
+  // CHECK-ALIAS-SAME: __equivalent_func_args__ = [0]
+  return %arg0 : tensor<?xf32>
+}
+
+// CHECK-ALIAS: func @bar(%[[arg0:.*]]: tensor<?xf32>
+func.func @bar(%arg0: tensor<?xf32>) -> tensor<?xf32> {
+  // CHECK-ALIAS: %[[call:.*]] = call @foo(%[[arg0]])
+  // CHECK-ALIAS-SAME: {__inplace_operands_attr__ = ["true"], __opresult_alias_set_attr__ = [{{\[}}"%[[call]]", "%[[arg0]]"]]}
+  %x = call @foo(%arg0) : (tensor<?xf32>) -> tensor<?xf32>
+  // CHECK-ALIAS: return
+  // CHECK-ALIAS-SAME: __equivalent_func_args__ = [0]
+  return %x : tensor<?xf32>
+}

[christopherbate]

LGTM, thanks!