TheRock/rocm-systems/emulation/rocjitsu/tests/shared_infra_test.cpp

// Copyright (c) 2026 Advanced Micro Devices, Inc.
// SPDX-License-Identifier: MIT

/// @file shared_infra_test.cpp
/// @brief Phase B unit tests: addr_calc, MMA execution, wavefront context, CU factory.

#include "rocjitsu/isa/arch/amdgpu/cdna1/isa.h"
#include "rocjitsu/isa/arch/amdgpu/cdna1/machine_insts.h"
#include "rocjitsu/isa/arch/amdgpu/cdna1/vop1.h"
#include "rocjitsu/isa/arch/amdgpu/cdna1/vopc.h"
#include "rocjitsu/isa/arch/amdgpu/cdna2/isa.h"
#include "rocjitsu/isa/arch/amdgpu/cdna2/machine_insts.h"
#include "rocjitsu/isa/arch/amdgpu/cdna2/vop1.h"
#include "rocjitsu/isa/arch/amdgpu/cdna2/vopc.h"
#include "rocjitsu/isa/arch/amdgpu/cdna3/isa.h"
#include "rocjitsu/isa/arch/amdgpu/cdna3/machine_insts.h"
#include "rocjitsu/isa/arch/amdgpu/cdna3/vop1.h"
#include "rocjitsu/isa/arch/amdgpu/cdna3/vopc.h"
#include "rocjitsu/isa/arch/amdgpu/cdna4/isa.h"
#include "rocjitsu/isa/arch/amdgpu/cdna4/machine_insts.h"
#include "rocjitsu/isa/arch/amdgpu/cdna4/opcodes.h"
#include "rocjitsu/isa/arch/amdgpu/cdna4/vop1.h"
#include "rocjitsu/isa/arch/amdgpu/cdna4/vopc.h"
#include "rocjitsu/isa/arch/amdgpu/gfx1250/addr_calc.h"
#include "rocjitsu/isa/arch/amdgpu/gfx1250/execution_backend.h"
#include "rocjitsu/isa/arch/amdgpu/gfx1250/isa.h"
#include "rocjitsu/isa/arch/amdgpu/gfx1250/machine_insts.h"
#include "rocjitsu/isa/arch/amdgpu/gfx1250/operand.h"
#include "rocjitsu/isa/arch/amdgpu/gfx1250/operand_types.h"
#include "rocjitsu/isa/arch/amdgpu/gfx1250/vop1.h"
#include "rocjitsu/isa/arch/amdgpu/gfx1250/vopc.h"
#include "rocjitsu/isa/arch/amdgpu/isa_properties.h"
#include "rocjitsu/isa/arch/amdgpu/rdna1/isa.h"
#include "rocjitsu/isa/arch/amdgpu/rdna1/machine_insts.h"
#include "rocjitsu/isa/arch/amdgpu/rdna1/vop1.h"
#include "rocjitsu/isa/arch/amdgpu/rdna1/vopc.h"
#include "rocjitsu/isa/arch/amdgpu/rdna2/isa.h"
#include "rocjitsu/isa/arch/amdgpu/rdna2/machine_insts.h"
#include "rocjitsu/isa/arch/amdgpu/rdna2/vop1.h"
#include "rocjitsu/isa/arch/amdgpu/rdna2/vopc.h"
#include "rocjitsu/isa/arch/amdgpu/rdna3/addr_calc.h"
#include "rocjitsu/isa/arch/amdgpu/rdna3/isa.h"
#include "rocjitsu/isa/arch/amdgpu/rdna3/machine_insts.h"
#include "rocjitsu/isa/arch/amdgpu/rdna3/operand_types.h"
#include "rocjitsu/isa/arch/amdgpu/rdna3/sopk.h"
#include "rocjitsu/isa/arch/amdgpu/rdna3/vop1.h"
#include "rocjitsu/isa/arch/amdgpu/rdna3/vopc.h"
#include "rocjitsu/isa/arch/amdgpu/rdna3_5/isa.h"
#include "rocjitsu/isa/arch/amdgpu/rdna3_5/machine_insts.h"
#include "rocjitsu/isa/arch/amdgpu/rdna3_5/vop1.h"
#include "rocjitsu/isa/arch/amdgpu/rdna3_5/vopc.h"
#include "rocjitsu/isa/arch/amdgpu/rdna4/addr_calc.h"
#include "rocjitsu/isa/arch/amdgpu/rdna4/isa.h"
#include "rocjitsu/isa/arch/amdgpu/rdna4/machine_insts.h"
#include "rocjitsu/isa/arch/amdgpu/rdna4/opcodes.h"
#include "rocjitsu/isa/arch/amdgpu/rdna4/operand_types.h"
#include "rocjitsu/isa/arch/amdgpu/rdna4/vop1.h"
#include "rocjitsu/isa/arch/amdgpu/rdna4/vop3.h"
#include "rocjitsu/isa/arch/amdgpu/rdna4/vop3p.h"
#include "rocjitsu/isa/arch/amdgpu/rdna4/vopc.h"
#include "rocjitsu/isa/arch/amdgpu/shared/addr_calc_flat.h"
#include "rocjitsu/isa/arch/amdgpu/shared/addr_calc_scalar.h"
#include "rocjitsu/isa/arch/amdgpu/shared/dpp_sdwa_ops.h"
#include "rocjitsu/isa/arch/amdgpu/shared/mma_exec.h"
#include "rocjitsu/isa/decoder.h"
#include "rocjitsu/isa/instruction.h"
#include "rocjitsu/isa/isa_traits.h"
#include "rocjitsu/vm/amdgpu/compute_unit.h"
#include "rocjitsu/vm/amdgpu/gpu_memory.h"
#include "rocjitsu/vm/amdgpu/l1_scalar_cache.h"
#include "rocjitsu/vm/amdgpu/l1_vector_cache.h"
#include "rocjitsu/vm/amdgpu/l2_cache.h"

#include "simdojo/sim/simulation.h"
#include "util/bit.h"
#include "util/data_types.h"
#include "util/except.h"

#include <gtest/gtest.h>

#include "rocjitsu/vm/amdgpu/register_access.h"
#include <algorithm>
#include <array>
#include <bit>
#include <cstddef>
#include <cstdint>
#include <cstring>
#include <memory>
#include <mutex>
#include <shared_mutex>
#include <span>
#include <string>
#include <utility>

