[llvm] [LLVM][OpenMPOpt] Fix a crash when associated function is nullptr (PR #66274)

[Shilei Tian via llvm-commits]

https://github.com/shiltian created https://github.com/llvm/llvm-project/pull/66274:

The associated function can be a nullptr if it is an indirect call.
This causes a crash in `CheckCallee` which always assumes the callee
is a valid pointer.

Fix #66904.


>From e628b71adbe5cccefccfbc77411d2fd5afc9a73a Mon Sep 17 00:00:00 2001
From: Shilei Tian <i at tianshilei.me>
Date: Wed, 13 Sep 2023 14:37:51 -0400
Subject: [PATCH] [LLVM][OpenMPOpt] Fix a crash when associated function is
 nullptr

The associated function can be a nullptr if it is an indirect call.
This causes a crash in `CheckCallee` which always assumes the callee
is a valid pointer.

Fix #66904.
---
 llvm/lib/Transforms/IPO/OpenMPOpt.cpp         |  4 +-
 .../OpenMP/indirect_call_kernel_info_crash.ll | 42 +++++++++++++++++++
 2 files changed, 45 insertions(+), 1 deletion(-)
 create mode 100644 llvm/test/Transforms/OpenMP/indirect_call_kernel_info_crash.ll

diff --git a/llvm/lib/Transforms/IPO/OpenMPOpt.cpp b/llvm/lib/Transforms/IPO/OpenMPOpt.cpp
index a18730ab35621ef..16bb1d37c23d4f5 100644
--- a/llvm/lib/Transforms/IPO/OpenMPOpt.cpp
+++ b/llvm/lib/Transforms/IPO/OpenMPOpt.cpp
@@ -5028,7 +5028,9 @@ struct AAKernelInfoCallSite : AAKernelInfo {
     const auto *AACE =
         A.getAAFor<AACallEdges>(*this, getIRPosition(), DepClassTy::OPTIONAL);
     if (!AACE || !AACE->getState().isValidState() || AACE->hasUnknownCallee()) {
-      CheckCallee(getAssociatedFunction(), /*NumCallees=*/1);
+      Function *F = getAssociatedFunction();
+      if (F)
+        CheckCallee(getAssociatedFunction(), /*NumCallees=*/1);
     } else {
       const auto &OptimisticEdges = AACE->getOptimisticEdges();
       for (auto *Callee : OptimisticEdges) {
diff --git a/llvm/test/Transforms/OpenMP/indirect_call_kernel_info_crash.ll b/llvm/test/Transforms/OpenMP/indirect_call_kernel_info_crash.ll
new file mode 100644
index 000000000000000..03bc31bac2034e6
--- /dev/null
+++ b/llvm/test/Transforms/OpenMP/indirect_call_kernel_info_crash.ll
@@ -0,0 +1,42 @@
+; RUN: opt -S -passes=openmp-opt < %s
+target datalayout = "e-p:64:64-p1:64:64-p2:32:32-p3:32:32-p4:64:64-p5:32:32-p6:32:32-p7:160:256:256:32-p8:128:128-i64:64-v16:16-v24:32-v32:32-v48:64-v96:128-v192:256-v256:256-v512:512-v1024:1024-v2048:2048-n32:64-S32-A5-G1-ni:7:8"
+target triple = "amdgcn-amd-amdhsa"
+
+%"struct.ompx::state::TeamStateTy" = type { %"struct.ompx::state::ICVStateTy", i32, i32, ptr }
+%"struct.ompx::state::ICVStateTy" = type { i32, i32, i32, i32, i32, i32, i32 }
+
+ at _ZN4ompx5state9TeamStateE = internal addrspace(3) global %"struct.ompx::state::TeamStateTy" undef
+
+define amdgpu_kernel void @__omp_offloading_32_70c2e76c_main_l24() {
+  %1 = tail call i32 @__kmpc_target_init(ptr null)
+  call void @__kmpc_parallel_51(ptr null, i32 0, i32 0, i32 0, i32 0, ptr @__omp_offloading_32_70c2e76c_main_l24_omp_outlined, ptr null, ptr null, i64 0)
+  ret void
+}
+
+define void @__omp_offloading_32_70c2e76c_main_l24_omp_outlined(ptr %0) {
+  call void @__kmpc_for_static_init_4()
+  br label %2
+
+2:                                                ; preds = %2, %1
+  %3 = load ptr, ptr addrspace(1) null, align 4294967296
+  %4 = call i32 %3(i32 0)
+  store i32 %4, ptr %0, align 4
+  br label %2
+}
+
+define internal i32 @__kmpc_target_init(ptr %0) {
+  store i32 0, ptr addrspace(3) @_ZN4ompx5state9TeamStateE, align 16
+  ret i32 0
+}
+
+declare void @__kmpc_parallel_51(ptr, i32, i32, i32, i32, ptr, ptr, ptr, i64)
+
+define void @__kmpc_for_static_init_4() {
+  %1 = load i32, ptr addrspace(3) @_ZN4ompx5state9TeamStateE, align 8
+  ret void
+}
+
+!llvm.module.flags = !{!0, !1}
+
+!0 = !{i32 7, !"openmp", i32 51}
+!1 = !{i32 7, !"openmp-device", i32 51}