[llvm] [OpenMPOpt] Don't assert when __kmpc_target_init has a null kernel environment argument (PR #196249)

[via llvm-commits]

https://github.com/IamYJLee updated https://github.com/llvm/llvm-project/pull/196249

>From e33daf90fe29f167a0d27be051a7e34f6efb5ab0 Mon Sep 17 00:00:00 2001
From: LeeYoungJoon <dog3hk.dev at gmail.com>
Date: Thu, 7 May 2026 15:42:03 +0900
Subject: [PATCH 1/2] [OpenMPOpt] Don't assert when __kmpc_target_init has a
 null kernel environment argument

---
 llvm/lib/Transforms/IPO/OpenMPOpt.cpp | 34 ++++++++++++++++++++-------
 1 file changed, 26 insertions(+), 8 deletions(-)

diff --git a/llvm/lib/Transforms/IPO/OpenMPOpt.cpp b/llvm/lib/Transforms/IPO/OpenMPOpt.cpp
index 31e9d41ca3410..dd8c598fd49c7 100644
--- a/llvm/lib/Transforms/IPO/OpenMPOpt.cpp
+++ b/llvm/lib/Transforms/IPO/OpenMPOpt.cpp
@@ -259,7 +259,7 @@ KERNEL_ENVIRONMENT_CONFIGURATION_GETTER(MaxTeams)
 GlobalVariable *
 getKernelEnvironementGVFromKernelInitCB(CallBase *KernelInitCB) {
   constexpr int InitKernelEnvironmentArgNo = 0;
-  return cast<GlobalVariable>(
+  return dyn_cast<GlobalVariable>(
       KernelInitCB->getArgOperand(InitKernelEnvironmentArgNo)
           ->stripPointerCasts());
 }
@@ -267,7 +267,9 @@ getKernelEnvironementGVFromKernelInitCB(CallBase *KernelInitCB) {
 ConstantStruct *getKernelEnvironementFromKernelInitCB(CallBase *KernelInitCB) {
   GlobalVariable *KernelEnvGV =
       getKernelEnvironementGVFromKernelInitCB(KernelInitCB);
-  return cast<ConstantStruct>(KernelEnvGV->getInitializer());
+  if (!KernelEnvGV || !KernelEnvGV->hasInitializer())
+    return nullptr;
+  return dyn_cast<ConstantStruct>(KernelEnvGV->getInitializer());
 }
 } // namespace KernelInfo
 
@@ -2947,6 +2949,8 @@ struct AAExecutionDomainFunction : public AAExecutionDomain {
         return false;
       ConstantStruct *KernelEnvC =
           KernelInfo::getKernelEnvironementFromKernelInitCB(CB);
+      if (!KernelEnvC)
+        return false;
       ConstantInt *ExecModeC =
           KernelInfo::getExecModeFromKernelEnvironment(KernelEnvC);
       return ExecModeC->getSExtValue() & OMP_TGT_EXEC_MODE_GENERIC;
@@ -3744,14 +3748,16 @@ struct AAKernelInfoFunction : AAKernelInfo {
     if (!KernelInitCB || !KernelDeinitCB)
       return;
 
-    // Add itself to the reaching kernel and set IsKernelEntry.
-    ReachingKernelEntries.insert(Fn);
-    IsKernelEntry = true;
-
     KernelEnvC =
         KernelInfo::getKernelEnvironementFromKernelInitCB(KernelInitCB);
     GlobalVariable *KernelEnvGV =
         KernelInfo::getKernelEnvironementGVFromKernelInitCB(KernelInitCB);
+    if (!KernelEnvC || !KernelEnvGV)
+      return;
+
+    // Add itself to the reaching kernel and set IsKernelEntry.
+    ReachingKernelEntries.insert(Fn);
+    IsKernelEntry = true;
 
     Attributor::GlobalVariableSimplifictionCallbackTy
         KernelConfigurationSimplifyCB =
@@ -3920,6 +3926,14 @@ struct AAKernelInfoFunction : AAKernelInfo {
     if (!KernelInitCB || !KernelDeinitCB)
       return ChangeStatus::UNCHANGED;
 
+    // Likewise, if the kernel-environment global is missing (e.g. the IR
+    // passes a null pointer as the first argument of __kmpc_target_init)
+    // there is no environment to update.
+    GlobalVariable *KernelEnvGV =
+        KernelInfo::getKernelEnvironementGVFromKernelInitCB(KernelInitCB);
+    if (!KernelEnvGV || !KernelEnvC)
+      return ChangeStatus::UNCHANGED;
+
     ChangeStatus Changed = ChangeStatus::UNCHANGED;
 
     bool HasBuiltStateMachine = true;
@@ -3941,8 +3955,6 @@ struct AAKernelInfoFunction : AAKernelInfo {
           OldUseGenericStateMachineVal);
 
     // At last, update the KernelEnvc
-    GlobalVariable *KernelEnvGV =
-        KernelInfo::getKernelEnvironementGVFromKernelInitCB(KernelInitCB);
     if (KernelEnvGV->getInitializer() != KernelEnvC) {
       KernelEnvGV->setInitializer(KernelEnvC);
       Changed = ChangeStatus::CHANGED;
@@ -4271,6 +4283,8 @@ struct AAKernelInfoFunction : AAKernelInfo {
     // Check if the kernel is already in SPMD mode, if so, return success.
     ConstantStruct *ExistingKernelEnvC =
         KernelInfo::getKernelEnvironementFromKernelInitCB(KernelInitCB);
+    if (!ExistingKernelEnvC)
+      return false;
     auto *ExecModeC =
         KernelInfo::getExecModeFromKernelEnvironment(ExistingKernelEnvC);
     const int8_t ExecModeVal = ExecModeC->getSExtValue();
@@ -4322,6 +4336,8 @@ struct AAKernelInfoFunction : AAKernelInfo {
 
     ConstantStruct *ExistingKernelEnvC =
         KernelInfo::getKernelEnvironementFromKernelInitCB(KernelInitCB);
+    if (!ExistingKernelEnvC)
+      return false;
 
     // Check if the current configuration is non-SPMD and generic state machine.
     // If we already have SPMD mode or a custom state machine we do not need to
@@ -4650,6 +4666,8 @@ struct AAKernelInfoFunction : AAKernelInfo {
 
         ConstantStruct *ExistingKernelEnvC =
             KernelInfo::getKernelEnvironementFromKernelInitCB(AA.KernelInitCB);
+        if (!ExistingKernelEnvC)
+          return;
 
         if (!AA.isValidState()) {
           AA.KernelEnvC = ExistingKernelEnvC;

>From abe257e68930e61a0a921419887b4f90fa4e9f81 Mon Sep 17 00:00:00 2001
From: LeeYoungJoon <dog3hk.dev at gmail.com>
Date: Thu, 7 May 2026 16:12:21 +0900
Subject: [PATCH 2/2] [OpenMPOpt] Add regression test for null kernel
 environment

---
 .../OpenMP/kernel_env_null_init_arg.ll        | 46 +++++++++++++++++++
 1 file changed, 46 insertions(+)
 create mode 100644 llvm/test/Transforms/OpenMP/kernel_env_null_init_arg.ll

diff --git a/llvm/test/Transforms/OpenMP/kernel_env_null_init_arg.ll b/llvm/test/Transforms/OpenMP/kernel_env_null_init_arg.ll
new file mode 100644
index 0000000000000..2d1c23b92a7cd
--- /dev/null
+++ b/llvm/test/Transforms/OpenMP/kernel_env_null_init_arg.ll
@@ -0,0 +1,46 @@
+; RUN: opt -S -passes=openmp-opt < %s | FileCheck %s
+
+target datalayout = "e-m:e-p270:32:32-p271:32:32-p272:64:64-i64:64-i128:128-f80:128-n8:16:32:64-S128"
+target triple = "x86_64-unknown-linux-gnu"
+
+; CHECK-LABEL: define ptx_kernel void @kernel_no_deinit
+; CHECK:         %{{.*}} = call i32 @__kmpc_target_init(ptr null, ptr null)
+define ptx_kernel void @kernel_no_deinit() {
+entry:
+  %0 = call i32 @__kmpc_target_init(ptr null, ptr null)
+  %exec_user_code = icmp eq i32 %0, -1
+  br i1 %exec_user_code, label %user_code.entry, label %worker.exit
+
+user_code.entry:
+  ret void
+
+worker.exit:
+  ret void
+}
+
+; CHECK-LABEL: define ptx_kernel void @kernel_with_deinit
+; CHECK:         %{{.*}} = call i32 @__kmpc_target_init(ptr null, ptr null)
+; CHECK:         call void @__kmpc_target_deinit()
+define ptx_kernel void @kernel_with_deinit() #0 {
+entry:
+  %0 = call i32 @__kmpc_target_init(ptr null, ptr null)
+  %exec_user_code = icmp eq i32 %0, -1
+  br i1 %exec_user_code, label %user_code.entry, label %worker.exit
+
+user_code.entry:
+  call void @__kmpc_target_deinit()
+  ret void
+
+worker.exit:
+  ret void
+}
+
+declare i32 @__kmpc_target_init(ptr, ptr)
+declare void @__kmpc_target_deinit()
+
+attributes #0 = { "kernel" }
+
+!llvm.module.flags = !{!0, !1}
+
+!0 = !{i32 7, !"openmp", i32 50}
+!1 = !{i32 7, !"openmp-device", i32 50}