namespace {

using namespace rocjitsu;

// ---------------------------------------------------------------------------
// Concept and trait verification (compile-time)
// ---------------------------------------------------------------------------

/*
 * \NPI new ISA family: add GpuIsa<<isa>::Isa> plus the relevant trait \
 * static_asserts (HasAccVgpr, HasMonolithicWaitcnt, ...) for it here.
 */
static_assert(GpuIsa<cdna3::Isa>);
static_assert(GpuIsa<gfx1250::Isa>);
static_assert(GpuIsa<rdna4::Isa>);
static_assert(HasAccVgpr<cdna2::Isa>);
static_assert(HasAccVgpr<cdna3::Isa>);
static_assert(HasAccVgpr<cdna4::Isa>);
static_assert(!HasAccVgpr<gfx1250::Isa>);
static_assert(!HasAccVgpr<rdna4::Isa>);
static_assert(HasMonolithicWaitcnt<cdna3::Isa>);
static_assert(!HasMonolithicWaitcnt<gfx1250::Isa>);
static_assert(!HasMonolithicWaitcnt<rdna4::Isa>);
static_assert(!isa_properties(ROCJITSU_CODE_ARCH_CDNA3).supports_wgp_mode);
static_assert(isa_properties(ROCJITSU_CODE_ARCH_RDNA4).supports_wgp_mode);
static_assert(!isa_properties(ROCJITSU_CODE_ARCH_GFX1250).supports_wgp_mode);
static_assert(!isa_properties(ROCJITSU_CODE_ARCH_CDNA3).uses_ttmp_workgroup_ids);
static_assert(isa_properties(ROCJITSU_CODE_ARCH_RDNA4).uses_ttmp_workgroup_ids);
static_assert(!isa_properties(ROCJITSU_CODE_ARCH_RDNA4).uses_cluster_ttmp_workgroup_ids);
static_assert(isa_properties(ROCJITSU_CODE_ARCH_GFX1250).uses_ttmp_workgroup_ids);
static_assert(isa_properties(ROCJITSU_CODE_ARCH_GFX1250).uses_cluster_ttmp_workgroup_ids);

// RDNA3/3.5 retain monolithic S_WAITCNT (GFX11 layout).
static_assert(HasMonolithicWaitcnt<rdna3::Isa>);

// RDNA2 supports Wave64 (WF_SIZE_MAX inherited as 64).
static_assert(rdna2::Isa::WF_SIZE_MAX == 64);
static_assert(gfx1250::Isa::WF_SIZE_MAX == 32);

// CDNA1 and GFX1250 have no AccVGPRs; CDNA2/3/4 have 256.
constexpr uint32_t kCdnaAccVgprsPerWf = cdna2::Isa::MAX_ACC_VGPRS_PER_WF;
static_assert(cdna1::Isa::MAX_ACC_VGPRS_PER_WF == 0);
static_assert(kCdnaAccVgprsPerWf == 256);
static_assert(cdna3::Isa::MAX_ACC_VGPRS_PER_WF == kCdnaAccVgprsPerWf);
static_assert(cdna4::Isa::MAX_ACC_VGPRS_PER_WF == kCdnaAccVgprsPerWf);
static_assert(gfx1250::Isa::MAX_ACC_VGPRS_PER_WF == 0);

class Rdna3MemoryTestCu
    : public amdgpu::IsaExecComputeUnit<simdojo::ExecMode::FUNCTIONAL, rdna3::Isa> {
public:
  using Base = amdgpu::IsaExecComputeUnit<simdojo::ExecMode::FUNCTIONAL, rdna3::Isa>;

  Rdna3MemoryTestCu(std::string name, const amdgpu::ComputeUnitCore::Config &config,
                    amdgpu::GpuMemory *memory, amdgpu::L2Cache *l2)
      : Base(std::move(name), config, memory, l2) {
    if (l2)
      l2->set_backing_memory(memory);
    set_memory(memory);
    set_l2(l2);
  }

  void execute_and_route(Instruction *inst, amdgpu::Wavefront &wf) {
    execute_instruction(inst, wf);
    if (inst->is_memory_op())
      route_memory_inst(inst, wf);
    else
      delete inst;
  }
};

TEST(RdnaWaitcntTest, Rdna3NamedSopkWaitcntsSetFineGrainedTargets) {
  amdgpu::GpuMemory mem("rdna3_waitcnt_mem");
  amdgpu::L2Cache l2("rdna3_waitcnt_l2");

  amdgpu::ComputeUnitCore::Config cfg{};
  cfg.arch = ROCJITSU_CODE_ARCH_RDNA3;
  cfg.num_wf_slots = 1;
  cfg.sgprs_per_wf = 128;
  cfg.vgprs_per_wf = 32;
  cfg.lds_size_kb = 64;

  auto cu = amdgpu::ComputeUnitCore::create("rdna3_waitcnt_cu", cfg, &mem, &l2);
  ASSERT_NE(cu, nullptr);

  auto dispatch = [&]() -> amdgpu::Wavefront * {
    // Recycle the single scratch slot: terminate a resident wave (freeing it) so
    // the next dispatch has a free slot, mirroring s_endpgm on real hardware.
    if (cu->wf(0) && !cu->wf(0)->is_halted())
      cu->wf(0)->halt();
    auto *wf = cu->dispatch_wf(0, 0, cfg.sgprs_per_wf, cfg.vgprs_per_wf);
    EXPECT_NE(wf, nullptr);
    return wf;
  };

  {
    auto *wf = dispatch();
    ASSERT_NE(wf, nullptr);
    wf->wait_counters().increment(amdgpu::WaitCounterType::STORECNT);
    const std::array<rdna3::MachineInst, 2> words = {0xBC000000u, 0x00000000u};
    rdna3::SWaitcntVscntSopk inst(words.data());

    EXPECT_NO_THROW(inst.execute_impl(*wf));
    EXPECT_EQ(wf->wait_target().vscnt, 0);
    EXPECT_EQ(wf->state(), amdgpu::WfState::WAITCNT);
  }

  {
    auto *wf = dispatch();
    ASSERT_NE(wf, nullptr);
    wf->wait_counters().increment(amdgpu::WaitCounterType::LOADCNT);
    const std::array<rdna3::MachineInst, 2> words = {0xBC800000u, 0x00000000u};
    rdna3::SWaitcntVmcntSopk inst(words.data());

    EXPECT_NO_THROW(inst.execute_impl(*wf));
    EXPECT_EQ(wf->wait_target().vmcnt, 0);
    EXPECT_EQ(wf->state(), amdgpu::WfState::WAITCNT);
  }

  {
    auto *wf = dispatch();
    ASSERT_NE(wf, nullptr);
    wf->wait_counters().increment(amdgpu::WaitCounterType::EXPCNT);
    const std::array<rdna3::MachineInst, 2> words = {0xBD000000u, 0x00000000u};
    rdna3::SWaitcntExpcntSopk inst(words.data());

    EXPECT_NO_THROW(inst.execute_impl(*wf));
    EXPECT_EQ(wf->wait_target().expcnt, 0);
    EXPECT_EQ(wf->state(), amdgpu::WfState::WAITCNT);
  }

  {
    auto *wf = dispatch();
    ASSERT_NE(wf, nullptr);
    wf->wait_counters().increment(amdgpu::WaitCounterType::DSCNT);
    const std::array<rdna3::MachineInst, 2> words = {0xBD800000u, 0x00000000u};
    rdna3::SWaitcntLgkmcntSopk inst(words.data());

    EXPECT_NO_THROW(inst.execute_impl(*wf));
    EXPECT_EQ(wf->wait_target().lgkmcnt, 0);
    EXPECT_EQ(wf->state(), amdgpu::WfState::WAITCNT);
  }
}

TEST(UtilBitTest, IsAlignedChecksPowerOfTwoAlignment) {
  EXPECT_TRUE(util::is_aligned<uint64_t>(0x1000u, 4u));
  EXPECT_TRUE(util::is_aligned<uint32_t>(0u, 8u));
  EXPECT_FALSE(util::is_aligned<uint64_t>(0x1003u, 4u));
}

// ---------------------------------------------------------------------------
// MFMA register layout tests
// ---------------------------------------------------------------------------

TEST(MfmaExecTest, InputLocF16_4x4x4) {
  // v_mfma_f32_4x4x4f16:
  // lanes_per_block = 64 / (4 * 4) = 4, elems_per_group = 4/4 = 1.
  auto loc = amdgpu::input_loc(4, 4, 4, /*i=*/2, /*k=*/0, /*b=*/0, 32);
  EXPECT_EQ(loc.vgpr_offset, 0u);
  EXPECT_EQ(loc.lane, 2u); // b*dim + ... = 0*4 + (0/1)*4*4 + 2 = 2
  EXPECT_EQ(loc.sub_element, 0u);
}

TEST(MfmaExecTest, InputLocF16_16x16) {
  // 16x16x16 with 1 block, f16 inputs: each lane holds 16 * 2B = 32B = 8 dwords.
  // lanes_per_block = 64 / (16 * 1) = 4
  // elems_per_group = 16 / 4 = 4
  // For i=0, k=0, b=0: local=0%4=0, lane=0*16+0*16*1+0=0, per_dword=2
  // vgpr_offset = 0/2 = 0, sub_element = 0%2 = 0
  auto loc = amdgpu::input_loc(16, 16, 1, 0, 0, 0, 16);
  EXPECT_EQ(loc.vgpr_offset, 0u);
  EXPECT_EQ(loc.lane, 0u);
  EXPECT_EQ(loc.sub_element, 0u);

  // k=1: local=1, vgpr_offset = 1/2 = 0, sub_element = 1
  auto loc1 = amdgpu::input_loc(16, 16, 1, 0, 1, 0, 16);
  EXPECT_EQ(loc1.vgpr_offset, 0u);
  EXPECT_EQ(loc1.sub_element, 1u);
}

TEST(MfmaExecTest, Gfx11WmmaIu8InputLocReplicatesKAcrossHalfwaves) {
  auto g0k0 = amdgpu::gfx11_wmma_input_loc(16, 16, /*i=*/3, /*k=*/0, 8,
                                           /*lane_group=*/0);
  EXPECT_EQ(g0k0.vgpr_offset, 0u);
  EXPECT_EQ(g0k0.lane, 3u);
  EXPECT_EQ(g0k0.sub_element, 0u);

  auto g0k15 = amdgpu::gfx11_wmma_input_loc(16, 16, /*i=*/3, /*k=*/15, 8,
                                            /*lane_group=*/0);
  EXPECT_EQ(g0k15.vgpr_offset, 3u);
  EXPECT_EQ(g0k15.lane, 3u);
  EXPECT_EQ(g0k15.sub_element, 3u);

  auto g1k15 = amdgpu::gfx11_wmma_input_loc(16, 16, /*i=*/3, /*k=*/15, 8,
                                            /*lane_group=*/1);
  EXPECT_EQ(g1k15.vgpr_offset, 3u);
  EXPECT_EQ(g1k15.lane, 19u);
  EXPECT_EQ(g1k15.sub_element, 3u);
}

TEST(MfmaExecTest, Gfx11WmmaOutputLoc32PairsRowsAcrossHalfwaves) {
  auto r0 = amdgpu::gfx11_wmma_output_loc_32(amdgpu::WMMA_WAVE32, 16, 16, /*row=*/0, /*col=*/5);
  EXPECT_EQ(r0.reg, 0u);
  EXPECT_EQ(r0.lane, 5u);

  auto r1 = amdgpu::gfx11_wmma_output_loc_32(amdgpu::WMMA_WAVE32, 16, 16, /*row=*/1, /*col=*/5);
  EXPECT_EQ(r1.reg, 0u);
  EXPECT_EQ(r1.lane, 21u);

  auto r2 = amdgpu::gfx11_wmma_output_loc_32(amdgpu::WMMA_WAVE32, 16, 16, /*row=*/2, /*col=*/5);
  EXPECT_EQ(r2.reg, 1u);
  EXPECT_EQ(r2.lane, 5u);

  auto r15 = amdgpu::gfx11_wmma_output_loc_32(amdgpu::WMMA_WAVE32, 16, 16, /*row=*/15, /*col=*/5);
  EXPECT_EQ(r15.reg, 7u);
  EXPECT_EQ(r15.lane, 21u);
}

TEST(MfmaExecTest, Gfx11WmmaOutputLoc32UsesFourLaneGroupsForWave64) {
  auto r0 = amdgpu::gfx11_wmma_output_loc_32(amdgpu::WMMA_WAVE64, 16, 16, /*row=*/0, /*col=*/5);
  EXPECT_EQ(r0.reg, 0u);
  EXPECT_EQ(r0.lane, 5u);

  auto r1 = amdgpu::gfx11_wmma_output_loc_32(amdgpu::WMMA_WAVE64, 16, 16, /*row=*/1, /*col=*/5);
  EXPECT_EQ(r1.reg, 0u);
  EXPECT_EQ(r1.lane, 21u);

  auto r2 = amdgpu::gfx11_wmma_output_loc_32(amdgpu::WMMA_WAVE64, 16, 16, /*row=*/2, /*col=*/5);
  EXPECT_EQ(r2.reg, 0u);
  EXPECT_EQ(r2.lane, 37u);

  auto r3 = amdgpu::gfx11_wmma_output_loc_32(amdgpu::WMMA_WAVE64, 16, 16, /*row=*/3, /*col=*/5);
  EXPECT_EQ(r3.reg, 0u);
  EXPECT_EQ(r3.lane, 53u);

  auto r15 = amdgpu::gfx11_wmma_output_loc_32(amdgpu::WMMA_WAVE64, 16, 16, /*row=*/15, /*col=*/5);
  EXPECT_EQ(r15.reg, 3u);
  EXPECT_EQ(r15.lane, 53u);
}

TEST(MfmaExecTest, SwmmacK32InputLocUsesSparseHardwareLayout) {
  // RDNA4 K=32 SWMMAC is not the dense WMMA K=16/K=32 register layout.
  // These positions match the gfx12 builtin layout used by the silicon tests.
  auto a_h_g2s0 = amdgpu::swmmac_a_input_loc(16, 32, /*row=*/3, /*compressed_k=*/4, 16);
  EXPECT_EQ(a_h_g2s0.lane, 19u);
  EXPECT_EQ(a_h_g2s0.vgpr_offset, 0u);
  EXPECT_EQ(a_h_g2s0.sub_element, 0u);

  auto a_h_g4s1 = amdgpu::swmmac_a_input_loc(16, 32, /*row=*/3, /*compressed_k=*/9, 16);
  EXPECT_EQ(a_h_g4s1.lane, 3u);
  EXPECT_EQ(a_h_g4s1.vgpr_offset, 2u);
  EXPECT_EQ(a_h_g4s1.sub_element, 1u);

  auto b_h_k10 = amdgpu::swmmac_b_input_loc(16, 32, /*col=*/5, /*dense_k=*/10, 16);
  EXPECT_EQ(b_h_k10.lane, 21u);
  EXPECT_EQ(b_h_k10.vgpr_offset, 1u);
  EXPECT_EQ(b_h_k10.sub_element, 0u);

  auto a_fp8_g4s1 = amdgpu::swmmac_a_input_loc(16, 32, /*row=*/3, /*compressed_k=*/9, 8);
  EXPECT_EQ(a_fp8_g4s1.lane, 19u);
  EXPECT_EQ(a_fp8_g4s1.vgpr_offset, 0u);
  EXPECT_EQ(a_fp8_g4s1.sub_element, 1u);

  auto b_fp8_k10 = amdgpu::swmmac_b_input_loc(16, 32, /*col=*/5, /*dense_k=*/10, 8);
  EXPECT_EQ(b_fp8_k10.lane, 5u);
  EXPECT_EQ(b_fp8_k10.vgpr_offset, 2u);
  EXPECT_EQ(b_fp8_k10.sub_element, 2u);

  auto h_idx_g2s0 = amdgpu::swmmac_index_loc(16, 32, 16, /*row=*/3, /*compressed_k=*/4, 16);
  EXPECT_EQ(h_idx_g2s0.lane, 19u);
  EXPECT_EQ(h_idx_g2s0.local_compressed_k, 0u);

  auto fp8_idx_g4s0 = amdgpu::swmmac_index_loc(16, 32, 8, /*row=*/3, /*compressed_k=*/8, 16);
  EXPECT_EQ(fp8_idx_g4s0.lane, 19u);
  EXPECT_EQ(fp8_idx_g4s0.local_compressed_k, 0u);
}

TEST(MfmaExecTest, Gfx12Wave64WmmaLocSplitsKAcrossFourLaneGroups) {
  auto a_k8 = amdgpu::gfx12_wmma_input_loc(amdgpu::WMMA_WAVE64, 16, 16, /*i=*/3, /*k=*/8, 16);
  EXPECT_EQ(a_k8.lane, 35u);
  EXPECT_EQ(a_k8.vgpr_offset, 0u);
  EXPECT_EQ(a_k8.sub_element, 0u);

  auto a_k15 = amdgpu::gfx12_wmma_input_loc(amdgpu::WMMA_WAVE64, 16, 16, /*i=*/3, /*k=*/15, 8);
  EXPECT_EQ(a_k15.lane, 51u);
  EXPECT_EQ(a_k15.vgpr_offset, 0u);
  EXPECT_EQ(a_k15.sub_element, 3u);

  auto out32 = amdgpu::gfx12_wmma_output_loc_32(amdgpu::WMMA_WAVE64, 16, 16, /*row=*/13, /*col=*/5);
  EXPECT_EQ(out32.lane, 53u);
  EXPECT_EQ(out32.reg, 1u);

  auto out16 = amdgpu::gfx12_wmma_output_loc_16(amdgpu::WMMA_WAVE64, 16, 16, /*row=*/13, /*col=*/5);
  EXPECT_EQ(out16.lane, 53u);
  EXPECT_EQ(out16.reg, 0u);
  EXPECT_EQ(out16.sub_element, 1u);
}

TEST(MfmaExecTest, Gfx12Wave64SwmmacK32LocUsesSparseHardwareLayout) {
  auto a_h_g4s1 =
      amdgpu::swmmac_a_input_loc(amdgpu::WMMA_WAVE64, 16, 32, /*row=*/3, /*compressed_k=*/9, 16);
  EXPECT_EQ(a_h_g4s1.lane, 35u);
  EXPECT_EQ(a_h_g4s1.vgpr_offset, 0u);
  EXPECT_EQ(a_h_g4s1.sub_element, 1u);

  auto b_h_k20 =
      amdgpu::swmmac_b_input_loc(amdgpu::WMMA_WAVE64, 16, 32, /*col=*/5, /*dense_k=*/20, 16);
  EXPECT_EQ(b_h_k20.lane, 37u);
  EXPECT_EQ(b_h_k20.vgpr_offset, 2u);
  EXPECT_EQ(b_h_k20.sub_element, 0u);

  auto a_fp8_g2s0 =
      amdgpu::swmmac_a_input_loc(amdgpu::WMMA_WAVE64, 16, 32, /*row=*/3, /*compressed_k=*/4, 8);
  EXPECT_EQ(a_fp8_g2s0.lane, 35u);
  EXPECT_EQ(a_fp8_g2s0.vgpr_offset, 0u);
  EXPECT_EQ(a_fp8_g2s0.sub_element, 0u);

  auto b_fp8_k10 =
      amdgpu::swmmac_b_input_loc(amdgpu::WMMA_WAVE64, 16, 32, /*col=*/5, /*dense_k=*/10, 8);
  EXPECT_EQ(b_fp8_k10.lane, 37u);
  EXPECT_EQ(b_fp8_k10.vgpr_offset, 0u);
  EXPECT_EQ(b_fp8_k10.sub_element, 2u);

  auto h_idx_g4s1 = amdgpu::swmmac_index_loc(amdgpu::WMMA_WAVE64, 16, 32, 16, /*row=*/3,
                                             /*compressed_k=*/9, 16);
  EXPECT_EQ(h_idx_g4s1.lane, 35u);
  EXPECT_EQ(h_idx_g4s1.local_compressed_k, 1u);

  auto fp8_idx_g2s0 = amdgpu::swmmac_index_loc(amdgpu::WMMA_WAVE64, 16, 32, 8, /*row=*/3,
                                               /*compressed_k=*/4, 16);
  EXPECT_EQ(fp8_idx_g2s0.lane, 35u);
  EXPECT_EQ(fp8_idx_g2s0.local_compressed_k, 0u);
}

TEST(MfmaExecTest, Cdna3CvtFp8Bf8UsesSameFnuzDecodeAsMfma) {
  amdgpu::GpuMemory mem("cdna3_cvt_fnuz_mem");
  amdgpu::L2Cache l2("cdna3_cvt_fnuz_l2");

  amdgpu::ComputeUnitCore::Config cfg{};
  cfg.arch = ROCJITSU_CODE_ARCH_CDNA3;
  cfg.num_wf_slots = 1;
  cfg.sgprs_per_wf = 104;
  cfg.vgprs_per_wf = 32;
  cfg.lds_size_kb = 64;

  auto cu = amdgpu::ComputeUnitCore::create("cdna3_cvt_fnuz_cu", cfg, &mem, &l2);
  ASSERT_NE(cu, nullptr);

  auto *wf = cu->dispatch_wf(0, 0, cfg.sgprs_per_wf, cfg.vgprs_per_wf);
  ASSERT_NE(wf, nullptr);
  ASSERT_EQ(wf->wf_size(), 64u);
  wf->set_exec(~0ULL);

  constexpr uint32_t kSrc = 4;
  constexpr uint32_t kDstFp8 = 8;
  constexpr uint32_t kDstBf8 = 9;
  uint32_t vbase = wf->vgpr_alloc().base;
  for (uint32_t lane = 0; lane < wf->wf_size(); ++lane) {
    uint8_t byte = (lane & 1u) ? 0xC0u : 0x40u;
    cu->write_vgpr(vbase + kSrc, lane, byte);
  }

  cdna3::Vop1MachineInst raw_fp8{};
  raw_fp8.src0 = 256 + kSrc;
  raw_fp8.vdst = kDstFp8;
  cdna3::VCvtF32Fp8Vop1 cvt_fp8(reinterpret_cast<const cdna3::MachineInst *>(&raw_fp8));
  cvt_fp8.execute_impl(*wf);

  cdna3::Vop1MachineInst raw_bf8{};
  raw_bf8.src0 = 256 + kSrc;
  raw_bf8.vdst = kDstBf8;
  cdna3::VCvtF32Bf8Vop1 cvt_bf8(reinterpret_cast<const cdna3::MachineInst *>(&raw_bf8));
  cvt_bf8.execute_impl(*wf);

  for (uint32_t lane = 0; lane < wf->wf_size(); ++lane) {
    amdgpu::InputLoc loc{/*vgpr_offset=*/0, lane, /*sub_element=*/0};
    float mfma_fp8 = amdgpu::extract_fp8_fnuz(*cu, vbase + kSrc, loc);
    float cvt_fp8_value = std::bit_cast<float>(cu->read_vgpr(vbase + kDstFp8, lane));
    EXPECT_EQ(cvt_fp8_value, mfma_fp8) << "fp8 lane=" << lane;

    float mfma_bf8 = amdgpu::extract_bf8_fnuz(*cu, vbase + kSrc, loc);
    float cvt_bf8_value = std::bit_cast<float>(cu->read_vgpr(vbase + kDstBf8, lane));
    EXPECT_EQ(cvt_bf8_value, mfma_bf8) << "bf8 lane=" << lane;
  }
}

void write_packed_byte(amdgpu::ComputeUnitCore &cu, uint32_t reg, uint32_t lane, uint32_t byte,
                       uint8_t value) {
  const uint32_t shift = 8u * byte;
  const uint32_t old = cu.read_vgpr(reg, lane);
  cu.write_vgpr(reg, lane, (old & ~(0xFFu << shift)) | (static_cast<uint32_t>(value) << shift));
}

void fill_vgprs(amdgpu::ComputeUnitCore &cu, uint32_t base, uint32_t regs, uint32_t lanes,
                uint32_t value) {
  for (uint32_t reg = 0; reg < regs; ++reg)
    for (uint32_t lane = 0; lane < lanes; ++lane)
      cu.write_vgpr(base + reg, lane, value);
}

TEST(MfmaExecTest, SwmmacF32K32Fp8MatchesSparseReference) {
  amdgpu::GpuMemory gpu_mem("rdna4_swmmac_fp8_exec_mem");
  amdgpu::L2Cache l2("rdna4_swmmac_fp8_exec_l2");

  amdgpu::ComputeUnitCore::Config cfg{};
  cfg.arch = ROCJITSU_CODE_ARCH_RDNA4;
  cfg.num_wf_slots = 1;
  cfg.sgprs_per_wf = 106;
  cfg.vgprs_per_wf = 256;
  cfg.lds_size_kb = 64;

  auto cu = amdgpu::ComputeUnitCore::create("rdna4_swmmac_fp8_exec", cfg, &gpu_mem, &l2);
  ASSERT_NE(cu, nullptr);

  auto *wf = cu->dispatch_wf(0, 0, cfg.sgprs_per_wf, cfg.vgprs_per_wf);
  ASSERT_NE(wf, nullptr);
  ASSERT_EQ(wf->wf_size(), 32u);
  wf->set_exec((1ULL << wf->wf_size()) - 1ULL);

  const uint32_t vb = wf->vgpr_alloc().base;
  const uint32_t a_base = vb + 0;
  const uint32_t b_base = vb + 4;
  const uint32_t c_base = vb + 12;
  const uint32_t d_base = vb + 20;
  const uint32_t index_base = vb + 32;

  fill_vgprs(*cu, a_base, 40, wf->wf_size(), 0);

  constexpr uint32_t M = 16;
  constexpr uint32_t N = 16;
  constexpr uint32_t K = 32;
  constexpr uint32_t Groups = K / 4;
  const std::array<std::array<uint32_t, 2>, 4> pairs = {{{0u, 2u}, {1u, 3u}, {0u, 1u}, {2u, 3u}}};

  auto a_value = [](uint32_t row, uint32_t group, uint32_t slot) {
    return static_cast<float>(static_cast<int>((row + 2u * group + slot) % 5u) - 2);
  };
  auto b_value = [](uint32_t k, uint32_t col) {
    return static_cast<float>(static_cast<int>((col + 3u * k) % 5u) - 2);
  };
  auto c_value = [](uint32_t row, uint32_t col) {
    return static_cast<float>(static_cast<int>((2u * row + col) % 7u) - 3);
  };

  for (uint32_t row = 0; row < M; ++row) {
    for (uint32_t group = 0; group < Groups; ++group) {
      const auto pair = pairs[(row + group) % pairs.size()];
      const uint32_t index_lane = row + 16u * (group / 4u);
      const uint32_t field = pair[0] | (pair[1] << 2u);
      const uint32_t index_shift = 4u * (group & 3u);
      const uint32_t old_index = cu->read_vgpr(index_base, index_lane);
      cu->write_vgpr(index_base, index_lane, old_index | (field << index_shift));
      for (uint32_t slot = 0; slot < 2; ++slot) {
        const uint32_t compressed_k = 2u * group + slot;
        const auto loc = amdgpu::swmmac_a_input_loc(M, K, row, compressed_k, 8);
        write_packed_byte(*cu, a_base + loc.vgpr_offset, loc.lane, loc.sub_element,
                          util::f32_to_fp8_e4m3_rne(a_value(row, group, slot)));
      }
    }
  }

  for (uint32_t k = 0; k < K; ++k) {
    for (uint32_t col = 0; col < N; ++col) {
      const auto loc = amdgpu::swmmac_b_input_loc(N, K, col, k, 8);
      write_packed_byte(*cu, b_base + loc.vgpr_offset, loc.lane, loc.sub_element,
                        util::f32_to_fp8_e4m3_rne(b_value(k, col)));
    }
  }

  for (uint32_t row = 0; row < M; ++row) {
    for (uint32_t col = 0; col < N; ++col) {
      const auto out = amdgpu::wmma_output_loc_32(M, N, row, col);
      cu->write_vgpr(c_base + out.reg, out.lane, std::bit_cast<uint32_t>(c_value(row, col)));
    }
  }

  amdgpu::exec_swmmac_f32(*cu, M, N, K, 8, d_base, a_base, b_base, c_base, index_base, 16,
                          /*index_key=*/0, amdgpu::extract_fp8, amdgpu::extract_fp8);

  for (uint32_t row = 0; row < M; ++row) {
    for (uint32_t col = 0; col < N; ++col) {
      float ref = c_value(row, col);
      for (uint32_t group = 0; group < Groups; ++group) {
        const auto pair = pairs[(row + group) % pairs.size()];
        for (uint32_t slot = 0; slot < 2; ++slot) {
          const uint32_t k = 4u * group + pair[slot];
          const float av =
              util::fp8_e4m3_to_f32(util::f32_to_fp8_e4m3_rne(a_value(row, group, slot)));
          const float bv = util::fp8_e4m3_to_f32(util::f32_to_fp8_e4m3_rne(b_value(k, col)));
          ref += av * bv;
        }
      }
      const auto out = amdgpu::wmma_output_loc_32(M, N, row, col);
      EXPECT_EQ(cu->read_vgpr(d_base + out.reg, out.lane), std::bit_cast<uint32_t>(ref))
          << "row=" << row << " col=" << col;
    }
  }

  wf->halt();
}

TEST(MfmaExecTest, WmmaF8f6f4K128InputLocUsesPairAwareSubbyteLayouts) {
  auto ab6 = amdgpu::wmma_f8f6f4_input_loc(16, 128, /*i=*/3, /*k=*/4, 6,
                                           /*mixed_subbyte=*/false);
  EXPECT_EQ(ab6.lane, 19u);
  EXPECT_EQ(ab6.vgpr_offset, 0u);
  EXPECT_EQ(ab6.bit_offset, 0u);
  EXPECT_EQ(ab6.data_bits, 6u);

  auto ab4 = amdgpu::wmma_f8f6f4_input_loc(16, 128, /*i=*/3, /*k=*/4, 4,
                                           /*mixed_subbyte=*/false);
  EXPECT_EQ(ab4.lane, 19u);
  EXPECT_EQ(ab4.vgpr_offset, 0u);
  EXPECT_EQ(ab4.bit_offset, 0u);
  EXPECT_EQ(ab4.data_bits, 4u);

  auto mixed6 = amdgpu::wmma_f8f6f4_input_loc(16, 128, /*i=*/3, /*k=*/4, 6,
                                              /*mixed_subbyte=*/true);
  EXPECT_EQ(mixed6.lane, 3u);
  EXPECT_EQ(mixed6.vgpr_offset, 3u);
  EXPECT_EQ(mixed6.bit_offset, 0u);

  auto mixed4 = amdgpu::wmma_f8f6f4_input_loc(16, 128, /*i=*/3, /*k=*/31, 4,
                                              /*mixed_subbyte=*/true);
  EXPECT_EQ(mixed4.lane, 3u);
  EXPECT_EQ(mixed4.vgpr_offset, 3u);
  EXPECT_EQ(mixed4.bit_offset, 28u);

  auto ab8 = amdgpu::wmma_f8f6f4_input_loc(16, 128, /*i=*/3, /*k=*/4, 8,
                                           /*mixed_subbyte=*/false);
  EXPECT_EQ(ab8.lane, 19u);
  EXPECT_EQ(ab8.vgpr_offset, 0u);
  EXPECT_EQ(ab8.sub_element, 0u);

  EXPECT_EQ(amdgpu::wmma_f8f6f4_scale_byte(/*k=*/4, /*data_bits=*/6,
                                           /*mixed_pair=*/false, /*scale16=*/false),
            1u);
  EXPECT_EQ(amdgpu::wmma_f8f6f4_scale_byte(/*k=*/32, /*data_bits=*/6,
                                           /*mixed_pair=*/false, /*scale16=*/true),
            1u);
  EXPECT_EQ(amdgpu::wmma_f8f6f4_scale_byte(/*k=*/64, /*data_bits=*/4,
                                           /*mixed_pair=*/false, /*scale16=*/true),
            4u);
  EXPECT_EQ(amdgpu::wmma_f8f6f4_scale_byte(/*k=*/32, /*data_bits=*/4,
                                           /*mixed_pair=*/true, /*scale16=*/false),
            1u);
  EXPECT_EQ(amdgpu::wmma_f8f6f4_scale_byte(/*k=*/36, /*data_bits=*/4,
                                           /*mixed_pair=*/true, /*scale16=*/true),
            3u);
}

TEST(MfmaExecTest, WmmaF4_32x16x128UsesSiliconGroundedALayoutAndScaleLane) {
  auto row0 = amdgpu::wmma_a_input_loc(32, 128, /*row=*/0, /*k=*/0, 4, 4);
  EXPECT_EQ(row0.lane, 0u);
  EXPECT_EQ(row0.vgpr_offset, 0u);
  EXPECT_EQ(row0.bit_offset, 0u);

  auto row8 = amdgpu::wmma_a_input_loc(32, 128, /*row=*/8, /*k=*/0, 4, 4);
  EXPECT_EQ(row8.lane, 0u);
  EXPECT_EQ(row8.vgpr_offset, 8u);
  EXPECT_EQ(row8.bit_offset, 0u);

  auto row16 = amdgpu::wmma_a_input_loc(32, 128, /*row=*/16, /*k=*/4, 4, 4);
  EXPECT_EQ(row16.lane, 24u);
  EXPECT_EQ(row16.vgpr_offset, 0u);
  EXPECT_EQ(row16.bit_offset, 0u);

  auto row24 = amdgpu::wmma_a_input_loc(32, 128, /*row=*/24, /*k=*/7, 4, 4);
  EXPECT_EQ(row24.lane, 24u);
  EXPECT_EQ(row24.vgpr_offset, 8u);
  EXPECT_EQ(row24.bit_offset, 12u);

  EXPECT_EQ(amdgpu::wmma_a_scale_lane(32, 128, /*row=*/0, 0, 4, 4), 0u);
  EXPECT_EQ(amdgpu::wmma_a_scale_lane(32, 128, /*row=*/8, 0, 4, 4), 16u);
  EXPECT_EQ(amdgpu::wmma_a_scale_lane(32, 128, /*row=*/16, 0, 4, 4), 8u);
  EXPECT_EQ(amdgpu::wmma_a_scale_lane(32, 128, /*row=*/24, 0, 4, 4), 24u);
}

TEST(MfmaExecTest, OutputLoc32_4x4) {
  // 4x4 matrix, block 0: reg = column index, lane = row index.
  auto loc = amdgpu::output_loc_32(4, 4, /*col=*/2, /*row=*/1, /*b=*/0);
  EXPECT_EQ(loc.reg, 2u);
  EXPECT_EQ(loc.lane, 1u);
}

TEST(MfmaExecTest, ResolveAccConstant) {
  // Encoding value 0-255 = inline constant. The callback should be invoked.
  uint32_t const_acc = 0;
  uint32_t result = amdgpu::resolve_acc<amdgpu::AccMode::Unified>(
      /*vb=*/100, /*dst=*/200, /*src2_ev=*/128, const_acc, [&]() -> uint32_t { return 42u; });
  EXPECT_EQ(const_acc, 42u);
  EXPECT_EQ(result, 200u); // Returns dst when constant.
}

TEST(MfmaExecTest, ResolveAccVgpr) {
  // Encoding value 256-511 = VGPR.
  uint32_t const_acc = 0;
  uint32_t result = amdgpu::resolve_acc<amdgpu::AccMode::Unified>(
      /*vb=*/100, /*dst=*/200, /*src2_ev=*/260, const_acc, [&]() -> uint32_t { return 99u; });
  EXPECT_EQ(const_acc, amdgpu::ACC_FROM_VGPR);
  EXPECT_EQ(result, 100u + 4u); // vb + (260 - 256)
}

TEST(MfmaExecTest, ResolveAccAccVgpr) {
  // Encoding value 768-1023 = AccVGPR (unified alias).
  uint32_t const_acc = 0;
  uint32_t result = amdgpu::resolve_acc<amdgpu::AccMode::Unified>(
      /*vb=*/100, /*dst=*/200, /*src2_ev=*/770, const_acc, [&]() -> uint32_t { return 99u; });
  EXPECT_EQ(const_acc, amdgpu::ACC_FROM_VGPR);
  EXPECT_EQ(result, 100u + amdgpu::ACC_VGPR_OFFSET + 2u);
}

// ---------------------------------------------------------------------------
// L2 cache tests
// ---------------------------------------------------------------------------

TEST(L2CacheTest, UcStoreInvalidatesResidentLineBeforeAtomicRmw) {
  amdgpu::GpuMemory mem("test_mem");
  amdgpu::L2Cache l2("test_l2");
  l2.set_backing_memory(&mem);

  constexpr uint64_t kAddr = 0x2000;
  mem.write32(kAddr, 1);

  uint32_t first_old = 0;
  l2.atomic_rmw(kAddr, sizeof(uint32_t), [&](uint8_t *line, uint32_t offset) {
    std::memcpy(&first_old, line + offset, sizeof(first_old));
    std::memcpy(line + offset, &first_old, sizeof(first_old));
  });
  ASSERT_EQ(first_old, 1u);

  const uint32_t unlocked = 0;
  l2.write(kAddr, reinterpret_cast<const uint8_t *>(&unlocked), sizeof(unlocked),
           amdgpu::Mtype::UC);

  uint32_t second_old = 1;
  l2.atomic_rmw(kAddr, sizeof(uint32_t), [&](uint8_t *line, uint32_t offset) {
    std::memcpy(&second_old, line + offset, sizeof(second_old));
    std::memcpy(line + offset, &second_old, sizeof(second_old));
  });
  EXPECT_EQ(second_old, 0u);
}

TEST(L2CacheTest, UcReadFlushesDirtyResidentLine) {
  amdgpu::GpuMemory mem("test_mem");
  amdgpu::L2Cache l2("test_l2");
  l2.set_backing_memory(&mem);

  constexpr uint64_t kAddr = 0x2080;
  mem.write32(kAddr, 7);

  uint8_t line[amdgpu::L2Cache::LINE_SIZE] = {};
  uint32_t dirty_value = 9;
  std::memcpy(line, &dirty_value, sizeof(dirty_value));
  l2.writeback_line(kAddr, line, amdgpu::Mtype::RW);

  uint32_t read_value = 0;
  l2.read(kAddr, reinterpret_cast<uint8_t *>(&read_value), sizeof(read_value), amdgpu::Mtype::UC);
  EXPECT_EQ(read_value, 9u);
}

TEST(L2CacheTest, UcReadCrossingLineBoundaryFlushesBothDirtyResidentLines) {
  amdgpu::GpuMemory mem("test_mem");
  amdgpu::L2Cache l2("test_l2");
  l2.set_backing_memory(&mem);

  constexpr uint64_t kBase = 0x3000;
  constexpr uint64_t kAddr = kBase + amdgpu::L2Cache::LINE_SIZE - 4;
  static_assert((kBase & (amdgpu::L2Cache::LINE_SIZE - 1)) == 0);

  std::array<uint8_t, amdgpu::L2Cache::LINE_SIZE> first_line{};
  std::array<uint8_t, amdgpu::L2Cache::LINE_SIZE> second_line{};
  for (uint32_t i = 0; i < amdgpu::L2Cache::LINE_SIZE; ++i) {
    first_line[i] = static_cast<uint8_t>(0x10u + i);
    second_line[i] = static_cast<uint8_t>(0x80u + i);
    mem.write8(kBase + i, 0xA0u);
    mem.write8(kBase + amdgpu::L2Cache::LINE_SIZE + i, 0xB0u);
  }

  l2.writeback_line(kBase, first_line.data(), amdgpu::Mtype::RW);
  l2.writeback_line(kBase + amdgpu::L2Cache::LINE_SIZE, second_line.data(), amdgpu::Mtype::RW);

  std::array<uint8_t, 8> read_value{};
  l2.read(kAddr, read_value.data(), read_value.size(), amdgpu::Mtype::UC);

  for (uint32_t i = 0; i < 4; ++i) {
    EXPECT_EQ(read_value[i], first_line[amdgpu::L2Cache::LINE_SIZE - 4 + i])
        << "first line byte " << i;
    EXPECT_EQ(read_value[4 + i], second_line[i]) << "second line byte " << i;
  }
}

TEST(L2CacheTest, UcWriteCrossingLineBoundaryFlushesBothDirtyResidentLines) {
  amdgpu::GpuMemory mem("test_mem");
  amdgpu::L2Cache l2("test_l2");
  l2.set_backing_memory(&mem);

  constexpr uint64_t kBase = 0x4000;
  constexpr uint64_t kAddr = kBase + amdgpu::L2Cache::LINE_SIZE - 4;
  static_assert((kBase & (amdgpu::L2Cache::LINE_SIZE - 1)) == 0);

  std::array<uint8_t, amdgpu::L2Cache::LINE_SIZE> first_line{};
  std::array<uint8_t, amdgpu::L2Cache::LINE_SIZE> second_line{};
  for (uint32_t i = 0; i < amdgpu::L2Cache::LINE_SIZE; ++i) {
    first_line[i] = static_cast<uint8_t>(0x20u + i);
    second_line[i] = static_cast<uint8_t>(0x90u + i);
    mem.write8(kBase + i, 0xC0u);
    mem.write8(kBase + amdgpu::L2Cache::LINE_SIZE + i, 0xD0u);
  }

  l2.writeback_line(kBase, first_line.data(), amdgpu::Mtype::RW);
  l2.writeback_line(kBase + amdgpu::L2Cache::LINE_SIZE, second_line.data(), amdgpu::Mtype::RW);

  const std::array<uint8_t, 8> store_value = {0x01u, 0x02u, 0x03u, 0x04u,
                                              0x05u, 0x06u, 0x07u, 0x08u};
  l2.write(kAddr, store_value.data(), store_value.size(), amdgpu::Mtype::UC);

  for (uint32_t i = 0; i < 4; ++i) {
    EXPECT_EQ(mem.read8(kAddr + i), store_value[i]) << "first line store byte " << i;
    EXPECT_EQ(mem.read8(kAddr + 4 + i), store_value[4 + i]) << "second line store byte " << i;
  }
  EXPECT_EQ(mem.read8(kBase), first_line[0])
      << "dirty byte outside the first-line UC store should be preserved";
  EXPECT_EQ(mem.read8(kBase + amdgpu::L2Cache::LINE_SIZE + 4), second_line[4])
      << "dirty byte outside the second-line UC store should be preserved";
}

TEST(L1ScalarCacheTest, UcReadInvalidatesResidentWriteThroughLine) {
  amdgpu::GpuMemory mem("test_mem");
  amdgpu::L2Cache l2("test_l2");
  amdgpu::L1ScalarCache l1(&l2);
  l2.set_backing_memory(&mem);
  l1.set_memory(&mem);

  constexpr uint32_t kVmid = 1;
  constexpr uint64_t kAddr = 0x5000;
  constexpr uint32_t kBackingValue = 0x11111111;
  constexpr uint32_t kStoredValue = 0x22222222;
  constexpr uint32_t kReloadValue = 0x33333333;

  std::array<uint8_t, KfdProcess::kPageSize> backing{};
  KfdProcess::PageTable page_table;
  std::shared_mutex page_table_mutex;
  page_table[kAddr >> KfdProcess::kPageShift] = {backing.data(), amdgpu::Mtype::RW};
  mem.register_process(kVmid, &page_table, &page_table_mutex);

  mem.write32(kAddr, kBackingValue, kVmid);
  l1.store(kAddr, /*num_dwords=*/1, &kStoredValue, kVmid);

  {
    std::unique_lock lock(page_table_mutex);
    page_table[kAddr >> KfdProcess::kPageShift].mtype = amdgpu::Mtype::UC;
  }
  uint32_t read_value = 0;
  l1.load(kAddr, /*num_dwords=*/1, &read_value, kVmid);

  EXPECT_EQ(read_value, kStoredValue);
  EXPECT_EQ(mem.read32(kAddr, kVmid), kStoredValue);

  l2.write(kAddr, reinterpret_cast<const uint8_t *>(&kReloadValue), sizeof(kReloadValue),
           amdgpu::Mtype::RW, kVmid);
  {
    std::unique_lock lock(page_table_mutex);
    page_table[kAddr >> KfdProcess::kPageShift].mtype = amdgpu::Mtype::RW;
  }

  read_value = 0;
  l1.load(kAddr, /*num_dwords=*/1, &read_value, kVmid);
  EXPECT_EQ(read_value, kReloadValue);
}

TEST(L1ScalarCacheTest, UcLoadBytesInvalidatesResidentWriteThroughLine) {
  amdgpu::GpuMemory mem("test_mem");
  amdgpu::L2Cache l2("test_l2");
  amdgpu::L1ScalarCache l1(&l2);
  l2.set_backing_memory(&mem);
  l1.set_memory(&mem);

  constexpr uint32_t kVmid = 4;
  constexpr uint64_t kAddr = 0x5400;
  constexpr uint32_t kBackingValue = 0x11111111;
  constexpr uint32_t kStoredValue = 0x44332211;
  constexpr uint32_t kReloadValue = 0x88776655;

  std::array<uint8_t, KfdProcess::kPageSize> backing{};
  KfdProcess::PageTable page_table;
  std::shared_mutex page_table_mutex;
  page_table[kAddr >> KfdProcess::kPageShift] = {backing.data(), amdgpu::Mtype::RW};
  mem.register_process(kVmid, &page_table, &page_table_mutex);

  mem.write32(kAddr, kBackingValue, kVmid);
  l1.store(kAddr, /*num_dwords=*/1, &kStoredValue, kVmid);

  {
    std::unique_lock lock(page_table_mutex);
    page_table[kAddr >> KfdProcess::kPageShift].mtype = amdgpu::Mtype::UC;
  }

  std::array<uint8_t, 2> read_bytes{};
  l1.load_bytes(kAddr + 1, read_bytes.size(), read_bytes.data(), kVmid);

  EXPECT_EQ(read_bytes[0], 0x22);
  EXPECT_EQ(read_bytes[1], 0x33);
  EXPECT_EQ(mem.read32(kAddr, kVmid), kStoredValue);

  l2.write(kAddr, reinterpret_cast<const uint8_t *>(&kReloadValue), sizeof(kReloadValue),
           amdgpu::Mtype::RW, kVmid);
  {
    std::unique_lock lock(page_table_mutex);
    page_table[kAddr >> KfdProcess::kPageShift].mtype = amdgpu::Mtype::RW;
  }

  read_bytes.fill(0);
  l1.load_bytes(kAddr + 1, read_bytes.size(), read_bytes.data(), kVmid);
  EXPECT_EQ(read_bytes[0], 0x66);
  EXPECT_EQ(read_bytes[1], 0x77);
}

TEST(L1ScalarCacheTest, CcReadInvalidatesResidentWriteThroughLine) {
  amdgpu::GpuMemory mem("test_mem");
  amdgpu::L2Cache l2("test_l2");
  amdgpu::L1ScalarCache l1(&l2);
  l2.set_backing_memory(&mem);
  l1.set_memory(&mem);

  constexpr uint32_t kVmid = 5;
  constexpr uint64_t kAddr = 0x5800;
  constexpr uint32_t kBackingValue = 0x11111111;
  constexpr uint32_t kStoredValue = 0x22222222;
  constexpr uint32_t kReloadValue = 0x33333333;

  std::array<uint8_t, KfdProcess::kPageSize> backing{};
  KfdProcess::PageTable page_table;
  std::shared_mutex page_table_mutex;
  page_table[kAddr >> KfdProcess::kPageShift] = {backing.data(), amdgpu::Mtype::RW};
  mem.register_process(kVmid, &page_table, &page_table_mutex);

  mem.write32(kAddr, kBackingValue, kVmid);
  l1.store(kAddr, /*num_dwords=*/1, &kStoredValue, kVmid);

  {
    std::unique_lock lock(page_table_mutex);
    page_table[kAddr >> KfdProcess::kPageShift].mtype = amdgpu::Mtype::CC;
  }
  uint32_t read_value = 0;
  l1.load(kAddr, /*num_dwords=*/1, &read_value, kVmid);

  EXPECT_EQ(read_value, kStoredValue);
  EXPECT_EQ(mem.read32(kAddr, kVmid), kStoredValue);

  l2.write(kAddr, reinterpret_cast<const uint8_t *>(&kReloadValue), sizeof(kReloadValue),
           amdgpu::Mtype::RW, kVmid);
  {
    std::unique_lock lock(page_table_mutex);
    page_table[kAddr >> KfdProcess::kPageShift].mtype = amdgpu::Mtype::RW;
  }

  read_value = 0;
  l1.load(kAddr, /*num_dwords=*/1, &read_value, kVmid);
  EXPECT_EQ(read_value, kReloadValue);
}

TEST(L1ScalarCacheTest, CcLoadBytesInvalidatesResidentWriteThroughLine) {
  amdgpu::GpuMemory mem("test_mem");
  amdgpu::L2Cache l2("test_l2");
  amdgpu::L1ScalarCache l1(&l2);
  l2.set_backing_memory(&mem);
  l1.set_memory(&mem);

  constexpr uint32_t kVmid = 6;
  constexpr uint64_t kAddr = 0x5C00;
  constexpr uint32_t kBackingValue = 0x11111111;
  constexpr uint32_t kStoredValue = 0x44332211;
  constexpr uint32_t kReloadValue = 0x88776655;

  std::array<uint8_t, KfdProcess::kPageSize> backing{};
  KfdProcess::PageTable page_table;
  std::shared_mutex page_table_mutex;
  page_table[kAddr >> KfdProcess::kPageShift] = {backing.data(), amdgpu::Mtype::RW};
  mem.register_process(kVmid, &page_table, &page_table_mutex);

  mem.write32(kAddr, kBackingValue, kVmid);
  l1.store(kAddr, /*num_dwords=*/1, &kStoredValue, kVmid);

  {
    std::unique_lock lock(page_table_mutex);
    page_table[kAddr >> KfdProcess::kPageShift].mtype = amdgpu::Mtype::CC;
  }

  std::array<uint8_t, 2> read_bytes{};
  l1.load_bytes(kAddr + 1, read_bytes.size(), read_bytes.data(), kVmid);

  EXPECT_EQ(read_bytes[0], 0x22);
  EXPECT_EQ(read_bytes[1], 0x33);
  EXPECT_EQ(mem.read32(kAddr, kVmid), kStoredValue);

  l2.write(kAddr, reinterpret_cast<const uint8_t *>(&kReloadValue), sizeof(kReloadValue),
           amdgpu::Mtype::RW, kVmid);
  {
    std::unique_lock lock(page_table_mutex);
    page_table[kAddr >> KfdProcess::kPageShift].mtype = amdgpu::Mtype::RW;
  }

  read_bytes.fill(0);
  l1.load_bytes(kAddr + 1, read_bytes.size(), read_bytes.data(), kVmid);
  EXPECT_EQ(read_bytes[0], 0x66);
  EXPECT_EQ(read_bytes[1], 0x77);
}

TEST(L1ScalarCacheTest, UcWriteInvalidatesResidentLineBeforeBypassStore) {
  amdgpu::GpuMemory mem("test_mem");
  amdgpu::L2Cache l2("test_l2");
  amdgpu::L1ScalarCache l1(&l2);
  l2.set_backing_memory(&mem);
  l1.set_memory(&mem);

  constexpr uint32_t kVmid = 2;
  constexpr uint64_t kBase = 0x6000;
  constexpr uint64_t kStoreAddr = kBase + 4;
  constexpr uint32_t kStoredOutsideValue = 0x11111111;
  constexpr uint32_t kStoredTargetValue = 0x22222222;
  constexpr uint32_t kUcStoreValue = 0x33333333;

  std::array<uint8_t, KfdProcess::kPageSize> backing{};
  KfdProcess::PageTable page_table;
  std::shared_mutex page_table_mutex;
  page_table[kBase >> KfdProcess::kPageShift] = {backing.data(), amdgpu::Mtype::RW};
  mem.register_process(kVmid, &page_table, &page_table_mutex);

  const uint32_t stored_values[] = {kStoredOutsideValue, kStoredTargetValue};
  l1.store(kBase, /*num_dwords=*/2, stored_values, kVmid);

  {
    std::unique_lock lock(page_table_mutex);
    page_table[kBase >> KfdProcess::kPageShift].mtype = amdgpu::Mtype::UC;
  }
  l1.store(kStoreAddr, /*num_dwords=*/1, &kUcStoreValue, kVmid);
  l1.writeback_all(kVmid);

  EXPECT_EQ(mem.read32(kBase, kVmid), kStoredOutsideValue);
  EXPECT_EQ(mem.read32(kStoreAddr, kVmid), kUcStoreValue);
}

TEST(L1ScalarCacheTest, CcWriteInvalidatesResidentLineBeforeBypassStore) {
  amdgpu::GpuMemory mem("test_mem");
  amdgpu::L2Cache l2("test_l2");
  amdgpu::L1ScalarCache l1(&l2);
  l2.set_backing_memory(&mem);
  l1.set_memory(&mem);

  constexpr uint32_t kVmid = 7;
  constexpr uint64_t kBase = 0x6400;
  constexpr uint64_t kStoreAddr = kBase + 4;
  constexpr uint32_t kBackingOutsideValue = 0x01010101;
  constexpr uint32_t kBackingTargetValue = 0x02020202;
  constexpr uint32_t kStoredOutsideValue = 0x11111111;
  constexpr uint32_t kStoredTargetValue = 0x22222222;
  constexpr uint32_t kCcStoreValue = 0x33333333;

  std::array<uint8_t, KfdProcess::kPageSize> backing{};
  KfdProcess::PageTable page_table;
  std::shared_mutex page_table_mutex;
  page_table[kBase >> KfdProcess::kPageShift] = {backing.data(), amdgpu::Mtype::RW};
  mem.register_process(kVmid, &page_table, &page_table_mutex);

  mem.write32(kBase, kBackingOutsideValue, kVmid);
  mem.write32(kStoreAddr, kBackingTargetValue, kVmid);
  const uint32_t stored_values[] = {kStoredOutsideValue, kStoredTargetValue};
  l1.store(kBase, /*num_dwords=*/2, stored_values, kVmid);

  {
    std::unique_lock lock(page_table_mutex);
    page_table[kBase >> KfdProcess::kPageShift].mtype = amdgpu::Mtype::CC;
  }
  l1.store(kStoreAddr, /*num_dwords=*/1, &kCcStoreValue, kVmid);

  EXPECT_EQ(mem.read32(kBase, kVmid), kStoredOutsideValue);
  EXPECT_EQ(mem.read32(kStoreAddr, kVmid), kCcStoreValue);
}

TEST(L1ScalarCacheTest, WriteThroughStoresAndCleanEvictionReachBacking) {
  amdgpu::GpuMemory mem("test_mem");
  amdgpu::L2Cache l2("test_l2");
  amdgpu::L1ScalarCache l1(&l2);
  l2.set_backing_memory(&mem);

  constexpr uint64_t kBase = 0x9000;
  constexpr uint64_t kSetStride = uint64_t{1}
                                  << (amdgpu::L1ScalarCache::LINE_SIZE_BITS +
                                      std::bit_width(amdgpu::L1ScalarCache::NUM_SETS - 1));
  static_assert(amdgpu::L1ScalarCache::ASSOCIATIVITY == 4);

  std::array<uint64_t, 5> addrs{};
  std::array<uint32_t, 5> values{};
  for (uint32_t i = 0; i < addrs.size(); ++i) {
    addrs[i] = kBase + i * kSetStride;
    values[i] = 0x11110000u + i;
  }

  for (uint32_t i = 0; i < 4; ++i) {
    l1.store(addrs[i], /*num_dwords=*/1, &values[i]);
    EXPECT_EQ(mem.read32(addrs[i]), values[i]) << "line " << i;
  }

  l1.store(addrs[4], /*num_dwords=*/1, &values[4]);
  EXPECT_EQ(mem.read32(addrs[4]), values[4]);

  l1.writeback_all();
  for (uint32_t i = 0; i < addrs.size(); ++i)
    EXPECT_EQ(mem.read32(addrs[i]), values[i]) << "line " << i;
}

TEST(L1ScalarCacheTest, CacheableStoresWriteThroughBeforeWriteback) {
  constexpr uint32_t kVmid = 9;
  constexpr uint64_t kAddr = 0x6800;

  for (const auto mtype : {amdgpu::Mtype::RW, amdgpu::Mtype::WB, amdgpu::Mtype::NT}) {
    SCOPED_TRACE(static_cast<int>(mtype));
    amdgpu::GpuMemory mem("test_mem");
    amdgpu::L2Cache l2("test_l2");
    amdgpu::L1ScalarCache l1(&l2);
    l2.set_backing_memory(&mem);
    l1.set_memory(&mem);

    std::array<uint8_t, KfdProcess::kPageSize> backing{};
    KfdProcess::PageTable page_table;
    std::shared_mutex page_table_mutex;
    page_table[kAddr >> KfdProcess::kPageShift] = {backing.data(), mtype};
    mem.register_process(kVmid, &page_table, &page_table_mutex);

    const uint32_t value = 0xCAFE0000u + static_cast<uint32_t>(mtype);
    l1.store(kAddr, /*num_dwords=*/1, &value, kVmid);
    EXPECT_EQ(mem.read32(kAddr, kVmid), value);

    l1.writeback_all(kVmid);
    EXPECT_EQ(mem.read32(kAddr, kVmid), value);
    mem.unregister_process(kVmid);
  }
}

TEST(L1ScalarCacheTest, UnalignedStoreCrossingLineWritesThroughExactBytes) {
  amdgpu::GpuMemory mem("test_mem");
  amdgpu::L2Cache l2("test_l2");
  amdgpu::L1ScalarCache l1(&l2);
  l2.set_backing_memory(&mem);
  l1.set_memory(&mem);

  constexpr uint32_t kVmid = 10;
  constexpr uint64_t kPageBase = 0x7000;
  constexpr uint64_t kAddr = kPageBase + (uint64_t{1} << amdgpu::L1ScalarCache::LINE_SIZE_BITS) - 2;
  constexpr uint32_t kValue = 0x44332211;

  std::array<uint8_t, KfdProcess::kPageSize> backing{};
  backing.fill(0xA5);
  KfdProcess::PageTable page_table;
  std::shared_mutex page_table_mutex;
  page_table[kPageBase >> KfdProcess::kPageShift] = {backing.data(), amdgpu::Mtype::RW};
  mem.register_process(kVmid, &page_table, &page_table_mutex);

  l1.store(kAddr, /*num_dwords=*/1, &kValue, kVmid);

  std::array<uint8_t, sizeof(kValue)> expected{};
  std::memcpy(expected.data(), &kValue, sizeof(kValue));
  const size_t backing_offset = kAddr - kPageBase;
  EXPECT_EQ(backing[backing_offset - 1], 0xA5);
  for (size_t i = 0; i < expected.size(); ++i)
    EXPECT_EQ(backing[backing_offset + i], expected[i]) << "byte " << i;
  EXPECT_EQ(backing[backing_offset + expected.size()], 0xA5);

  l1.writeback_all(kVmid);
  for (size_t i = 0; i < expected.size(); ++i)
    EXPECT_EQ(backing[backing_offset + i], expected[i]) << "byte after writeback " << i;
  mem.unregister_process(kVmid);
}

TEST(L1ScalarCacheTest, ScalarWritebackDoesNotClobberAtomicAtDisjointAddress) {
  amdgpu::GpuMemory mem("test_mem");
  amdgpu::L2Cache l2a("l2a");
  amdgpu::L2Cache l2b("l2b");
  l2a.set_backing_memory(&mem);
  l2b.set_backing_memory(&mem);
  amdgpu::L1ScalarCache l1(&l2a);
  l1.set_memory(&mem);

  constexpr uint32_t kVmid = 17;
  constexpr uint64_t kVa = 0x500000;
  constexpr uint32_t kScalarValue = 0x5A5A5A5A;
  std::array<uint8_t, KfdProcess::kPageSize> backing{};
  KfdProcess::PageTable page_table;
  std::shared_mutex page_table_mutex;
  page_table[kVa >> KfdProcess::kPageShift] = {backing.data(), amdgpu::Mtype::RW};
  mem.register_process(kVmid, &page_table, &page_table_mutex);

  l1.store(kVa + sizeof(uint32_t), /*num_dwords=*/1, &kScalarValue, kVmid);
  l2b.atomic_rmw(
      kVa, sizeof(uint32_t),
      [](uint8_t *storage, uint32_t offset) {
        uint32_t value = 0;
        std::memcpy(&value, storage + offset, sizeof(value));
        ++value;
        std::memcpy(storage + offset, &value, sizeof(value));
      },
      kVmid);

  l1.writeback_all(kVmid);
  EXPECT_EQ(mem.read32(kVa, kVmid), 1u);
  EXPECT_EQ(mem.read32(kVa + sizeof(uint32_t), kVmid), kScalarValue);
  mem.unregister_process(kVmid);
}

TEST(L1ScalarCacheTest, CleanEvictionDoesNotClobberAtomicAtDisjointAddress) {
  amdgpu::GpuMemory mem("test_mem");
  amdgpu::L2Cache l2a("l2a");
  amdgpu::L2Cache l2b("l2b");
  l2a.set_backing_memory(&mem);
  l2b.set_backing_memory(&mem);
  amdgpu::L1ScalarCache l1(&l2a);

  constexpr uint64_t kBase = 0x600000;
  constexpr uint64_t kSetStride = uint64_t{1}
                                  << (amdgpu::L1ScalarCache::LINE_SIZE_BITS +
                                      std::bit_width(amdgpu::L1ScalarCache::NUM_SETS - 1));
  constexpr uint32_t kScalarValue = 0x6B6B6B6B;
  l1.store(kBase + sizeof(uint32_t), /*num_dwords=*/1, &kScalarValue);
  l2b.atomic_rmw(kBase, sizeof(uint32_t), [](uint8_t *storage, uint32_t offset) {
    uint32_t value = 0;
    std::memcpy(&value, storage + offset, sizeof(value));
    ++value;
    std::memcpy(storage + offset, &value, sizeof(value));
  });

  for (uint32_t i = 1; i <= amdgpu::L1ScalarCache::ASSOCIATIVITY; ++i) {
    uint32_t ignored = 0;
    l1.load(kBase + i * kSetStride, /*num_dwords=*/1, &ignored);
  }

  EXPECT_EQ(mem.read32(kBase), 1u);
  EXPECT_EQ(mem.read32(kBase + sizeof(uint32_t)), kScalarValue);
}

TEST(L1ScalarCacheTest, UcAndCcFlushDoNotClobberAtomicAtDisjointAddress) {
  constexpr uint32_t kVmid = 18;
  constexpr uint64_t kVa = 0x700000;
  constexpr uint32_t kScalarValue = 0x7C7C7C7C;

  for (const auto mtype : {amdgpu::Mtype::UC, amdgpu::Mtype::CC}) {
    SCOPED_TRACE(static_cast<int>(mtype));
    amdgpu::GpuMemory mem("test_mem");
    amdgpu::L2Cache l2a("l2a");
    amdgpu::L2Cache l2b("l2b");
    l2a.set_backing_memory(&mem);
    l2b.set_backing_memory(&mem);
    amdgpu::L1ScalarCache l1(&l2a);
    l1.set_memory(&mem);

    std::array<uint8_t, KfdProcess::kPageSize> backing{};
    KfdProcess::PageTable page_table;
    std::shared_mutex page_table_mutex;
    page_table[kVa >> KfdProcess::kPageShift] = {backing.data(), amdgpu::Mtype::RW};
    mem.register_process(kVmid, &page_table, &page_table_mutex);

    l1.store(kVa + sizeof(uint32_t), /*num_dwords=*/1, &kScalarValue, kVmid);
    l2b.atomic_rmw(
        kVa, sizeof(uint32_t),
        [](uint8_t *storage, uint32_t offset) {
          uint32_t value = 0;
          std::memcpy(&value, storage + offset, sizeof(value));
          ++value;
          std::memcpy(storage + offset, &value, sizeof(value));
        },
        kVmid);

    {
      std::unique_lock lock(page_table_mutex);
      page_table[kVa >> KfdProcess::kPageShift].mtype = mtype;
    }
    uint32_t ignored = 0;
    l1.load(kVa + 2 * sizeof(uint32_t), /*num_dwords=*/1, &ignored, kVmid);

    EXPECT_EQ(mem.read32(kVa, kVmid), 1u);
    EXPECT_EQ(mem.read32(kVa + sizeof(uint32_t), kVmid), kScalarValue);
    mem.unregister_process(kVmid);
  }
}

TEST(DeviceCacheCoherenceTest, ScalarWriteThroughCannotClobberRemoteAtomic) {
  amdgpu::GpuMemory mem("test_mem");
  amdgpu::L2Cache scalar_l2("scalar_l2");
  amdgpu::L2Cache atomic_l2("atomic_l2");
  amdgpu::L1ScalarCache scalar_l1(&scalar_l2);
  scalar_l2.set_backing_memory(&mem);
  atomic_l2.set_backing_memory(&mem);

  constexpr uint64_t kLine = 0xA000;
  constexpr uint64_t kAtomicAddr = kLine;
  constexpr uint64_t kScalarAddr = kLine + sizeof(uint32_t);
  constexpr uint32_t kScalarValue = 0xA5A5A5A5;
  mem.write32(kAtomicAddr, 0);
  mem.write32(kScalarAddr, 0);

  // The scalar store read-allocates the atomic dword's old value into the same
  // K$ line. The store must update only its target bytes, and the atomic must
  // invalidate the stale clean snapshot.
  scalar_l1.store(kScalarAddr, /*num_dwords=*/1, &kScalarValue);
  atomic_l2.atomic_rmw(kAtomicAddr, sizeof(uint32_t), [](uint8_t *line, uint32_t offset) {
    uint32_t value = 0;
    std::memcpy(&value, line + offset, sizeof(value));
    ++value;
    std::memcpy(line + offset, &value, sizeof(value));
  });
  scalar_l1.writeback_all();

  EXPECT_EQ(mem.read32(kAtomicAddr), 1u);
  EXPECT_EQ(mem.read32(kScalarAddr), kScalarValue);
}

TEST(DeviceCacheCoherenceTest, DisjointScalarWriteThroughStoresSurviveRemoteAtomic) {
  amdgpu::GpuMemory mem("test_mem");
  amdgpu::L2Cache first_scalar_l2("first_scalar_l2");
  amdgpu::L2Cache second_scalar_l2("second_scalar_l2");
  amdgpu::L2Cache atomic_l2("atomic_l2");
  amdgpu::L1ScalarCache first_scalar_l1(&first_scalar_l2);
  amdgpu::L1ScalarCache second_scalar_l1(&second_scalar_l2);
  first_scalar_l2.set_backing_memory(&mem);
  second_scalar_l2.set_backing_memory(&mem);
  atomic_l2.set_backing_memory(&mem);

  constexpr uint64_t kLine = 0xA080;
  constexpr uint64_t kAtomicAddr = kLine;
  constexpr uint64_t kFirstScalarAddr = kLine + sizeof(uint32_t);
  constexpr uint64_t kSecondScalarAddr = kLine + 2 * sizeof(uint32_t);
  constexpr uint32_t kFirstScalarValue = 0x11112222;
  constexpr uint32_t kSecondScalarValue = 0x33334444;
  mem.write32(kAtomicAddr, 0);
  mem.write32(kFirstScalarAddr, 0);
  mem.write32(kSecondScalarAddr, 0);

  // Both K$ instances fill the same old line. Their disjoint write-through
  // stores must merge without either cached snapshot replacing the other.
  first_scalar_l1.store(kFirstScalarAddr, /*num_dwords=*/1, &kFirstScalarValue);
  second_scalar_l1.store(kSecondScalarAddr, /*num_dwords=*/1, &kSecondScalarValue);
  atomic_l2.atomic_rmw(kAtomicAddr, sizeof(uint32_t), [](uint8_t *line, uint32_t offset) {
    uint32_t value = 0;
    std::memcpy(&value, line + offset, sizeof(value));
    ++value;
    std::memcpy(line + offset, &value, sizeof(value));
  });
  first_scalar_l1.writeback_all();
  second_scalar_l1.writeback_all();

  EXPECT_EQ(mem.read32(kAtomicAddr), 1u);
  EXPECT_EQ(mem.read32(kFirstScalarAddr), kFirstScalarValue);
  EXPECT_EQ(mem.read32(kSecondScalarAddr), kSecondScalarValue);
}

TEST(DeviceCacheCoherenceTest, RemoteAtomicInvalidatesScalarCachedRead) {
  amdgpu::GpuMemory mem("test_mem");
  amdgpu::L2Cache scalar_l2("scalar_l2");
  amdgpu::L2Cache atomic_l2("atomic_l2");
  amdgpu::L1ScalarCache scalar_l1(&scalar_l2);
  scalar_l2.set_backing_memory(&mem);
  atomic_l2.set_backing_memory(&mem);

  constexpr uint64_t kAddr = 0xA100;
  mem.write32(kAddr, 10);
  uint32_t value = 0;
  scalar_l1.load(kAddr, /*num_dwords=*/1, &value);
  ASSERT_EQ(value, 10u);

  atomic_l2.atomic_rmw(kAddr, sizeof(uint32_t), [](uint8_t *line, uint32_t offset) {
    uint32_t current = 0;
    std::memcpy(&current, line + offset, sizeof(current));
    ++current;
    std::memcpy(line + offset, &current, sizeof(current));
  });

  value = 0;
  scalar_l1.load(kAddr, /*num_dwords=*/1, &value);
  EXPECT_EQ(value, 11u);
}

TEST(DeviceCacheCoherenceTest, RemoteAtomicInvalidatesVectorCachedRead) {
  amdgpu::GpuMemory mem("test_mem");
  amdgpu::L2Cache vector_l2("vector_l2");
  amdgpu::L2Cache atomic_l2("atomic_l2");
  amdgpu::L1VectorCache vector_l1(&vector_l2);
  vector_l2.set_backing_memory(&mem);
  atomic_l2.set_backing_memory(&mem);

  constexpr uint64_t kAddr = 0xA200;
  mem.write32(kAddr, 10);
  uint64_t addrs[cdna3::Isa::WF_SIZE] = {};
  addrs[0] = kAddr;
  std::array<uint8_t, cdna3::Isa::WF_SIZE * sizeof(uint32_t)> bytes{};
  auto load_value = [&] {
    bytes.fill(0);
    vector_l1.load(addrs, /*lane_mask=*/1, /*elem_size=*/sizeof(uint32_t),
                   /*num_elems=*/1, bytes.data(), amdgpu::Mtype::RW,
                   /*non_temporal=*/false, /*request_l1_bypass=*/false, cdna3::Isa::WF_SIZE);
    uint32_t value = 0;
    std::memcpy(&value, bytes.data(), sizeof(value));
    return value;
  };
  ASSERT_EQ(load_value(), 10u);

  atomic_l2.atomic_rmw(kAddr, sizeof(uint32_t), [](uint8_t *line, uint32_t offset) {
    uint32_t current = 0;
    std::memcpy(&current, line + offset, sizeof(current));
    ++current;
    std::memcpy(line + offset, &current, sizeof(current));
  });

  EXPECT_EQ(load_value(), 11u);
}

TEST(DeviceCacheCoherenceTest, AtomicConsumesScalarWriteThroughTarget) {
  amdgpu::GpuMemory mem("test_mem");
  amdgpu::L2Cache scalar_l2("scalar_l2");
  amdgpu::L2Cache atomic_l2("atomic_l2");
  amdgpu::L1ScalarCache scalar_l1(&scalar_l2);
  scalar_l2.set_backing_memory(&mem);
  atomic_l2.set_backing_memory(&mem);

  constexpr uint64_t kAddr = 0xA300;
  constexpr uint32_t kStoredValue = 40;
  mem.write32(kAddr, 0);
  scalar_l1.store(kAddr, /*num_dwords=*/1, &kStoredValue);

  atomic_l2.atomic_rmw(kAddr, sizeof(uint32_t), [](uint8_t *line, uint32_t offset) {
    uint32_t current = 0;
    std::memcpy(&current, line + offset, sizeof(current));
    ++current;
    std::memcpy(line + offset, &current, sizeof(current));
  });

  EXPECT_EQ(mem.read32(kAddr), 41u);
  uint32_t reloaded = 0;
  scalar_l1.load(kAddr, /*num_dwords=*/1, &reloaded);
  EXPECT_EQ(reloaded, 41u);
}

TEST(DeviceCacheCoherenceTest, DestroyedCachesAreRemovedFromRegistry) {
  amdgpu::GpuMemory mem("test_mem");
  constexpr uint64_t kAddr = 0xA400;

  struct alignas(amdgpu::L2Cache) L2Storage {
    std::byte data[sizeof(amdgpu::L2Cache)];
  };
  struct alignas(amdgpu::L1ScalarCache) ScalarStorage {
    std::byte data[sizeof(amdgpu::L1ScalarCache)];
  };
  struct alignas(amdgpu::L1VectorCache) VectorStorage {
    std::byte data[sizeof(amdgpu::L1VectorCache)];
  };

  auto l2_storage = std::make_unique<L2Storage>();
  auto scalar_storage = std::make_unique<ScalarStorage>();
  auto vector_storage = std::make_unique<VectorStorage>();
  auto *transient_l2 =
      std::construct_at(reinterpret_cast<amdgpu::L2Cache *>(l2_storage->data), "transient_l2");
  auto *transient_scalar = std::construct_at(
      reinterpret_cast<amdgpu::L1ScalarCache *>(scalar_storage->data), transient_l2);
  auto *transient_vector = std::construct_at(
      reinterpret_cast<amdgpu::L1VectorCache *>(vector_storage->data), transient_l2);
  transient_l2->set_backing_memory(&mem);

  constexpr uint32_t kDirty = 7;
  transient_scalar->store(kAddr + sizeof(uint32_t), /*num_dwords=*/1, &kDirty);
  uint64_t addrs[cdna3::Isa::WF_SIZE] = {};
  addrs[0] = kAddr;
  std::array<uint8_t, cdna3::Isa::WF_SIZE * sizeof(uint32_t)> bytes{};
  transient_vector->load(addrs, /*lane_mask=*/1, /*elem_size=*/sizeof(uint32_t),
                         /*num_elems=*/1, bytes.data(), amdgpu::Mtype::RW,
                         /*non_temporal=*/false, /*request_l1_bypass=*/false, cdna3::Isa::WF_SIZE);

  std::destroy_at(transient_vector);
  std::destroy_at(transient_scalar);
  std::destroy_at(transient_l2);
  std::memset(vector_storage->data, 0xA5, sizeof(vector_storage->data));
  std::memset(scalar_storage->data, 0xA5, sizeof(scalar_storage->data));
  std::memset(l2_storage->data, 0xA5, sizeof(l2_storage->data));

  amdgpu::L2Cache survivor("survivor");
  EXPECT_NE(static_cast<const void *>(&survivor), static_cast<const void *>(transient_l2));
  survivor.set_backing_memory(&mem);
  mem.write32(kAddr, 0);
  survivor.atomic_rmw(kAddr, sizeof(uint32_t), [](uint8_t *line, uint32_t offset) {
    uint32_t value = 0;
    std::memcpy(&value, line + offset, sizeof(value));
    ++value;
    std::memcpy(line + offset, &value, sizeof(value));
  });
  EXPECT_EQ(mem.read32(kAddr), 1u);
}

TEST(L1VectorCacheTest, UcReadInvalidatesResidentLine) {
  amdgpu::GpuMemory mem("test_mem");
  amdgpu::L2Cache l2("test_l2");
  amdgpu::L1VectorCache l1(&l2);
  l2.set_backing_memory(&mem);

  constexpr uint64_t kAddr = 0x7000;
  constexpr uint32_t kOldValue = 0x11111111;
  constexpr uint32_t kNewValue = 0x22222222;

  uint64_t addrs[64] = {};
  addrs[0] = kAddr;
  std::array<uint8_t, 64 * sizeof(uint32_t)> bytes{};

  mem.write32(kAddr, kOldValue);
  l1.load(addrs, /*lane_mask=*/0x1, /*elem_size=*/4, /*num_elems=*/1, bytes.data(),
          amdgpu::Mtype::RW, /*non_temporal=*/false, /*request_l1_bypass=*/false,
          cdna3::Isa::WF_SIZE);
  uint32_t value = 0;
  std::memcpy(&value, bytes.data(), sizeof(value));
  ASSERT_EQ(value, kOldValue);

  mem.write32(kAddr, kNewValue);
  bytes.fill(0);
  l1.load(addrs, /*lane_mask=*/0x1, /*elem_size=*/4, /*num_elems=*/1, bytes.data(),
          amdgpu::Mtype::UC, /*non_temporal=*/false, /*request_l1_bypass=*/false,
          cdna3::Isa::WF_SIZE);
  std::memcpy(&value, bytes.data(), sizeof(value));
  ASSERT_EQ(value, kNewValue);

  bytes.fill(0);
  l1.load(addrs, /*lane_mask=*/0x1, /*elem_size=*/4, /*num_elems=*/1, bytes.data(),
          amdgpu::Mtype::RW, /*non_temporal=*/false, /*request_l1_bypass=*/false,
          cdna3::Isa::WF_SIZE);
  std::memcpy(&value, bytes.data(), sizeof(value));
  EXPECT_EQ(value, kNewValue);
}

TEST(GpuMemoryTest, BlockAccessSpansSparseFallbackPages) {
  amdgpu::GpuMemory mem("test_mem");
  constexpr uint64_t kAddr = amdgpu::GpuMemory::PAGE_SIZE - 8;
  constexpr std::array<uint8_t, 16> kData = {0x00, 0x11, 0x22, 0x33, 0x44, 0x55, 0x66, 0x77,
                                             0x88, 0x99, 0xAA, 0xBB, 0xCC, 0xDD, 0xEE, 0xFF};

  mem.write_block(kAddr, std::span<const uint8_t>(kData));
  std::array<uint8_t, kData.size()> result{};
  mem.read_block(kAddr, std::span<uint8_t>(result));

  EXPECT_EQ(result, kData);
}

TEST(GpuMemoryTest, BlockAccessSpansMappedPages) {
  amdgpu::GpuMemory mem("test_mem");
  constexpr uint32_t kVmid = 8;
  constexpr uint64_t kAddr = KfdProcess::kPageSize - 8;
  constexpr std::array<uint8_t, 16> kData = {0x00, 0x11, 0x22, 0x33, 0x44, 0x55, 0x66, 0x77,
                                             0x88, 0x99, 0xAA, 0xBB, 0xCC, 0xDD, 0xEE, 0xFF};

  std::array<uint8_t, KfdProcess::kPageSize> first_page{};
  std::array<uint8_t, KfdProcess::kPageSize> second_page{};
  KfdProcess::PageTable page_table;
  std::shared_mutex page_table_mutex;
  page_table[0] = {first_page.data(), amdgpu::Mtype::RW};
  page_table[1] = {second_page.data(), amdgpu::Mtype::RW};
  mem.register_process(kVmid, &page_table, &page_table_mutex);

  mem.write_block(kAddr, std::span<const uint8_t>(kData), kVmid);
  std::array<uint8_t, kData.size()> result{};
  mem.read_block(kAddr, std::span<uint8_t>(result), kVmid);

  EXPECT_EQ(result, kData);
}

TEST(GpuMemoryTest, BlockAccessRechecksTranslationAfterSparseFallbackPage) {
  amdgpu::GpuMemory mem("test_mem");
  constexpr uint32_t kVmid = 9;
  constexpr uint64_t kAddr = KfdProcess::kPageSize - 8;
  constexpr std::array<uint8_t, 16> kReadData = {0x10, 0x11, 0x12, 0x13, 0x14, 0x15, 0x16, 0x17,
                                                 0x20, 0x21, 0x22, 0x23, 0x24, 0x25, 0x26, 0x27};
  constexpr std::array<uint8_t, 16> kWriteData = {0x30, 0x31, 0x32, 0x33, 0x34, 0x35, 0x36, 0x37,
                                                  0x40, 0x41, 0x42, 0x43, 0x44, 0x45, 0x46, 0x47};

  std::array<uint8_t, KfdProcess::kPageSize> second_page{};
  std::copy(kReadData.begin() + 8, kReadData.end(), second_page.begin());
  KfdProcess::PageTable page_table;
  std::shared_mutex page_table_mutex;
  page_table[1] = {second_page.data(), amdgpu::Mtype::RW};
  mem.register_process(kVmid, &page_table, &page_table_mutex);
  for (size_t i = 0; i < 8; ++i)
    mem.write8(kAddr + i, kReadData[i], kVmid);

  std::array<uint8_t, kReadData.size()> result{};
  mem.read_block(kAddr, std::span<uint8_t>(result), kVmid);
  EXPECT_EQ(result, kReadData);

  mem.write_block(kAddr, std::span<const uint8_t>(kWriteData), kVmid);
  for (size_t i = 0; i < 8; ++i)
    EXPECT_EQ(mem.read8(kAddr + i, kVmid), kWriteData[i]);
  EXPECT_TRUE(std::equal(kWriteData.begin() + 8, kWriteData.end(), second_page.begin()));
}

TEST(L1VectorCacheTest, UcDwordx4RoundTripPreservesVectorTransaction) {
  amdgpu::GpuMemory mem("test_mem");
  amdgpu::L2Cache l2("test_l2");
  amdgpu::L1VectorCache l1(&l2);
  l2.set_backing_memory(&mem);

  constexpr uint64_t kAddr = 0x7800;
  constexpr std::array<uint32_t, 4> kLane0 = {0x01234567u, 0x89ABCDEFu, 0x76543210u, 0xFEDCBA98u};
  constexpr std::array<uint32_t, 4> kLane1 = {0x11112222u, 0x33334444u, 0x55556666u, 0x77778888u};
  uint64_t addrs[64] = {};
  addrs[0] = kAddr;
  addrs[1] = kAddr + 0x100;
  std::array<uint8_t, 64 * sizeof(kLane0)> store_bytes{};
  std::memcpy(store_bytes.data(), kLane0.data(), sizeof(kLane0));
  std::memcpy(store_bytes.data() + sizeof(kLane0), kLane1.data(), sizeof(kLane1));

  l1.store(addrs, /*lane_mask=*/0x3, /*elem_size=*/4, /*num_elems=*/4, store_bytes.data(),
           amdgpu::Mtype::UC, /*non_temporal=*/false, cdna3::Isa::WF_SIZE);
  EXPECT_EQ(l1.store_l2_writes(), 2u);
  EXPECT_EQ(l2.backing_write_transactions(), 2u);

  std::array<uint8_t, 64 * sizeof(kLane0)> load_bytes{};
  l1.load(addrs, /*lane_mask=*/0x3, /*elem_size=*/4, /*num_elems=*/4, load_bytes.data(),
          amdgpu::Mtype::UC, /*non_temporal=*/false, /*request_l1_bypass=*/false,
          cdna3::Isa::WF_SIZE);
  EXPECT_EQ(l2.backing_read_transactions(), 2u);
  EXPECT_EQ(std::memcmp(load_bytes.data(), kLane0.data(), sizeof(kLane0)), 0);
  EXPECT_EQ(std::memcmp(load_bytes.data() + sizeof(kLane0), kLane1.data(), sizeof(kLane1)), 0);
}

TEST(L1VectorCacheTest, UcDwordx4CoalescesAdjacentLanes) {
  amdgpu::GpuMemory mem("test_mem");
  amdgpu::L2Cache l2("test_l2");
  amdgpu::L1VectorCache l1(&l2);
  l2.set_backing_memory(&mem);

  constexpr uint64_t kAddr = 0x7900;
  constexpr std::array<uint32_t, 4> kLane0 = {0x01234567u, 0x89ABCDEFu, 0x76543210u, 0xFEDCBA98u};
  constexpr std::array<uint32_t, 4> kLane1 = {0x11112222u, 0x33334444u, 0x55556666u, 0x77778888u};
  uint64_t addrs[64] = {};
  addrs[0] = kAddr;
  addrs[1] = kAddr + sizeof(kLane0);
  std::array<uint8_t, 64 * sizeof(kLane0)> store_bytes{};
  std::memcpy(store_bytes.data(), kLane0.data(), sizeof(kLane0));
  std::memcpy(store_bytes.data() + sizeof(kLane0), kLane1.data(), sizeof(kLane1));

  l1.store(addrs, /*lane_mask=*/0x3, /*elem_size=*/4, /*num_elems=*/4, store_bytes.data(),
           amdgpu::Mtype::UC, /*non_temporal=*/false, cdna3::Isa::WF_SIZE);
  EXPECT_EQ(l1.store_l2_writes(), 1u);
  EXPECT_EQ(l2.backing_write_transactions(), 1u);

  std::array<uint8_t, 64 * sizeof(kLane0)> load_bytes{};
  l1.load(addrs, /*lane_mask=*/0x3, /*elem_size=*/4, /*num_elems=*/4, load_bytes.data(),
          amdgpu::Mtype::UC, /*non_temporal=*/false, /*request_l1_bypass=*/false,
          cdna3::Isa::WF_SIZE);
  EXPECT_EQ(l2.backing_read_transactions(), 1u);
  EXPECT_EQ(std::memcmp(load_bytes.data(), kLane0.data(), sizeof(kLane0)), 0);
  EXPECT_EQ(std::memcmp(load_bytes.data() + sizeof(kLane0), kLane1.data(), sizeof(kLane1)), 0);
}

TEST(L1VectorCacheTest, UcWriteInvalidatesResidentLine) {
  amdgpu::GpuMemory mem("test_mem");
  amdgpu::L2Cache l2("test_l2");
  amdgpu::L1VectorCache l1(&l2);
  l2.set_backing_memory(&mem);

  constexpr uint64_t kAddr = 0x8000;
  constexpr uint32_t kOldValue = 0x11111111;
  constexpr uint32_t kNewValue = 0x22222222;

  uint64_t addrs[64] = {};
  addrs[0] = kAddr;
  std::array<uint8_t, 64 * sizeof(uint32_t)> bytes{};

  mem.write32(kAddr, kOldValue);
  l1.load(addrs, /*lane_mask=*/0x1, /*elem_size=*/4, /*num_elems=*/1, bytes.data(),
          amdgpu::Mtype::RW, /*non_temporal=*/false, /*request_l1_bypass=*/false,
          cdna3::Isa::WF_SIZE);
  uint32_t value = 0;
  std::memcpy(&value, bytes.data(), sizeof(value));
  ASSERT_EQ(value, kOldValue);

  std::array<uint8_t, 64 * sizeof(uint32_t)> store_bytes{};
  std::memcpy(store_bytes.data(), &kNewValue, sizeof(kNewValue));
  l1.store(addrs, /*lane_mask=*/0x1, /*elem_size=*/4, /*num_elems=*/1, store_bytes.data(),
           amdgpu::Mtype::UC, /*non_temporal=*/false, cdna3::Isa::WF_SIZE);

  bytes.fill(0);
  l1.load(addrs, /*lane_mask=*/0x1, /*elem_size=*/4, /*num_elems=*/1, bytes.data(),
          amdgpu::Mtype::RW, /*non_temporal=*/false, /*request_l1_bypass=*/false,
          cdna3::Isa::WF_SIZE);
  std::memcpy(&value, bytes.data(), sizeof(value));
  EXPECT_EQ(value, kNewValue);
}

TEST(L1VectorCacheTest, NonTemporalReadInvalidatesResidentLine) {
  amdgpu::GpuMemory mem("test_mem");
  amdgpu::L2Cache l2("test_l2");
  amdgpu::L1VectorCache l1(&l2);
  l2.set_backing_memory(&mem);

  constexpr uint64_t kAddr = 0x8100;
  constexpr uint32_t kOldValue = 0x11111111;
  constexpr uint32_t kNewValue = 0x22222222;

  uint64_t addrs[64] = {};
  addrs[0] = kAddr;
  std::array<uint8_t, 64 * sizeof(uint32_t)> bytes{};

  mem.write32(kAddr, kOldValue);
  l1.load(addrs, /*lane_mask=*/0x1, /*elem_size=*/4, /*num_elems=*/1, bytes.data(),
          amdgpu::Mtype::RW, /*non_temporal=*/false, /*request_l1_bypass=*/false,
          cdna3::Isa::WF_SIZE);
  uint32_t value = 0;
  std::memcpy(&value, bytes.data(), sizeof(value));
  ASSERT_EQ(value, kOldValue);

  l2.write(kAddr, reinterpret_cast<const uint8_t *>(&kNewValue), sizeof(kNewValue),
           amdgpu::Mtype::RW);

  bytes.fill(0);
  l1.load(addrs, /*lane_mask=*/0x1, /*elem_size=*/4, /*num_elems=*/1, bytes.data(),
          amdgpu::Mtype::RW, /*non_temporal=*/true, /*request_l1_bypass=*/false,
          cdna3::Isa::WF_SIZE);
  std::memcpy(&value, bytes.data(), sizeof(value));
  ASSERT_EQ(value, kNewValue);

  bytes.fill(0);
  l1.load(addrs, /*lane_mask=*/0x1, /*elem_size=*/4, /*num_elems=*/1, bytes.data(),
          amdgpu::Mtype::RW, /*non_temporal=*/false, /*request_l1_bypass=*/false,
          cdna3::Isa::WF_SIZE);
  std::memcpy(&value, bytes.data(), sizeof(value));
  EXPECT_EQ(value, kNewValue);
}

TEST(L1VectorCacheTest, L1BypassReadInvalidatesResidentLine) {
  amdgpu::GpuMemory mem("test_mem");
  amdgpu::L2Cache l2("test_l2");
  amdgpu::L1VectorCache l1(&l2);
  l2.set_backing_memory(&mem);

  constexpr uint64_t kAddr = 0x8200;
  constexpr uint32_t kOldValue = 0x11111111;
  constexpr uint32_t kNewValue = 0x22222222;

  uint64_t addrs[64] = {};
  addrs[0] = kAddr;
  std::array<uint8_t, 64 * sizeof(uint32_t)> bytes{};

  mem.write32(kAddr, kOldValue);
  l1.load(addrs, /*lane_mask=*/0x1, /*elem_size=*/4, /*num_elems=*/1, bytes.data(),
          amdgpu::Mtype::RW, /*non_temporal=*/false, /*request_l1_bypass=*/false,
          cdna3::Isa::WF_SIZE);
  uint32_t value = 0;
  std::memcpy(&value, bytes.data(), sizeof(value));
  ASSERT_EQ(value, kOldValue);

  l2.write(kAddr, reinterpret_cast<const uint8_t *>(&kNewValue), sizeof(kNewValue),
           amdgpu::Mtype::RW);

  bytes.fill(0);
  l1.load(addrs, /*lane_mask=*/0x1, /*elem_size=*/4, /*num_elems=*/1, bytes.data(),
          amdgpu::Mtype::RW, /*non_temporal=*/false, /*request_l1_bypass=*/true,
          cdna3::Isa::WF_SIZE);
  std::memcpy(&value, bytes.data(), sizeof(value));
  ASSERT_EQ(value, kNewValue);

  bytes.fill(0);
  l1.load(addrs, /*lane_mask=*/0x1, /*elem_size=*/4, /*num_elems=*/1, bytes.data(),
          amdgpu::Mtype::RW, /*non_temporal=*/false, /*request_l1_bypass=*/false,
          cdna3::Isa::WF_SIZE);
  std::memcpy(&value, bytes.data(), sizeof(value));
  EXPECT_EQ(value, kNewValue);
}

TEST(L1VectorCacheTest, NonTemporalWriteInvalidatesResidentLine) {
  amdgpu::GpuMemory mem("test_mem");
  amdgpu::L2Cache l2("test_l2");
  amdgpu::L1VectorCache l1(&l2);
  l2.set_backing_memory(&mem);

  constexpr uint64_t kAddr = 0x8300;
  constexpr uint32_t kOldValue = 0x11111111;
  constexpr uint32_t kNewValue = 0x22222222;

  uint64_t addrs[64] = {};
  addrs[0] = kAddr;
  std::array<uint8_t, 64 * sizeof(uint32_t)> bytes{};

  mem.write32(kAddr, kOldValue);
  l1.load(addrs, /*lane_mask=*/0x1, /*elem_size=*/4, /*num_elems=*/1, bytes.data(),
          amdgpu::Mtype::RW, /*non_temporal=*/false, /*request_l1_bypass=*/false,
          cdna3::Isa::WF_SIZE);
  uint32_t value = 0;
  std::memcpy(&value, bytes.data(), sizeof(value));
  ASSERT_EQ(value, kOldValue);

  std::array<uint8_t, 64 * sizeof(uint32_t)> store_bytes{};
  std::memcpy(store_bytes.data(), &kNewValue, sizeof(kNewValue));
  l1.store(addrs, /*lane_mask=*/0x1, /*elem_size=*/4, /*num_elems=*/1, store_bytes.data(),
           amdgpu::Mtype::RW, /*non_temporal=*/true, cdna3::Isa::WF_SIZE);

  bytes.fill(0);
  l1.load(addrs, /*lane_mask=*/0x1, /*elem_size=*/4, /*num_elems=*/1, bytes.data(),
          amdgpu::Mtype::RW, /*non_temporal=*/false, /*request_l1_bypass=*/false,
          cdna3::Isa::WF_SIZE);
  std::memcpy(&value, bytes.data(), sizeof(value));
  EXPECT_EQ(value, kNewValue);
}

// ---------------------------------------------------------------------------
// CU factory tests — verify all 9 ISAs can be instantiated
// ---------------------------------------------------------------------------

class CuFactoryTest : public ::testing::TestWithParam<rj_code_arch_t> {};

TEST_P(CuFactoryTest, CreatesSuccessfully) {
  auto arch = GetParam();
  amdgpu::GpuMemory mem("test_mem");
  amdgpu::L2Cache l2("test_l2");

  amdgpu::ComputeUnitCore::Config cfg{};
  cfg.arch = arch;
  cfg.num_wf_slots = 2;
  cfg.sgprs_per_wf = 102;
  cfg.vgprs_per_wf = 256;
  cfg.lds_size_kb = 64;

  auto cu = amdgpu::ComputeUnitCore::create("test_cu", cfg, &mem, &l2);
  ASSERT_NE(cu, nullptr);
  EXPECT_EQ(cu->arch(), arch);
}

TEST(CuFactoryTest, CdnaAccVgprsDoNotAliasNextWaveSlot) {
  for (rj_code_arch_t arch :
       {ROCJITSU_CODE_ARCH_CDNA2, ROCJITSU_CODE_ARCH_CDNA3, ROCJITSU_CODE_ARCH_CDNA4}) {
    SCOPED_TRACE(::testing::Message() << "arch=" << static_cast<int>(arch));
    amdgpu::GpuMemory mem("test_mem");
    amdgpu::L2Cache l2("test_l2");

    amdgpu::ComputeUnitCore::Config cfg{};
    cfg.arch = arch;
    cfg.num_wf_slots = 2;
    cfg.sgprs_per_wf = 104;
    cfg.vgprs_per_wf = 256;
    cfg.lds_size_kb = 64;

    auto cu = amdgpu::ComputeUnitCore::create("test_cu", cfg, &mem, &l2);
    ASSERT_NE(cu, nullptr);

    auto *wf0 = cu->dispatch_wf(0, 0, cfg.sgprs_per_wf, cfg.vgprs_per_wf);
    ASSERT_NE(wf0, nullptr);
    auto *wf1 = cu->dispatch_wf(1, 0, cfg.sgprs_per_wf, cfg.vgprs_per_wf);
    ASSERT_NE(wf1, nullptr);

    const uint32_t wf0_acc0 = wf0->vgpr_alloc().base + amdgpu::ACC_VGPR_OFFSET;
    const uint32_t wf0_acc_last = wf0_acc0 + kCdnaAccVgprsPerWf - 1;
    const uint32_t wf1_v0 = wf1->vgpr_alloc().base;
    EXPECT_NE(wf0_acc0, wf1_v0);
    EXPECT_LT(wf0_acc_last, wf1_v0);

    cu->write_vgpr(wf0_acc0, 0, 0xA55A0001u);
    cu->write_vgpr(wf0_acc_last, 0, 0xDEADBEEFu);
    cu->write_vgpr(wf1_v0, 0, 0x5AA50002u);

    EXPECT_EQ(cu->read_vgpr(wf0_acc0, 0), 0xA55A0001u);
    EXPECT_EQ(cu->read_vgpr(wf0_acc_last, 0), 0xDEADBEEFu);
    EXPECT_EQ(cu->read_vgpr(wf1_v0, 0), 0x5AA50002u);
  }
}

TEST(CuFactoryTest, CdnaAccVgprsAreClearedOnRedispatch) {
  for (rj_code_arch_t arch :
       {ROCJITSU_CODE_ARCH_CDNA2, ROCJITSU_CODE_ARCH_CDNA3, ROCJITSU_CODE_ARCH_CDNA4}) {
    SCOPED_TRACE(::testing::Message() << "arch=" << static_cast<int>(arch));
    amdgpu::GpuMemory mem("test_mem");
    amdgpu::L2Cache l2("test_l2");

    amdgpu::ComputeUnitCore::Config cfg{};
    cfg.arch = arch;
    cfg.num_wf_slots = 1;
    cfg.sgprs_per_wf = 104;
    cfg.vgprs_per_wf = 256;
    cfg.lds_size_kb = 64;

    auto cu = amdgpu::ComputeUnitCore::create("test_cu", cfg, &mem, &l2);
    ASSERT_NE(cu, nullptr);

    auto *wf = cu->dispatch_wf(0, 0, cfg.sgprs_per_wf, cfg.vgprs_per_wf);
    ASSERT_NE(wf, nullptr);
    const uint32_t acc0 = wf->vgpr_alloc().base + amdgpu::ACC_VGPR_OFFSET;
    const uint32_t acc_last = acc0 + kCdnaAccVgprsPerWf - 1;
    cu->write_vgpr(acc0, 0, 0xFFFFFFFFu);
    cu->write_vgpr(acc_last, 0, 0xDEADBEEFu);

    wf->halt();
    wf = cu->dispatch_wf(0, 0, cfg.sgprs_per_wf, cfg.vgprs_per_wf);
    ASSERT_NE(wf, nullptr);

    EXPECT_EQ(cu->read_vgpr(wf->vgpr_alloc().base + amdgpu::ACC_VGPR_OFFSET, 0), 0u);
    EXPECT_EQ(
        cu->read_vgpr(wf->vgpr_alloc().base + amdgpu::ACC_VGPR_OFFSET + kCdnaAccVgprsPerWf - 1, 0),
        0u);
  }
}

INSTANTIATE_TEST_SUITE_P(AllIsas, CuFactoryTest,
                         ::testing::Values(ROCJITSU_CODE_ARCH_CDNA1, ROCJITSU_CODE_ARCH_CDNA2,
                                           ROCJITSU_CODE_ARCH_CDNA3, ROCJITSU_CODE_ARCH_CDNA4,
                                           ROCJITSU_CODE_ARCH_GFX1250, ROCJITSU_CODE_ARCH_RDNA1,
                                           ROCJITSU_CODE_ARCH_RDNA2, ROCJITSU_CODE_ARCH_RDNA3,
                                           ROCJITSU_CODE_ARCH_RDNA3_5, ROCJITSU_CODE_ARCH_RDNA4));

// ---------------------------------------------------------------------------
// DPP permutation tests
// ---------------------------------------------------------------------------

TEST(DppPermuteTest, QuadPerm) {
  using namespace amdgpu::dpp;
  // quad_perm(1,0,3,2) = swap pairs within each quad
  // Encoding: lane0->1, lane1->0, lane2->3, lane3->2
  // = (1 << 0) | (0 << 2) | (3 << 4) | (2 << 6) = 0xB1
  bool oob = false;
  EXPECT_EQ(dpp_permute(0xB1, 0, 64, oob), 1);
  EXPECT_FALSE(oob);
  EXPECT_EQ(dpp_permute(0xB1, 1, 64, oob), 0);
  EXPECT_FALSE(oob);
  EXPECT_EQ(dpp_permute(0xB1, 2, 64, oob), 3);
  EXPECT_FALSE(oob);
  EXPECT_EQ(dpp_permute(0xB1, 3, 64, oob), 2);
  EXPECT_FALSE(oob);
  // Quad boundary: lane 4 starts a new quad, same permutation.
  EXPECT_EQ(dpp_permute(0xB1, 4, 64, oob), 5);
  EXPECT_FALSE(oob);
  EXPECT_EQ(dpp_permute(0xB1, 5, 64, oob), 4);
  EXPECT_FALSE(oob);
}

TEST(DppPermuteTest, RowShr1) {
  using namespace amdgpu::dpp;
  bool oob = false;
  // row_shr 1 = 0x111: data shifts right, so lane K reads from lane K-1.
  EXPECT_EQ(dpp_permute(0x111, 1, 64, oob), 0);
  EXPECT_FALSE(oob);
  EXPECT_EQ(dpp_permute(0x111, 15, 64, oob), 14);
  EXPECT_FALSE(oob);
  // Lane 0 (first in row) goes OOB (no lane -1).
  oob = false;
  dpp_permute(0x111, 0, 64, oob);
  EXPECT_TRUE(oob);
}

TEST(DppPermuteTest, RowShl1) {
  using namespace amdgpu::dpp;
  bool oob = false;
  // row_shl 1 = 0x101: data shifts left, so lane K reads from lane K+1.
  EXPECT_EQ(dpp_permute(0x101, 0, 64, oob), 1);
  EXPECT_FALSE(oob);
  EXPECT_EQ(dpp_permute(0x101, 14, 64, oob), 15);
  EXPECT_FALSE(oob);
  // Lane 15 (last in row) goes OOB (no lane 16 in this row).
  oob = false;
  dpp_permute(0x101, 15, 64, oob);
  EXPECT_TRUE(oob);
}

TEST(DppPermuteTest, RowRor1) {
  using namespace amdgpu::dpp;
  bool oob = false;
  // row_ror 1 = 0x121: lane K reads from lane K-1, wrapping within the row.
  EXPECT_EQ(dpp_permute(0x121, 0, 64, oob), 15);
  EXPECT_FALSE(oob);
  EXPECT_EQ(dpp_permute(0x121, 1, 64, oob), 0);
  EXPECT_FALSE(oob);
  EXPECT_EQ(dpp_permute(0x121, 16, 64, oob), 31);
  EXPECT_FALSE(oob);
}

TEST(DppPermuteTest, WaveShiftAndRotate) {
  using namespace amdgpu::dpp;
  bool oob = false;

  EXPECT_EQ(dpp_permute(WF_SHL1, 0, 64, oob), 1);
  EXPECT_FALSE(oob);
  oob = false;
  dpp_permute(WF_SHL1, 63, 64, oob);
  EXPECT_TRUE(oob);

  oob = false;
  EXPECT_EQ(dpp_permute(WF_ROL1, 0, 64, oob), 1);
  EXPECT_FALSE(oob);
  EXPECT_EQ(dpp_permute(WF_ROL1, 63, 64, oob), 0);
  EXPECT_FALSE(oob);
  EXPECT_EQ(dpp_permute(WF_ROL1, 31, 32, oob), 0);
  EXPECT_FALSE(oob);

  EXPECT_EQ(dpp_permute(WF_SRL1, 1, 64, oob), 0);
  EXPECT_FALSE(oob);
  oob = false;
  dpp_permute(WF_SRL1, 0, 64, oob);
  EXPECT_TRUE(oob);

  oob = false;
  EXPECT_EQ(dpp_permute(WF_ROR1, 0, 64, oob), 63);
  EXPECT_FALSE(oob);
  EXPECT_EQ(dpp_permute(WF_ROR1, 1, 64, oob), 0);
  EXPECT_FALSE(oob);
}

TEST(DppPermuteTest, RowMirrors) {
  using namespace amdgpu::dpp;
  bool oob = false;
  // row_mirror = 0x140: reverse lane order within a row.
  EXPECT_EQ(dpp_permute(0x140, 0, 64, oob), 15);
  EXPECT_EQ(dpp_permute(0x140, 15, 64, oob), 0);
  EXPECT_EQ(dpp_permute(0x140, 7, 64, oob), 8);
  // Second row.
  EXPECT_EQ(dpp_permute(0x140, 16, 64, oob), 31);

  // row_half_mirror = 0x141: reverse lane order within each 8-lane half-row.
  EXPECT_EQ(dpp_permute(0x141, 0, 64, oob), 7);
  EXPECT_EQ(dpp_permute(0x141, 7, 64, oob), 0);
  EXPECT_EQ(dpp_permute(0x141, 8, 64, oob), 15);
  EXPECT_EQ(dpp_permute(0x141, 16, 64, oob), 23);
}

TEST(DppPermuteTest, RowShareAndXmask) {
  using namespace amdgpu::dpp;
  bool oob = false;
  // row_share with lane_sel=1 = 0x151: every lane in a row reads row lane 1.
  EXPECT_EQ(dpp_permute(0x151, 0, 64, oob), 1);
  EXPECT_EQ(dpp_permute(0x151, 15, 64, oob), 1);
  EXPECT_EQ(dpp_permute(0x151, 16, 64, oob), 17);
  EXPECT_EQ(dpp_permute(0x151, 31, 64, oob), 17);

  // row_xmask with mask=1 = 0x161: XOR lane offset with 1 (swap adjacent pairs).
  EXPECT_EQ(dpp_permute(0x161, 0, 64, oob), 1);
  EXPECT_EQ(dpp_permute(0x161, 1, 64, oob), 0);
  EXPECT_EQ(dpp_permute(0x161, 2, 64, oob), 3);
  EXPECT_EQ(dpp_permute(0x161, 3, 64, oob), 2);
  EXPECT_EQ(dpp_permute(0x161, 16, 64, oob), 17);
  EXPECT_EQ(dpp_permute(0x161, 17, 64, oob), 16);
}

TEST(DppPermuteTest, RowBroadcasts) {
  using namespace amdgpu::dpp;
  bool oob = false;

  EXPECT_EQ(dpp_permute(ROW_BCAST15, 16, 64, oob), 15);
  EXPECT_FALSE(oob);
  EXPECT_EQ(dpp_permute(ROW_BCAST15, 31, 64, oob), 15);
  EXPECT_FALSE(oob);
  EXPECT_EQ(dpp_permute(ROW_BCAST15, 32, 64, oob), 31);
  EXPECT_FALSE(oob);
  EXPECT_EQ(dpp_permute(ROW_BCAST15, 47, 64, oob), 31);
  EXPECT_FALSE(oob);
  EXPECT_EQ(dpp_permute(ROW_BCAST15, 48, 64, oob), 47);
  EXPECT_FALSE(oob);
  EXPECT_EQ(dpp_permute(ROW_BCAST15, 63, 64, oob), 47);
  EXPECT_FALSE(oob);

  oob = false;
  dpp_permute(ROW_BCAST15, 15, 64, oob);
  EXPECT_TRUE(oob);

  oob = false;
  EXPECT_EQ(dpp_permute(ROW_BCAST31, 32, 64, oob), 31);
  EXPECT_FALSE(oob);
  EXPECT_EQ(dpp_permute(ROW_BCAST31, 63, 64, oob), 31);
  EXPECT_FALSE(oob);

  oob = false;
  dpp_permute(ROW_BCAST31, 31, 64, oob);
  EXPECT_TRUE(oob);
}

TEST(DppPermuteTest, Dpp8SelectsWithinGroupsOfEight) {
  using namespace amdgpu::dpp;
  const uint32_t lane_sel = (7u << 0u) | (0u << 3u) | (3u << 6u) | (2u << 9u) | (5u << 12u) |
                            (4u << 15u) | (1u << 18u) | (6u << 21u);

  EXPECT_EQ(dpp8_src_lane(0, lane_sel), 7u);
  EXPECT_EQ(dpp8_src_lane(1, lane_sel), 0u);
  EXPECT_EQ(dpp8_src_lane(6, lane_sel), 1u);
  EXPECT_EQ(dpp8_src_lane(7, lane_sel), 6u);
  EXPECT_EQ(dpp8_src_lane(8, lane_sel), 15u);
  EXPECT_EQ(dpp8_src_lane(15, lane_sel), 14u);
}

TEST(DppPermuteTest, True16SourceByteMaskFollowsOpSel) {
  using namespace amdgpu::dpp;

  EXPECT_EQ(true16_source_byte_mask(/*opsel=*/0b0000, /*source_index=*/0),
            ExecutionPlugin::kLowHalfByteMask);
  EXPECT_EQ(true16_source_byte_mask(/*opsel=*/0b0001, /*source_index=*/0),
            ExecutionPlugin::kHighHalfByteMask);
  EXPECT_EQ(true16_source_byte_mask(/*opsel=*/0b0000, /*source_index=*/1),
            ExecutionPlugin::kLowHalfByteMask);
  EXPECT_EQ(true16_source_byte_mask(/*opsel=*/0b0010, /*source_index=*/1),
            ExecutionPlugin::kHighHalfByteMask);
}

TEST(DppPermuteTest, AccessPlanTracksSelectedSources) {
  using namespace amdgpu::dpp;
  constexpr uint64_t kAllLanesActive = ~0ULL;
  const auto plan =
      make_dpp_access_plan(64, ROW_SHR1, 0xF, 0xF, /*bound_ctrl=*/1, /*fi=*/1, kAllLanesActive);

  EXPECT_EQ(plan.source_lane_for_destination[0], DppAccessPlan::kNoSourceLane);
  EXPECT_EQ(plan.source_lane_for_destination[1], 0);
  EXPECT_EQ(plan.source_lane_for_destination[5], 4);
  EXPECT_EQ(plan.source_lane_for_destination[17], 16);
  EXPECT_TRUE(plan.source_lane_mask & (uint64_t{1} << 0));
  EXPECT_TRUE(plan.source_lane_mask & (uint64_t{1} << 4));
  EXPECT_TRUE(plan.source_lane_mask & (uint64_t{1} << 16));

  const auto row_masked =
      make_dpp_access_plan(64, ROW_SHR1, 0xE, 0xF, /*bound_ctrl=*/1, /*fi=*/1, kAllLanesActive);
  EXPECT_EQ(row_masked.source_lane_for_destination[5], DppAccessPlan::kNoSourceLane);

  const auto bank_masked =
      make_dpp_access_plan(64, ROW_SHR1, 0xF, 0xD, /*bound_ctrl=*/1, /*fi=*/1, kAllLanesActive);
  EXPECT_EQ(bank_masked.source_lane_for_destination[5], DppAccessPlan::kNoSourceLane);
}

TEST(DppPermuteTest, FetchInactiveControlsInactiveSourceReads) {
  using namespace amdgpu::dpp;
  constexpr uint64_t kLane0Inactive = ~1ULL;

  const auto dpp_fi0 =
      make_dpp_access_plan(64, ROW_SHR1, 0xF, 0xF, /*bound_ctrl=*/1, /*fi=*/0, kLane0Inactive);
  EXPECT_EQ(dpp_fi0.source_lane_for_destination[1], DppAccessPlan::kNoSourceLane);
  EXPECT_EQ(dpp_fi0.source_lane_mask & 1u, 0u);

  const auto dpp_fi1 =
      make_dpp_access_plan(64, ROW_SHR1, 0xF, 0xF, /*bound_ctrl=*/1, /*fi=*/1, kLane0Inactive);
  EXPECT_EQ(dpp_fi1.source_lane_for_destination[1], 0);
  EXPECT_EQ(dpp_fi1.source_lane_mask & 1u, 1u);

  constexpr uint32_t kAllLanesSelectLane0 = 0;
  const auto dpp8_fi0 = make_dpp8_access_plan(32, kAllLanesSelectLane0, /*fi=*/0, kLane0Inactive);
  EXPECT_EQ(dpp8_fi0.source_lane_for_destination[1], DppAccessPlan::kNoSourceLane);
  EXPECT_EQ(dpp8_fi0.source_lane_mask & 1u, 0u);

  const auto dpp8_fi1 = make_dpp8_access_plan(32, kAllLanesSelectLane0, /*fi=*/1, kLane0Inactive);
  EXPECT_EQ(dpp8_fi1.source_lane_for_destination[1], 0);
  EXPECT_EQ(dpp8_fi1.source_lane_mask & 1u, 1u);

  EXPECT_EQ(src_dpp8_fi(amdgpu::SRC_DPP8_FI_0), 0u);
  EXPECT_EQ(src_dpp8_fi(amdgpu::SRC_DPP8_FI_1), 1u);
}

TEST(DppPermuteTest, WriteMaskHonorsBoundCtrlAndBroadcastValidity) {
  using namespace amdgpu::dpp;

  EXPECT_FALSE(dpp_lane_write_enabled(0, 64, ROW_SHR1, 0xF, 0xF, 0));
  EXPECT_TRUE(dpp_lane_write_enabled(0, 64, ROW_SHR1, 0xF, 0xF, 1));
  EXPECT_TRUE(dpp_lane_write_enabled(1, 64, ROW_SHR1, 0xF, 0xF, 0));

  EXPECT_FALSE(dpp_lane_write_enabled(15, 64, ROW_BCAST15, 0xF, 0xF, 0));
  EXPECT_TRUE(dpp_lane_write_enabled(16, 64, ROW_BCAST15, 0xF, 0xF, 0));
  EXPECT_TRUE(dpp_lane_write_enabled(32, 64, ROW_BCAST15, 0xF, 0xF, 0));
  EXPECT_FALSE(dpp_lane_write_enabled(31, 64, ROW_BCAST31, 0xF, 0xF, 0));
  EXPECT_TRUE(dpp_lane_write_enabled(32, 64, ROW_BCAST31, 0xF, 0xF, 0));

  EXPECT_FALSE(dpp_lane_write_enabled(16, 64, ROW_BCAST15, 0xD, 0xF, 0));
  EXPECT_FALSE(dpp_lane_write_enabled(20, 64, ROW_BCAST15, 0xF, 0xD, 0));

  uint64_t mask = dpp_write_mask(64, ROW_BCAST15, 0xF, 0xF, 0);
  EXPECT_EQ(mask, 0xFFFFFFFFFFFF0000ULL);

  mask = dpp_write_mask(64, ROW_BCAST31, 0xF, 0xF, 0);
  EXPECT_EQ(mask, 0xFFFFFFFF00000000ULL);

  mask = dpp_write_mask(32, ROW_BCAST15, 0xF, 0xF, 0);
  EXPECT_EQ(mask, 0xFFFF0000ULL);
}

struct Cdna1DppTraits {
  static constexpr const char *name = "cdna1";
  static constexpr rj_code_arch_t arch = ROCJITSU_CODE_ARCH_CDNA1;
  using MachineInst = cdna1::MachineInst;
  using Vop1VopDppMachineInst = cdna1::Vop1VopDppMachineInst;
  using VMovB32Vop1 = cdna1::VMovB32Vop1;
  using VCmpEqU32Vopc = cdna1::VCmpEqU32Vopc;
  using VCmpxEqU32Vopc = cdna1::VCmpxEqU32Vopc;
};

struct Cdna2DppTraits {
  static constexpr const char *name = "cdna2";
  static constexpr rj_code_arch_t arch = ROCJITSU_CODE_ARCH_CDNA2;
  static constexpr uint32_t wf_size = 64;
  using MachineInst = cdna2::MachineInst;
  using Vop1VopDppMachineInst = cdna2::Vop1VopDppMachineInst;
  using Vop1F64DppMachineInst = cdna2::Vop1VopDppMachineInst;
  using VMovB32Vop1 = cdna2::VMovB32Vop1;
  using VCvtF64I32Vop1 = cdna2::VCvtF64I32Vop1;
  using VCmpEqU32Vopc = cdna2::VCmpEqU32Vopc;
  using VCmpxEqU32Vopc = cdna2::VCmpxEqU32Vopc;
};

struct Cdna3DppTraits {
  static constexpr const char *name = "cdna3";
  static constexpr rj_code_arch_t arch = ROCJITSU_CODE_ARCH_CDNA3;
  static constexpr uint32_t wf_size = 64;
  using MachineInst = cdna3::MachineInst;
  using Vop1VopDppMachineInst = cdna3::Vop1VopDppMachineInst;
  using Vop1F64DppMachineInst = cdna3::Vop1VopDppMachineInst;
  using VMovB32Vop1 = cdna3::VMovB32Vop1;
  using VCvtF64I32Vop1 = cdna3::VCvtF64I32Vop1;
  using VCmpEqU32Vopc = cdna3::VCmpEqU32Vopc;
  using VCmpxEqU32Vopc = cdna3::VCmpxEqU32Vopc;
};

struct Cdna4DppTraits {
  static constexpr const char *name = "cdna4";
  static constexpr rj_code_arch_t arch = ROCJITSU_CODE_ARCH_CDNA4;
  static constexpr uint32_t wf_size = 64;
  using MachineInst = cdna4::MachineInst;
  using Vop1VopDppMachineInst = cdna4::Vop1VopDppMachineInst;
  using Vop1F64DppMachineInst = cdna4::Vop1VopDppMachineInst;
  using VMovB32Vop1 = cdna4::VMovB32Vop1;
  using VCvtF64I32Vop1 = cdna4::VCvtF64I32Vop1;
  using VCmpEqU32Vopc = cdna4::VCmpEqU32Vopc;
  using VCmpxEqU32Vopc = cdna4::VCmpxEqU32Vopc;
};

struct Rdna1DppTraits {
  static constexpr const char *name = "rdna1";
  static constexpr rj_code_arch_t arch = ROCJITSU_CODE_ARCH_RDNA1;
  using MachineInst = rdna1::MachineInst;
  using VopcMachineInst = rdna1::VopcMachineInst;
  using Vop1VopDpp16MachineInst = rdna1::Vop1VopDpp16MachineInst;
  using Vop1VopDpp8MachineInst = rdna1::Vop1VopDpp8MachineInst;
  using VMovB32Vop1 = rdna1::VMovB32Vop1;
  using VCmpEqU32Vopc = rdna1::VCmpEqU32Vopc;
  using VCmpxEqU32Vopc = rdna1::VCmpxEqU32Vopc;
};

struct Rdna2DppTraits {
  static constexpr const char *name = "rdna2";
  static constexpr rj_code_arch_t arch = ROCJITSU_CODE_ARCH_RDNA2;
  using MachineInst = rdna2::MachineInst;
  using VopcMachineInst = rdna2::VopcMachineInst;
  using Vop1VopDpp16MachineInst = rdna2::Vop1VopDpp16MachineInst;
  using Vop1VopDpp8MachineInst = rdna2::Vop1VopDpp8MachineInst;
  using VMovB32Vop1 = rdna2::VMovB32Vop1;
  using VCmpEqU32Vopc = rdna2::VCmpEqU32Vopc;
  using VCmpxEqU32Vopc = rdna2::VCmpxEqU32Vopc;
};

struct Rdna4DppTraits {
  static constexpr const char *name = "rdna4";
  static constexpr rj_code_arch_t arch = ROCJITSU_CODE_ARCH_RDNA4;
  using MachineInst = rdna4::MachineInst;
  using VopcMachineInst = rdna4::VopcMachineInst;
  using Vop1VopDpp16MachineInst = rdna4::Vop1VopDpp16MachineInst;
  using Vop1VopDpp8MachineInst = rdna4::Vop1VopDpp8MachineInst;
  using Vop1F64DppMachineInst = rdna4::Vop1VopDpp16MachineInst;
  using VopcVopDpp16MachineInst = rdna4::VopcVopDpp16MachineInst;
  using VMovB32Vop1 = rdna4::VMovB32Vop1;
  using VCvtF64I32Vop1 = rdna4::VCvtF64I32Vop1;
  using VCmpEqU32Vopc = rdna4::VCmpEqU32Vopc;
  using VCmpxEqU32Vopc = rdna4::VCmpxEqU32Vopc;
};

struct Rdna3DppTraits {
  static constexpr const char *name = "rdna3";
  static constexpr rj_code_arch_t arch = ROCJITSU_CODE_ARCH_RDNA3;
  using MachineInst = rdna3::MachineInst;
  using VopcMachineInst = rdna3::VopcMachineInst;
  using Vop1VopDpp16MachineInst = rdna3::Vop1VopDpp16MachineInst;
  using Vop1VopDpp8MachineInst = rdna3::Vop1VopDpp8MachineInst;
  using VopcVopDpp16MachineInst = rdna3::VopcVopDpp16MachineInst;
  using VMovB32Vop1 = rdna3::VMovB32Vop1;
  using VCmpEqU32Vopc = rdna3::VCmpEqU32Vopc;
  using VCmpxEqU32Vopc = rdna3::VCmpxEqU32Vopc;
};

struct Rdna3_5DppTraits {
  static constexpr const char *name = "rdna3_5";
  static constexpr rj_code_arch_t arch = ROCJITSU_CODE_ARCH_RDNA3_5;
  using MachineInst = rdna3_5::MachineInst;
  using VopcMachineInst = rdna3_5::VopcMachineInst;
  using Vop1VopDpp16MachineInst = rdna3_5::Vop1VopDpp16MachineInst;
  using Vop1VopDpp8MachineInst = rdna3_5::Vop1VopDpp8MachineInst;
  using VopcVopDpp16MachineInst = rdna3_5::VopcVopDpp16MachineInst;
  using VMovB32Vop1 = rdna3_5::VMovB32Vop1;
  using VCmpEqU32Vopc = rdna3_5::VCmpEqU32Vopc;
  using VCmpxEqU32Vopc = rdna3_5::VCmpxEqU32Vopc;
};

struct Gfx1250DppTraits {
  static constexpr const char *name = "gfx1250";
  static constexpr rj_code_arch_t arch = ROCJITSU_CODE_ARCH_GFX1250;
  using MachineInst = gfx1250::MachineInst;
  using VopcMachineInst = gfx1250::VopcMachineInst;
  using Vop1VopDpp16MachineInst = gfx1250::Vop1VopDpp16MachineInst;
  using Vop1VopDpp8MachineInst = gfx1250::Vop1VopDpp8MachineInst;
  using Vop1F64DppMachineInst = gfx1250::Vop1VopDpp16MachineInst;
  using VopcVopDpp16MachineInst = gfx1250::VopcVopDpp16MachineInst;
  using VMovB32Vop1 = gfx1250::VMovB32Vop1;
  using VCvtF64I32Vop1 = gfx1250::VCvtF64I32Vop1;
  using VCmpEqU32Vopc = gfx1250::VCmpEqU32Vopc;
  using VCmpxEqU32Vopc = gfx1250::VCmpxEqU32Vopc;
};

// Test-only RDNA4 trait for exercising 64-bit EXEC handling. The production
// RDNA factory still constructs the ISA default Wave32 form.
struct Rdna4Wave64DppTestIsa : rdna4::Isa {
  static constexpr uint32_t WF_SIZE = 64;
  static constexpr uint32_t WF_SIZE_MAX = 64;
};

static_assert(GpuIsa<Rdna4Wave64DppTestIsa>);

template <typename Traits> void cdna_generated_vop1_uses_shared_row_broadcast() {
  SCOPED_TRACE(Traits::name);
  amdgpu::GpuMemory mem(std::string(Traits::name) + "_dpp_broadcast_mem");
  amdgpu::L2Cache l2(std::string(Traits::name) + "_dpp_broadcast_l2");

  amdgpu::ComputeUnitCore::Config cfg{};
  cfg.arch = Traits::arch;
  cfg.num_wf_slots = 1;
  cfg.sgprs_per_wf = 104;
  cfg.vgprs_per_wf = 32;
  cfg.lds_size_kb = 64;

  auto cu = amdgpu::ComputeUnitCore::create(std::string(Traits::name) + "_dpp_broadcast_cu", cfg,
                                            &mem, &l2);
  ASSERT_NE(cu, nullptr);

  auto *wf = cu->dispatch_wf(0, 0, cfg.sgprs_per_wf, cfg.vgprs_per_wf);
  ASSERT_NE(wf, nullptr);
  ASSERT_EQ(wf->wf_size(), 64u);
  wf->set_exec(~0ULL);

  constexpr uint32_t kSrc = 4;
  constexpr uint32_t kDst = 8;
  uint32_t vbase = wf->vgpr_alloc().base;
  for (uint32_t lane = 0; lane < wf->wf_size(); ++lane) {
    cu->write_vgpr(vbase + kSrc, lane, 0x1000u + lane);
    cu->write_vgpr(vbase + kDst, lane, 0xDEAD0000u + lane);
  }

  typename Traits::Vop1VopDppMachineInst raw{};
  raw.src0 = amdgpu::SRC_DPP;
  raw.vsrc0 = kSrc;
  raw.vdst = kDst;
  raw.dpp_ctrl = amdgpu::dpp::ROW_BCAST15;
  raw.bound_ctrl = 1;
  raw.bank_mask = 0xF;
  raw.row_mask = 0xF;

  typename Traits::VMovB32Vop1 inst(reinterpret_cast<const typename Traits::MachineInst *>(&raw));
  inst.execute_impl(*wf);

  EXPECT_EQ(cu->read_vgpr(vbase + kDst, 0), 0u);
  EXPECT_EQ(cu->read_vgpr(vbase + kDst, 15), 0u);
  EXPECT_EQ(cu->read_vgpr(vbase + kDst, 16), 0x100Fu);
  EXPECT_EQ(cu->read_vgpr(vbase + kDst, 31), 0x100Fu);
  EXPECT_EQ(cu->read_vgpr(vbase + kDst, 32), 0x101Fu);
  EXPECT_EQ(cu->read_vgpr(vbase + kDst, 47), 0x101Fu);
  EXPECT_EQ(cu->read_vgpr(vbase + kDst, 48), 0x102Fu);
  EXPECT_EQ(cu->read_vgpr(vbase + kDst, 63), 0x102Fu);
}

template <typename Traits> void cdna_generated_vop1_dpp_write_mask_honors_bound_ctrl() {
  SCOPED_TRACE(Traits::name);
  amdgpu::GpuMemory mem(std::string(Traits::name) + "_dpp_write_mask_mem");
  amdgpu::L2Cache l2(std::string(Traits::name) + "_dpp_write_mask_l2");

  amdgpu::ComputeUnitCore::Config cfg{};
  cfg.arch = Traits::arch;
  cfg.num_wf_slots = 1;
  cfg.sgprs_per_wf = 104;
  cfg.vgprs_per_wf = 32;
  cfg.lds_size_kb = 64;

  auto cu = amdgpu::ComputeUnitCore::create(std::string(Traits::name) + "_dpp_write_mask_cu", cfg,
                                            &mem, &l2);
  ASSERT_NE(cu, nullptr);

  auto *wf = cu->dispatch_wf(0, 0, cfg.sgprs_per_wf, cfg.vgprs_per_wf);
  ASSERT_NE(wf, nullptr);
  ASSERT_EQ(wf->wf_size(), 64u);
  wf->set_exec(~0ULL);

  constexpr uint32_t kSrc = 4;
  constexpr uint32_t kDst = 8;
  uint32_t vbase = wf->vgpr_alloc().base;
  for (uint32_t lane = 0; lane < wf->wf_size(); ++lane) {
    cu->write_vgpr(vbase + kSrc, lane, 0x1000u + lane);
    cu->write_vgpr(vbase + kDst, lane, 0xDEAD0000u + lane);
  }

  typename Traits::Vop1VopDppMachineInst raw{};
  raw.src0 = amdgpu::SRC_DPP;
  raw.vsrc0 = kSrc;
  raw.vdst = kDst;
  raw.dpp_ctrl = amdgpu::dpp::ROW_BCAST15;
  raw.bound_ctrl = 0;
  raw.bank_mask = 0xF;
  raw.row_mask = 0xF;

  typename Traits::VMovB32Vop1 inst(reinterpret_cast<const typename Traits::MachineInst *>(&raw));
  inst.execute_impl(*wf);

  EXPECT_EQ(cu->read_vgpr(vbase + kDst, 0), 0xDEAD0000u);
  EXPECT_EQ(cu->read_vgpr(vbase + kDst, 15), 0xDEAD000Fu);
  EXPECT_EQ(cu->read_vgpr(vbase + kDst, 16), 0x100Fu);
  EXPECT_EQ(cu->read_vgpr(vbase + kDst, 31), 0x100Fu);
  EXPECT_EQ(cu->read_vgpr(vbase + kDst, 32), 0x101Fu);
  EXPECT_EQ(cu->read_vgpr(vbase + kDst, 47), 0x101Fu);
  EXPECT_EQ(cu->read_vgpr(vbase + kDst, 48), 0x102Fu);
  EXPECT_EQ(cu->read_vgpr(vbase + kDst, 63), 0x102Fu);
}

template <typename Traits> void wave32_generated_vop1_dpp_write_mask_honors_bound_ctrl() {
  SCOPED_TRACE(Traits::name);
  amdgpu::GpuMemory mem(std::string(Traits::name) + "_dpp_vop1_wave32_write_mask_mem");
  amdgpu::L2Cache l2(std::string(Traits::name) + "_dpp_vop1_wave32_write_mask_l2");

  amdgpu::ComputeUnitCore::Config cfg{};
  cfg.arch = Traits::arch;
  cfg.num_wf_slots = 1;
  cfg.sgprs_per_wf = 106;
  cfg.vgprs_per_wf = 32;
  cfg.lds_size_kb = 64;

  auto cu = amdgpu::ComputeUnitCore::create(
      std::string(Traits::name) + "_dpp_vop1_wave32_write_mask_cu", cfg, &mem, &l2);
  ASSERT_NE(cu, nullptr);

  auto *wf = cu->dispatch_wf(0, 0, cfg.sgprs_per_wf, cfg.vgprs_per_wf);
  ASSERT_NE(wf, nullptr);
  ASSERT_EQ(wf->wf_size(), 32u);
  wf->set_exec(0xFFFFFFFFULL);

  constexpr uint32_t kSrc = 4;
  constexpr uint32_t kDst = 8;
  uint32_t vbase = wf->vgpr_alloc().base;
  for (uint32_t lane = 0; lane < wf->wf_size(); ++lane) {
    cu->write_vgpr(vbase + kSrc, lane, 0x1000u + lane);
    cu->write_vgpr(vbase + kDst, lane, 0xDEAD0000u + lane);
  }

  typename Traits::Vop1VopDpp16MachineInst raw{};
  raw.src0 = amdgpu::SRC_DPP;
  raw.vsrc0 = kSrc;
  raw.vdst = kDst;
  raw.dpp_ctrl = amdgpu::dpp::ROW_BCAST15;
  raw.bound_ctrl = 0;
  raw.bank_mask = 0xF;
  raw.row_mask = 0xF;

  typename Traits::VMovB32Vop1 inst(reinterpret_cast<const typename Traits::MachineInst *>(&raw));
  inst.execute_impl(*wf);

  for (uint32_t lane = 0; lane < 16; ++lane)
    EXPECT_EQ(cu->read_vgpr(vbase + kDst, lane), 0xDEAD0000u + lane);
  for (uint32_t lane = 16; lane < 32; ++lane)
    EXPECT_EQ(cu->read_vgpr(vbase + kDst, lane), 0x100Fu);
}

template <typename Traits> void generated_vop1_dpp64_preserves_masked_destination() {
  SCOPED_TRACE(Traits::name);
  amdgpu::GpuMemory mem(std::string(Traits::name) + "_dpp_vop1_f64_mask_mem");
  amdgpu::L2Cache l2(std::string(Traits::name) + "_dpp_vop1_f64_mask_l2");

  amdgpu::ComputeUnitCore::Config cfg{};
  cfg.arch = Traits::arch;
  cfg.num_wf_slots = 1;
  cfg.sgprs_per_wf = 106;
  cfg.vgprs_per_wf = 32;
  cfg.lds_size_kb = 64;

  auto cu = amdgpu::ComputeUnitCore::create(std::string(Traits::name) + "_dpp_vop1_f64_mask_cu",
                                            cfg, &mem, &l2);
  ASSERT_NE(cu, nullptr);

  auto *wf = cu->dispatch_wf(0, 0, cfg.sgprs_per_wf, cfg.vgprs_per_wf);
  ASSERT_NE(wf, nullptr);
  ASSERT_EQ(wf->wf_size(), Traits::wf_size);
  wf->set_exec(Traits::wf_size == 64 ? ~0ULL : 0xFFFFFFFFULL);

  constexpr uint32_t kSrc = 4;
  constexpr uint32_t kDst = 8;
  constexpr uint64_t kOldDstBase = 0xA5A500005A5A0000ULL;
  uint32_t vbase = wf->vgpr_alloc().base;
  for (uint32_t lane = 0; lane < wf->wf_size(); ++lane) {
    const uint64_t old_dst = kOldDstBase + lane;
    cu->write_vgpr(vbase + kSrc, lane, 100u + lane);
    cu->write_vgpr(vbase + kDst, lane, static_cast<uint32_t>(old_dst));
    cu->write_vgpr(vbase + kDst + 1, lane, static_cast<uint32_t>(old_dst >> 32));
  }

  typename Traits::Vop1F64DppMachineInst raw{};
  raw.src0 = amdgpu::SRC_DPP;
  raw.vsrc0 = kSrc;
  raw.vdst = kDst;
  raw.dpp_ctrl = amdgpu::dpp::ROW_SHARE_BASE + 1; // row_newbcast:1
  raw.bound_ctrl = 1;
  raw.bank_mask = 0xF;
  raw.row_mask = 0x1;

  typename Traits::VCvtF64I32Vop1 inst(
      reinterpret_cast<const typename Traits::MachineInst *>(&raw));
  inst.execute_impl(*wf);

  for (uint32_t lane = 0; lane < wf->wf_size(); ++lane) {
    const uint64_t actual = static_cast<uint64_t>(cu->read_vgpr(vbase + kDst, lane)) |
                            (static_cast<uint64_t>(cu->read_vgpr(vbase + kDst + 1, lane)) << 32);
    const uint64_t expected =
        lane < 16 ? std::bit_cast<uint64_t>(static_cast<double>(101u)) : kOldDstBase + lane;
    EXPECT_EQ(actual, expected) << "lane " << lane;
  }
}

void rdna4_generated_vop1_64_preserves_inactive_destination() {
  amdgpu::GpuMemory mem("rdna4_vop1_f64_exec_mask_mem");
  amdgpu::L2Cache l2("rdna4_vop1_f64_exec_mask_l2");

  amdgpu::ComputeUnitCore::Config cfg{};
  cfg.arch = ROCJITSU_CODE_ARCH_RDNA4;
  cfg.num_wf_slots = 1;
  cfg.sgprs_per_wf = 106;
  cfg.vgprs_per_wf = 32;
  cfg.lds_size_kb = 64;

  auto cu = amdgpu::ComputeUnitCore::create("rdna4_vop1_f64_exec_mask_cu", cfg, &mem, &l2);
  ASSERT_NE(cu, nullptr);

  auto *wf = cu->dispatch_wf(0, 0, cfg.sgprs_per_wf, cfg.vgprs_per_wf);
  ASSERT_NE(wf, nullptr);
  ASSERT_EQ(wf->wf_size(), 32u);
  constexpr uint64_t kExec = 0x55555555ULL;
  wf->set_exec(kExec);

  constexpr uint32_t kSrc = 4;
  constexpr uint32_t kDst = 8;
  constexpr uint64_t kOldDstBase = 0xA5A500005A5A0000ULL;
  uint32_t vbase = wf->vgpr_alloc().base;
  for (uint32_t lane = 0; lane < wf->wf_size(); ++lane) {
    const uint64_t old_dst = kOldDstBase + lane;
    cu->write_vgpr(vbase + kSrc, lane, 100u + lane);
    cu->write_vgpr(vbase + kDst, lane, static_cast<uint32_t>(old_dst));
    cu->write_vgpr(vbase + kDst + 1, lane, static_cast<uint32_t>(old_dst >> 32));
  }

  rdna4::Vop1MachineInst raw{};
  raw.src0 = 256 + kSrc;
  raw.vdst = kDst;

  rdna4::VCvtF64I32Vop1 inst(reinterpret_cast<const rdna4::MachineInst *>(&raw));
  inst.execute_impl(*wf);

  for (uint32_t lane = 0; lane < wf->wf_size(); ++lane) {
    const uint64_t expected = (kExec & (1ULL << lane))
                                  ? std::bit_cast<uint64_t>(static_cast<double>(100u + lane))
                                  : kOldDstBase + lane;
    EXPECT_EQ(cu->read_vgpr(vbase + kDst, lane), static_cast<uint32_t>(expected))
        << "low dword, lane " << lane;
    EXPECT_EQ(cu->read_vgpr(vbase + kDst + 1, lane), static_cast<uint32_t>(expected >> 32))
        << "high dword, lane " << lane;
  }
}

template <typename Traits> void wave32_generated_vop1_dpp16_fetch_inactive_uses_fi() {
  SCOPED_TRACE(Traits::name);
  amdgpu::GpuMemory mem(std::string(Traits::name) + "_dpp16_fi_mem");
  amdgpu::L2Cache l2(std::string(Traits::name) + "_dpp16_fi_l2");

  amdgpu::ComputeUnitCore::Config cfg{};
  cfg.arch = Traits::arch;
  cfg.num_wf_slots = 1;
  cfg.sgprs_per_wf = 106;
  cfg.vgprs_per_wf = 32;
  cfg.lds_size_kb = 64;

  auto cu =
      amdgpu::ComputeUnitCore::create(std::string(Traits::name) + "_dpp16_fi_cu", cfg, &mem, &l2);
  ASSERT_NE(cu, nullptr);

  auto *wf = cu->dispatch_wf(0, 0, cfg.sgprs_per_wf, cfg.vgprs_per_wf);
  ASSERT_NE(wf, nullptr);
  ASSERT_EQ(wf->wf_size(), 32u);

  constexpr uint32_t kSrc = 4;
  constexpr uint32_t kDst = 8;
  constexpr uint32_t kSrcLane0Value = 0xA5A50000u;
  uint32_t vbase = wf->vgpr_alloc().base;

  auto run = [&](uint32_t fi) {
    wf->set_exec(0xFFFFFFFEULL);
    for (uint32_t lane = 0; lane < wf->wf_size(); ++lane) {
      cu->write_vgpr(vbase + kSrc, lane, kSrcLane0Value + lane);
      cu->write_vgpr(vbase + kDst, lane, 0xDEAD0000u + lane);
    }

    typename Traits::Vop1VopDpp16MachineInst raw{};
    raw.src0 = amdgpu::SRC_DPP;
    raw.vsrc0 = kSrc;
    raw.vdst = kDst;
    raw.dpp_ctrl = amdgpu::dpp::ROW_SHR1;
    raw.fi = fi;
    raw.bound_ctrl = 1;
    raw.bank_mask = 0xF;
    raw.row_mask = 0xF;

    typename Traits::VMovB32Vop1 inst(reinterpret_cast<const typename Traits::MachineInst *>(&raw));
    inst.execute_impl(*wf);
    return cu->read_vgpr(vbase + kDst, 1);
  };

  EXPECT_EQ(run(0), 0u);
  EXPECT_EQ(run(1), kSrcLane0Value);
}

template <typename Traits> void wave32_generated_vop1_dpp16_fi_zero_precedes_bound_ctrl() {
  SCOPED_TRACE(Traits::name);
  amdgpu::GpuMemory mem(std::string(Traits::name) + "_dpp16_fi_bound_mem");
  amdgpu::L2Cache l2(std::string(Traits::name) + "_dpp16_fi_bound_l2");

  amdgpu::ComputeUnitCore::Config cfg{};
  cfg.arch = Traits::arch;
  cfg.num_wf_slots = 1;
  cfg.sgprs_per_wf = 106;
  cfg.vgprs_per_wf = 32;
  cfg.lds_size_kb = 64;

  auto cu = amdgpu::ComputeUnitCore::create(std::string(Traits::name) + "_dpp16_fi_bound_cu", cfg,
                                            &mem, &l2);
  ASSERT_NE(cu, nullptr);

  auto *wf = cu->dispatch_wf(0, 0, cfg.sgprs_per_wf, cfg.vgprs_per_wf);
  ASSERT_NE(wf, nullptr);
  ASSERT_EQ(wf->wf_size(), 32u);
  wf->set_exec(0xFFFFFFFEULL);

  constexpr uint32_t kSrc = 4;
  constexpr uint32_t kDst = 8;
  constexpr uint32_t kSrcLane0Value = 0xA5A50000u;
  uint32_t vbase = wf->vgpr_alloc().base;
  for (uint32_t lane = 0; lane < wf->wf_size(); ++lane) {
    cu->write_vgpr(vbase + kSrc, lane, kSrcLane0Value + lane);
    cu->write_vgpr(vbase + kDst, lane, 0xDEAD0000u + lane);
  }

  typename Traits::Vop1VopDpp16MachineInst raw{};
  raw.src0 = amdgpu::SRC_DPP;
  raw.vsrc0 = kSrc;
  raw.vdst = kDst;
  raw.dpp_ctrl = amdgpu::dpp::ROW_SHR1;
  raw.fi = 0;
  raw.bound_ctrl = 0;
  raw.bank_mask = 0xF;
  raw.row_mask = 0xF;

  typename Traits::VMovB32Vop1 inst(reinterpret_cast<const typename Traits::MachineInst *>(&raw));
  inst.execute_impl(*wf);

  EXPECT_EQ(cu->read_vgpr(vbase + kDst, 0), 0xDEAD0000u);
  EXPECT_EQ(cu->read_vgpr(vbase + kDst, 1), 0u);
}

void rdna4_wave64_generated_vop1_dpp16_fetch_inactive_uses_upper_exec_bit() {
  amdgpu::GpuMemory mem("rdna4_wave64_dpp16_fi_mem");
  amdgpu::L2Cache l2("rdna4_wave64_dpp16_fi_l2");

  amdgpu::ComputeUnitCore::Config cfg{};
  cfg.arch = ROCJITSU_CODE_ARCH_RDNA4;
  cfg.num_wf_slots = 1;
  cfg.sgprs_per_wf = rdna4::Isa::MAX_SGPRS_PER_WF;
  cfg.vgprs_per_wf = 32;
  cfg.lds_size_kb = 64;

  auto cu = std::make_unique<
      amdgpu::IsaExecComputeUnit<simdojo::ExecMode::FUNCTIONAL, Rdna4Wave64DppTestIsa>>(
      "rdna4_wave64_dpp16_fi_cu", cfg, &mem, &l2);

  auto *wf = cu->dispatch_wf(0, 0, cfg.sgprs_per_wf, cfg.vgprs_per_wf);
  ASSERT_NE(wf, nullptr);
  ASSERT_EQ(wf->wf_size(), 64u);

  constexpr uint32_t kSrc = 4;
  constexpr uint32_t kDst = 8;
  constexpr uint32_t kSrcLane32Value = 0x5A5A0020u;
  uint32_t vbase = wf->vgpr_alloc().base;

  auto run = [&](uint32_t fi) {
    wf->set_exec(~(1ULL << 32u));
    for (uint32_t lane = 0; lane < wf->wf_size(); ++lane) {
      cu->write_vgpr(vbase + kSrc, lane, 0x5A5A0000u + lane);
      cu->write_vgpr(vbase + kDst, lane, 0xDEAD0000u + lane);
    }

    rdna4::Vop1VopDpp16MachineInst raw{};
    raw.src0 = amdgpu::SRC_DPP;
    raw.vsrc0 = kSrc;
    raw.vdst = kDst;
    raw.dpp_ctrl = amdgpu::dpp::ROW_SHR1;
    raw.fi = fi;
    raw.bound_ctrl = 0;
    raw.bank_mask = 0xF;
    raw.row_mask = 0xF;

    rdna4::VMovB32Vop1 inst(reinterpret_cast<const rdna4::MachineInst *>(&raw));
    inst.execute_impl(*wf);
    return cu->read_vgpr(vbase + kDst, 33);
  };

  EXPECT_EQ(run(0), 0u);
  EXPECT_EQ(run(1), kSrcLane32Value);
}

void rdna4_wave64_generated_vop1_dpp8_fetch_inactive_uses_upper_exec_bit() {
  amdgpu::GpuMemory mem("rdna4_wave64_dpp8_fi_mem");
  amdgpu::L2Cache l2("rdna4_wave64_dpp8_fi_l2");

  amdgpu::ComputeUnitCore::Config cfg{};
  cfg.arch = ROCJITSU_CODE_ARCH_RDNA4;
  cfg.num_wf_slots = 1;
  cfg.sgprs_per_wf = rdna4::Isa::MAX_SGPRS_PER_WF;
  cfg.vgprs_per_wf = 32;
  cfg.lds_size_kb = 64;

  auto cu = std::make_unique<
      amdgpu::IsaExecComputeUnit<simdojo::ExecMode::FUNCTIONAL, Rdna4Wave64DppTestIsa>>(
      "rdna4_wave64_dpp8_fi_cu", cfg, &mem, &l2);

  auto *wf = cu->dispatch_wf(0, 0, cfg.sgprs_per_wf, cfg.vgprs_per_wf);
  ASSERT_NE(wf, nullptr);
  ASSERT_EQ(wf->wf_size(), 64u);

  constexpr uint32_t kSrc = 4;
  constexpr uint32_t kDst = 8;
  uint32_t vbase = wf->vgpr_alloc().base;

  struct Dpp8Case {
    uint32_t dst_lane;
    uint32_t src_lane;
  };
  constexpr std::array<Dpp8Case, 2> cases{{
      {33, 32},
      {35, 37},
  }};

  auto set_lane_sel = [](rdna4::Vop1VopDpp8MachineInst &raw, uint32_t dst_lane, uint32_t src_lane) {
    uint32_t selector = src_lane & 7u;
    switch (dst_lane & 7u) {
    case 0:
      raw.lane_sel_0 = selector;
      break;
    case 1:
      raw.lane_sel_1 = selector;
      break;
    case 2:
      raw.lane_sel_2 = selector;
      break;
    case 3:
      raw.lane_sel_3 = selector;
      break;
    case 4:
      raw.lane_sel_4 = selector;
      break;
    case 5:
      raw.lane_sel_5 = selector;
      break;
    case 6:
      raw.lane_sel_6 = selector;
      break;
    case 7:
      raw.lane_sel_7 = selector;
      break;
    }
  };

  auto run = [&](const Dpp8Case &test_case, uint32_t src0_marker) {
    wf->set_exec(~(1ULL << test_case.src_lane));
    for (uint32_t lane = 0; lane < wf->wf_size(); ++lane) {
      cu->write_vgpr(vbase + kSrc, lane, 0x6B6B0000u + lane);
      cu->write_vgpr(vbase + kDst, lane, 0xDEAD0000u + lane);
    }

    rdna4::Vop1VopDpp8MachineInst raw{};
    raw.src0 = src0_marker;
    raw.vsrc0 = kSrc;
    raw.vdst = kDst;
    set_lane_sel(raw, test_case.dst_lane, test_case.src_lane);

    rdna4::VMovB32Vop1 inst(reinterpret_cast<const rdna4::MachineInst *>(&raw));
    inst.execute_impl(*wf);
    return cu->read_vgpr(vbase + kDst, test_case.dst_lane);
  };

  for (const auto &test_case : cases) {
    SCOPED_TRACE("dst lane " + std::to_string(test_case.dst_lane) + " src lane " +
                 std::to_string(test_case.src_lane));
    EXPECT_EQ(run(test_case, amdgpu::SRC_DPP8_FI_0), 0u);
    EXPECT_EQ(run(test_case, amdgpu::SRC_DPP8_FI_1), 0x6B6B0000u + test_case.src_lane);
  }
}

template <typename Traits> void wave32_generated_vop1_dpp8_fetch_inactive_uses_fi() {
  SCOPED_TRACE(Traits::name);
  amdgpu::GpuMemory mem(std::string(Traits::name) + "_dpp8_fi_mem");
  amdgpu::L2Cache l2(std::string(Traits::name) + "_dpp8_fi_l2");

  amdgpu::ComputeUnitCore::Config cfg{};
  cfg.arch = Traits::arch;
  cfg.num_wf_slots = 1;
  cfg.sgprs_per_wf = 106;
  cfg.vgprs_per_wf = 32;
  cfg.lds_size_kb = 64;

  auto cu =
      amdgpu::ComputeUnitCore::create(std::string(Traits::name) + "_dpp8_fi_cu", cfg, &mem, &l2);
  ASSERT_NE(cu, nullptr);

  auto *wf = cu->dispatch_wf(0, 0, cfg.sgprs_per_wf, cfg.vgprs_per_wf);
  ASSERT_NE(wf, nullptr);
  ASSERT_EQ(wf->wf_size(), 32u);

  constexpr uint32_t kSrc = 4;
  constexpr uint32_t kDst = 8;
  constexpr uint32_t kSrcLane0Value = 0x5A5A0000u;
  uint32_t vbase = wf->vgpr_alloc().base;

  auto run = [&](uint32_t src0_marker) {
    wf->set_exec(0xFFFFFFFEULL);
    for (uint32_t lane = 0; lane < wf->wf_size(); ++lane) {
      cu->write_vgpr(vbase + kSrc, lane, kSrcLane0Value + lane);
      cu->write_vgpr(vbase + kDst, lane, 0xDEAD0000u + lane);
    }

    typename Traits::Vop1VopDpp8MachineInst raw{};
    raw.src0 = src0_marker;
    raw.vsrc0 = kSrc;
    raw.vdst = kDst;

    typename Traits::VMovB32Vop1 inst(reinterpret_cast<const typename Traits::MachineInst *>(&raw));
    inst.execute_impl(*wf);
    return cu->read_vgpr(vbase + kDst, 1);
  };

  EXPECT_EQ(run(amdgpu::SRC_DPP8_FI_0), 0u);
  EXPECT_EQ(run(amdgpu::SRC_DPP8_FI_1), kSrcLane0Value);
}

template <typename Traits> void wave32_generated_vopc_dpp_write_mask_honors_bound_ctrl() {
  SCOPED_TRACE(Traits::name);
  amdgpu::GpuMemory mem(std::string(Traits::name) + "_dpp_vopc_wave32_write_mask_mem");
  amdgpu::L2Cache l2(std::string(Traits::name) + "_dpp_vopc_wave32_write_mask_l2");

  amdgpu::ComputeUnitCore::Config cfg{};
  cfg.arch = Traits::arch;
  cfg.num_wf_slots = 1;
  cfg.sgprs_per_wf = 106;
  cfg.vgprs_per_wf = 32;
  cfg.lds_size_kb = 64;

  auto cu = amdgpu::ComputeUnitCore::create(
      std::string(Traits::name) + "_dpp_vopc_wave32_write_mask_cu", cfg, &mem, &l2);
  ASSERT_NE(cu, nullptr);

  auto *wf = cu->dispatch_wf(0, 0, cfg.sgprs_per_wf, cfg.vgprs_per_wf);
  ASSERT_NE(wf, nullptr);
  ASSERT_EQ(wf->wf_size(), 32u);
  wf->set_exec(0xFFFFFFFFULL);
  wf->set_vcc(0x00000000001000A5ULL);

  constexpr uint32_t kSrc0 = 4;
  constexpr uint32_t kSrc1 = 8;
  uint32_t vbase = wf->vgpr_alloc().base;
  for (uint32_t lane = 0; lane < wf->wf_size(); ++lane) {
    uint32_t src = 0x1000u + lane;
    uint32_t cmp = lane >= 16 ? 0x100Fu : src;
    if (lane == 20)
      cmp = 0xDEAD0020u;
    cu->write_vgpr(vbase + kSrc0, lane, src);
    cu->write_vgpr(vbase + kSrc1, lane, cmp);
  }

  typename Traits::VopcVopDpp16MachineInst raw{};
  raw.src0 = amdgpu::SRC_DPP;
  raw.vsrc1 = kSrc1;
  raw.vsrc0 = kSrc0;
  raw.dpp_ctrl = amdgpu::dpp::ROW_BCAST15;
  raw.fi = 1;
  raw.bound_ctrl = 0;
  raw.bank_mask = 0xF;
  raw.row_mask = 0xF;

  typename Traits::VCmpEqU32Vopc inst(reinterpret_cast<const typename Traits::MachineInst *>(&raw));
  inst.execute_impl(*wf);

  EXPECT_EQ(wf->vcc(), 0x00000000FFEF00A5ULL);
}

template <typename Traits> void wave32_generated_vcmpx_dpp_write_mask_preserves_exec() {
  SCOPED_TRACE(Traits::name);
  amdgpu::GpuMemory mem(std::string(Traits::name) + "_dpp_vcmpx_wave32_exec_mask_mem");
  amdgpu::L2Cache l2(std::string(Traits::name) + "_dpp_vcmpx_wave32_exec_mask_l2");

  amdgpu::ComputeUnitCore::Config cfg{};
  cfg.arch = Traits::arch;
  cfg.num_wf_slots = 1;
  cfg.sgprs_per_wf = 106;
  cfg.vgprs_per_wf = 32;
  cfg.lds_size_kb = 64;

  auto cu = amdgpu::ComputeUnitCore::create(
      std::string(Traits::name) + "_dpp_vcmpx_wave32_exec_mask_cu", cfg, &mem, &l2);
  ASSERT_NE(cu, nullptr);

  auto *wf = cu->dispatch_wf(0, 0, cfg.sgprs_per_wf, cfg.vgprs_per_wf);
  ASSERT_NE(wf, nullptr);
  ASSERT_EQ(wf->wf_size(), 32u);

  constexpr uint64_t kOldExec = 0xA5A55A5AFFFF005AULL;
  constexpr uint64_t kOldVcc = 0x00000000000000A5ULL;
  wf->set_exec_raw(kOldExec);
  wf->set_vcc(kOldVcc);

  constexpr uint32_t kSrc0 = 4;
  constexpr uint32_t kSrc1 = 8;
  uint32_t vbase = wf->vgpr_alloc().base;
  for (uint32_t lane = 0; lane < wf->wf_size(); ++lane) {
    uint32_t src = 0x1000u + lane;
    uint32_t cmp = lane >= 16 ? 0x100Fu : 0xDEAD0000u + lane;
    if (lane == 20)
      cmp = 0xDEAD0020u;
    cu->write_vgpr(vbase + kSrc0, lane, src);
    cu->write_vgpr(vbase + kSrc1, lane, cmp);
  }

  typename Traits::VopcVopDpp16MachineInst raw{};
  raw.src0 = amdgpu::SRC_DPP;
  raw.vsrc1 = kSrc1;
  raw.vsrc0 = kSrc0;
  raw.dpp_ctrl = amdgpu::dpp::ROW_BCAST15;
  raw.fi = 1;
  raw.bound_ctrl = 0;
  raw.bank_mask = 0xF;
  raw.row_mask = 0xF;

  typename Traits::VCmpxEqU32Vopc inst(
      reinterpret_cast<const typename Traits::MachineInst *>(&raw));
  inst.execute_impl(*wf);

  EXPECT_EQ(wf->vcc(), kOldVcc);
  EXPECT_EQ(wf->exec(), 0x00000000FFEF005AULL);
  EXPECT_EQ(wf->exec_raw(), 0xA5A55A5AFFEF005AULL);
}

template <typename Traits> void wave32_generated_vcmpx_preserves_exec_hi() {
  SCOPED_TRACE(Traits::name);
  amdgpu::GpuMemory mem(std::string(Traits::name) + "_vcmpx_wave32_exec_hi_mem");
  amdgpu::L2Cache l2(std::string(Traits::name) + "_vcmpx_wave32_exec_hi_l2");

  amdgpu::ComputeUnitCore::Config cfg{};
  cfg.arch = Traits::arch;
  cfg.num_wf_slots = 1;
  cfg.sgprs_per_wf = 106;
  cfg.vgprs_per_wf = 32;
  cfg.lds_size_kb = 64;

  auto cu = amdgpu::ComputeUnitCore::create(std::string(Traits::name) + "_vcmpx_wave32_exec_hi_cu",
                                            cfg, &mem, &l2);
  ASSERT_NE(cu, nullptr);

  auto *wf = cu->dispatch_wf(0, 0, cfg.sgprs_per_wf, cfg.vgprs_per_wf);
  ASSERT_NE(wf, nullptr);
  ASSERT_EQ(wf->wf_size(), 32u);
  wf->set_exec_raw(0xA5A55A5AFFFFFFFFULL);

  constexpr uint32_t kSrc0 = 4;
  constexpr uint32_t kSrc1 = 8;
  uint32_t vbase = wf->vgpr_alloc().base;
  for (uint32_t lane = 0; lane < wf->wf_size(); ++lane) {
    cu->write_vgpr(vbase + kSrc0, lane, lane);
    cu->write_vgpr(vbase + kSrc1, lane, (lane & 4u) == 0 ? lane : lane + 1);
  }

  constexpr uint32_t kVgprSrcEncodingBase = 256;
  typename Traits::VopcMachineInst raw{};
  raw.src0 = kVgprSrcEncodingBase + kSrc0;
  raw.vsrc1 = kSrc1;

  typename Traits::VCmpxEqU32Vopc inst(
      reinterpret_cast<const typename Traits::MachineInst *>(&raw));
  inst.execute_impl(*wf);

  EXPECT_EQ(wf->exec(), 0x0F0F0F0FULL);
  EXPECT_EQ(wf->exec_raw(), 0xA5A55A5A0F0F0F0FULL);
}

template <typename Traits> void unsupported_rdna_vopc_dpp_throws() {
  SCOPED_TRACE(Traits::name);

  auto expect_throws = [](uint32_t src0) {
    typename Traits::Vop1VopDpp16MachineInst raw{};
    raw.src0 = src0;

    EXPECT_THROW(typename Traits::VCmpEqU32Vopc(
                     reinterpret_cast<const typename Traits::MachineInst *>(&raw)),
                 util::InvalidInst);
  };

  expect_throws(amdgpu::SRC_DPP);
  expect_throws(amdgpu::SRC_DPP8_FI_0);
  expect_throws(amdgpu::SRC_DPP8_FI_1);
}

template <typename Traits> void unsupported_cdna_vopc_dpp_throws() {
  SCOPED_TRACE(Traits::name);

  typename Traits::Vop1VopDppMachineInst raw{};
  raw.src0 = amdgpu::SRC_DPP;

  EXPECT_THROW(
      typename Traits::VCmpEqU32Vopc(reinterpret_cast<const typename Traits::MachineInst *>(&raw)),
      util::InvalidInst);
}

void cdna4_vop1_sdwa_availability_is_instruction_specific() {
  cdna4::Vop1VopSdwaMachineInst raw{};
  raw.src0 = amdgpu::SRC_SDWA;
  raw.vsrc0 = 4;
  raw.vdst = 8;
  raw.src0_sel = amdgpu::sdwa::DWORD;
  raw.dst_sel = amdgpu::sdwa::DWORD;
  raw.dst_unused = amdgpu::sdwa::UNUSED_PAD;

  EXPECT_NO_THROW(cdna4::VMovB32Vop1(reinterpret_cast<const cdna4::MachineInst *>(&raw)));
  EXPECT_THROW(cdna4::VCvtF64I32Vop1(reinterpret_cast<const cdna4::MachineInst *>(&raw)),
               util::InvalidInst);
}

void cdna4_unsupported_sdwa_decode_halts_wave() {
  cdna4::Vop1VopSdwaMachineInst raw{};
  raw.src0 = amdgpu::SRC_SDWA;
  raw.op = cdna4::kVCvtF64I32Vop1;
  raw.vsrc0 = 4;
  raw.vdst = 8;
  raw.encoding = cdna4::encoding::kVop1 >> 2;
  raw.src0_sel = amdgpu::sdwa::DWORD;
  raw.dst_sel = amdgpu::sdwa::DWORD;
  raw.dst_unused = amdgpu::sdwa::UNUSED_PAD;

  static_assert(sizeof(raw) == 2 * sizeof(uint32_t));
  const auto encoded = std::bit_cast<std::array<uint32_t, 2>>(raw);
  const std::array<uint32_t, 4> words{encoded[0], encoded[1], 0, 0};

  auto decoder = Decoder::create(ROCJITSU_CODE_ARCH_CDNA4);
  ASSERT_NE(decoder, nullptr);
  EXPECT_THROW(
      { std::unique_ptr<Instruction> inst(decoder->decode(words.data())); }, util::InvalidInst);

  amdgpu::GpuMemory mem("cdna4_unsupported_sdwa_mem");
  amdgpu::L2Cache l2("cdna4_unsupported_sdwa_l2");
  amdgpu::ComputeUnitCore::Config cfg{};
  cfg.arch = ROCJITSU_CODE_ARCH_CDNA4;
  cfg.num_wf_slots = 1;
  cfg.sgprs_per_wf = 106;
  cfg.vgprs_per_wf = 32;
  cfg.lds_size_kb = 64;

  auto cu = amdgpu::ComputeUnitCore::create("cdna4_unsupported_sdwa_cu", cfg, &mem, &l2);
  ASSERT_NE(cu, nullptr);
  auto *wf = cu->dispatch_wf(0, 0, cfg.sgprs_per_wf, cfg.vgprs_per_wf);
  ASSERT_NE(wf, nullptr);
  for (uint32_t index = 0; index < words.size(); ++index)
    mem.write32(index * sizeof(uint32_t), words[index]);

  EXPECT_NO_THROW(static_cast<void>(cu->step()));
  EXPECT_TRUE(wf->is_halted());
}

template <typename Raw> void rdna4_vop3_dpp_marker_is_instruction_specific(uint32_t marker) {
  Raw raw{};
  raw.src0 = marker;
  raw.op = rdna4::kVAddF32Vop3;
  raw.encoding = rdna4::encoding::kVop3 >> 3;
  raw.vsrc0 = 4;
  raw.src1 = 5;
  raw.vdst = 8;

  EXPECT_NO_THROW(rdna4::VAddF32Vop3(reinterpret_cast<const rdna4::MachineInst *>(&raw)));
  raw.op = rdna4::kVAddF64Vop3;
  EXPECT_THROW(rdna4::VAddF64Vop3(reinterpret_cast<const rdna4::MachineInst *>(&raw)),
               util::InvalidInst);
}

template <typename Raw> void rdna4_vop3p_dpp_marker_is_unsupported(uint32_t marker) {
  Raw raw{};
  raw.src0 = marker;
  raw.op = rdna4::kVPkAddF16Vop3p;
  raw.encoding = rdna4::encoding::kVop3p >> 1;
  raw.vsrc0 = 4;
  raw.src1 = 5;
  raw.vdst = 8;
  EXPECT_THROW(rdna4::VPkAddF16Vop3p(reinterpret_cast<const rdna4::MachineInst *>(&raw)),
               util::InvalidInst);
}

void rdna4_vop3_dpp_availability_is_instruction_specific() {
  rdna4_vop3_dpp_marker_is_instruction_specific<rdna4::Vop3VopDpp16MachineInst>(amdgpu::SRC_DPP);
  rdna4_vop3_dpp_marker_is_instruction_specific<rdna4::Vop3VopDpp8MachineInst>(
      amdgpu::SRC_DPP8_FI_0);
  rdna4_vop3p_dpp_marker_is_unsupported<rdna4::Vop3pVopDpp16MachineInst>(amdgpu::SRC_DPP);
  rdna4_vop3p_dpp_marker_is_unsupported<rdna4::Vop3pVopDpp8MachineInst>(amdgpu::SRC_DPP8_FI_0);
}

TEST(DppPermuteTest, CdnaGeneratedVop1UsesSharedRowBroadcast) {
  cdna_generated_vop1_uses_shared_row_broadcast<Cdna1DppTraits>();
  cdna_generated_vop1_uses_shared_row_broadcast<Cdna2DppTraits>();
  cdna_generated_vop1_uses_shared_row_broadcast<Cdna3DppTraits>();
  cdna_generated_vop1_uses_shared_row_broadcast<Cdna4DppTraits>();
}

TEST(DppPermuteTest, CdnaGeneratedVop1DppWriteMaskHonorsBoundCtrl) {
  cdna_generated_vop1_dpp_write_mask_honors_bound_ctrl<Cdna1DppTraits>();
  cdna_generated_vop1_dpp_write_mask_honors_bound_ctrl<Cdna2DppTraits>();
  cdna_generated_vop1_dpp_write_mask_honors_bound_ctrl<Cdna3DppTraits>();
  cdna_generated_vop1_dpp_write_mask_honors_bound_ctrl<Cdna4DppTraits>();
}

TEST(DppPermuteTest, CdnaVopcDppThrowsUnsupported) {
  unsupported_cdna_vopc_dpp_throws<Cdna1DppTraits>();
  unsupported_cdna_vopc_dpp_throws<Cdna2DppTraits>();
  unsupported_cdna_vopc_dpp_throws<Cdna3DppTraits>();
  unsupported_cdna_vopc_dpp_throws<Cdna4DppTraits>();
}

TEST(DppPermuteTest, Cdna4Vop1SdwaAvailabilityIsInstructionSpecific) {
  cdna4_vop1_sdwa_availability_is_instruction_specific();
}

TEST(DppPermuteTest, Cdna4UnsupportedSdwaDecodeHaltsWave) {
  cdna4_unsupported_sdwa_decode_halts_wave();
}

TEST(DppPermuteTest, Rdna4Vop3DppAvailabilityIsInstructionSpecific) {
  rdna4_vop3_dpp_availability_is_instruction_specific();
}

TEST(DppPermuteTest, RdnaGeneratedVop1DppWriteMaskHonorsBoundCtrl) {
  wave32_generated_vop1_dpp_write_mask_honors_bound_ctrl<Rdna1DppTraits>();
  wave32_generated_vop1_dpp_write_mask_honors_bound_ctrl<Rdna2DppTraits>();
  wave32_generated_vop1_dpp_write_mask_honors_bound_ctrl<Rdna3DppTraits>();
  wave32_generated_vop1_dpp_write_mask_honors_bound_ctrl<Rdna3_5DppTraits>();
  wave32_generated_vop1_dpp_write_mask_honors_bound_ctrl<Rdna4DppTraits>();
}

TEST(DppPermuteTest, CdnaGeneratedVop1Dpp64PreservesMaskedDestination) {
  generated_vop1_dpp64_preserves_masked_destination<Cdna2DppTraits>();
  generated_vop1_dpp64_preserves_masked_destination<Cdna3DppTraits>();
  generated_vop1_dpp64_preserves_masked_destination<Cdna4DppTraits>();
}

TEST(DppPermuteTest, Rdna4GeneratedVop1F64PreservesInactiveDestination) {
  rdna4_generated_vop1_64_preserves_inactive_destination();
}

TEST(DppPermuteTest, Gfx1250GeneratedVop1DppWriteMaskHonorsBoundCtrl) {
  ScopedIsaExecutionBackend execution_backend_scope{&gfx1250::execution_backend()};
  wave32_generated_vop1_dpp_write_mask_honors_bound_ctrl<Gfx1250DppTraits>();
}

TEST(DppPermuteTest, RdnaGeneratedVop1Dpp16FetchInactiveUsesFi) {
  wave32_generated_vop1_dpp16_fetch_inactive_uses_fi<Rdna1DppTraits>();
  wave32_generated_vop1_dpp16_fetch_inactive_uses_fi<Rdna2DppTraits>();
  wave32_generated_vop1_dpp16_fetch_inactive_uses_fi<Rdna3DppTraits>();
  wave32_generated_vop1_dpp16_fetch_inactive_uses_fi<Rdna3_5DppTraits>();
  wave32_generated_vop1_dpp16_fetch_inactive_uses_fi<Rdna4DppTraits>();
}

TEST(DppPermuteTest, Gfx1250GeneratedVop1Dpp16FetchInactiveUsesFi) {
  ScopedIsaExecutionBackend execution_backend_scope{&gfx1250::execution_backend()};
  wave32_generated_vop1_dpp16_fetch_inactive_uses_fi<Gfx1250DppTraits>();
}

TEST(DppPermuteTest, RdnaGeneratedVop1Dpp16FiZeroPrecedesBoundCtrl) {
  wave32_generated_vop1_dpp16_fi_zero_precedes_bound_ctrl<Rdna1DppTraits>();
  wave32_generated_vop1_dpp16_fi_zero_precedes_bound_ctrl<Rdna2DppTraits>();
  wave32_generated_vop1_dpp16_fi_zero_precedes_bound_ctrl<Rdna3DppTraits>();
  wave32_generated_vop1_dpp16_fi_zero_precedes_bound_ctrl<Rdna3_5DppTraits>();
  wave32_generated_vop1_dpp16_fi_zero_precedes_bound_ctrl<Rdna4DppTraits>();
}

TEST(DppPermuteTest, Gfx1250GeneratedVop1Dpp16FiZeroPrecedesBoundCtrl) {
  ScopedIsaExecutionBackend execution_backend_scope{&gfx1250::execution_backend()};
  wave32_generated_vop1_dpp16_fi_zero_precedes_bound_ctrl<Gfx1250DppTraits>();
}

TEST(DppPermuteTest, Rdna4GeneratedVop1Dpp16Wave64FetchInactiveUsesUpperExecBit) {
  rdna4_wave64_generated_vop1_dpp16_fetch_inactive_uses_upper_exec_bit();
}

TEST(DppPermuteTest, Rdna4GeneratedVop1Dpp8Wave64FetchInactiveUsesUpperExecBit) {
  rdna4_wave64_generated_vop1_dpp8_fetch_inactive_uses_upper_exec_bit();
}

TEST(DppPermuteTest, RdnaGeneratedVop1Dpp8FetchInactiveUsesFi) {
  wave32_generated_vop1_dpp8_fetch_inactive_uses_fi<Rdna1DppTraits>();
  wave32_generated_vop1_dpp8_fetch_inactive_uses_fi<Rdna2DppTraits>();
  wave32_generated_vop1_dpp8_fetch_inactive_uses_fi<Rdna3DppTraits>();
  wave32_generated_vop1_dpp8_fetch_inactive_uses_fi<Rdna3_5DppTraits>();
  wave32_generated_vop1_dpp8_fetch_inactive_uses_fi<Rdna4DppTraits>();
}

TEST(DppPermuteTest, Gfx1250GeneratedVop1Dpp8FetchInactiveUsesFi) {
  ScopedIsaExecutionBackend execution_backend_scope{&gfx1250::execution_backend()};
  wave32_generated_vop1_dpp8_fetch_inactive_uses_fi<Gfx1250DppTraits>();
}

TEST(DppPermuteTest, RdnaGeneratedVopcDppWriteMaskHonorsBoundCtrl) {
  wave32_generated_vopc_dpp_write_mask_honors_bound_ctrl<Rdna3DppTraits>();
  wave32_generated_vopc_dpp_write_mask_honors_bound_ctrl<Rdna3_5DppTraits>();
  wave32_generated_vopc_dpp_write_mask_honors_bound_ctrl<Rdna4DppTraits>();
}

TEST(DppPermuteTest, Gfx1250GeneratedVopcDppWriteMaskHonorsBoundCtrl) {
  ScopedIsaExecutionBackend execution_backend_scope{&gfx1250::execution_backend()};
  wave32_generated_vopc_dpp_write_mask_honors_bound_ctrl<Gfx1250DppTraits>();
}

TEST(DppPermuteTest, RdnaGeneratedVcmpxDppWave32WriteMaskPreservesExec) {
  wave32_generated_vcmpx_dpp_write_mask_preserves_exec<Rdna3DppTraits>();
  wave32_generated_vcmpx_dpp_write_mask_preserves_exec<Rdna3_5DppTraits>();
  wave32_generated_vcmpx_dpp_write_mask_preserves_exec<Rdna4DppTraits>();
}

TEST(DppPermuteTest, Gfx1250GeneratedVcmpxDppWave32WriteMaskPreservesExec) {
  ScopedIsaExecutionBackend execution_backend_scope{&gfx1250::execution_backend()};
  wave32_generated_vcmpx_dpp_write_mask_preserves_exec<Gfx1250DppTraits>();
}

TEST(ExecMaskTest, RdnaGeneratedVcmpxWave32PreservesExecHi) {
  ScopedIsaExecutionBackend execution_backend_scope{&gfx1250::execution_backend()};
  wave32_generated_vcmpx_preserves_exec_hi<Rdna1DppTraits>();
  wave32_generated_vcmpx_preserves_exec_hi<Rdna2DppTraits>();
  wave32_generated_vcmpx_preserves_exec_hi<Rdna3DppTraits>();
  wave32_generated_vcmpx_preserves_exec_hi<Rdna3_5DppTraits>();
  wave32_generated_vcmpx_preserves_exec_hi<Rdna4DppTraits>();
  wave32_generated_vcmpx_preserves_exec_hi<Gfx1250DppTraits>();
}

TEST(DppPermuteTest, Rdna1VopcDppThrowsUnsupported) {
  unsupported_rdna_vopc_dpp_throws<Rdna1DppTraits>();
}

TEST(DppPermuteTest, Rdna2VopcDppThrowsUnsupported) {
  unsupported_rdna_vopc_dpp_throws<Rdna2DppTraits>();
}

// ---------------------------------------------------------------------------
// SDWA tests
// ---------------------------------------------------------------------------

// These helper tests pin byte placement independently of instruction decode.
// End-to-end callback tests can then distinguish a selector/merge bug here
// from a generated wrapper passing the wrong read or write mask.
TEST(SdwaTest, SrcSelect) {
  using namespace amdgpu::sdwa;
  uint32_t val = 0xDEADBEEF;

  EXPECT_EQ(sdwa_src_select(val, BYTE_0, false), 0xEFu);
  EXPECT_EQ(sdwa_src_select(val, BYTE_1, false), 0xBEu);
  EXPECT_EQ(sdwa_src_select(val, BYTE_2, false), 0xADu);
  EXPECT_EQ(sdwa_src_select(val, BYTE_3, false), 0xDEu);
  EXPECT_EQ(sdwa_src_select(val, WORD_0, false), 0xBEEFu);
  EXPECT_EQ(sdwa_src_select(val, WORD_1, false), 0xDEADu);
  EXPECT_EQ(sdwa_src_select(val, DWORD, false), val);

  // Sign extension.
  EXPECT_EQ(sdwa_src_select(0x00000080, BYTE_0, true), 0xFFFFFF80u);
  EXPECT_EQ(sdwa_src_select(0x00000080, BYTE_0, false), 0x80u);
  EXPECT_EQ(sdwa_src_select(0x00008000, WORD_0, true), 0xFFFF8000u);
}

TEST(SdwaTest, DstMerge) {
  using namespace amdgpu::sdwa;
  // Write result byte 0x42 into BYTE_1, zero-pad rest.
  uint32_t merged = sdwa_dst_merge(0x42, 0xAAAAAAAA, BYTE_1, UNUSED_PAD);
  EXPECT_EQ(merged, 0x00004200u);

  // Preserve unused bytes.
  merged = sdwa_dst_merge(0x42, 0xAABBCCDD, BYTE_1, UNUSED_PRESERVE);
  EXPECT_EQ(merged, 0xAABB42DDu);

  // Full dword: just return result.
  merged = sdwa_dst_merge(0x12345678, 0xAAAAAAAA, DWORD, UNUSED_PAD);
  EXPECT_EQ(merged, 0x12345678u);

  // Sign-extension fills bytes above the selected byte/word and zeroes bytes
  // below it.
  merged = sdwa_dst_merge(0x80, 0xAABBCCDD, BYTE_1, UNUSED_SEXT);
  EXPECT_EQ(merged, 0xFFFF8000u);
  merged = sdwa_dst_merge(0x7F, 0xAABBCCDD, BYTE_3, UNUSED_SEXT);
  EXPECT_EQ(merged, 0x7F000000u);

  // Word destinations exercise both preservation and sign-extension.
  merged = sdwa_dst_merge(0x12345678, 0xAABBCCDD, WORD_0, UNUSED_PRESERVE);
  EXPECT_EQ(merged, 0xAABB5678u);
  merged = sdwa_dst_merge(0x00008000, 0xAABBCCDD, WORD_1, UNUSED_SEXT);
  EXPECT_EQ(merged, 0x80000000u);

  // The source selector and preserve destination mask use the same byte
  // windows; pad/sext destinations are full-dword writes.
  EXPECT_EQ(sdwa_src_byte_mask(BYTE_0), 0b0001);
  EXPECT_EQ(sdwa_src_byte_mask(BYTE_3), 0b1000);
  EXPECT_EQ(sdwa_src_byte_mask(WORD_0), 0b0011);
  EXPECT_EQ(sdwa_src_byte_mask(WORD_1), 0b1100);
  EXPECT_EQ(sdwa_src_byte_mask(DWORD), 0b1111);
  EXPECT_EQ(sdwa_dst_byte_mask(BYTE_1, UNUSED_PRESERVE), 0b0010);
  EXPECT_EQ(sdwa_dst_byte_mask(WORD_1, UNUSED_PRESERVE), 0b1100);
  EXPECT_EQ(sdwa_dst_byte_mask(BYTE_1, UNUSED_PAD), 0b1111);
  EXPECT_EQ(sdwa_dst_byte_mask(BYTE_1, UNUSED_SEXT), 0b1111);
}

// ---------------------------------------------------------------------------
// Scratch address calculation tests
// ---------------------------------------------------------------------------

TEST(ScratchAddrCalcTest, FlatScratchUsesWavefrontBase) {
  // Verify that FLAT with seg==1 (SCRATCH) computes:
  //   address = scratch_base + VGPR[lane] + offset
  amdgpu::GpuMemory mem("test_mem");
  amdgpu::L2Cache l2("test_l2");
  amdgpu::ComputeUnitCore::Config cfg{};
  cfg.arch = ROCJITSU_CODE_ARCH_CDNA4;
  cfg.num_wf_slots = 1;
  cfg.sgprs_per_wf = 104;
  cfg.vgprs_per_wf = 16;
  cfg.lds_size_kb = 64;
  auto cu = amdgpu::ComputeUnitCore::create("scratch_cu", cfg, &mem, &l2);
  ASSERT_NE(cu, nullptr);

  auto *wf = cu->dispatch_wf(0, 0, 104, 16);
  ASSERT_NE(wf, nullptr);

  // Set scratch base via the wavefront's dedicated FLAT_SCRATCH register.
  // On CDNA4 this is an architected HW register, not SGPRs s[102:103].
  constexpr uint64_t SCRATCH_BASE = 0x1'0000'0000ULL;
  wf->set_scratch_base(SCRATCH_BASE);

  // Write a 32-bit offset into VGPR[0] lane 0.
  uint32_t vbase = wf->vgpr_alloc().base;
  cu->write_vgpr(vbase, 0, 0x100); // lane 0: offset 0x100

  // Set EXEC so only lane 0 is active.
  wf->set_exec(1ULL);

  // Build a FlatScratchMachineInst with seg=1 (SCRATCH), sve=1 (VADDR enabled),
  // saddr=0x7F (no SADDR), offset=0x10.
  cdna4::FlatScratchMachineInst inst{};
  inst.seg = 1;       // SCRATCH
  inst.sve = 1;       // VADDR enabled
  inst.saddr = 0x7F;  // No SADDR
  inst.addr = 0;      // VGPR index 0
  inst.offset = 0x10; // 13-bit signed offset

  amdgpu::VectorMemState d(amdgpu::GLOBAL_MEM);
  amdgpu::addr_calc::flat_calculate_addresses(inst, *wf, d);

  EXPECT_EQ(d.lane_mask, 1ULL);
  // scratch_base (0x1_0000_0000) + VGPR (0x100) + offset (0x10) = 0x1_0000_0110
  EXPECT_EQ(d.per_lane_addr[0], SCRATCH_BASE + 0x100 + 0x10);
}

TEST(ScratchAddrCalcTest, FlatGlobalDoesNotUseScratchBase) {
  // Verify that FLAT with seg==2 (GLOBAL) does NOT add scratch_base.
  amdgpu::GpuMemory mem("test_mem");
  amdgpu::L2Cache l2("test_l2");
  amdgpu::ComputeUnitCore::Config cfg{};
  cfg.arch = ROCJITSU_CODE_ARCH_CDNA4;
  cfg.num_wf_slots = 1;
  cfg.sgprs_per_wf = 104;
  cfg.vgprs_per_wf = 16;
  cfg.lds_size_kb = 64;
  auto cu = amdgpu::ComputeUnitCore::create("global_cu", cfg, &mem, &l2);
  ASSERT_NE(cu, nullptr);

  auto *wf = cu->dispatch_wf(0, 0, 104, 16);
  ASSERT_NE(wf, nullptr);

  // Set scratch base — should be ignored for GLOBAL.
  wf->set_scratch_base(0xDEAD'0000ULL);

  // Write a 64-bit address into VGPR[0:1] lane 0.
  uint32_t vbase = wf->vgpr_alloc().base;
  cu->write_vgpr(vbase, 0, 0x2000);     // low 32
  cu->write_vgpr(vbase + 1, 0, 0x0001); // high 32 → addr = 0x1_0000_2000
  wf->set_exec(1ULL);

  cdna4::FlatMachineInst inst{};
  inst.seg = 2;      // GLOBAL
  inst.saddr = 0x7F; // No SADDR → use 64-bit VGPR pair
  inst.addr = 0;
  inst.offset = 0;
  inst.pad_12 = 0;

  amdgpu::VectorMemState d(amdgpu::GLOBAL_MEM);
  amdgpu::addr_calc::flat_calculate_addresses(inst, *wf, d);

  EXPECT_EQ(d.per_lane_addr[0], 0x1'0000'2000ULL); // No scratch_base added.
}

TEST(RdnaAddrCalcTest, Rdna3Saddr7cUsesVgprPair) {
  amdgpu::GpuMemory mem("rdna3_addr_mem");
  amdgpu::L2Cache l2("rdna3_addr_l2");
  amdgpu::ComputeUnitCore::Config cfg{};
  cfg.arch = ROCJITSU_CODE_ARCH_RDNA3;
  cfg.num_wf_slots = 1;
  cfg.sgprs_per_wf = 128;
  cfg.vgprs_per_wf = 16;
  cfg.lds_size_kb = 64;
  auto cu = amdgpu::ComputeUnitCore::create("rdna3_addr_cu", cfg, &mem, &l2);
  ASSERT_NE(cu, nullptr);

  auto *wf = cu->dispatch_wf(0, 0, 128, 16);
  ASSERT_NE(wf, nullptr);
  wf->set_exec(1ULL);

  uint32_t vbase = wf->vgpr_alloc().base;
  cu->write_vgpr(vbase, 0, 0x2000);
  cu->write_vgpr(vbase + 1, 0, 0x0001);
  cu->write_sgpr(wf->sgpr_alloc().base + 0x7C, 0xDEAD0000);
  cu->write_sgpr(wf->sgpr_alloc().base + 0x7D, 0xDEAD0001);

  rdna3::FlatMachineInst inst{};
  inst.saddr = 0x7C;
  inst.addr = 0;
  inst.offset = 0x20;

  amdgpu::VectorMemState d(amdgpu::GLOBAL_MEM);
  rdna3::flat_calculate_addresses(inst, *wf, d);
  EXPECT_EQ(d.per_lane_addr[0], 0x1'0000'2020ULL);

  cu->write_sgpr(wf->sgpr_alloc().base + 4, 0x3000);
  cu->write_sgpr(wf->sgpr_alloc().base + 5, 0x0002);
  cu->write_vgpr(vbase, 0, 0x40);
  cu->write_vgpr(vbase + 1, 0, 0xDEAD);

  inst.saddr = 4;
  inst.offset = 0x10;
  rdna3::flat_calculate_addresses(inst, *wf, d);
  EXPECT_EQ(d.per_lane_addr[0], 0x2'0000'3050ULL);
}

TEST(RdnaAddrCalcTest, Rdna3ScratchOffUsesScratchBaseAndLaneStride) {
  amdgpu::GpuMemory mem("rdna3_scratch_addr_mem");
  amdgpu::L2Cache l2("rdna3_scratch_addr_l2");
  amdgpu::ComputeUnitCore::Config cfg{};
  cfg.arch = ROCJITSU_CODE_ARCH_RDNA3;
  cfg.num_wf_slots = 1;
  cfg.sgprs_per_wf = 128;
  cfg.vgprs_per_wf = 16;
  cfg.lds_size_kb = 64;
  auto cu = amdgpu::ComputeUnitCore::create("rdna3_scratch_addr_cu", cfg, &mem, &l2);
  ASSERT_NE(cu, nullptr);

  auto *wf = cu->dispatch_wf(0, 0, 128, 16);
  ASSERT_NE(wf, nullptr);
  wf->set_exec(0x3ULL);
  wf->set_scratch_base(0x6'0000'0000ULL);
  wf->set_scratch_lane_size(0x100);

  uint32_t sbase = wf->sgpr_alloc().base;
  uint32_t vbase = wf->vgpr_alloc().base;
  cu->write_vgpr(vbase, 0, 0x20);
  cu->write_vgpr(vbase, 1, 0x30);
  cu->write_vgpr(vbase + 1, 0, 0xBAD0);
  cu->write_vgpr(vbase + 1, 1, 0xBAD1);
  cu->write_sgpr(sbase + 0x7C, 0xDEAD0000);
  cu->write_sgpr(sbase + 0x7D, 0xDEAD0001);

  rdna3::FlatMachineInst inst{};
  inst.seg = 1;
  inst.saddr = 0x7C;
  inst.sve = 1;
  inst.addr = 0;
  inst.offset = 0x10;

  amdgpu::VectorMemState d(amdgpu::GLOBAL_MEM);
  rdna3::flat_calculate_addresses(inst, *wf, d);
  EXPECT_EQ(d.per_lane_addr[0], 0x6'0000'0030ULL);
  EXPECT_EQ(d.per_lane_addr[1], 0x6'0000'0140ULL);

  cu->write_sgpr(sbase + 4, 0x70);
  cu->write_vgpr(vbase, 0, 0x200);
  cu->write_vgpr(vbase, 1, 0x300);
  inst.saddr = 4;
  inst.sve = 0;
  rdna3::flat_calculate_addresses(inst, *wf, d);
  EXPECT_EQ(d.per_lane_addr[0], 0x6'0000'0080ULL);
  EXPECT_EQ(d.per_lane_addr[1], 0x6'0000'0180ULL);
}

TEST(RdnaScratchExecutionTest, Rdna3B128StoreFeedsB32LoadsWithVectorOffsets) {
  amdgpu::GpuMemory mem("rdna3_scratch_exec_mem");
  amdgpu::L2Cache l2("rdna3_scratch_exec_l2");
  l2.set_backing_memory(&mem);

  amdgpu::ComputeUnitCore::Config cfg{};
  cfg.arch = ROCJITSU_CODE_ARCH_RDNA3;
  cfg.num_wf_slots = 1;
  cfg.sgprs_per_wf = 128;
  cfg.vgprs_per_wf = 32;
  cfg.lds_size_kb = 64;
  auto cu = std::make_unique<Rdna3MemoryTestCu>("rdna3_scratch_exec_cu", cfg, &mem, &l2);

  auto *wf = cu->dispatch_wf(0, 0, 128, 32);
  ASSERT_NE(wf, nullptr);
  wf->set_exec(0x1FULL);
  wf->set_scratch_base(0x6'0000'0000ULL);
  wf->set_scratch_lane_size(0x100);

  const uint32_t vbase = wf->vgpr_alloc().base;
  for (uint32_t lane = 0; lane < 5; ++lane) {
    cu->write_vgpr(vbase + 10, lane, 0x1000u + lane);
    cu->write_vgpr(vbase + 11, lane, 0x1100u + lane);
    cu->write_vgpr(vbase + 12, lane, 0x1200u + lane);
    cu->write_vgpr(vbase + 13, lane, 0x1300u + lane);
    cu->write_vgpr(vbase + 2, lane, lane * 4);
  }

  auto decoder = Decoder::create(ROCJITSU_CODE_ARCH_RDNA3);
  ASSERT_NE(decoder, nullptr);
  auto execute_mem = [&](std::array<uint32_t, 2> words) {
    std::unique_ptr<Instruction> inst(decoder->decode(words.data()));
    ASSERT_NE(inst, nullptr);
    cu->execute_and_route(inst.release(), *wf);
  };

  // scratch_store_b128 off, v[10:13], off
  execute_mem({0xDC750000u, 0x007C0A00u});
  // scratch_load_b32 v1, v2, off
  execute_mem({0xDC510000u, 0x01FC0002u});

  EXPECT_EQ(cu->read_vgpr(vbase + 1, 0), 0x1000u);
  EXPECT_EQ(cu->read_vgpr(vbase + 1, 1), 0x1101u);
  EXPECT_EQ(cu->read_vgpr(vbase + 1, 2), 0x1202u);
  EXPECT_EQ(cu->read_vgpr(vbase + 1, 3), 0x1303u);
  EXPECT_EQ(cu->read_vgpr(vbase + 1, 4), 0u);
}

TEST(RdnaAddrCalcTest, Rdna3SmemSoffsetHandlesNullM0AndSgprSelectors) {
  amdgpu::GpuMemory mem("rdna3_smem_addr_mem");
  amdgpu::L2Cache l2("rdna3_smem_addr_l2");
  amdgpu::ComputeUnitCore::Config cfg{};
  cfg.arch = ROCJITSU_CODE_ARCH_RDNA3;
  cfg.num_wf_slots = 1;
  cfg.sgprs_per_wf = 128;
  cfg.vgprs_per_wf = 16;
  cfg.lds_size_kb = 64;
  auto cu = amdgpu::ComputeUnitCore::create("rdna3_smem_addr_cu", cfg, &mem, &l2);
  ASSERT_NE(cu, nullptr);

  auto *wf = cu->dispatch_wf(0, 0, 128, 16);
  ASSERT_NE(wf, nullptr);

  constexpr uint64_t kBase = 0x1'0000'1000ULL;
  uint32_t sbase = wf->sgpr_alloc().base;
  cu->write_sgpr(sbase, static_cast<uint32_t>(kBase));
  cu->write_sgpr(sbase + 1, static_cast<uint32_t>(kBase >> 32));
  cu->write_sgpr(sbase + rdna3::OPR_SMEM_OFFSET_NULL, 0x100);
  cu->write_sgpr(sbase + rdna3::OPR_SMEM_OFFSET_M0, 0x200);
  cu->write_sgpr(sbase + 8, 0x80);
  wf->set_m0(0x40);

  rdna3::SmemMachineInst inst{};
  inst.sbase = 0;
  inst.offset = 0x20;

  inst.soffset = rdna3::OPR_SMEM_OFFSET_NULL;
  EXPECT_EQ(rdna3::smem_calculate_address(inst, *wf), kBase + 0x20);

  inst.soffset = rdna3::OPR_SMEM_OFFSET_M0;
  EXPECT_EQ(rdna3::smem_calculate_address(inst, *wf), kBase + 0x20 + 0x40);

  inst.soffset = 8;
  EXPECT_EQ(rdna3::smem_calculate_address(inst, *wf), kBase + 0x20 + 0x80);
}

TEST(RdnaAddrCalcTest, Rdna3MubufWrapsOffsetPartBeforeBoundsCheck) {
  amdgpu::GpuMemory mem("rdna3_mubuf_wrap_mem");
  amdgpu::L2Cache l2("rdna3_mubuf_wrap_l2");
  amdgpu::ComputeUnitCore::Config cfg{};
  cfg.arch = ROCJITSU_CODE_ARCH_RDNA3;
  cfg.num_wf_slots = 1;
  cfg.sgprs_per_wf = 128;
  cfg.vgprs_per_wf = 16;
  cfg.lds_size_kb = 64;
  auto cu = amdgpu::ComputeUnitCore::create("rdna3_mubuf_wrap_cu", cfg, &mem, &l2);
  ASSERT_NE(cu, nullptr);

  auto *wf = cu->dispatch_wf(0, 0, 128, 16);
  ASSERT_NE(wf, nullptr);
  wf->set_exec(1ULL);

  constexpr uint64_t kBase = 0x2'0000'1000ULL;
  uint32_t sbase = wf->sgpr_alloc().base;
  uint32_t vbase = wf->vgpr_alloc().base;
  cu->write_sgpr(sbase, static_cast<uint32_t>(kBase));
  cu->write_sgpr(sbase + 1, static_cast<uint32_t>(kBase >> 32));
  cu->write_sgpr(sbase + 2, 0x1000);
  cu->write_sgpr(sbase + 3, 0);
  cu->write_vgpr(vbase + 4, 0, 0xFFFF'FFF0u);

  rdna3::MubufMachineInst inst{};
  inst.srsrc = 0;
  inst.soffset = 0x80;
  inst.offen = 1;
  inst.idxen = 0;
  inst.vaddr = 4;
  inst.offset = 0x10;

  amdgpu::VectorMemState d(amdgpu::GLOBAL_MEM);
  rdna3::mubuf_calculate_addresses(inst, *wf, d);
  EXPECT_EQ(d.lane_mask, 1ULL);
  EXPECT_EQ(d.per_lane_addr[0], kBase);
}

TEST(RdnaAddrCalcTest, Rdna4Saddr7cCoversGlobalFlatAndScratch) {
  amdgpu::GpuMemory mem("rdna4_addr_mem");
  amdgpu::L2Cache l2("rdna4_addr_l2");
  amdgpu::ComputeUnitCore::Config cfg{};
  cfg.arch = ROCJITSU_CODE_ARCH_RDNA4;
  cfg.num_wf_slots = 1;
  cfg.sgprs_per_wf = 128;
  cfg.vgprs_per_wf = 16;
  cfg.lds_size_kb = 64;
  auto cu = amdgpu::ComputeUnitCore::create("rdna4_addr_cu", cfg, &mem, &l2);
  ASSERT_NE(cu, nullptr);

  auto *wf = cu->dispatch_wf(0, 0, 128, 16);
  ASSERT_NE(wf, nullptr);
  wf->set_exec(1ULL);

  uint32_t vbase = wf->vgpr_alloc().base;
  cu->write_vgpr(vbase, 0, 0x4000);
  cu->write_vgpr(vbase + 1, 0, 0x0003);
  cu->write_sgpr(wf->sgpr_alloc().base + 0x7C, 0xDEAD0000);
  cu->write_sgpr(wf->sgpr_alloc().base + 0x7D, 0xDEAD0001);

  rdna4::VglobalMachineInst global_inst{};
  global_inst.saddr = 0x7C;
  global_inst.vaddr = 0;
  global_inst.ioffset = 0x20;

  amdgpu::VectorMemState d(amdgpu::GLOBAL_MEM);
  rdna4::flat_calculate_addresses(global_inst, *wf, d);
  EXPECT_EQ(d.per_lane_addr[0], 0x3'0000'4020ULL);

  cu->write_sgpr(wf->sgpr_alloc().base + 4, 0x5000);
  cu->write_sgpr(wf->sgpr_alloc().base + 5, 0x0004);
  cu->write_vgpr(vbase, 0, 0x60);
  cu->write_vgpr(vbase + 1, 0, 0xDEAD);

  global_inst.saddr = 4;
  global_inst.ioffset = 0x10;
  rdna4::flat_calculate_addresses(global_inst, *wf, d);
  EXPECT_EQ(d.per_lane_addr[0], 0x4'0000'5070ULL);

  cu->write_vgpr(vbase, 0, 0x8000);
  cu->write_vgpr(vbase + 1, 0, 0x0005);
  cu->write_sgpr(wf->sgpr_alloc().base + 0x7C, 0xBAD00000);
  cu->write_sgpr(wf->sgpr_alloc().base + 0x7D, 0xBAD00001);

  rdna4::VflatMachineInst flat_inst{};
  flat_inst.saddr = 0x7C;
  flat_inst.vaddr = 0;
  flat_inst.ioffset = 0x30;
  rdna4::flat_calculate_addresses(flat_inst, *wf, d);
  EXPECT_EQ(d.per_lane_addr[0], 0x5'0000'8030ULL);

  wf->set_exec(0x3ULL);
  wf->set_scratch_base(0x6'0000'0000ULL);
  wf->set_scratch_lane_size(0x100);
  cu->write_vgpr(vbase, 0, 0x80);
  cu->write_vgpr(vbase, 1, 0x90);
  cu->write_vgpr(vbase + 1, 0, 0xBAD);
  cu->write_sgpr(wf->sgpr_alloc().base + 0x7C, 0xBAD00000);

  rdna4::VscratchMachineInst scratch_inst{};
  scratch_inst.saddr = 0x7C;
  scratch_inst.sve = 1;
  scratch_inst.vaddr = 0;
  scratch_inst.ioffset = 0x24;
  rdna4::flat_calculate_addresses(scratch_inst, *wf, d);
  EXPECT_EQ(d.per_lane_addr[0], 0x6'0000'00A4ULL);
  EXPECT_EQ(d.per_lane_addr[1], 0x6'0000'01B4ULL);

  cu->write_sgpr(wf->sgpr_alloc().base + 4, 0x7000);
  cu->write_vgpr(vbase, 0, 0x60);
  cu->write_vgpr(vbase, 1, 0x70);
  scratch_inst.saddr = 4;
  scratch_inst.ioffset = 0x10;
  rdna4::flat_calculate_addresses(scratch_inst, *wf, d);
  EXPECT_EQ(d.per_lane_addr[0], 0x6'0000'7070ULL);
  EXPECT_EQ(d.per_lane_addr[1], 0x6'0000'7180ULL);

  scratch_inst.sve = 0;
  rdna4::flat_calculate_addresses(scratch_inst, *wf, d);
  EXPECT_EQ(d.per_lane_addr[0], 0x6'0000'7010ULL);
  EXPECT_EQ(d.per_lane_addr[1], 0x6'0000'7110ULL);
}

TEST(Gfx1250AddrCalcTest, FlatPrivateScratchDecodesLaneBits) {
  ScopedIsaExecutionBackend execution_backend_scope{&gfx1250::execution_backend()};
  amdgpu::GpuMemory mem("gfx1250_flat_private_mem");
  amdgpu::L2Cache l2("gfx1250_flat_private_l2");
  amdgpu::ComputeUnitCore::Config cfg{};
  cfg.arch = ROCJITSU_CODE_ARCH_GFX1250;
  cfg.num_wf_slots = 1;
  cfg.sgprs_per_wf = 128;
  cfg.vgprs_per_wf = 32;
  cfg.lds_size_kb = 64;
  auto cu = amdgpu::ComputeUnitCore::create("gfx1250_flat_private_cu", cfg, &mem, &l2);
  ASSERT_NE(cu, nullptr);

  auto *wf = cu->dispatch_wf(0, 0, 128, 32);
  ASSERT_NE(wf, nullptr);
  ASSERT_EQ(wf->wf_size(), 32u);
  wf->set_exec(0x7ULL);

  constexpr uint64_t kScratchBase = 0x0002'0000'0000'0000ULL;
  constexpr uint64_t kPrivateBase = 0x0007'0000'0000'0000ULL;
  constexpr uint32_t kPrivateSegmentSize = 0x80;
  wf->set_scratch_base(kScratchBase);
  wf->set_scratch_lane_size(kPrivateSegmentSize);
  wf->set_apertures(0x0001'0000'0000'0000ULL, 0x0001'0000'ffff'ffffULL, kPrivateBase,
                    kPrivateBase + 0xffff'ffffULL);

  gfx1250::Operand flat_scratch_base(
      64, gfx1250::OperandType::OPR_SRC,
      static_cast<int>(gfx1250::OpSelSrc::OPR_SRC_SRC_FLAT_SCRATCH_BASE_LO));
  ASSERT_EQ(amdgpu::RegisterAccess(*wf).read_scalar64(flat_scratch_base), kScratchBase);

  const uint64_t private_offsets[] = {0x10, 0x14, kPrivateSegmentSize + 0x20};
  uint32_t vbase = wf->vgpr_alloc().base;
  for (uint32_t lane = 0; lane < 3; ++lane) {
    uint64_t private_offset = private_offsets[lane];
    uint64_t flat_private_addr =
        kScratchBase + (static_cast<uint64_t>(lane) << 52) + private_offset;
    cu->write_vgpr(vbase, lane, static_cast<uint32_t>(flat_private_addr));
    cu->write_vgpr(vbase + 1, lane, static_cast<uint32_t>(flat_private_addr >> 32));
  }

  gfx1250::VflatMachineInst inst{};
  inst.saddr = gfx1250::OPR_SREG_NULL;
  inst.vaddr = 0;
  inst.ioffset = 4;

  amdgpu::VectorMemState d(amdgpu::GLOBAL_MEM);
  gfx1250::flat_calculate_addresses(inst, *wf, d);

  for (uint32_t lane = 0; lane < 3; ++lane) {
    uint64_t private_offset = private_offsets[lane] + inst.ioffset;
    uint64_t expected =
        kScratchBase + static_cast<uint64_t>(lane) * kPrivateSegmentSize + private_offset;
    EXPECT_EQ(d.per_lane_addr[lane], expected) << "lane " << lane;
  }
}

TEST(RdnaAddrCalcTest, Rdna4SmemSoffsetHandlesNullM0AndSgprSelectors) {
  amdgpu::GpuMemory mem("rdna4_smem_addr_mem");
  amdgpu::L2Cache l2("rdna4_smem_addr_l2");
  amdgpu::ComputeUnitCore::Config cfg{};
  cfg.arch = ROCJITSU_CODE_ARCH_RDNA4;
  cfg.num_wf_slots = 1;
  cfg.sgprs_per_wf = 128;
  cfg.vgprs_per_wf = 16;
  cfg.lds_size_kb = 64;
  auto cu = amdgpu::ComputeUnitCore::create("rdna4_smem_addr_cu", cfg, &mem, &l2);
  ASSERT_NE(cu, nullptr);

  auto *wf = cu->dispatch_wf(0, 0, 128, 16);
  ASSERT_NE(wf, nullptr);

  constexpr uint64_t kBase = 0x2'0000'1000ULL;
  uint32_t sbase = wf->sgpr_alloc().base;
  cu->write_sgpr(sbase, static_cast<uint32_t>(kBase));
  cu->write_sgpr(sbase + 1, static_cast<uint32_t>(kBase >> 32));
  cu->write_sgpr(sbase + rdna4::OPR_SMEM_OFFSET_NULL, 0x100);
  cu->write_sgpr(sbase + rdna4::OPR_SMEM_OFFSET_M0, 0x200);
  cu->write_sgpr(sbase + 8, 0x80);
  wf->set_m0(0x40);

  rdna4::SmemMachineInst inst{};
  inst.sbase = 0;
  inst.ioffset = 0x20;

  inst.soffset = rdna4::OPR_SMEM_OFFSET_NULL;
  EXPECT_EQ(rdna4::smem_calculate_address(inst, *wf), kBase + 0x20);

  inst.soffset = rdna4::OPR_SMEM_OFFSET_M0;
  EXPECT_EQ(rdna4::smem_calculate_address(inst, *wf), kBase + 0x20 + 0x40);

  inst.soffset = 8;
  EXPECT_EQ(rdna4::smem_calculate_address(inst, *wf), kBase + 0x20 + 0x80);
}

TEST(RdnaAddrCalcTest, Rdna4VbufferUsesDecodedRsrcAndOptionalSoffset) {
  amdgpu::GpuMemory mem("rdna4_vbuffer_addr_mem");
  amdgpu::L2Cache l2("rdna4_vbuffer_addr_l2");
  amdgpu::ComputeUnitCore::Config cfg{};
  cfg.arch = ROCJITSU_CODE_ARCH_RDNA4;
  cfg.num_wf_slots = 1;
  cfg.sgprs_per_wf = 128;
  cfg.vgprs_per_wf = 16;
  cfg.lds_size_kb = 64;
  auto cu = amdgpu::ComputeUnitCore::create("rdna4_vbuffer_addr_cu", cfg, &mem, &l2);
  ASSERT_NE(cu, nullptr);

  auto *wf = cu->dispatch_wf(0, 0, 128, 16);
  ASSERT_NE(wf, nullptr);
  wf->set_exec(0x3ULL);

  uint32_t sbase = wf->sgpr_alloc().base;
  uint32_t vbase = wf->vgpr_alloc().base;
  constexpr uint64_t kBase = 0x2'0000'1000ULL;
  cu->write_sgpr(sbase + 4, 0x1000);
  cu->write_sgpr(sbase + 5, 0x0002);
  cu->write_sgpr(sbase + 6, 0xDEAD0006);
  cu->write_sgpr(sbase + 7, 0xDEAD0007);

  // A previous implementation incorrectly treated rsrc as a descriptor index and
  // read s[rsrc * 4:rsrc * 4 + 3]. Keep that range distinct so the test fails
  // if the decoded first-SGPR selector is scaled again.
  cu->write_sgpr(sbase + 16, 0xBAD00010);
  cu->write_sgpr(sbase + 17, 0xBAD00011);

  cu->write_vgpr(vbase + 4, 0, 0x20);
  cu->write_vgpr(vbase + 4, 1, 0x30);

  rdna4::VbufferMachineInst inst{};
  inst.rsrc = 4;
  inst.soffset = rdna4::OPR_SREG_M0_NULL;
  inst.offen = 1;
  inst.idxen = 0;
  inst.vaddr = 4;
  inst.ioffset = 0x10;

  amdgpu::VectorMemState d(amdgpu::GLOBAL_MEM);
  rdna4::mubuf_calculate_addresses(inst, *wf, d);
  EXPECT_EQ(d.lane_mask, 0x3ULL);
  EXPECT_EQ(d.per_lane_addr[0], kBase + 0x20 + 0x10);
  EXPECT_EQ(d.per_lane_addr[1], kBase + 0x30 + 0x10);

  cu->write_sgpr(sbase + 8, 0x80);
  inst.soffset = 8;
  rdna4::mubuf_calculate_addresses(inst, *wf, d);
  EXPECT_EQ(d.per_lane_addr[0], kBase + 0x20 + 0x10 + 0x80);

  wf->set_m0(0x40);
  inst.soffset = rdna4::OPR_SREG_M0_M0;
  rdna4::mubuf_calculate_addresses(inst, *wf, d);
  EXPECT_EQ(d.per_lane_addr[0], kBase + 0x20 + 0x10 + 0x40);
}

TEST(RdnaAddrCalcTest, Rdna4VbufferWrapsOffsetPartBeforeBaseAddition) {
  amdgpu::GpuMemory mem("rdna4_vbuffer_wrap_mem");
  amdgpu::L2Cache l2("rdna4_vbuffer_wrap_l2");
  amdgpu::ComputeUnitCore::Config cfg{};
  cfg.arch = ROCJITSU_CODE_ARCH_RDNA4;
  cfg.num_wf_slots = 1;
  cfg.sgprs_per_wf = 128;
  cfg.vgprs_per_wf = 16;
  cfg.lds_size_kb = 64;
  auto cu = amdgpu::ComputeUnitCore::create("rdna4_vbuffer_wrap_cu", cfg, &mem, &l2);
  ASSERT_NE(cu, nullptr);

  auto *wf = cu->dispatch_wf(0, 0, 128, 16);
  ASSERT_NE(wf, nullptr);
  wf->set_exec(1ULL);

  constexpr uint64_t kBase = 0x2'0000'1000ULL;
  uint32_t sbase = wf->sgpr_alloc().base;
  uint32_t vbase = wf->vgpr_alloc().base;
  cu->write_sgpr(sbase + 4, static_cast<uint32_t>(kBase));
  cu->write_sgpr(sbase + 5, static_cast<uint32_t>(kBase >> 32));
  cu->write_sgpr(sbase + 6, 0x1000);
  cu->write_sgpr(sbase + 7, 0);
  cu->write_vgpr(vbase + 4, 0, 0xFFFF'8200u);

  rdna4::VbufferMachineInst inst{};
  inst.rsrc = 4;
  inst.soffset = rdna4::OPR_SREG_M0_NULL;
  inst.offen = 1;
  inst.idxen = 0;
  inst.vaddr = 4;
  inst.ioffset = 0x7E00;

  amdgpu::VectorMemState d(amdgpu::GLOBAL_MEM);
  rdna4::mubuf_calculate_addresses(inst, *wf, d);
  EXPECT_EQ(d.lane_mask, 1ULL);
  EXPECT_EQ(d.per_lane_addr[0], kBase);
}

TEST(RdnaAddrCalcTest, Gfx1250VbufferWrapsOffsetPartBeforeBoundsCheck) {
  amdgpu::GpuMemory mem("gfx1250_vbuffer_wrap_mem");
  amdgpu::L2Cache l2("gfx1250_vbuffer_wrap_l2");
  amdgpu::ComputeUnitCore::Config cfg{};
  cfg.arch = ROCJITSU_CODE_ARCH_GFX1250;
  cfg.num_wf_slots = 1;
  cfg.sgprs_per_wf = 128;
  cfg.vgprs_per_wf = 16;
  cfg.lds_size_kb = 64;
  auto cu = amdgpu::ComputeUnitCore::create("gfx1250_vbuffer_wrap_cu", cfg, &mem, &l2);
  ASSERT_NE(cu, nullptr);

  auto *wf = cu->dispatch_wf(0, 0, 128, 16);
  ASSERT_NE(wf, nullptr);
  wf->set_exec(1ULL);

  constexpr uint64_t kBase = 0x2'0000'1000ULL;
  uint32_t sbase = wf->sgpr_alloc().base;
  uint32_t vbase = wf->vgpr_alloc().base;
  cu->write_sgpr(sbase + 4, static_cast<uint32_t>(kBase));
  cu->write_sgpr(sbase + 5, static_cast<uint32_t>(kBase >> 32));
  cu->write_sgpr(sbase + 6, 0x100);
  cu->write_sgpr(sbase + 7, 0);
  cu->write_vgpr(vbase + 4, 0, 0xFFFF'8200u);

  gfx1250::VbufferMachineInst inst{};
  inst.rsrc = 4;
  inst.soffset = gfx1250::OPR_SREG_NULL;
  inst.offen = 1;
  inst.idxen = 0;
  inst.vaddr = 4;
  inst.ioffset = 0x7E00;

  amdgpu::VectorMemState d(amdgpu::GLOBAL_MEM);
  gfx1250::mubuf_calculate_addresses(inst, *wf, d);
  EXPECT_EQ(d.lane_mask, 1ULL);
  EXPECT_EQ(d.per_lane_addr[0], kBase);
}

TEST(Gfx1250AddrCalcTest, VbufferZeroNumRecordsMasksAllLanes) {
  amdgpu::GpuMemory mem("gfx1250_vbuffer_zero_records_mem");
  amdgpu::L2Cache l2("gfx1250_vbuffer_zero_records_l2");
  amdgpu::ComputeUnitCore::Config cfg{};
  cfg.arch = ROCJITSU_CODE_ARCH_GFX1250;
  cfg.num_wf_slots = 1;
  cfg.sgprs_per_wf = 128;
  cfg.vgprs_per_wf = 16;
  cfg.lds_size_kb = 64;
  auto cu = amdgpu::ComputeUnitCore::create("gfx1250_vbuffer_zero_records_cu", cfg, &mem, &l2);
  ASSERT_NE(cu, nullptr);

  auto *wf = cu->dispatch_wf(0, 0, 128, 16);
  ASSERT_NE(wf, nullptr);
  wf->set_exec(0x3ULL);

  // A null optional pointer is represented by an all-zero descriptor. Its
  // zero-sized range makes every access out of bounds, so loads return zero
  // and stores are dropped without touching address zero.
  uint32_t vbase = wf->vgpr_alloc().base;
  cu->write_vgpr(vbase + 4, 0, 0);
  cu->write_vgpr(vbase + 4, 1, 0x1000);

  gfx1250::VbufferMachineInst inst{};
  inst.rsrc = 40;
  inst.soffset = gfx1250::OPR_SREG_NULL;
  inst.offen = 1;
  inst.idxen = 0;
  inst.vaddr = 4;

  amdgpu::VectorMemState d(amdgpu::GLOBAL_MEM);
  gfx1250::mubuf_calculate_addresses(inst, *wf, d);
  EXPECT_EQ(d.exec_mask, 0x3ULL);
  EXPECT_EQ(d.lane_mask, 0ULL);
  EXPECT_EQ(d.per_lane_addr[0], 0ULL);
  EXPECT_EQ(d.per_lane_addr[1], 0ULL);
}

TEST(Gfx1250AddrCalcTest, VbufferStructuredStrideChecksIndexBounds) {
  amdgpu::GpuMemory mem("gfx1250_vbuffer_structured_stride_mem");
  amdgpu::L2Cache l2("gfx1250_vbuffer_structured_stride_l2");
  amdgpu::ComputeUnitCore::Config cfg{};
  cfg.arch = ROCJITSU_CODE_ARCH_GFX1250;
  cfg.num_wf_slots = 1;
  cfg.sgprs_per_wf = 128;
  cfg.vgprs_per_wf = 16;
  cfg.lds_size_kb = 64;
  auto cu = amdgpu::ComputeUnitCore::create("gfx1250_vbuffer_structured_stride_cu", cfg, &mem, &l2);
  ASSERT_NE(cu, nullptr);

  auto *wf = cu->dispatch_wf(0, 0, 128, 16);
  ASSERT_NE(wf, nullptr);
  wf->set_exec(0x3ULL);

  constexpr uint64_t kBase = 0x2'0000'1000ULL;
  constexpr uint32_t kNumRecords = 2;
  constexpr uint32_t kStride = 16;
  uint32_t sbase = wf->sgpr_alloc().base;
  uint32_t vbase = wf->vgpr_alloc().base;
  cu->write_sgpr(sbase + 40, static_cast<uint32_t>(kBase));
  cu->write_sgpr(sbase + 41, static_cast<uint32_t>(kBase >> 32) | (kNumRecords << 25));
  cu->write_sgpr(sbase + 42, 0);
  cu->write_sgpr(sbase + 43, kStride << 12);
  cu->write_vgpr(vbase + 4, 0, kNumRecords - 1);
  cu->write_vgpr(vbase + 4, 1, kNumRecords);

  gfx1250::VbufferMachineInst inst{};
  inst.rsrc = 40;
  inst.soffset = gfx1250::OPR_SREG_NULL;
  inst.offen = 0;
  inst.idxen = 1;
  inst.vaddr = 4;

  amdgpu::VectorMemState d(amdgpu::GLOBAL_MEM);
  gfx1250::mubuf_calculate_addresses(inst, *wf, d);
  EXPECT_EQ(d.exec_mask, 0x3ULL);
  EXPECT_EQ(d.lane_mask, 0x1ULL);
  EXPECT_EQ(d.per_lane_addr[0], kBase + (kNumRecords - 1) * kStride);
  EXPECT_EQ(d.per_lane_addr[1], 0ULL);
}

void expect_vector_lane_reads_use_own_wave_vgprs(rj_code_arch_t arch) {
  amdgpu::GpuMemory mem("rdna_lane_read_mem");
  amdgpu::L2Cache l2("rdna_lane_read_l2");
  amdgpu::ComputeUnitCore::Config cfg{};
  cfg.arch = arch;
  cfg.num_wf_slots = 4;
  cfg.sgprs_per_wf = 128;
  cfg.vgprs_per_wf = 256;
  cfg.lds_size_kb = 64;

  auto cu = amdgpu::ComputeUnitCore::create("rdna_lane_read_cu", cfg, &mem, &l2);
  ASSERT_NE(cu, nullptr);

  std::array<amdgpu::Wavefront *, 4> wfs{};
  for (uint32_t i = 0; i < wfs.size(); ++i) {
    wfs[i] = cu->dispatch_wf(0, 0, cfg.sgprs_per_wf, cfg.vgprs_per_wf);
    ASSERT_NE(wfs[i], nullptr);
    const uint32_t vbase = wfs[i]->vgpr_alloc().base;
    cu->write_vgpr(vbase + 1, 0, 0x100u + i);
    cu->write_vgpr(vbase + 4, 0, 0xdead0000u + i);
    cu->write_vgpr(vbase + 4, 31, 0x300u + i);
    cu->write_vgpr(vbase + 141, 0, 0xbeef0000u + i);
    cu->write_vgpr(vbase + 141, 2, 0x200u + i);
    cu->write_vgpr(vbase + 141, 31, 0x400u + i);
  }

  auto decoder = Decoder::create(arch);
  ASSERT_NE(decoder, nullptr);

  constexpr std::array<uint32_t, 2> kReadfirstlaneS24V1 = {
      0x7e300501u, // v_readfirstlane_b32_e32 s24, v1
      0u,
  };
  constexpr std::array<uint32_t, 2> kReadlaneS4V141Lane2 = {
      0xd7600004u, // v_readlane_b32 s4, v141, 2
      0x0201058du,
  };
  constexpr std::array<uint32_t, 2> kReadlaneS4V4Lane31 = {
      0xd7600004u, // v_readlane_b32 s4, v4, 31
      0x02013f04u,
  };
  constexpr std::array<uint32_t, 2> kReadlaneS4V141S2 = {
      0xd7600004u, // v_readlane_b32 s4, v141, s2
      0x0200058du,
  };
  constexpr std::array<uint32_t, 2> kWritelaneV141S4S2 = {
      0xd761008du, // v_writelane_b32 v141, s4, s2
      0x02000404u,
  };

  for (uint32_t i = 0; i < wfs.size(); ++i) {
    const uint32_t sbase = wfs[i]->sgpr_alloc().base;
    std::unique_ptr<Instruction> inst(decoder->decode(kReadfirstlaneS24V1.data()));
    ASSERT_NE(inst, nullptr);
    cu->write_sgpr(sbase + 24, 0);
    cu->execute_instruction(inst.get(), *wfs[i]);
    EXPECT_EQ(cu->read_sgpr(sbase + 24), 0x100u + i);
  }

  for (uint32_t i = 0; i < wfs.size(); ++i) {
    const uint32_t sbase = wfs[i]->sgpr_alloc().base;
    std::unique_ptr<Instruction> inst(decoder->decode(kReadlaneS4V141Lane2.data()));
    ASSERT_NE(inst, nullptr);
    cu->write_sgpr(sbase + 2, 0);
    cu->write_sgpr(sbase + 4, 0);
    cu->execute_instruction(inst.get(), *wfs[i]);
    EXPECT_EQ(cu->read_sgpr(sbase + 4), 0x200u + i);
  }

  for (uint32_t i = 0; i < wfs.size(); ++i) {
    const uint32_t sbase = wfs[i]->sgpr_alloc().base;
    std::unique_ptr<Instruction> inst(decoder->decode(kReadlaneS4V4Lane31.data()));
    ASSERT_NE(inst, nullptr);
    cu->write_sgpr(sbase + 4, 0);
    cu->write_sgpr(sbase + 31, 0);
    cu->execute_instruction(inst.get(), *wfs[i]);
    EXPECT_EQ(cu->read_sgpr(sbase + 4), 0x300u + i);
  }

  for (uint32_t i = 0; i < wfs.size(); ++i) {
    const uint32_t sbase = wfs[i]->sgpr_alloc().base;
    std::unique_ptr<Instruction> inst(decoder->decode(kReadlaneS4V141S2.data()));
    ASSERT_NE(inst, nullptr);
    cu->write_sgpr(sbase + 2, 31);
    cu->write_sgpr(sbase + 4, 0);
    cu->execute_instruction(inst.get(), *wfs[i]);
    EXPECT_EQ(cu->read_sgpr(sbase + 4), 0x400u + i);
  }

  for (uint32_t i = 0; i < wfs.size(); ++i) {
    const uint32_t sbase = wfs[i]->sgpr_alloc().base;
    const uint32_t vbase = wfs[i]->vgpr_alloc().base;
    std::unique_ptr<Instruction> inst(decoder->decode(kWritelaneV141S4S2.data()));
    ASSERT_NE(inst, nullptr);
    cu->write_sgpr(sbase + 2, 31);
    cu->write_sgpr(sbase + 4, 0x500u + i);
    cu->execute_instruction(inst.get(), *wfs[i]);
    EXPECT_EQ(cu->read_vgpr(vbase + 141, 2), 0x200u + i);
    EXPECT_EQ(cu->read_vgpr(vbase + 141, 31), 0x500u + i);
  }
}

TEST(RdnaVectorLaneReadTest, ReadlaneFamilyUsesDecodedSourceVgprPerWave) {
  {
    SCOPED_TRACE("rdna3");
    expect_vector_lane_reads_use_own_wave_vgprs(ROCJITSU_CODE_ARCH_RDNA3);
  }
  {
    SCOPED_TRACE("rdna4");
    expect_vector_lane_reads_use_own_wave_vgprs(ROCJITSU_CODE_ARCH_RDNA4);
  }
}

} // namespace