[llvm] [offload] Verify replay config of teams/threads is allowed (PR #192784)

Kevin Sala Penades via llvm-commits llvm-commits at lists.llvm.org
Sat Apr 18 03:41:54 PDT 2026


https://github.com/kevinsala created https://github.com/llvm/llvm-project/pull/192784

None

>From a66f51e5d2e5ca0bfbde5fc1192d2b468ba2aa2a Mon Sep 17 00:00:00 2001
From: Kevin Sala <salapenades1 at llnl.gov>
Date: Sat, 18 Apr 2026 01:42:13 -0700
Subject: [PATCH 1/2] [offload] Cleanup of llvm-omp-kernel-replay tool

---
 .../kernelreplay/llvm-omp-kernel-replay.cpp   | 323 +++++++++++-------
 1 file changed, 204 insertions(+), 119 deletions(-)

diff --git a/offload/tools/kernelreplay/llvm-omp-kernel-replay.cpp b/offload/tools/kernelreplay/llvm-omp-kernel-replay.cpp
index 3f22c7c223fa4..9548a3c33e25c 100644
--- a/offload/tools/kernelreplay/llvm-omp-kernel-replay.cpp
+++ b/offload/tools/kernelreplay/llvm-omp-kernel-replay.cpp
@@ -16,6 +16,7 @@
 
 #include "llvm/Frontend/Offloading/Utility.h"
 #include "llvm/Support/CommandLine.h"
+#include "llvm/Support/Error.h"
 #include "llvm/Support/ErrorHandling.h"
 #include "llvm/Support/JSON.h"
 #include "llvm/Support/MemoryBuffer.h"
@@ -26,17 +27,19 @@
 
 using namespace llvm;
 
-cl::OptionCategory ReplayOptions("llvm-omp-kernel-replay Options");
+#define TOOL_NAME "llvm-omp-kernel-replay"
+#define TOOL_PREFIX "[" TOOL_NAME "]"
 
-// InputFilename - The filename to read the json description of the kernel.
-static cl::opt<std::string> InputFilename(cl::Positional,
-                                          cl::desc("<input kernel json file>"),
-                                          cl::Required);
+cl::OptionCategory ReplayOptions(TOOL_NAME " Options");
+
+/// The filename to read the JSON kernel description.
+static cl::opt<std::string> JsonFilename(cl::Positional,
+                                         cl::desc("<input kernel JSON file>"),
+                                         cl::Required);
 
 static cl::opt<bool> VerifyOpt(
     "verify",
-    cl::desc(
-        "Verify device memory post execution against the original output."),
+    cl::desc("Verify device memory after replaying against the record output."),
     cl::init(false), cl::cat(ReplayOptions));
 
 static cl::opt<bool> SaveOutputOpt(
@@ -55,88 +58,159 @@ static cl::opt<uint32_t> NumThreadsOpt("num-threads",
 static cl::opt<int32_t> DeviceIdOpt("device-id", cl::desc("Set the device id."),
                                     cl::init(-1), cl::cat(ReplayOptions));
 
-int main(int argc, char **argv) {
-  cl::HideUnrelatedOptions(ReplayOptions);
-  cl::ParseCommandLineOptions(argc, argv, "llvm-omp-kernel-replay\n");
+template <typename... ArgsTy>
+Error createErr(const char *ErrFmt, ArgsTy &&...Args) {
+  return llvm::createStringError(llvm::inconvertibleErrorCode(), ErrFmt,
+                                 std::forward<ArgsTy>(Args)...);
+}
+
+template <typename T>
+Error getInteger(const json::Object *Obj, StringRef Key, T &Result) {
+  auto OptInt = Obj->getInteger(Key);
+  if (!OptInt)
+    return createErr("failed to read JSON integer %s", Key.data());
+  Result = static_cast<T>(*OptInt);
+  return Error::success();
+}
+
+Error getPointer(const json::Object *Obj, StringRef Key, void *&Result) {
+  auto OptInt = Obj->getInteger(Key);
+  if (!OptInt)
+    return createErr("failed to read JSON integer %s", Key.data());
+  Result = reinterpret_cast<void *>(*OptInt);
+  return Error::success();
+}
 
-  ErrorOr<std::unique_ptr<MemoryBuffer>> KernelInfoMB =
-      MemoryBuffer::getFile(InputFilename, /*isText=*/true,
+Error getString(const json::Object *Obj, StringRef Key, StringRef &Result) {
+  auto OptStr = Obj->getString(Key);
+  if (!OptStr)
+    return createErr("failed to read JSON string %s", Key.data());
+  Result = *OptStr;
+  return Error::success();
+}
+
+template <typename Func>
+Error processIntegerArray(const json::Object *Obj, StringRef Key,
+                          Func ProcessFunc) {
+  auto Array = Obj->getArray(Key);
+  if (!Array)
+    return createErr("failed to read JSON array %s", Key.data());
+
+  for (const auto &Val : *Array) {
+    if (auto OptInt = Val.getAsInteger())
+      ProcessFunc(*OptInt);
+    else
+      return createErr("failed to read an integer from JSON array %s",
+                       Key.data());
+  }
+  return Error::success();
+}
+
+/// Replay the kernel and return whether verification occurred.
+Error replayKernel() {
+  // Load the kernel descriptor JSON file.
+  auto KernelDescrBufferOrErr =
+      MemoryBuffer::getFile(JsonFilename, /*isText=*/true,
                             /*RequiresNullTerminator=*/true);
-  if (!KernelInfoMB)
-    reportFatalUsageError("Error reading the kernel info json file");
-  Expected<json::Value> JsonKernelInfo =
-      json::parse(KernelInfoMB.get()->getBuffer());
-  if (auto Err = JsonKernelInfo.takeError())
-    reportFatalUsageError("Cannot parse the kernel info json file");
-
-  auto NumTeamsJson =
-      JsonKernelInfo->getAsObject()->getInteger("NumTeamsClause");
-  uint32_t NumTeams = (NumTeamsOpt > 0 ? NumTeamsOpt : NumTeamsJson.value());
-  auto NumThreadsJson =
-      JsonKernelInfo->getAsObject()->getInteger("ThreadLimitClause");
-  uint32_t NumThreads =
-      (NumThreadsOpt > 0 ? NumThreadsOpt : NumThreadsJson.value());
-  uint32_t SharedMemorySize =
-      JsonKernelInfo->getAsObject()->getInteger("SharedMemorySize").value();
-  // TODO: Print a warning if number of teams/threads is explicitly set in the
-  // kernel info but overridden through command line options.
-  auto LoopTripCount =
-      JsonKernelInfo->getAsObject()->getInteger("LoopTripCount");
-  auto KernelFunc = JsonKernelInfo->getAsObject()->getString("Name");
-  std::string KernelName = KernelFunc.value().str();
+  if (!KernelDescrBufferOrErr)
+    return createErr("failed read the kernel info JSON file");
+
+  // Parse the JSON file.
+  auto JsonDescrOrErr = json::parse(KernelDescrBufferOrErr.get()->getBuffer());
+  if (!JsonDescrOrErr)
+    return JsonDescrOrErr.takeError();
+
+  auto JsonObj = JsonDescrOrErr->getAsObject();
+  if (!JsonObj)
+    return createErr("invalid JSON file");
+
+  // Retrieve the values from the JSON file.
+  uint32_t NumTeams, NumThreads, SharedMemorySize, DeviceId, NumArgs;
+  if (auto Err = getInteger(JsonObj, "NumTeamsClause", NumTeams))
+    return Err;
+  if (auto Err = getInteger(JsonObj, "ThreadLimitClause", NumThreads))
+    return Err;
+  if (auto Err = getInteger(JsonObj, "SharedMemorySize", SharedMemorySize))
+    return Err;
+  if (auto Err = getInteger(JsonObj, "DeviceId", DeviceId))
+    return Err;
+  if (auto Err = getInteger(JsonObj, "NumArgs", NumArgs))
+    return Err;
+
+  uint64_t LoopTripCount, VAllocSize;
+  if (auto Err = getInteger(JsonObj, "VAllocSize", VAllocSize))
+    return Err;
+  if (auto Err = getInteger(JsonObj, "LoopTripCount", LoopTripCount))
+    return Err;
+
+  void *VAllocAddr;
+  if (auto Err = getPointer(JsonObj, "VAllocAddr", VAllocAddr))
+    return Err;
+
+  StringRef KernelName;
+  if (auto Err = getString(JsonObj, "Name", KernelName))
+    return Err;
+
+  // If needed, adjust number of teams and threads, and the device identifier.
+  NumTeams = NumTeamsOpt > 0 ? NumTeamsOpt : NumTeams;
+  NumThreads = NumThreadsOpt > 0 ? NumThreadsOpt : NumThreads;
+  DeviceId = DeviceIdOpt >= 0 ? DeviceIdOpt : DeviceId;
 
   SmallVector<void *> TgtArgs;
+  auto Err = processIntegerArray(JsonObj, "ArgPtrs", [&](uint64_t Val) {
+    TgtArgs.push_back(reinterpret_cast<void *>(Val));
+  });
+  if (Err)
+    return Err;
+
   SmallVector<ptrdiff_t> TgtArgOffsets;
-  auto NumArgs = JsonKernelInfo->getAsObject()->getInteger("NumArgs");
-  auto *TgtArgsArray = JsonKernelInfo->getAsObject()->getArray("ArgPtrs");
-  for (auto It : *TgtArgsArray)
-    TgtArgs.push_back(reinterpret_cast<void *>(It.getAsInteger().value()));
-  auto *TgtArgOffsetsArray =
-      JsonKernelInfo->getAsObject()->getArray("ArgOffsets");
-  for (auto It : *TgtArgOffsetsArray)
-    TgtArgOffsets.push_back(static_cast<ptrdiff_t>(It.getAsInteger().value()));
-
-  void *VAllocAddr = reinterpret_cast<void *>(
-      JsonKernelInfo->getAsObject()->getInteger("VAllocAddr").value());
-  uint64_t VAllocSize =
-      JsonKernelInfo->getAsObject()->getInteger("VAllocSize").value();
-
-  auto Filepath = std::filesystem::path(InputFilename.getValue());
+  Err = processIntegerArray(JsonObj, "ArgOffsets", [&](uint64_t Val) {
+    TgtArgOffsets.push_back(static_cast<ptrdiff_t>(Val));
+  });
+  if (Err)
+    return Err;
+
+  // Keep the filepath and directory for future use.
+  auto Filepath = std::filesystem::path(JsonFilename.getValue());
   auto Directory = Filepath.parent_path();
 
+  // Load the recorded globals file.
   Filepath.replace_extension("globals");
-  ErrorOr<std::unique_ptr<MemoryBuffer>> GlobalsMB =
+  auto GlobalsBufferOrErr =
       MemoryBuffer::getFile(Filepath.string(), /*isText=*/false,
                             /*RequiresNullTerminator=*/false);
-
-  if (!GlobalsMB)
-    reportFatalUsageError("Error reading the globals file");
+  if (!GlobalsBufferOrErr)
+    return createErr("failed to read the globals file");
+  auto GlobalsBuffer = std::move(GlobalsBufferOrErr.get());
 
   // On AMD for currently unknown reasons we cannot copy memory mapped data to
   // device. This is a work-around.
-  uint8_t *RecordedGlobals = new uint8_t[GlobalsMB.get()->getBufferSize()];
+  uint8_t *RecordedGlobals = new uint8_t[GlobalsBuffer->getBufferSize()];
   std::memcpy(RecordedGlobals,
-              const_cast<char *>(GlobalsMB.get()->getBuffer().data()),
-              GlobalsMB.get()->getBufferSize());
+              const_cast<char *>(GlobalsBuffer->getBuffer().data()),
+              GlobalsBuffer->getBufferSize());
 
   void *BufferPtr = (void *)RecordedGlobals;
   uint32_t NumGlobals = *((uint32_t *)(BufferPtr));
   BufferPtr = utils::advancePtr(BufferPtr, sizeof(uint32_t));
 
-  llvm::SmallVector<llvm::offloading::EntryTy> OffloadEntries(
+  SmallVector<llvm::offloading::EntryTy> OffloadEntries(
       NumGlobals + 1, {0x0, 0x1, object::OffloadKind::OFK_OpenMP, 0, nullptr,
                        nullptr, 0, 0, nullptr});
 
-  OffloadEntries[0].SymbolName = const_cast<char *>(KernelName.c_str());
-  // Anything non-zero works to uniquely identify the kernel.
+  // The first offload entry corresponds to the kernel function.
+  OffloadEntries[0].SymbolName = const_cast<char *>(KernelName.data());
+  // Use a unique identifier.
   OffloadEntries[0].Address = (void *)0x1;
 
+  // The rest of entries correspond to the recorded global variables.
   for (uint32_t I = 0; I < NumGlobals; ++I) {
     auto &Global = OffloadEntries[I + 1];
 
     // Use a unique identifier.
     Global.Address = static_cast<char *>(OffloadEntries[0].Address) + I + 1;
 
+    // Setup the offload entry using the information from the file.
     uint32_t NameSize = *((uint32_t *)(BufferPtr));
     BufferPtr = utils::advancePtr(BufferPtr, sizeof(uint32_t));
     uint64_t Size = *((uint64_t *)(BufferPtr));
@@ -148,16 +222,19 @@ int main(int argc, char **argv) {
     BufferPtr = utils::advancePtr(BufferPtr, Size);
   }
 
+  // Load the device image file.
   Filepath.replace_extension("image");
-  ErrorOr<std::unique_ptr<MemoryBuffer>> ImageMB =
+  auto ImageBufferOrErr =
       MemoryBuffer::getFile(Filepath.string(), /*isText=*/false,
                             /*RequiresNullTerminator=*/false);
-  if (!ImageMB)
-    reportFatalUsageError("Error reading the kernel image file");
+  if (!ImageBufferOrErr)
+    return createErr("failed to read the kernel image file");
+  auto ImageBuffer = std::move(ImageBufferOrErr.get());
 
+  // Prepare the device image and binary descriptor.
   __tgt_device_image DeviceImage;
-  DeviceImage.ImageStart = const_cast<char *>(ImageMB.get()->getBufferStart());
-  DeviceImage.ImageEnd = const_cast<char *>(ImageMB.get()->getBufferEnd());
+  DeviceImage.ImageStart = const_cast<char *>(ImageBuffer->getBufferStart());
+  DeviceImage.ImageEnd = const_cast<char *>(ImageBuffer->getBufferEnd());
   DeviceImage.EntriesBegin = &OffloadEntries[0];
   DeviceImage.EntriesEnd = &OffloadEntries[OffloadEntries.size() - 1] + 1;
 
@@ -167,81 +244,89 @@ int main(int argc, char **argv) {
   Desc.HostEntriesEnd = &OffloadEntries[OffloadEntries.size() - 1] + 1;
   Desc.DeviceImages = &DeviceImage;
 
-  auto DeviceIdJson = JsonKernelInfo->getAsObject()->getInteger("DeviceId");
-  // TODO: Print warning if the user overrides the device id in the json file.
-  int32_t DeviceId = (DeviceIdOpt > -1 ? DeviceIdOpt : DeviceIdJson.value());
-
-  // TODO: do we need requires?
-  //__tgt_register_requires(/*Flags=*/1);
-
+  // Register the image and the offload entries.
   __tgt_register_lib(&Desc);
 
   int Rc = __tgt_activate_record_replay(
-      DeviceId, VAllocSize, VAllocAddr, /*IsRecord=*/false, VerifyOpt,
+      DeviceId, VAllocSize, VAllocAddr, /*IsRecord=*/false,
+      VerifyOpt || SaveOutputOpt,
       /*EmitReport=*/false, Directory.c_str());
   if (Rc != OMP_TGT_SUCCESS)
-    reportFatalUsageError("Error activating record replay");
+    return createErr("failed to activate record replay");
 
+  // Load the record input file.
   Filepath.replace_extension("record_input");
-  ErrorOr<std::unique_ptr<MemoryBuffer>> DeviceMemoryMB =
+  auto RecordInputBufferOrErr =
       MemoryBuffer::getFile(Filepath.string(), /*isText=*/false,
                             /*RequiresNullTerminator=*/false);
-
-  if (!DeviceMemoryMB)
-    reportFatalUsageError("Error reading the kernel record input file");
+  if (!RecordInputBufferOrErr)
+    return createErr("failed to read the kernel record input file");
+  auto RecordInputBuffer = std::move(RecordInputBufferOrErr.get());
 
   // On AMD for currently unknown reasons we cannot copy memory mapped data to
   // device. This is a work-around.
-  uint8_t *RecordedData = new uint8_t[DeviceMemoryMB.get()->getBufferSize()];
+  uint8_t *RecordedData = new uint8_t[RecordInputBuffer->getBufferSize()];
   std::memcpy(RecordedData,
-              const_cast<char *>(DeviceMemoryMB.get()->getBuffer().data()),
-              DeviceMemoryMB.get()->getBufferSize());
+              const_cast<char *>(RecordInputBuffer->getBuffer().data()),
+              RecordInputBuffer->getBufferSize());
 
   KernelReplayOutcomeTy ReplayOutcome;
-
   Rc = __tgt_target_kernel_replay(
       /*Loc=*/nullptr, DeviceId, OffloadEntries[0].Address,
-      (char *)RecordedData, DeviceMemoryMB.get()->getBufferSize(),
+      (char *)RecordedData, RecordInputBuffer->getBufferSize(),
       NumGlobals ? &OffloadEntries[1] : nullptr, NumGlobals, TgtArgs.data(),
-      TgtArgOffsets.data(), NumArgs.value(), NumTeams, NumThreads,
-      SharedMemorySize, LoopTripCount.value(), &ReplayOutcome);
+      TgtArgOffsets.data(), NumArgs, NumTeams, NumThreads, SharedMemorySize,
+      LoopTripCount, &ReplayOutcome);
   if (Rc != OMP_TGT_SUCCESS)
-    reportFatalUsageError("Error replaying kernel");
-
-  int ErrorDetected = 0;
-  if (VerifyOpt) {
-    if (ReplayOutcome.OutputFilepath.empty())
-      reportFatalUsageError("Replay output file was not generated");
-
-    Filepath.replace_extension("record_output");
-    ErrorOr<std::unique_ptr<MemoryBuffer>> OriginalOutputMB =
-        MemoryBuffer::getFile(Filepath.string(),
-                              /*isText=*/false,
-                              /*RequiresNullTerminator=*/false);
-    if (!OriginalOutputMB)
-      reportFatalUsageError(
-          "Error reading the kernel record output file. Make sure "
-          "LIBOMPTARGET_RECORD_OUTPUT is set when recording");
-
-    ErrorOr<std::unique_ptr<MemoryBuffer>> ReplayOutputMB =
-        MemoryBuffer::getFile(ReplayOutcome.OutputFilepath.c_str(),
-                              /*isText=*/false,
-                              /*RequiresNullTerminator=*/false);
-    if (!ReplayOutputMB)
-      reportFatalUsageError("Error reading the kernel replay output file");
-
-    StringRef OriginalOutput = OriginalOutputMB.get()->getBuffer();
-    StringRef ReplayOutput = ReplayOutputMB.get()->getBuffer();
-    if (OriginalOutput == ReplayOutput) {
-      outs() << "[llvm-omp-kernel-replay] Replay device memory verified!\n";
-    } else {
-      ErrorDetected = 1;
-      outs() << "[llvm-omp-kernel-replay] Replay device memory failed to "
-                "verify!\n";
-    }
-  }
+    return createErr("failed to replay kernel");
 
   delete[] RecordedData;
 
-  return ErrorDetected;
+  // Nothing else to do if no verification was required.
+  if (!VerifyOpt) {
+    outs() << TOOL_PREFIX << "Replay finished (verification skipped)\n";
+    return Error::success();
+  }
+
+  if (ReplayOutcome.OutputFilepath.empty())
+    return createErr("replay output file was not generated");
+
+  // Load the record output file.
+  Filepath.replace_extension("record_output");
+  auto RecordOutputBufferOrErr =
+      MemoryBuffer::getFile(Filepath.string(),
+                            /*isText=*/false,
+                            /*RequiresNullTerminator=*/false);
+  if (!RecordOutputBufferOrErr)
+    return createErr("failed to read the kernel record output file");
+
+  // Load the replay output file.
+  auto ReplayOutputBufferOrErr =
+      MemoryBuffer::getFile(ReplayOutcome.OutputFilepath.c_str(),
+                            /*isText=*/false,
+                            /*RequiresNullTerminator=*/false);
+  if (!ReplayOutputBufferOrErr)
+    return createErr("failed to read the kernel replay output file");
+
+  // Compare record and replay outputs to verify they match.
+  StringRef RecordOutput = RecordOutputBufferOrErr.get()->getBuffer();
+  StringRef ReplayOutput = ReplayOutputBufferOrErr.get()->getBuffer();
+  if (RecordOutput != ReplayOutput)
+    return createErr("replay device memory failed to verify");
+
+  // Sucessfully verified.
+  outs() << TOOL_PREFIX << "Replay device memory verified\n";
+  return Error::success();
+}
+
+int main(int Argc, char **Argv) {
+  cl::HideUnrelatedOptions(ReplayOptions);
+  cl::ParseCommandLineOptions(Argc, Argv, TOOL_NAME "\n");
+
+  if (auto Err = replayKernel()) {
+    errs() << TOOL_PREFIX << " Error: " << llvm::toString(std::move(Err))
+           << "\n";
+    return 1;
+  }
+  return 0;
 }

>From 7151bb81ba8dc758896df66571fc2500ae1bc472 Mon Sep 17 00:00:00 2001
From: Kevin Sala <salapenades1 at llnl.gov>
Date: Sat, 18 Apr 2026 03:36:48 -0700
Subject: [PATCH 2/2] [offload] Verify replay config of teams/threads is
 allowed

---
 .../common/src/RecordReplay.cpp               | 14 ++++++--
 .../kernelreplay/llvm-omp-kernel-replay.cpp   | 33 +++++++++++++++++--
 2 files changed, 42 insertions(+), 5 deletions(-)

diff --git a/offload/plugins-nextgen/common/src/RecordReplay.cpp b/offload/plugins-nextgen/common/src/RecordReplay.cpp
index d4ae33c040344..6d95ed1debe54 100644
--- a/offload/plugins-nextgen/common/src/RecordReplay.cpp
+++ b/offload/plugins-nextgen/common/src/RecordReplay.cpp
@@ -206,14 +206,24 @@ Error NativeRecordReplayTy::recordDescImpl(
   json::Object JsonKernelInfo;
   JsonKernelInfo["Name"] = Kernel.getName();
   JsonKernelInfo["NumArgs"] = KernelArgs.NumArgs;
-  JsonKernelInfo["NumTeamsClause"] = Instance.NumTeams;
-  JsonKernelInfo["ThreadLimitClause"] = Instance.NumThreads;
+  JsonKernelInfo["NumTeams"] = Instance.NumTeams;
+  JsonKernelInfo["NumThreads"] = Instance.NumThreads;
   JsonKernelInfo["SharedMemorySize"] = Instance.SharedMemorySize;
   JsonKernelInfo["LoopTripCount"] = KernelArgs.Tripcount;
   JsonKernelInfo["DeviceId"] = Device.getDeviceId();
   JsonKernelInfo["VAllocAddr"] = (intptr_t)StartAddr;
   JsonKernelInfo["VAllocSize"] = TotalSize;
 
+  json::Array JsonTeamsLimits;
+  JsonTeamsLimits.push_back(KernelArgs.NumTeams[0]);
+  JsonTeamsLimits.push_back(KernelArgs.NumTeams[0]);
+  JsonKernelInfo["TeamsLimits"] = json::Value(std::move(JsonTeamsLimits));
+
+  json::Array JsonThreadsLimits;
+  JsonThreadsLimits.push_back(KernelArgs.ThreadLimit[0] > 0);
+  JsonThreadsLimits.push_back(KernelArgs.ThreadLimit[0]);
+  JsonKernelInfo["ThreadsLimits"] = json::Value(std::move(JsonThreadsLimits));
+
   json::Array JsonArgPtrs;
   for (uint32_t I = 0; I < KernelArgs.NumArgs; ++I)
     JsonArgPtrs.push_back((intptr_t)(*(void **)LaunchParams.Ptrs[I]));
diff --git a/offload/tools/kernelreplay/llvm-omp-kernel-replay.cpp b/offload/tools/kernelreplay/llvm-omp-kernel-replay.cpp
index 9548a3c33e25c..80b39637c5882 100644
--- a/offload/tools/kernelreplay/llvm-omp-kernel-replay.cpp
+++ b/offload/tools/kernelreplay/llvm-omp-kernel-replay.cpp
@@ -126,9 +126,9 @@ Error replayKernel() {
 
   // Retrieve the values from the JSON file.
   uint32_t NumTeams, NumThreads, SharedMemorySize, DeviceId, NumArgs;
-  if (auto Err = getInteger(JsonObj, "NumTeamsClause", NumTeams))
+  if (auto Err = getInteger(JsonObj, "NumTeams", NumTeams))
     return Err;
-  if (auto Err = getInteger(JsonObj, "ThreadLimitClause", NumThreads))
+  if (auto Err = getInteger(JsonObj, "NumThreads", NumThreads))
     return Err;
   if (auto Err = getInteger(JsonObj, "SharedMemorySize", SharedMemorySize))
     return Err;
@@ -156,8 +156,35 @@ Error replayKernel() {
   NumThreads = NumThreadsOpt > 0 ? NumThreadsOpt : NumThreads;
   DeviceId = DeviceIdOpt >= 0 ? DeviceIdOpt : DeviceId;
 
+  // Retrieve the teams and threads limits (min and max).
+  SmallVector<uint32_t> TeamsLimits;
+  auto Err = processIntegerArray(JsonObj, "TeamsLimits", [&](uint64_t Val) {
+    TeamsLimits.push_back(static_cast<uint32_t>(Val));
+  });
+  if (Err)
+    return Err;
+
+  SmallVector<uint32_t> ThreadsLimits;
+  Err = processIntegerArray(JsonObj, "ThreadsLimits", [&](uint64_t Val) {
+    ThreadsLimits.push_back(static_cast<uint32_t>(Val));
+  });
+  if (Err)
+    return Err;
+
+  if (TeamsLimits.size() != 2 || ThreadsLimits.size() != 2)
+    return createErr("TeamsLimits and ThreadsLimits must have a min and max");
+
+  // If the limits were specified, verify the selected values are valid.
+  if (TeamsLimits[0] > 0 &&
+      (NumTeams < TeamsLimits[0] || NumTeams > TeamsLimits[1]))
+    return createErr("number of teams is out of the allowed limits");
+  if (ThreadsLimits[0] > 0 &&
+      (NumThreads < ThreadsLimits[0] || NumThreads > ThreadsLimits[1]))
+    return createErr("number of threads is out of the allowed limits");
+
+  // Retrieve the arguments of the kernel.
   SmallVector<void *> TgtArgs;
-  auto Err = processIntegerArray(JsonObj, "ArgPtrs", [&](uint64_t Val) {
+  Err = processIntegerArray(JsonObj, "ArgPtrs", [&](uint64_t Val) {
     TgtArgs.push_back(reinterpret_cast<void *>(Val));
   });
   if (Err)



More information about the llvm-commits mailing list