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

[llvmlistbot at llvm.org]

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

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.

>From ad47238a96f5d65b6584ac4fa369046421d41ab2 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   | 15 ++++++++++-----
 .../Analysis/DataFlow/test-next-access.mlir    | 18 ++++++++++++++++++
 2 files changed, 28 insertions(+), 5 deletions(-)

diff --git a/mlir/lib/Analysis/DataFlow/DenseAnalysis.cpp b/mlir/lib/Analysis/DataFlow/DenseAnalysis.cpp
index 08d89d6db788c8..33056ff7cd25a5 100644
--- a/mlir/lib/Analysis/DataFlow/DenseAnalysis.cpp
+++ b/mlir/lib/Analysis/DataFlow/DenseAnalysis.cpp
@@ -283,17 +283,22 @@ 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
+}