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

[Kevin Sala Penades via llvm-commits]

================
@@ -167,81 +271,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();
----------------
kevinsala wrote:

Will apply it to #192784

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