[mlir] support non-interprocedural dataflow analyses (#75583)

The core implementation of the dataflow anlysis framework is
interpocedural by design. While this offers better analysis precision,
it also comes with additional cost as it takes longer for the analysis
to reach the fixpoint state. Add a configuration mechanism to the
dataflow solver to control whether it operates inteprocedurally or not
to offer clients a choice.

As a positive side effect, this change also adds hooks for explicitly
processing external/opaque function calls in the dataflow analyses,
e.g., based off of attributes present in the the function declaration or
call operation such as alias scopes and modref available in the LLVM
dialect.

This change should not affect existing analyses and the default solver
configuration remains interprocedural.

Co-authored-by: Jacob Peng <jacobmpeng@gmail.com>

3 files changed, 286 insertions, 69 deletions

diff --git a/mlir/test/Analysis/DataFlow/test-last-modified-callgraph.mlir b/mlir/test/Analysis/DataFlow/test-last-modified-callgraph.mlir
index 709d787..a5eba43 100644
--- a/mlir/test/Analysis/DataFlow/test-last-modified-callgraph.mlir
+++ b/mlir/test/Analysis/DataFlow/test-last-modified-callgraph.mlir
@@ -1,8 +1,32 @@
-// RUN: mlir-opt -test-last-modified --split-input-file %s 2>&1 | FileCheck %s
+// RUN: mlir-opt -test-last-modified --split-input-file %s 2>&1 |\
+// RUN:          FileCheck %s --check-prefixes=CHECK,IP,IP_ONLY
+// RUN: mlir-opt -test-last-modified='assume-func-writes=true' \
+// RUN:          --split-input-file %s 2>&1 |\
+// RUN:          FileCheck %s --check-prefixes=CHECK,IP,IP_AW
+// RUN: mlir-opt -test-last-modified='interprocedural=false' \
+// RUN:          --split-input-file %s 2>&1 |\
+// RUN:          FileCheck %s --check-prefixes=CHECK,LOCAL
+// RUN: mlir-opt \
+// RUN:    -test-last-modified='interprocedural=false assume-func-writes=true' \
+// RUN:    --split-input-file %s 2>&1 |\
+// RUN:    FileCheck %s --check-prefixes=CHECK,LC_AW
+
+// Check prefixes are as follows:
+// 'check': common for all runs;
+// 'ip': interprocedural runs;
+// 'ip_aw': interpocedural runs assuming calls to external functions write to
+//          all arguments;
+// 'ip_only': interprocedural runs not assuming calls writing;
+// 'local': local (non-interprocedural) analysis not assuming calls writing;
+// 'lc_aw': local analysis assuming external calls writing to all arguments.
 
 // CHECK-LABEL: test_tag: test_callsite
-// CHECK: operand #0
-// CHECK-NEXT: - a
+// IP:    operand #0
+// IP-NEXT: - a
+// LOCAL: operand #0
+// LOCAL-NEXT: - <unknown>
+// LC_AW: operand #0
+// LC_AW-NEXT: - <unknown>
 func.func private @single_callsite_fn(%ptr: memref<i32>) -> memref<i32> {
   return {tag = "test_callsite"} %ptr : memref<i32>
 }
@@ -16,8 +40,12 @@ func.func @test_callsite() {
 }
 
 // CHECK-LABEL: test_tag: test_return_site
-// CHECK: operand #0
-// CHECK-NEXT: - b
+// IP:    operand #0
+// IP-NEXT:    - b
+// LOCAL: operand #0
+// LOCAL-NEXT: - <unknown>
+// LC_AW: operand #0
+// LC_AW-NEXT: - <unknown>
 func.func private @single_return_site_fn(%ptr: memref<i32>) -> memref<i32> {
   %c0 = arith.constant 0 : i32
   memref.store %c0, %ptr[] {tag_name = "b"} : memref<i32>
@@ -25,9 +53,13 @@ func.func private @single_return_site_fn(%ptr: memref<i32>) -> memref<i32> {
 }
 
 // CHECK-LABEL: test_tag: test_multiple_callsites
-// CHECK: operand #0
-// CHECK-NEXT: write0
-// CHECK-NEXT: write1
+// IP:    operand #0
+// IP-NEXT:    write0
+// IP-NEXT:    write1
+// LOCAL: operand #0
+// LOCAL-NEXT: - <unknown>
+// LC_AW: operand #0
+// LC_AW-NEXT: - <unknown>
 func.func @test_return_site(%ptr: memref<i32>) -> memref<i32> {
   %0 = func.call @single_return_site_fn(%ptr) : (memref<i32>) -> memref<i32>
   return {tag = "test_return_site"} %0 : memref<i32>
@@ -46,9 +78,13 @@ func.func @test_multiple_callsites(%a: i32, %ptr: memref<i32>) -> memref<i32> {
 }
 
 // CHECK-LABEL: test_tag: test_multiple_return_sites
-// CHECK: operand #0
-// CHECK-NEXT: return0
-// CHECK-NEXT: return1
+// IP:    operand #0
+// IP-NEXT:    return0
+// IP-NEXT:    return1
+// LOCAL: operand #0
+// LOCAL-NEXT: - <unknown>
+// LC_AW: operand #0
+// LC_AW-NEXT: - <unknown>
 func.func private @multiple_return_site_fn(%cond: i1, %a: i32, %ptr: memref<i32>) -> memref<i32> {
   cf.cond_br %cond, ^a, ^b
 
@@ -69,8 +105,12 @@ func.func @test_multiple_return_sites(%cond: i1, %a: i32, %ptr: memref<i32>) ->
 // -----
 
 // CHECK-LABEL: test_tag: after_call
-// CHECK: operand #0
-// CHECK-NEXT: - write0
+// IP:    operand #0
+// IP-NEXT:    - write0
+// LOCAL: operand #0
+// LOCAL-NEXT: - <unknown>
+// LC_AW: operand #0
+// LC_AW-NEXT: - func.call
 func.func private @void_return(%ptr: memref<i32>) {
   return
 }
@@ -98,17 +138,29 @@ func.func private @callee(%arg0: memref<f32>) -> memref<f32> {
 // "pre" -> "call" -> "callee" -> "post"
 
 // CHECK-LABEL: test_tag: call_and_store_before::enter_callee:
-// CHECK:  operand #0
-// CHECK:   - call
+// IP:     operand #0
+// IP:      - call
+// LOCAL:  operand #0
+// LOCAL:   - <unknown>
+// LC_AW:  operand #0
+// LC_AW:   - <unknown>
+
 // CHECK: test_tag: exit_callee:
 // CHECK:  operand #0
 // CHECK:   - callee
+
 // CHECK: test_tag: before_call:
 // CHECK:  operand #0
 // CHECK:   - pre
+
 // CHECK: test_tag: after_call:
-// CHECK:  operand #0
-// CHECK:   - callee
+// IP:     operand #0
+// IP:      - callee
+// LOCAL:  operand #0
+// LOCAL:   - <unknown>
+// LC_AW:  operand #0
+// LC_AW:   - call
+
 // CHECK: test_tag: return:
 // CHECK:  operand #0
 // CHECK:   - post
@@ -138,17 +190,29 @@ func.func private @callee(%arg0: memref<f32>) -> memref<f32> {
 // "pre" -> "callee" -> "call" -> "post"
 
 // CHECK-LABEL: test_tag: call_and_store_after::enter_callee:
-// CHECK:  operand #0
-// CHECK:   - pre
+// IP:     operand #0
+// IP:      - pre
+// LOCAL:  operand #0
+// LOCAL:   - <unknown>
+// LC_AW:  operand #0
+// LC_AW:   - <unknown>
+
 // CHECK: test_tag: exit_callee:
 // CHECK:  operand #0
 // CHECK:   - callee
+
 // CHECK: test_tag: before_call:
 // CHECK:  operand #0
 // CHECK:   - pre
-// CHECK: test_tag: after_call:
-// CHECK:  operand #0
-// CHECK:   - call
+
+// CHECK:    test_tag: after_call:
+// IP:     operand #0
+// IP:      - call
+// LOCAL:  operand #0
+// LOCAL:   - <unknown>
+// LC_AW:  operand #0
+// LC_AW:   - call
+
 // CHECK: test_tag: return:
 // CHECK:  operand #0
 // CHECK:   - post
@@ -162,3 +226,20 @@ func.func @call_and_store_after(%arg0: memref<f32>) -> memref<f32> {
   memref.store %1, %arg0[] {tag_name = "post"} : memref<f32>
   return {tag = "return"} %arg0 : memref<f32>
 }
+
+// -----
+
+func.func private @void_return(%ptr: memref<i32>)
+
+// CHECK-LABEL: test_tag: after_opaque_call:
+// CHECK:        operand #0
+// IP_ONLY:       - <unknown>
+// IP_AW:         - func.call
+func.func @test_opaque_call_return() {
+  %ptr = memref.alloc() : memref<i32>
+  %c0 = arith.constant 0 : i32
+  memref.store %c0, %ptr[] {tag_name = "write0"} : memref<i32>
+  func.call @void_return(%ptr) : (memref<i32>) -> ()
+  memref.load %ptr[] {tag = "after_opaque_call"} : memref<i32>
+  return
+}
diff --git a/mlir/test/Analysis/DataFlow/test-next-access.mlir b/mlir/test/Analysis/DataFlow/test-next-access.mlir
index 313a75c..de0788f 100644
--- a/mlir/test/Analysis/DataFlow/test-next-access.mlir
+++ b/mlir/test/Analysis/DataFlow/test-next-access.mlir
@@ -1,4 +1,22 @@
-// RUN: mlir-opt %s --test-next-access --split-input-file | FileCheck %s
+// RUN: mlir-opt %s --test-next-access --split-input-file |\
+// RUN:             FileCheck %s --check-prefixes=CHECK,IP
+// RUN: mlir-opt %s --test-next-access='interprocedural=false' \
+// RUN:             --split-input-file |\
+// RUN:             FileCheck %s --check-prefixes=CHECK,LOCAL
+// RUN: mlir-opt %s --test-next-access='assume-func-reads=true' \
+// RUN:             --split-input-file |\
+// RUN:             FileCheck %s --check-prefixes=CHECK,IP_AR
+// RUN: mlir-opt %s \
+// RUN:      --test-next-access='interprocedural=false assume-func-reads=true' \
+// RUN:      --split-input-file | FileCheck %s --check-prefixes=CHECK,LC_AR
+
+// Check prefixes are as follows:
+// 'check': common for all runs;
+// 'ip_ar': interpocedural runs assuming calls to external functions read
+//          all arguments;
+// 'ip': interprocedural runs not assuming function calls reading;
+// 'local': local (non-interprocedural) analysis not assuming calls reading;
+// 'lc_ar': local analysis assuming external calls reading all arguments.
 
 // CHECK-LABEL: @trivial
 func.func @trivial(%arg0: memref<f32>, %arg1: f32) -> f32 {
@@ -252,8 +270,10 @@ func.func @known_conditional_cf(%arg0: memref<f32>) {
 // -----
 
 func.func private @callee1(%arg0: memref<f32>) {
-  // CHECK:      name = "callee1"
-  // CHECK-SAME: next_access = {{\[}}["post"]]
+  // IP:         name = "callee1"
+  // IP-SAME:    next_access = {{\[}}["post"]]
+  // LOCAL:      name = "callee1"
+  // LOCAL-SAME: next_access = ["unknown"]
   memref.load %arg0[] {name = "callee1"} : memref<f32>
   return
 }
@@ -267,10 +287,14 @@ func.func private @callee2(%arg0: memref<f32>) {
 
 // CHECK-LABEL: @simple_call
 func.func @simple_call(%arg0: memref<f32>) {
-  // CHECK:      name = "caller"
-  // CHECK-SAME: next_access = {{\[}}["callee1"]]
+  // IP:         name = "caller"
+  // IP-SAME:    next_access = {{\[}}["callee1"]]
+  // LOCAL:      name = "caller"
+  // LOCAL-SAME: next_access = ["unknown"]
+  // LC_AR:      name = "caller"
+  // LC_AR-SAME: next_access = {{\[}}["call"]]
   memref.load %arg0[] {name = "caller"} : memref<f32>
-  func.call @callee1(%arg0) : (memref<f32>) -> ()
+  func.call @callee1(%arg0) {name = "call"} : (memref<f32>) -> ()
   memref.load %arg0[] {name = "post"} : memref<f32>
   return
 }
@@ -279,10 +303,14 @@ func.func @simple_call(%arg0: memref<f32>) {
 
 // CHECK-LABEL: @infinite_recursive_call
 func.func @infinite_recursive_call(%arg0: memref<f32>) {
-  // CHECK:      name = "pre"
-  // CHECK-SAME: next_access = {{\[}}["pre"]]
+  // IP:         name = "pre"
+  // IP-SAME:    next_access = {{\[}}["pre"]]
+  // 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 @infinite_recursive_call(%arg0) : (memref<f32>) -> ()
+  func.call @infinite_recursive_call(%arg0) {name = "call"} : (memref<f32>) -> ()
   memref.load %arg0[] {name = "post"} : memref<f32>
   return
 }
@@ -291,11 +319,15 @@ func.func @infinite_recursive_call(%arg0: memref<f32>) {
 
 // CHECK-LABEL: @recursive_call
 func.func @recursive_call(%arg0: memref<f32>, %cond: i1) {
-  // CHECK:      name = "pre"
-  // CHECK-SAME: next_access = {{\[}}["post", "pre"]]
+  // IP:         name = "pre"
+  // IP-SAME:    next_access = {{\[}}["post", "pre"]]
+  // LOCAL:      name = "pre"
+  // LOCAL-SAME: next_access = ["unknown"]
+  // LC_AR:      name = "pre"
+  // LC_AR-SAME: next_access = {{\[}}["post", "call"]]
   memref.load %arg0[] {name = "pre"} : memref<f32>
   scf.if %cond {
-    func.call @recursive_call(%arg0, %cond) : (memref<f32>, i1) -> ()
+    func.call @recursive_call(%arg0, %cond) {name = "call"} : (memref<f32>, i1) -> ()
   }
   memref.load %arg0[] {name = "post"} : memref<f32>
   return
@@ -305,12 +337,16 @@ func.func @recursive_call(%arg0: memref<f32>, %cond: i1) {
 
 // CHECK-LABEL: @recursive_call_cf
 func.func @recursive_call_cf(%arg0: memref<f32>, %cond: i1) {
-  // CHECK:      name = "pre"
-  // CHECK-SAME: next_access = {{\[}}["pre", "post"]]
+  // IP:         name = "pre"
+  // IP-SAME:    next_access = {{\[}}["pre", "post"]]
+  // LOCAL:      name = "pre"
+  // LOCAL-SAME: next_access = ["unknown"]
+  // LC_AR:      name = "pre"
+  // LC_AR-SAME: next_access = {{\[}}["call", "post"]]
   %0 = memref.load %arg0[] {name = "pre"} : memref<f32>
   cf.cond_br %cond, ^bb1, ^bb2
 ^bb1:
-  call @recursive_call_cf(%arg0, %cond) : (memref<f32>, i1) -> ()
+  call @recursive_call_cf(%arg0, %cond) {name = "call"} : (memref<f32>, i1) -> ()
   cf.br ^bb2
 ^bb2:
   %2 = memref.load %arg0[] {name = "post"} : memref<f32>
@@ -320,27 +356,35 @@ func.func @recursive_call_cf(%arg0: memref<f32>, %cond: i1) {
 // -----
 
 func.func private @callee1(%arg0: memref<f32>) {
-  // CHECK:      name = "callee1"
-  // CHECK-SAME: next_access = {{\[}}["post"]]
+  // IP:         name = "callee1"
+  // IP-SAME:    next_access = {{\[}}["post"]]
+  // LOCAL:      name = "callee1"
+  // LOCAL-SAME: next_access = ["unknown"]
   memref.load %arg0[] {name = "callee1"} : memref<f32>
   return
 }
 
 func.func private @callee2(%arg0: memref<f32>) {
-  // CHECK:      name = "callee2"
-  // CHECK-SAME: next_access = {{\[}}["post"]]
+  // IP:         name = "callee2"
+  // IP-SAME:    next_access = {{\[}}["post"]]
+  // LOCAL:      name = "callee2"
+  // LOCAL-SAME: next_access = ["unknown"]
   memref.load %arg0[] {name = "callee2"} : memref<f32>
   return
 }
 
 func.func @conditonal_call(%arg0: memref<f32>, %cond: i1) {
-  // CHECK:      name = "pre"
-  // CHECK-SAME: next_access = {{\[}}["callee1", "callee2"]]
+  // IP:         name = "pre"
+  // IP-SAME:    next_access = {{\[}}["callee1", "callee2"]]
+  // LOCAL:      name = "pre"
+  // LOCAL-SAME: next_access = ["unknown"]
+  // LC_AR:      name = "pre"
+  // LC_AR-SAME: next_access = {{\[}}["call1", "call2"]]
   memref.load %arg0[] {name = "pre"} : memref<f32>
   scf.if %cond {
-    func.call @callee1(%arg0) : (memref<f32>) -> ()
+    func.call @callee1(%arg0) {name = "call1"} : (memref<f32>) -> ()
   } else {
-    func.call @callee2(%arg0) : (memref<f32>) -> ()
+    func.call @callee2(%arg0) {name = "call2"} : (memref<f32>) -> ()
   }
   memref.load %arg0[] {name = "post"} : memref<f32>
   return
@@ -354,16 +398,22 @@ func.func @conditonal_call(%arg0: memref<f32>, %cond: i1) {
 // "caller" -> "call" -> "callee" -> "post"
 
 func.func private @callee(%arg0: memref<f32>) {
-  // CHECK:              name = "callee"
-  // CHECK-SAME-LITERAL: next_access = [["post"]]
+  // IP:              name = "callee"
+  // IP-SAME-LITERAL: next_access = [["post"]]
+  // LOCAL:           name = "callee"
+  // LOCAL-SAME:      next_access = ["unknown"]
   memref.load %arg0[] {name = "callee"} : memref<f32>
   return
 }
 
 // CHECK-LABEL: @call_and_store_before
 func.func @call_and_store_before(%arg0: memref<f32>) {
-  // CHECK:              name = "caller"
-  // CHECK-SAME-LITERAL: next_access = [["call"]]
+  // IP:              name = "caller"
+  // IP-SAME-LITERAL: next_access = [["call"]]
+  // LOCAL:           name = "caller"
+  // LOCAL-SAME:      next_access = ["unknown"]
+  // LC_AR:           name = "caller"
+  // LC_AR-SAME:      next_access = {{\[}}["call"]]
   memref.load %arg0[] {name = "caller"} : memref<f32>
   // Note that the access after the entire call is "post".
   // CHECK:              name = "call"
@@ -382,20 +432,26 @@ func.func @call_and_store_before(%arg0: memref<f32>) {
 // "caller" -> "callee" -> "call" -> "post"
 
 func.func private @callee(%arg0: memref<f32>) {
-  // CHECK:              name = "callee"
-  // CHECK-SAME-LITERAL: next_access = [["call"]]
+  // IP:              name = "callee"
+  // IP-SAME-LITERAL: next_access = [["call"]]
+  // LOCAL:           name = "callee"
+  // LOCAL-SAME:      next_access = ["unknown"]
   memref.load %arg0[] {name = "callee"} : memref<f32>
   return
 }
 
 // CHECK-LABEL: @call_and_store_after
 func.func @call_and_store_after(%arg0: memref<f32>) {
-  // CHECK:              name = "caller"
-  // CHECK-SAME-LITERAL: next_access = [["callee"]]
+  // IP:              name = "caller"
+  // IP-SAME-LITERAL: next_access = [["callee"]]
+  // LOCAL:           name = "caller"
+  // LOCAL-SAME:      next_access = ["unknown"]
+  // LC_AR:           name = "caller"
+  // LC_AR-SAME:      next_access = {{\[}}["call"]]
   memref.load %arg0[] {name = "caller"} : memref<f32>
   // CHECK:              name = "call"
   // CHECK-SAME-LITERAL: next_access = [["post"], ["post"]]
-  test.call_and_store @callee(%arg0), %arg0 {name = "call", store_before_call = true} : (memref<f32>, memref<f32>) -> ()
+  test.call_and_store @callee(%arg0), %arg0 {name = "call", store_before_call = false} : (memref<f32>, memref<f32>) -> ()
   // CHECK:              name = "post"
   // CHECK-SAME-LITERAL: next_access = ["unknown"]
   memref.load %arg0[] {name = "post"} : memref<f32>
@@ -499,3 +555,23 @@ func.func @store_with_a_region_after_containing_a_load(%arg0: memref<f32>) {
   memref.load %arg0[] {name = "post"} : memref<f32>
   return
 }
+
+// -----
+
+func.func private @opaque_callee(%arg0: memref<f32>)
+
+// CHECK-LABEL: @call_opaque_callee
+func.func @call_opaque_callee(%arg0: memref<f32>) {
+  // IP:         name = "pre"
+  // IP-SAME:    next_access = ["unknown"]
+  // IP_AR:      name = "pre"
+  // IP_AR-SAME: next_access = {{\[}}["call"]]
+  // 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 @opaque_callee(%arg0) {name = "call"} : (memref<f32>) -> ()
+  memref.load %arg0[] {name = "post"} : memref<f32>
+  return
+}
diff --git a/mlir/test/Analysis/DataFlow/test-written-to.mlir b/mlir/test/Analysis/DataFlow/test-written-to.mlir
index 82fe755..4fc9af1 100644
--- a/mlir/test/Analysis/DataFlow/test-written-to.mlir
+++ b/mlir/test/Analysis/DataFlow/test-written-to.mlir
@@ -1,4 +1,28 @@
-// RUN: mlir-opt -split-input-file -test-written-to %s 2>&1 | FileCheck %s
+// RUN: mlir-opt -split-input-file -test-written-to %s 2>&1 |\
+// RUN:          FileCheck %s --check-prefixes=CHECK,IP
+// RUN: mlir-opt -split-input-file -test-written-to='interprocedural=false' %s \
+// RUN:          2>&1 | FileCheck %s --check-prefixes=CHECK,LOCAL
+// RUN: mlir-opt -split-input-file \
+// RUN:          -test-written-to='assume-func-writes=true' %s 2>&1 |\
+// RUN:          FileCheck %s --check-prefixes=CHECK,IP_AW
+// RUN: mlir-opt -split-input-file \
+// RUN:       -test-written-to='interprocedural=false assume-func-writes=true' \
+// RUN:       %s 2>&1 | FileCheck %s --check-prefixes=CHECK,LC_AW
+
+// Check prefixes are as follows:
+// 'check': common for all runs;
+// 'ip': interprocedural runs;
+// 'ip_aw': interpocedural runs assuming calls to external functions write to
+//          all arguments;
+// 'local': local (non-interprocedural) analysis not assuming calls writing;
+// 'lc_aw': local analysis assuming external calls writing to all arguments.
+
+// Note that despite the name of the test analysis being "written to", it is set
+// up in a peculiar way where passing a value through a block or region argument
+// (via visitCall/BranchOperand) is considered as "writing" that value to the
+// corresponding operand, which is itself a value and not necessarily "memory".
+// This is arguably okay for testing purposes, but may be surprising for readers
+// trying to interpret this test using their intuition.
 
 // CHECK-LABEL: test_tag: constant0
 // CHECK: result #0: [a]
@@ -105,7 +129,9 @@ func.func @test_switch(%flag: i32, %m0: memref<i32>) {
 // -----
 
 // CHECK-LABEL: test_tag: add
-// CHECK: result #0: [a]
+// IP:    result #0: [a]
+// LOCAL: result #0: [callarg0]
+// LC_AW: result #0: [func.call]
 func.func @test_caller(%m0: memref<f32>, %arg: f32) {
   %0 = arith.addf %arg, %arg {tag = "add"} : f32
   %1 = func.call @callee(%0) : (f32) -> f32
@@ -130,7 +156,9 @@ func.func private @callee(%0 : f32) -> f32 {
 }
 
 // CHECK-LABEL: test_tag: sub
-// CHECK: result #0: [a]
+// IP:    result #0: [a]
+// LOCAL: result #0: [callarg0]
+// LC_AW: result #0: [func.call]
 func.func @test_caller_below_callee(%m0: memref<f32>, %arg: f32) {
   %0 = arith.subf %arg, %arg {tag = "sub"} : f32
   %1 = func.call @callee(%0) : (f32) -> f32
@@ -155,7 +183,9 @@ func.func private @callee3(%0 : f32) -> f32 {
 }
 
 // CHECK-LABEL: test_tag: mul
-// CHECK: result #0: [a]
+// IP:    result #0: [a]
+// LOCAL: result #0: [callarg0]
+// LC_AW: result #0: [func.call]
 func.func @test_callchain(%m0: memref<f32>, %arg: f32) {
   %0 = arith.mulf %arg, %arg {tag = "mul"} : f32
   %1 = func.call @callee1(%0) : (f32) -> f32
@@ -239,19 +269,19 @@ func.func @test_for(%m0: memref<i32>) {
 // -----
 
 // CHECK-LABEL: test_tag: default_a
-// CHECK-LABEL: result #0: [a]
+// CHECK:       result #0: [a]
 // CHECK-LABEL: test_tag: default_b
-// CHECK-LABEL: result #0: [b]
+// CHECK:       result #0: [b]
 // CHECK-LABEL: test_tag: 1a
-// CHECK-LABEL: result #0: [a]
+// CHECK:       result #0: [a]
 // CHECK-LABEL: test_tag: 1b
-// CHECK-LABEL: result #0: [b]
+// CHECK:       result #0: [b]
 // CHECK-LABEL: test_tag: 2a
-// CHECK-LABEL: result #0: [a]
+// CHECK:       result #0: [a]
 // CHECK-LABEL: test_tag: 2b
-// CHECK-LABEL: result #0: [b]
+// CHECK:       result #0: [b]
 // CHECK-LABEL: test_tag: switch
-// CHECK-LABEL: operand #0: [brancharg0]
+// CHECK:       operand #0: [brancharg0]
 func.func @test_switch(%arg0 : index, %m0: memref<i32>) {
   %0, %1 = scf.index_switch %arg0 {tag="switch"} -> i32, i32
   case 1 {
@@ -276,6 +306,9 @@ func.func @test_switch(%arg0 : index, %m0: memref<i32>) {
 
 // -----
 
+// The point of this test is to ensure the analysis doesn't crash in presence of
+// external functions.
+
 // CHECK-LABEL: llvm.func @decl(i64)
 // CHECK-LABEL: llvm.func @func(%arg0: i64) {
 // CHECK-NEXT:  llvm.call @decl(%arg0) : (i64) -> ()
@@ -295,12 +328,39 @@ func.func private @callee(%arg0 : i32, %arg1 : i32) -> i32 {
 }
 
 // CHECK-LABEL: test_tag: a
-// CHECK-LABEL:  operand #0: [b]
-// CHECK-LABEL:  operand #1: []
-// CHECK-LABEL:  operand #2: [callarg2]
-// CHECK-LABEL:  result #0: [b]
+
+// IP:           operand #0: [b]
+// LOCAL:        operand #0: [callarg0]
+// LC_AW:        operand #0: [test.call_on_device]
+
+// IP:           operand #1: []
+// LOCAL:        operand #1: [callarg1]
+// LC_AW:        operand #1: [test.call_on_device]
+
+// IP:           operand #2: [callarg2]
+// LOCAL:        operand #2: [callarg2]
+// LC_AW:        operand #2: [test.call_on_device]
+
+// CHECK:        result #0: [b]
 func.func @test_call_on_device(%arg0: i32, %arg1: i32, %device: i32, %m0: memref<i32>) {
   %0 = test.call_on_device @callee(%arg0, %arg1), %device {tag = "a"} : (i32, i32, i32) -> (i32)
   memref.store %0, %m0[] {tag_name = "b"} : memref<i32>
   return
 }
+
+// -----
+
+func.func private @external_callee(%arg0: i32) -> i32
+
+// CHECK-LABEL: test_tag: add_external
+// IP:    operand #0: [callarg0]
+// LOCAL: operand #0: [callarg0]
+// LC_AW: operand #0: [func.call]
+// IP_AW: operand #0: [func.call]
+
+func.func @test_external_callee(%arg0: i32, %m0: memref<i32>) {
+  %0 = arith.addi %arg0, %arg0 { tag = "add_external"}: i32
+  %1 = func.call @external_callee(%arg0) : (i32) -> i32
+  memref.store %1, %m0[] {tag_name = "a"} : memref<i32>
+  return
+}