[llvm] [Offload][OpenMP] Keep user Numblocks requested value for RecordReplay (PR #228345)

[Alex Duran via llvm-commits]

https://github.com/adurang created https://github.com/llvm/llvm-project/pull/228345

RecordReplay still uses the Numblocks value requested by the user so keep passing that value through the PluginInterface for now.

Fixes issue introduced by #222607 

>From abc0af639b15537e15ab0b9b86b43f8a40d3ae43 Mon Sep 17 00:00:00 2001
From: "Duran, Alex" <alejandro.duran at intel.com>
Date: Fri, 2 Oct 2026 00:10:00 -0700
Subject: [PATCH] [Offload][OpenMP] Keep user Numblocks requested value for
 RecordReplay

---
 offload/liboffload/src/OffloadImpl.cpp                   | 1 +
 offload/libompaccsupport/device.cpp                      | 3 +++
 offload/plugins-nextgen/common/include/PluginInterface.h | 3 +++
 offload/plugins-nextgen/common/src/RecordReplay.cpp      | 2 +-
 4 files changed, 8 insertions(+), 1 deletion(-)

diff --git a/offload/liboffload/src/OffloadImpl.cpp b/offload/liboffload/src/OffloadImpl.cpp
index 656a6bc85ab8d..515aecf70c4e6 100644
--- a/offload/liboffload/src/OffloadImpl.cpp
+++ b/offload/liboffload/src/OffloadImpl.cpp
@@ -1297,6 +1297,7 @@ Error olLaunchKernel_impl(ol_queue_handle_t Queue, ol_device_handle_t Device,
   LaunchArgs.UserNumBlocks[0] = LaunchSizeArgs->NumGroups.x;
   LaunchArgs.UserNumBlocks[1] = LaunchSizeArgs->NumGroups.y;
   LaunchArgs.UserNumBlocks[2] = LaunchSizeArgs->NumGroups.z;
+  LaunchArgs.KernelLaunchInfo.RequestedNumBlocks = LaunchSizeArgs->NumGroups.x;
   LaunchArgs.UserThreadLimit[0] = LaunchSizeArgs->GroupSize.x;
   LaunchArgs.UserThreadLimit[1] = LaunchSizeArgs->GroupSize.y;
   LaunchArgs.UserThreadLimit[2] = LaunchSizeArgs->GroupSize.z;
diff --git a/offload/libompaccsupport/device.cpp b/offload/libompaccsupport/device.cpp
index 097d2840b485a..ce26a62940ef5 100644
--- a/offload/libompaccsupport/device.cpp
+++ b/offload/libompaccsupport/device.cpp
@@ -536,6 +536,9 @@ int32_t DeviceTy::launchKernel(void *TgtEntryPtr, void **TgtVarsPtr,
   LaunchArgs.Tripcount = KernelArgs.Tripcount;
   LaunchArgs.DynCGroupMem = KernelArgs.DynCGroupMem;
   llvm::copy(KernelArgs.UserNumBlocks, LaunchArgs.UserNumBlocks);
+  // Save the requested value before computing the effective number of blocks so
+  // it can be used by record-replay mechanisms.
+  LaunchArgs.KernelLaunchInfo.RequestedNumBlocks = KernelArgs.UserNumBlocks[0];
   llvm::copy(KernelArgs.UserThreadLimit, LaunchArgs.UserThreadLimit);
   LaunchArgs.Flags.Cooperative = KernelArgs.Flags.Cooperative;
   LaunchArgs.Flags.DynCGroupMemFallback = KernelArgs.Flags.DynCGroupMemFallback;
diff --git a/offload/plugins-nextgen/common/include/PluginInterface.h b/offload/plugins-nextgen/common/include/PluginInterface.h
index bed5b550fc43f..33740c3af6e16 100644
--- a/offload/plugins-nextgen/common/include/PluginInterface.h
+++ b/offload/plugins-nextgen/common/include/PluginInterface.h
@@ -461,6 +461,9 @@ struct KernelLaunchArgsTy {
     uint32_t ReductionDataSize = 0;
     /// Maximum number of threads per block that this kernel may use.
     uint32_t MaxNumThreads = 0;
+    /// Number of blocks originally requested by the program for the first
+    /// dimension (e.g., num_teams clause), or 0 if none was requested.
+    uint32_t RequestedNumBlocks = 0;
   } KernelLaunchInfo;
   struct {
     uint64_t Cooperative : 1; // Was this kernel spawned as cooperative.
diff --git a/offload/plugins-nextgen/common/src/RecordReplay.cpp b/offload/plugins-nextgen/common/src/RecordReplay.cpp
index 73a298be2effa..5fc491060e1d1 100644
--- a/offload/plugins-nextgen/common/src/RecordReplay.cpp
+++ b/offload/plugins-nextgen/common/src/RecordReplay.cpp
@@ -271,7 +271,7 @@ Error NativeRecordReplayTy::recordDescImpl(
 
   // Export minimum and maximum for allowed number of teams. If zero, it means
   // there was no restriction provided by the program.
-  uint32_t MinMaxBlocks = std::max(LaunchArgs.UserNumBlocks[0], uint32_t(0));
+  uint32_t MinMaxBlocks = LaunchArgs.KernelLaunchInfo.RequestedNumBlocks;
   json::Array JsonTeamsLimits;
   JsonTeamsLimits.push_back(MinMaxBlocks);
   JsonTeamsLimits.push_back(MinMaxBlocks);