feat(rocjitsu): detect SGPR write-after-write races

emulation/rocjitsu/docs/race-detector.md

@@ -148,8 +148,10 @@ races. Some examples:
 - **VGPR races**: a vector register is read or overwritten by an instruction
   before a pending global or LDS load has completed (`s_waitcnt vmcnt` /
   `s_waitcnt lgkmcnt` insufficient).
-- **SGPR races**: a scalar register is read before a pending scalar load has
-  completed (`s_waitcnt lgkmcnt` insufficient).
+- **Scalar-register races**: an SGPR or TTMP is read or overwritten before a
+  pending scalar-memory load has completed. A later scalar load to the same
+  destination is also checked because scalar-memory results can complete out
+  of order.
 - **LDS races**: an LDS byte is read by one wave while another wave has an
   outstanding write to the same byte, or written by one wave while another
   wave has an outstanding read of the same byte, without an intervening
@@ -183,9 +185,15 @@ operations are in flight. When an instruction in the emulator accesses an LDS
 byte, there is a check to see what memory events are still in flight that
 read/write that byte, from the perspective of the accessing thread. In this way,
 RAW (read-after-write) and WAR (write-after-read) hazards can be detected.
-For VGPRs, the detector also flags WAW when an instruction write can be
-clobbered by a pending asynchronous load. Similar logic applies for VGPR and
-SGPR reads.
+For VGPRs and scalar registers, the detector also flags WAW when an instruction
+write can be clobbered by a pending asynchronous load. The detector also reports
+WAW between scalar loads targeting the same SGPR or TTMP.
+
+On architectures where scalar-memory and data-share operations use a combined
+`lgkmcnt`, a nonzero partial wait cannot identify which scalar destination has
+completed. Those scalar destinations remain pending until `lgkmcnt(0)`. The
+same partial wait can still retire older data-share operations that are provably
+complete from their in-order completion rule.
 
 **LDS race detection** uses coarse-grained counters (one per 16-byte chunk) for
 fast-path checks, with interval-based overlap scanning as a fallback. Live
@@ -298,8 +306,10 @@ ctest --test-dir $BUILD_DIR -R "RaceTest"
   races, races between dispatches, or host-device synchronization issues.
 
 - **Limited WAW detection**: VGPR WAW is limited to synchronous instruction
-  writes that overlap a pending asynchronous load. WAW between two asynchronous
-  VGPR writers and LDS WAW are not currently reported.
+  writes that overlap a pending asynchronous load. Scalar-register WAW covers
+  instruction writes and scalar loads that overlap pending scalar-memory
+  destinations. WAW between two asynchronous VGPR writers and LDS WAW are not
+  currently reported.
 
 - **Conservative DPP/SDWA write masks**: WAW precision depends on the execution
   plugin's instruction-write lane and byte masks. DPP destinations can currently

emulation/rocjitsu/lib/rocjitsu/src/rocjitsu/vm/plugins/race_detector/core/wave_race_state.cpp

