[rocjitsu] Fix ds_read_b64_tr_b16 lane layout in gfx950 simulator

emulation/rocjitsu/lib/rocjitsu/src/rocjitsu/isa/arch/amdgpu/shared/ds_transpose.h

@@ -12,7 +12,9 @@
 /// matrix instructions.
 ///
 /// TR_B8 uses byte-level transpose with groups of 4 consecutive lanes.
-/// TR_B16 transposes 16-bit elements within groups of 8 consecutive lanes.
+/// TR_B16 on a wave64 B64 load (CDNA4 ds_read_b64_tr_b16) is a 4x16-lane
+/// halfword transpose with stride-4 source lanes; other TR_B16 loads
+/// transpose 16-bit elements within groups of 8 consecutive lanes.
 
 #include "rocjitsu/vm/amdgpu/mem_state.h"
 
@@ -71,8 +73,13 @@ inline void transpose_b64(std::vector<uint8_t> &response_data, uint32_t num_elem
 
 /// @brief TR_B16: 16-bit element transpose.
 ///
-/// Groups of 8 source lanes, each reading 8 halfwords. Destination lane
-/// `group_start + halfword_index` receives that halfword from each source lane.
+/// Wave64 B64 (CDNA4 ds_read_b64_tr_b16): 4x16-lane halfword transpose within
+/// each 16-lane group. Destination lane `l` halfword `n` comes from source lane
+/// `((l & ~0xf) | ((l >> 2) & 3)) + 4 * n`, halfword `l & 3`.
+///
+/// Otherwise (B128): groups of 8 source lanes, each reading 8 halfwords.
+/// Destination lane `group_start + halfword_index` receives that halfword from
+/// each source lane.
 inline void transpose_b16(std::vector<uint8_t> &response_data, uint32_t num_elems,
                           uint32_t wf_size) {
   constexpr uint32_t lanes_per_half = 32;
@@ -84,6 +91,25 @@ inline void transpose_b16(std::vector<uint8_t> &response_data, uint32_t num_elem
 
   std::vector<uint8_t> output(response_data.size(), 0);
 
+  if (halfwords_per_source_lane == 4 && wf_size == 64) {
+    for (uint32_t dest_lane = 0; dest_lane < wf_size; ++dest_lane) {
+      const uint32_t source_halfword = dest_lane & 3;
+      const uint32_t source_base = (dest_lane & ~0xfu) | ((dest_lane >> 2) & 3);
+      for (uint32_t halfword_index = 0; halfword_index < halfwords_per_source_lane;
+           ++halfword_index) {
+        const uint32_t source_lane = source_base + 4 * halfword_index;
+        const uint32_t source_offset = source_lane * bytes_per_lane_total + source_halfword * 2;
+        const uint32_t dest_offset = dest_lane * bytes_per_lane_total + halfword_index * 2;
+
+        if (source_offset + 2 <= response_data.size() && dest_offset + 2 <= output.size())
+          std::memcpy(&output[dest_offset], &response_data[source_offset], 2);
+      }
+    }
+
+    response_data = std::move(output);
+    return;
+  }
+
   for (uint32_t half_index = 0; half_index < num_halves; ++half_index) {
     const uint32_t lane_base = half_index * lanes_per_half;
 

emulation/rocjitsu/tests/CMakeLists.txt

@@ -2511,6 +2511,7 @@ if(RJ_ENABLE_HIP_TESTS)
         # generic HIP smoke tests are built for a local RDNA host such as
         # gfx1201.
         rj_add_hip_test(hip_ds_transpose_acc_test ARCH gfx950)
+        rj_add_hip_test(hip_ds_read_tr_b16_sim_test ARCH gfx950)
 
         rj_add_hip_test(
             cdna4_to_cdna3_lds_hip_test

emulation/rocjitsu/tests/amdgpu_vm_test.cpp

@@ -3125,6 +3125,77 @@ TEST(DsTransposeTest, ReadB64TrB16_AccBit) {
                           << " should contain LDS data, not the VGPR sentinel";
 }
 
+constexpr uint32_t tr_b16_halfword_value(uint32_t lane, uint32_t halfword) {
+  return (0x1200u + lane * 0x11u + halfword) & 0xffffu;
+}
+
+constexpr uint32_t pack_u16_pair(uint32_t lo, uint32_t hi) {
+  return (lo & 0xffffu) | ((hi & 0xffffu) << 16);
+}
+
+// Verify the ds_read_b64_tr_b16 cross-lane layout: within each 16-lane group,
+// destination lane l halfword n comes from source lane
+// ((l & 0x30) | ((l & 0xc) >> 2)) + 4 * n, halfword (l & 3). Expected values use
+// the same reference formula as tests/dbt/cdna4_to_cdna3_lds_hip_test.cpp.
+TEST(DsTransposeTest, ReadB64TrB16_LaneLayout) {
+  VmFixture f("cdna4", 1, 10);
+  auto *snap = f.capture_halts();
+
+  constexpr uint32_t VDST = 4;
+  constexpr uint32_t TID_REG = 0;
+  constexpr uint32_t ADDR_REG = 1;
+  constexpr uint32_t DS_OP = 227; // ds_read_b64_tr_b16
+  constexpr uint32_t WAVE_SIZE = 64;
+  constexpr uint32_t BYTES_PER_LANE = 8;
+
+  // Kernel:
+  //   v_add_u32 v1, v0, v0     ; v1 = 2 * tid
+  //   v_add_u32 v1, v1, v1     ; v1 = 4 * tid
+  //   v_add_u32 v1, v1, v1     ; v1 = 8 * tid (per-lane LDS byte address)
+  //   ds_read_b64_tr_b16 v[4:5], v1
+  //   s_waitcnt lgkmcnt(0)
+  //   s_endpgm
+  using namespace enc;
+  const uint32_t code[] = {
+      v_add_u32(ADDR_REG, VGPR_SRC(TID_REG), TID_REG),
+      v_add_u32(ADDR_REG, VGPR_SRC(ADDR_REG), ADDR_REG),
+      v_add_u32(ADDR_REG, VGPR_SRC(ADDR_REG), ADDR_REG),
+      ds_lo(DS_OP),
+      ds_hi(VDST, /*data0=*/0, ADDR_REG),
+      S_WAITCNT_0,
+      S_ENDPGM,
+  };
+  uint64_t ko = f.write_kernel(0x1000, code, sizeof(code));
+
+  auto *cu = f.cu();
+  for (uint32_t lane = 0; lane < WAVE_SIZE; ++lane) {
+    const uint32_t source_lo =
+        pack_u16_pair(tr_b16_halfword_value(lane, 0), tr_b16_halfword_value(lane, 1));
+    const uint32_t source_hi =
+        pack_u16_pair(tr_b16_halfword_value(lane, 2), tr_b16_halfword_value(lane, 3));
+    cu->lds().write32(lane * BYTES_PER_LANE, source_lo);
+    cu->lds().write32(lane * BYTES_PER_LANE + 4, source_hi);
+  }
+
+  test::AqlQueue queue(f.mem(), f.cp());
+  queue.dispatch(ko, WAVE_SIZE);
+  f.engine->run();
+
+  ASSERT_EQ(snap->snapshots().size(), 1u);
+  const auto &wf = snap->snapshots().front();
+
+  for (uint32_t lane = 0; lane < WAVE_SIZE; ++lane) {
+    const uint32_t halfword = lane & 3;
+    const uint32_t source_base = (lane & 0x30) + ((lane & 0x0c) >> 2);
+    const uint32_t expected_lo = pack_u16_pair(tr_b16_halfword_value(source_base + 0, halfword),
+                                               tr_b16_halfword_value(source_base + 4, halfword));
+    const uint32_t expected_hi = pack_u16_pair(tr_b16_halfword_value(source_base + 8, halfword),
+                                               tr_b16_halfword_value(source_base + 12, halfword));
+    EXPECT_EQ(wf.vgpr(VDST, lane), expected_lo) << "lane " << lane << " v" << VDST;
+    EXPECT_EQ(wf.vgpr(VDST + 1, lane), expected_hi) << "lane " << lane << " v" << (VDST + 1);
+  }
+}
+
 // Write-through scalar stores must use the store's VMID. Two page tables map
 // the same GPU VA to different host pages; a store under VMID 7 followed by the
 // no-op writeback_all(8) must land only in VMID 7's backing.

emulation/rocjitsu/tests/hip_ds_read_tr_b16_sim_test.cpp

@@ -0,0 +1,95 @@
+// Copyright (c) 2026 Advanced Micro Devices, Inc.
+// SPDX-License-Identifier: MIT
+
+/// @file hip_ds_read_tr_b16_sim_test.cpp
+/// @brief Verify the ds_read_b64_tr_b16 cross-lane layout on gfx950.
+///
+/// One wave writes a known per-lane 4-halfword pattern to LDS with
+/// ds_write_b64, reads it back with ds_read_b64_tr_b16, and checks every lane
+/// against the 4x16-lane halfword transpose: destination lane l halfword n
+/// comes from source lane ((l & 0x30) | ((l & 0xc) >> 2)) + 4 * n, halfword
+/// (l & 3).
+
+#include <cstdint>
+#include <hip/hip_runtime.h>
+#include <vector>
+
+#include <gtest/gtest.h>
+
+int main(int argc, char **argv) {
+  ::testing::InitGoogleTest(&argc, argv);
+  int rc = RUN_ALL_TESTS();
+  (void)hipDeviceReset();
+  return rc;
+}
+
+#define HIP_ASSERT(call)                                                                           \
+  do {                                                                                             \
+    hipError_t err = (call);                                                                       \
+    ASSERT_EQ(err, hipSuccess) << "HIP error: " << hipGetErrorString(err);                         \
+  } while (0)
+
+namespace {
+
+constexpr int kWaveSize = 64;
+
+using U32x2 = uint32_t __attribute__((ext_vector_type(2)));
+
+constexpr uint32_t tr_b16_halfword_value(int thread, int halfword) {
+  return (0x1200u + static_cast<uint32_t>(thread) * 0x11u + static_cast<uint32_t>(halfword)) &
+         0xffffu;
+}
+
+constexpr uint32_t pack_u16_pair(uint32_t lo, uint32_t hi) {
+  return (lo & 0xffffu) | ((hi & 0xffffu) << 16);
+}
+
+__global__ __launch_bounds__(kWaveSize) void ds_read_tr_b16_kernel(uint32_t *out) {
+  __shared__ uint32_t lds[kWaveSize * 2];
+  const int tid = threadIdx.x;
+  const uint32_t addr = static_cast<uint32_t>(reinterpret_cast<uintptr_t>(lds)) + tid * 8;
+  U32x2 raw = {pack_u16_pair(tr_b16_halfword_value(tid, 0), tr_b16_halfword_value(tid, 1)),
+               pack_u16_pair(tr_b16_halfword_value(tid, 2), tr_b16_halfword_value(tid, 3))};
+  U32x2 value = {};
+  asm volatile("ds_write_b64 %0, %1\n"
+               "s_waitcnt lgkmcnt(0)\n"
+               :
+               : "v"(addr), "v"(raw)
+               : "memory");
+  __syncthreads();
+  asm volatile("ds_read_b64_tr_b16 %0, %1\n"
+               "s_waitcnt lgkmcnt(0)\n"
+               : "=v"(value)
+               : "v"(addr)
+               : "memory");
+  out[tid * 2 + 0] = value[0];
+  out[tid * 2 + 1] = value[1];
+}
+
+} // namespace
+
+TEST(DsReadTrB16SimTest, LaneLayout) {
+  constexpr int kOutWords = kWaveSize * 2;
+  uint32_t *device_out = nullptr;
+  HIP_ASSERT(hipMalloc(&device_out, kOutWords * sizeof(uint32_t)));
+  HIP_ASSERT(hipMemset(device_out, 0, kOutWords * sizeof(uint32_t)));
+
+  ds_read_tr_b16_kernel<<<1, kWaveSize>>>(device_out);
+  HIP_ASSERT(hipDeviceSynchronize());
+
+  std::vector<uint32_t> host_out(kOutWords);
+  HIP_ASSERT(
+      hipMemcpy(host_out.data(), device_out, kOutWords * sizeof(uint32_t), hipMemcpyDeviceToHost));
+  (void)hipFree(device_out);
+
+  for (int lane = 0; lane < kWaveSize; ++lane) {
+    const int halfword = lane & 3;
+    const int source_base = (lane & 0x30) + ((lane & 0x0c) >> 2);
+    const uint32_t expected_lo = pack_u16_pair(tr_b16_halfword_value(source_base + 0, halfword),
+                                               tr_b16_halfword_value(source_base + 4, halfword));
+    const uint32_t expected_hi = pack_u16_pair(tr_b16_halfword_value(source_base + 8, halfword),
+                                               tr_b16_halfword_value(source_base + 12, halfword));
+    EXPECT_EQ(host_out[lane * 2 + 0], expected_lo) << "lane " << lane << " low dword";
+    EXPECT_EQ(host_out[lane * 2 + 1], expected_hi) << "lane " << lane << " high dword";
+  }
+}