[Mlir-commits] [mlir] [mlir][dataflow]Fix dense backward dataflow intraprocedural hook (PR #76865)

[llvmlistbot at llvm.org]

https://github.com/drblallo updated https://github.com/llvm/llvm-project/pull/76865

>From 915a045f40032031209085785092bbde053dedd1 Mon Sep 17 00:00:00 2001
From: massimo <mo.fioravanti at gmail.com>
Date: Mon, 25 Dec 2023 21:18:05 +0100
Subject: [PATCH] [MLIR] Fix dense dataflow intraprocedural hook

The dataflow analysis framework within MLIR allows to customize the
transfer function when a `call-like` operation is encuntered.

The check to see if the analysis was executed in intraprocedural mode
was executed after the check to see if the callee had the
CallableOpInterface, and thus intraprocedural analyses would behave as
interpocedural ones when performing indirect calls.

This commit fixes the issue by performing the check for
intraprocedurality first.
---
 mlir/lib/Analysis/DataFlow/DenseAnalysis.cpp   | 16 +++++++++++-----
 .../Analysis/DataFlow/test-next-access.mlir    | 18 ++++++++++++++++++
 2 files changed, 29 insertions(+), 5 deletions(-)

diff --git a/mlir/lib/Analysis/DataFlow/DenseAnalysis.cpp b/mlir/lib/Analysis/DataFlow/DenseAnalysis.cpp
index 08d89d6db788c8..d4b9134ab9eea3 100644
--- a/mlir/lib/Analysis/DataFlow/DenseAnalysis.cpp
+++ b/mlir/lib/Analysis/DataFlow/DenseAnalysis.cpp
@@ -283,17 +283,23 @@ void AbstractDenseBackwardDataFlowAnalysis::visitCallOperation(
     AbstractDenseLattice *before) {
   // Find the callee.
   Operation *callee = call.resolveCallable(&symbolTable);
-  auto callable = dyn_cast_or_null<CallableOpInterface>(callee);
-  if (!callable)
-    return setToExitState(before);
 
+  auto callable = dyn_cast_or_null<CallableOpInterface>(callee);
   // No region means the callee is only declared in this module.
-  Region *region = callable.getCallableRegion();
-  if (!region || region->empty() || !getSolverConfig().isInterprocedural()) {
+  // If that is the case or if the solver is not interprocedural,
+  // let the hook handle it.
+  if (!getSolverConfig().isInterprocedural() ||
+      (callable && (!callable.getCallableRegion() ||
+                    callable.getCallableRegion()->empty()))) {
     return visitCallControlFlowTransfer(
         call, CallControlFlowAction::ExternalCallee, after, before);
   }
 
+  if (!callable)
+    return setToExitState(before);
+
+  Region *region = callable.getCallableRegion();
+
   // Call-level control flow specifies the data flow here.
   //
   //   func.func @callee() {
diff --git a/mlir/test/Analysis/DataFlow/test-next-access.mlir b/mlir/test/Analysis/DataFlow/test-next-access.mlir
index 70069b10a93983..8825c699dd1305 100644
--- a/mlir/test/Analysis/DataFlow/test-next-access.mlir
+++ b/mlir/test/Analysis/DataFlow/test-next-access.mlir
@@ -575,3 +575,21 @@ func.func @call_opaque_callee(%arg0: memref<f32>) {
   memref.load %arg0[] {name = "post"} : memref<f32>
   return
 }
+
+// -----
+
+// CHECK-LABEL: @indirect_call
+func.func @indirect_call(%arg0: memref<f32>, %arg1: (memref<f32>) -> ()) {
+  // IP:         name = "pre"
+  // IP-SAME:    next_access = ["unknown"]
+  // IP_AR:      name = "pre"
+  // IP_AR-SAME: next_access = ["unknown"] 
+  // LOCAL:      name = "pre"
+  // LOCAL-SAME: next_access = ["unknown"]
+  // LC_AR:      name = "pre"
+  // LC_AR-SAME: next_access = {{\[}}["call"]]
+  memref.load %arg0[] {name = "pre"} : memref<f32>
+  func.call_indirect %arg1(%arg0) {name = "call"} : (memref<f32>) -> ()
+  memref.load %arg0[] {name = "post"} : memref<f32>
+  return
+}