From c710a06ee03c2e82ebe8e201c22ae7f5c3230176 Mon Sep 17 00:00:00 2001
From: Jonathan Kim <jonathan.kim@amd.com>
Date: Fri, 21 Mar 2025 17:13:32 -0400
Subject: [PATCH] kfdtest: fix trap on wave start and end

The debugger override will set the initial request mask to the
previously set request mask so use a different mask to assert
enablement.
Trap on wave start and end also run back to back, so fix the
previous override mask check as well.

In addition, unlike instruction traps, trap on wave start and end
will not require a rewind of the program counter on wave exit.
---
 libhsakmt/tests/kfdtest/src/KFDDBGTest.cpp  | 9 ++++-----
 libhsakmt/tests/kfdtest/src/ShaderStore.cpp | 5 +++++
 2 files changed, 9 insertions(+), 5 deletions(-)

diff --git a/libhsakmt/tests/kfdtest/src/KFDDBGTest.cpp b/libhsakmt/tests/kfdtest/src/KFDDBGTest.cpp
index ab19dddc33..f9a00b9ab7 100644
--- a/libhsakmt/tests/kfdtest/src/KFDDBGTest.cpp
+++ b/libhsakmt/tests/kfdtest/src/KFDDBGTest.cpp
@@ -361,10 +361,8 @@ TEST_F(KFDDBGTest, HitTrapOnWaveStartEndEvent) {
         for (int i = 0; i < 2; i++) {
             uint32_t enableMask = !!!(i % 2) ? KFD_DBG_TRAP_MASK_TRAP_ON_WAVE_START :
                                                KFD_DBG_TRAP_MASK_TRAP_ON_WAVE_END;
-            uint32_t supportedMask = enableMask;
-            debug->SetWaveLaunchOverride(KFD_DBG_TRAP_OVERRIDE_OR,
-                                                    &enableMask,
-                                                    &supportedMask);
+            uint32_t supportedMask = enableMask, reqMask = enableMask;
+            debug->SetWaveLaunchOverride(KFD_DBG_TRAP_OVERRIDE_OR, &reqMask, &supportedMask);
 
             if (!!!(supportedMask & enableMask)) {
                 EXPECT_SUCCESS(queue.Destroy());
@@ -374,7 +372,8 @@ TEST_F(KFDDBGTest, HitTrapOnWaveStartEndEvent) {
                 goto exit;
             }
 
-            ASSERT_EQ(enableMask, 0); // previous set mask
+	    // previous set mask
+            ASSERT_EQ(reqMask, !!!(i % 2) ? 0 : KFD_DBG_TRAP_MASK_TRAP_ON_WAVE_START);
 
             Dispatch *dispatch;
             dispatch = new Dispatch(isaBuf);
diff --git a/libhsakmt/tests/kfdtest/src/ShaderStore.cpp b/libhsakmt/tests/kfdtest/src/ShaderStore.cpp
index ed3c602f21..946adbc967 100644
--- a/libhsakmt/tests/kfdtest/src/ShaderStore.cpp
+++ b/libhsakmt/tests/kfdtest/src/ShaderStore.cpp
@@ -1130,6 +1130,11 @@ const char *TrapHandlerIsa =
         s_waitcnt lgkmcnt(0)
         s_mov_b32 m0, ttmp2
         v_mov_b32 v4, ttmp1
+        .if (.amdgcn.gfx_generation_number >= 12)
+            s_getreg_b32 ttmp14, hwreg(HW_REG_EXCP_FLAG_PRIV)
+            s_and_b32 ttmp2, ttmp14, 0x300
+            s_cbranch_scc1 RESTORE_AND_EXIT
+         .endif
         /* restore and increment program counter to skip shader trap jump*/
         s_add_u32 ttmp0, ttmp0, 4
         s_addc_u32 ttmp1, ttmp1, 0