[llvm] [OpenMPOpt] Block SPMDization when omp_alloc/omp_free are present (PR #190875)

[David Pagan via llvm-commits]

https://github.com/ddpagan created https://github.com/llvm/llvm-project/pull/190875

Currently, OpenMPOpt's SPMDization pass converts generic-mode target
kernels containing omp_alloc()/omp_free() calls to SPMD mode, but it
does so incorrectly. SPMDization of omp_alloc/omp_free is currently
allowed because they are not registered in OMPKinds.def, so OpenMPOpt
doesn't recognize them as OpenMP runtime calls. The callee-level analysis
then finds no SPMD incompatibility (the functions themselves contain no
parallel regions), so SPMDization proceeds without guarding the calls.

The potential problem is that all threads independently call the device
memory allocator instead of only thread 0, producing duplicate allocations.
Whether this causes wrong results or not depends on the allocator
implementation that omp_alloc eventually resolves to after LTO inlining.
The current default allocator (trunk) ends up as an atomic bump-pointer,
which just happens to tolerate thread duplication (each thread gets a
unique valid pointer), which masks the bug. A different allocator, such as
__ockl_dm_alloc, does not. For this case, the concurrent calls from all
threads corrupt its internal state, producing wrong results at -O2. Note
that situation can be replicated in the trunk.

Fixing SPMDization to recognize omp_alloc/omp_free and to guard
appropriately would be the best solution. However, simply fixing
SPMDization to recognize the need for guarding and allowing it to proceed
through the existing guarding code doesn't always work. In particular,
there is at least one case where incorrect code is generated that
mishandles the omp_alloc/omp_free pair such that the omp_free uses a bad
pointer value. This then leads to an AMDGPU error.

A simple, conservative fix is to register omp_alloc and omp_free in
OMPKinds.def and then handle them in AAKernelInfoCallSite's SPMD
compatibility switch with a pessimistic fixpoint. This prevents
SPMDization of any kernel containing these calls, so it works regardless
of the ultimate allocator used.

>From 670316751d88cb57a830aaec92fce07a286518fd Mon Sep 17 00:00:00 2001
From: Dave Pagan <dave.pagan at amd.com>
Date: Thu, 19 Mar 2026 10:35:54 -0500
Subject: [PATCH] [OpenMPOpt] Block SPMDization when omp_alloc/omp_free are
 present

Currently, OpenMPOpt's SPMDization pass converts generic-mode target
kernels containing omp_alloc()/omp_free() calls to SPMD mode, but it
does so incorrectly. SPMDization of omp_alloc/omp_free is currently
allowed because they are not registered in OMPKinds.def, so OpenMPOpt
doesn't recognize them as OpenMP runtime calls. The callee-level analysis
then finds no SPMD incompatibility (the functions themselves contain no
parallel regions), so SPMDization proceeds without guarding the calls.

The potential problem is that all threads independently call the device
memory allocator instead of only thread 0, producing duplicate allocations.
Whether this causes wrong results or not depends on the allocator
implementation that omp_alloc eventually resolves to after LTO inlining.
The current default allocator (trunk) ends up as an atomic bump-pointer,
which just happens to tolerate thread duplication (each thread gets a
unique valid pointer), which masks the bug. A different allocator, such as
__ockl_dm_alloc, does not. For this case, the concurrent calls from all
threads corrupt its internal state, producing wrong results at -O2. Note
that situation can be replicated in the trunk.

Fixing SPMDization to recognize omp_alloc/omp_free and to guard
appropriately would be the best solution. However, simply fixing
SPMDization to recognize the need for guarding and allowing it to proceed
through the existing guarding code doesn't always work. In particular,
there is at least one case where incorrect code is generated that
mishandles the omp_alloc/omp_free pair such that the omp_free uses a bad
pointer value. This then leads to an AMDGPU error.

A simple, conservative fix is to register omp_alloc and omp_free in
OMPKinds.def and then handle them in AAKernelInfoCallSite's SPMD
compatibility switch with a pessimistic fixpoint. This prevents
SPMDization of any kernel containing these calls, so it works regardless
of the ultimate allocator used.
---
 llvm/include/llvm/Frontend/OpenMP/OMPKinds.def | 3 +++
 llvm/lib/Transforms/IPO/OpenMPOpt.cpp          | 7 +++++++
 2 files changed, 10 insertions(+)

diff --git a/llvm/include/llvm/Frontend/OpenMP/OMPKinds.def b/llvm/include/llvm/Frontend/OpenMP/OMPKinds.def
index ad44625f352fe..02aff351c0bf5 100644
--- a/llvm/include/llvm/Frontend/OpenMP/OMPKinds.def
+++ b/llvm/include/llvm/Frontend/OpenMP/OMPKinds.def
@@ -260,6 +260,9 @@ __OMP_RTL(omp_set_nested, false, Void, Int32)
 __OMP_RTL(omp_set_schedule, false, Void, Int32, Int32)
 __OMP_RTL(omp_set_max_active_levels, false, Void, Int32)
 
+__OMP_RTL(omp_alloc, false, VoidPtr, SizeTy, SizeTy)
+__OMP_RTL(omp_free, false, Void, VoidPtr, SizeTy)
+
 __OMP_RTL(__kmpc_master, false, Int32, IdentPtr, Int32)
 __OMP_RTL(__kmpc_end_master, false, Void, IdentPtr, Int32)
 __OMP_RTL(__kmpc_masked, false, Int32, IdentPtr, Int32, Int32)
diff --git a/llvm/lib/Transforms/IPO/OpenMPOpt.cpp b/llvm/lib/Transforms/IPO/OpenMPOpt.cpp
index 5266103d1ac24..3ab65c2561725 100644
--- a/llvm/lib/Transforms/IPO/OpenMPOpt.cpp
+++ b/llvm/lib/Transforms/IPO/OpenMPOpt.cpp
@@ -5017,6 +5017,13 @@ struct AAKernelInfoCallSite : AAKernelInfo {
         SPMDCompatibilityTracker.insert(&CB);
         ReachedUnknownParallelRegions.insert(&CB);
         break;
+      case OMPRTL_omp_alloc:
+      case OMPRTL_omp_free:
+        // omp_alloc/omp_free have side effects (allocator state) that are
+        // unsafe to duplicate across threads during SPMDization.
+        SPMDCompatibilityTracker.indicatePessimisticFixpoint();
+        SPMDCompatibilityTracker.insert(&CB);
+        break;
       case OMPRTL___kmpc_alloc_shared:
       case OMPRTL___kmpc_free_shared:
         // Return without setting a fixpoint, to be resolved in updateImpl.