@@ -60,6 +60,8 @@ void WaveRaceState::registerScalarLoad(uint64_t pc, RegisterRef destination, uin
   if (destination.index >= limit || destination.width > limit - destination.index)
     return;
 
+  checkScalarWrite(destination);
+
   std::vector<uint32_t> registers(destination.width);
   for (uint32_t i = 0; i < destination.width; ++i)
     registers[i] = destination.index + i;
@@ -297,7 +299,7 @@ bool WaveRaceState::isOutstandingFromVgpr(int lane, int reg) const {
   return false;
 }
 
-void WaveRaceState::checkScalarRead(RegisterRef ref) const {
+void WaveRaceState::checkScalarAccess(RegisterRef ref, bool isWrite) const {
   const std::vector<std::vector<EventId>> *events = nullptr;
   const std::vector<int> *counts = nullptr;
   RaceViolation::Space space;
@@ -318,16 +320,28 @@ void WaveRaceState::checkScalarRead(RegisterRef ref) const {
 
   if (ref.width == 0 || ref.index >= events->size() || ref.width > events->size() - ref.index)
     return;
+  std::vector<EventId> reportedEvents;
   for (uint32_t offset = 0; offset < ref.width; ++offset) {
     const uint32_t index = ref.index + offset;
     if ((*counts)[index] == 0)
       continue;
     for (EventId eid : (*events)[index]) {
-      if (detector->events().type(eid) == type)
-        detector->getRaceHandler()({space, static_cast<int>(index), waveId.value, -1, false,
-                                    detector->getWorkgroupId(), eid});
+      if (detector->events().type(eid) != type ||
+          std::find(reportedEvents.begin(), reportedEvents.end(), eid) != reportedEvents.end())
+        continue;
+      reportedEvents.push_back(eid);
+      detector->getRaceHandler()({space, static_cast<int>(index), waveId.value, -1, isWrite,
+                                  detector->getWorkgroupId(), eid});
     }
   }
 }
 
+void WaveRaceState::checkScalarRead(RegisterRef ref) const {
+  checkScalarAccess(ref, /*isWrite=*/false);
+}
+
+void WaveRaceState::checkScalarWrite(RegisterRef ref) const {
+  checkScalarAccess(ref, /*isWrite=*/true);
+}
+
 } // namespace rocjitsu::plugins::race_detector

emulation/rocjitsu/lib/rocjitsu/src/rocjitsu/vm/plugins/race_detector/core/wave_race_state.h

@@ -94,6 +94,10 @@ class WaveRaceState {
   /// violation (outstanding scalar load targeting the same register).
   void checkScalarRead(RegisterRef reg) const;
 
+  /// Check a scalar instruction write for conflicts with pending scalar loads.
+  /// A wide write reports each conflicting memory event at most once.
+  void checkScalarWrite(RegisterRef reg) const;
+
   /// True if any outstanding global/LDS store reads from the given VGPR lane.
   bool isOutstandingFromVgpr(int lane, int reg) const;
 
@@ -120,6 +124,7 @@ class WaveRaceState {
                                   uint64_t execMask, uint8_t byteMask, IntervalSet ldsIntervals,
                                   amdgpu::WaitCounterType waitCounterType);
   void retireEventRegisters(EventId);
+  void checkScalarAccess(RegisterRef reg, bool isWrite) const;
 
   template <typename Pred> void resolveWaitCnt(int limit, Pred isTargetType);
   void applyWaitCounter(amdgpu::WaitCounterType type, int threshold);

emulation/rocjitsu/lib/rocjitsu/src/rocjitsu/vm/plugins/race_detector/plugin.cpp

@@ -378,6 +378,12 @@ void RaceDetectorPlugin::onAmdgpuReadScalarRegister(const amdgpu::Wavefront *wf,
   s->race_state->checkScalarRead(reg);
 }
 
+void RaceDetectorPlugin::onAmdgpuWriteScalarRegister(const amdgpu::Wavefront *wf, RegisterRef reg) {
+  auto *s = get_state(wf);
+  assert(s && s->race_state);
+  s->race_state->checkScalarWrite(reg);
+}
+
 void RaceDetectorPlugin::onAmdgpuBeforeExecuteInstruction(uint64_t pc, const Instruction &inst,
                                                           amdgpu::Wavefront &wf) {
   auto *s = get_state(wf);

emulation/rocjitsu/lib/rocjitsu/src/rocjitsu/vm/plugins/race_detector/plugin.h

@@ -113,6 +113,8 @@ class RaceDetectorPlugin : public ExecutionPlugin {
 
   void onAmdgpuReadScalarRegister(const amdgpu::Wavefront *wf, RegisterRef reg) override;
 
+  void onAmdgpuWriteScalarRegister(const amdgpu::Wavefront *wf, RegisterRef reg) override;
+
   void onAmdgpuBeforeExecuteInstruction(uint64_t pc, const Instruction &inst,
                                         amdgpu::Wavefront &wf) override;
 

emulation/rocjitsu/tests/execution_plugin_test.cpp

@@ -2951,6 +2951,35 @@ TEST(RaceDetectorPluginTest, ScalarLoadToTtmpReportsReadBeforeWait) {
   EXPECT_NE(sink.str().find("type=TTMP"), std::string::npos);
 }
 
+TEST(RaceDetectorPluginTest, ScalarLoadToTtmpReportsWriteBeforeWait) {
+  PluginFixture f(/*num_wf_slots=*/1);
+  PluginSinkConfig sink_config;
+  StringSink &sink = sink_config.emplace<StringSink>();
+  f.plugin_group_ = std::make_shared<ExecutionPluginGroup>(std::move(sink_config));
+  ASSERT_TRUE(f.plugin_group_->add(std::make_unique<RaceDetectorPlugin>()));
+  f.soc->set_plugin_group(f.plugin_group_);
+  f.plugin_group_->onInit();
+
+  auto *cu = f.cu();
+  auto *wf = cu->dispatch_wf(/*wg_id=*/0, /*pc=*/0, /*sgprs=*/104, /*vgprs=*/256);
+  ASSERT_NE(wf, nullptr);
+  wf->set_exec(1u);
+  std::array<amdgpu::Wavefront *, 1> waves{wf};
+  f.plugin_group_->onAmdgpuWorkgroupDispatched(
+      /*dispatch_id=*/1, /*wg_id=*/0, /*physical_vgpr_count=*/256,
+      /*physical_sgpr_count=*/104, waves);
+
+  auto state = std::make_unique<ScalarMemState>();
+  state->dst_register = {ScalarRegisterStorage::TTMP, 0, 1};
+  state->is_load = true;
+  TestMemoryInstruction load(std::move(state));
+  f.plugin_group_->onAmdgpuRouteMemoryInstruction(load, *wf);
+
+  RegisterAccess(*wf).write_ttmp(0, 0x12345678u);
+  EXPECT_NE(sink.str().find("type=TTMP"), std::string::npos);
+  EXPECT_NE(sink.str().find("access=write"), std::string::npos);
+}
+
 TEST(RaceDetectorPluginTest, ScalarLoadToTtmpHonorsSplitKmcntWait) {
   PluginFixture f(/*num_wf_slots=*/1, /*arch=*/"rdna4", /*wavefront_size=*/32,
                   /*sgprs_per_wf=*/128);

emulation/rocjitsu/tests/race-detector/CMakeLists.txt

@@ -15,13 +15,23 @@ set(RJ_RACE_TEST_SUPPORT_HEADERS
     ${CMAKE_CURRENT_SOURCE_DIR}/race_log_expectation.hpp
     ${CMAKE_CURRENT_SOURCE_DIR}/race_test_support.hpp
 )
-rj_add_hip_test(
+
+function(rj_add_hip_race_test TEST_NAME)
+    rj_add_hip_test(${TEST_NAME} ${ARGN})
+    # The race plugin is loaded at runtime, so no link dependency pulls its
+    # shared object into focused builds of the HIP race-test targets.
+    if(TARGET rocjitsu_plugin_race_so)
+        add_dependencies(${TEST_NAME}_target rocjitsu_plugin_race_so)
+    endif()
+endfunction()
+
+rj_add_hip_race_test(
     hip_race_tests_gfx950
     ARCH gfx950
     SOURCE hip_race_gfx950_test.hip
     DEPENDENCIES ${RJ_RACE_TEST_SUPPORT_HEADERS}
 )
-rj_add_hip_test(
+rj_add_hip_race_test(
     hip_race_tests_gfx1151
     ARCH gfx1151
     SOURCE hip_race_gfx1151_test.hip
@@ -132,6 +142,10 @@ rj_add_gfx950_race_test_case(vgpr_waitcnt)
 rj_add_gfx950_race_test_case(vgpr_waitcnt_race)
 rj_add_gfx950_race_test_case(sgpr_waitcnt)
 rj_add_gfx950_race_test_case(sgpr_waitcnt_race)
+rj_add_gfx950_race_test_case(sgpr_waw_load_then_mov)
+rj_add_gfx950_race_test_case(sgpr_waw_load_then_mov_race)
+rj_add_gfx950_race_test_case(sgpr_waw_load_then_load)
+rj_add_gfx950_race_test_case(sgpr_waw_load_then_load_race)
 rj_add_gfx950_race_test_case(lds_cross_wave)
 rj_add_gfx950_race_test_case(lds_cross_wave_race)
 rj_add_gfx950_race_test_case(lds_same_wave_order)

emulation/rocjitsu/tests/race-detector/hip_race_gfx950_test.hip

@@ -128,6 +128,103 @@ TEST_F(Gfx950RaceTest, sgpr_waitcnt_race) {
   ExpectRace(kSgprWaitcntRaceExpectation);
 }
 
+// --- SGPR write-after-write ---
+
+__global__ void sgpr_waw_load_then_mov_kernel(const int *src, int *dst) {
+  int val;
+  asm volatile(
+      "s_load_dword %0, %1, 0x0\n"
+      "s_waitcnt lgkmcnt(0)\n"
+      "s_mov_b32 %0, 7\n"
+      : "=&s"(val) : "s"(src) : "memory"
+  );
+  dst[threadIdx.x] = val;
+}
+
+static constexpr RaceExpectation kSgprWawLoadThenMovExpectation{
+    .kernel = "sgpr_waw_load_then_mov_race_kernel",
+    .type = "SGPR",
+    .access = "write",
+    .wave = 0,
+    .producer = "s_load_dword",
+    .consumer = "s_mov_b32",
+};
+
+__global__ void sgpr_waw_load_then_mov_race_kernel(const int *src, int *dst) {
+  int val;
+  asm volatile(
+      "s_load_dword %0, %1, 0x0\n"
+      // BUG: no wait before a scalar instruction overwrites the load destination.
+      "s_mov_b32 %0, 7\n"
+      : "=&s"(val) : "s"(src)
+  );
+  asm volatile("s_waitcnt lgkmcnt(0)\n" ::: "memory");
+  dst[threadIdx.x] = val;
+}
+
+static constexpr RaceExpectation kSgprWawLoadThenLoadExpectation{
+    .kernel = "sgpr_waw_load_then_load_race_kernel",
+    .type = "SGPR",
+    .access = "write",
+    .wave = 0,
+    .producer = "s_load_dword",
+    .consumer = "s_load_dword",
+};
+
+__global__ void sgpr_waw_load_then_load_kernel(const int *src0, const int *src1, int *dst) {
+  int val;
+  asm volatile(
+      "s_load_dword %0, %1, 0x0\n"
+      "s_waitcnt lgkmcnt(0)\n"
+      "s_load_dword %0, %2, 0x0\n"
+      : "=&s"(val) : "s"(src0), "s"(src1) : "memory"
+  );
+  asm volatile("s_waitcnt lgkmcnt(0)\n" ::: "memory");
+  dst[threadIdx.x] = val;
+}
+
+__global__ void sgpr_waw_load_then_load_race_kernel(const int *src0, const int *src1, int *dst) {
+  int val;
+  asm volatile(
+      "s_load_dword %0, %1, 0x0\n"
+      // BUG: scalar-memory reads can return out of order.
+      "s_load_dword %0, %2, 0x0\n"
+      : "=&s"(val) : "s"(src0), "s"(src1)
+  );
+  asm volatile("s_waitcnt lgkmcnt(0)\n" ::: "memory");
+  dst[threadIdx.x] = val;
+}
+
+TEST_F(Gfx950RaceTest, sgpr_waw_load_then_mov) {
+  auto *src = allocWithData<int>(64), *dst = alloc<int>(64);
+  sgpr_waw_load_then_mov_kernel<<<1, 64>>>(src, dst);
+  sync();
+  ExpectNoRace();
+}
+
+TEST_F(Gfx950RaceTest, sgpr_waw_load_then_mov_race) {
+  auto *src = allocWithData<int>(64), *dst = alloc<int>(64);
+  sgpr_waw_load_then_mov_race_kernel<<<1, 64>>>(src, dst);
+  sync();
+  ExpectRace(kSgprWawLoadThenMovExpectation);
+}
+
+TEST_F(Gfx950RaceTest, sgpr_waw_load_then_load) {
+  auto *src0 = allocWithData<int>(64), *src1 = allocWithData<int>(64);
+  auto *dst = alloc<int>(64);
+  sgpr_waw_load_then_load_kernel<<<1, 64>>>(src0, src1, dst);
+  sync();
+  ExpectNoRace();
+}
+
+TEST_F(Gfx950RaceTest, sgpr_waw_load_then_load_race) {
+  auto *src0 = allocWithData<int>(64), *src1 = allocWithData<int>(64);
+  auto *dst = alloc<int>(64);
+  sgpr_waw_load_then_load_race_kernel<<<1, 64>>>(src0, src1, dst);
+  sync();
+  ExpectRace(kSgprWawLoadThenLoadExpectation);
+}
+
 // --- LDS cross-wave ---
 
 __attribute__((amdgpu_flat_work_group_size(128, 128)))