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

#include "aql_queue.h"
#include "halt_snapshot_plugin.h"

#include "embedded_schema.h"
#include "rocjitsu/code/amdgpu_elf.h"
#include "rocjitsu/code/kernel_symbol.h"
#include "rocjitsu/config/config_loader.h"
#include "rocjitsu/isa/arch/amdgpu/cdna4/opcodes.h"
#include "rocjitsu/isa/arch/amdgpu/cdna4/vop3p.h"
#include "rocjitsu/isa/arch/amdgpu/shared/mma_exec.h"
#include "rocjitsu/kmd/linux/kfd_process.h"
#include "rocjitsu/vm/amdgpu/dispatch_entry.h"
#include "rocjitsu/vm/amdgpu/gpu_memory.h"
#include "rocjitsu/vm/amdgpu/l1_scalar_cache.h"
#include "rocjitsu/vm/amdgpu/l2_cache.h"
#include "rocjitsu/vm/amdgpu/wavefront.h"
#include "rocjitsu/vm/rj_vm.h"
#include "rocjitsu/vm/soc.h"

#include "simdojo/sim/simulation.h"

#include "rocjitsu/base/rj_compiler.h"
RJ_DIAGNOSTIC_PUSH
RJ_DIAGNOSTIC_IGNORE_PEDANTIC
#include "hsa/AMDHSAKernelDescriptor.h"
RJ_DIAGNOSTIC_POP

#include <gtest/gtest.h>

#include <algorithm>
#include <array>
#include <atomic>
#include <barrier>
#include <bit>
#include <chrono>
#include <cstdint>
#include <cstring>
#include <limits>
#include <memory>
#include <new>
#include <span>
#include <stdexcept>
#include <string>
#include <string_view>
#include <sys/mman.h>
#include <thread>
#include <unistd.h>
#include <vector>

namespace {

// SOPP encoding: bits[31:23] = 0x17F (SOPP prefix), bits[22:16] = op.
constexpr uint32_t SOPP_S_NOP = 0xBF800000;
constexpr uint32_t SOPP_S_ENDPGM = 0xBF810000;

using namespace rocjitsu;

struct VmFixture {
  std::unique_ptr<simdojo::SimulationEngine> engine;
  SoC *soc_ptr = nullptr;
  amdgpu::GpuMemory *gpu_mem = nullptr;

  VmFixture(std::string_view arch = "cdna3", uint32_t num_cus = 1, uint32_t num_wf_slots = 10,
            uint32_t lds_size_kb = 64, uint32_t sgprs_per_wf = 104) {
    std::string cu_range = "cu[0:" + std::to_string(num_cus) + "]";
    std::string links;
    for (uint32_t i = 0; i < num_cus; ++i) {
      if (i > 0)
        links += ",";
      links += R"({"src":"xcd0.cp.req_)" + std::to_string(i) + R"(","dst":"xcd0.se0.cu)" +
               std::to_string(i) + R"(.cpl","latency":1,"weight":2})";
      links += R"(,{"src":"xcd0.se0.cu)" + std::to_string(i) + R"(.req","dst":"xcd0.l2.cpl_)" +
               std::to_string(i) + R"(","latency":1,"weight":10})";
    }

    std::string json = R"({"max_ticks":10000,"num_threads":1,"vm":{"arch":")" + std::string(arch) +
                       R"("},)"
                       R"("topology":{"root":{"name":"soc","type":"soc","children":[)"
                       R"({"name":"vram","type":"gpu_memory"},)"
                       R"({"name":"xcd0","type":"xcd","children":[)"
                       R"({"name":"l2","type":"l2_cache"},)"
                       R"({"name":"cp","type":"command_processor"},)"
                       R"({"name":"se0","type":"shader_engine","children":[)"
                       R"({"name":")" +
                       cu_range +
                       R"(","type":"compute_unit","config":[)"
                       R"({"key":"num_wf_slots","value":")" +
                       std::to_string(num_wf_slots) +
                       R"("},)"
                       R"({"key":"sgprs_per_wf","value":")" +
                       std::to_string(sgprs_per_wf) +
                       R"("},)"
                       R"({"key":"vgprs_per_wf","value":"256"},)"
                       R"({"key":"lds_size_kb","value":")" +
                       std::to_string(lds_size_kb) +
                       R"("})"
                       R"(]}]}]}]},"links":[)" +
                       links + R"(]}})";
    auto loaded = config::load_config_from_string(json, rocjitsu::kEmbeddedSchema);
    soc_ptr = loaded.soc();
    gpu_mem = loaded.memory();
    engine = std::make_unique<simdojo::SimulationEngine>(loaded.engine_config);
    engine->topology().set_root(loaded.take_root());
    loaded.wire_links(engine->topology());
    engine->create();
  }

  amdgpu::Xcd *xcd(uint32_t idx = 0) { return soc_ptr->xcd(idx); }
  amdgpu::ShaderEngine *se(uint32_t idx = 0) { return soc_ptr->xcd(0)->shader_engine(idx); }
  amdgpu::GpuMemory *mem() { return gpu_mem; }
  amdgpu::ComputeUnitCore *cu(uint32_t idx = 0) { return se()->compute_unit(idx); }
  amdgpu::CommandProcessor *cp(uint32_t idx = 0) { return xcd(idx)->command_processor(); }

  /// Install a HaltSnapshotPlugin that records each wavefront's final register
  /// state at s_endpgm (before resources are freed). Call before dispatching.
  test::HaltSnapshotPlugin *capture_halts() {
    plugin_group_ = test::make_halt_snapshot_group(&snapshot_plugin_);
    soc_ptr->set_plugin_group(plugin_group_);
    return snapshot_plugin_;
  }

  /// Place a single resident wavefront on CU 0 without running it to s_endpgm.
  /// Used by tests that drive individual instructions and then read the register
  /// file directly — the wave stays resident (its resources are not freed) until
  /// the fixture is destroyed.
  amdgpu::Wavefront *dispatch_scratch_wf(uint32_t num_sgprs = 104, uint32_t num_vgprs = 256) {
    return cu()->dispatch_wf(/*wg_id=*/0, /*pc=*/0x1040, num_sgprs, num_vgprs);
  }

  std::shared_ptr<ExecutionPluginGroup> plugin_group_;
  test::HaltSnapshotPlugin *snapshot_plugin_ = nullptr;

  /// Write a kernel descriptor + instructions to GPU memory per AMDHSA ABI.
  /// Returns the kernel_object address.
  uint64_t write_kernel(uint64_t addr, const void *code, size_t code_size, uint32_t sgprs = 104,
                        uint32_t vgprs = 256, uint32_t user_sgprs = 2,
                        uint32_t group_segment_fixed_size = 0, bool wgp_mode = false,
                        uint32_t enable_vgpr_workitem_id = 0,
                        uint32_t extra_compute_pgm_rsrc1 = 0) {
    using namespace rocr::llvm::amdhsa;
    kernel_descriptor_t kd{};
    kd.kernel_code_entry_byte_offset = sizeof(kernel_descriptor_t);
    AMDHSA_BITS_SET(kd.compute_pgm_rsrc1, COMPUTE_PGM_RSRC1_GRANULATED_WORKITEM_VGPR_COUNT,
                    ((vgprs / 8) - 1));
    AMDHSA_BITS_SET(kd.compute_pgm_rsrc1, COMPUTE_PGM_RSRC1_GRANULATED_WAVEFRONT_SGPR_COUNT,
                    ((sgprs / 8) - 1));
    AMDHSA_BITS_SET(kd.compute_pgm_rsrc2, COMPUTE_PGM_RSRC2_USER_SGPR_COUNT, user_sgprs);
    AMDHSA_BITS_SET(kd.compute_pgm_rsrc1, COMPUTE_PGM_RSRC1_WGP_MODE, (wgp_mode ? 1u : 0u));
    kd.group_segment_fixed_size = group_segment_fixed_size;
    AMDHSA_BITS_SET(kd.compute_pgm_rsrc2, COMPUTE_PGM_RSRC2_ENABLE_VGPR_WORKITEM_ID,
                    enable_vgpr_workitem_id);
    kd.compute_pgm_rsrc1 |= extra_compute_pgm_rsrc1;

    mem()->load_image(reinterpret_cast<const uint8_t *>(&kd), sizeof(kd), addr);
    mem()->load_image(static_cast<const uint8_t *>(code), code_size,
                      addr + sizeof(kernel_descriptor_t));
    return addr;
  }
};

// Drive the engine until the listed CUs have no resident wavefronts. A wavefront
// frees itself at s_endpgm (num_wfs()/has_active_wfs() drop as it halts), so the
// kernel is complete once every listed CU reports idle. Waves that need their final
// register state inspected should capture it via HaltSnapshotPlugin, which snapshots
// at halt regardless of where this loop stops.
void step_until_halted(simdojo::SimulationEngine &engine,
                       std::initializer_list<amdgpu::ComputeUnitCore *> cus,
                       uint32_t max_steps = 10000) {
  auto any_active = [&]() {
    for (auto *cu : cus)
      if (cu->has_active_wfs())
        return true;
    return false;
  };
  bool saw_work = false;
  for (uint32_t i = 0; i < max_steps && engine.step(); ++i) {
    if (any_active())
      saw_work = true;
    else if (saw_work)
      break;
  }
}

std::vector<uint8_t> make_loaded_kernel_symbol_elf(uint64_t kernel_descriptor_offset,
                                                   std::string_view symbol_name,
                                                   bool include_symbol_terminator = true) {
  constexpr uint64_t dyn_offset = 0x100;
  constexpr uint64_t symtab_offset = 0x200;
  constexpr uint64_t strtab_offset = 0x300;
  constexpr uint64_t hash_offset = 0x380;

  std::vector<uint8_t> image(4096, 0);

  Elf64_Ehdr ehdr{};
  std::memcpy(ehdr.e_ident, EI_MAGIC, EI_MAGIC_SIZE);
  ehdr.e_ident[EI_CLASS] = ELFCLASS64;
  ehdr.e_ident[EI_OSABI] = ELFOSABI_AMDGPU_HSA;
  ehdr.e_type = ET_DYN;
  ehdr.e_machine = EM_AMDGPU;
  ehdr.e_version = 1;
  ehdr.e_phoff = sizeof(Elf64_Ehdr);
  ehdr.e_ehsize = sizeof(Elf64_Ehdr);
  ehdr.e_phentsize = sizeof(Elf64_Phdr);
  ehdr.e_phnum = 1;
  std::memcpy(image.data(), &ehdr, sizeof(ehdr));

  Elf64_Phdr phdr{};
  phdr.p_type = PT_DYNAMIC;
  phdr.p_vaddr = dyn_offset;
  phdr.p_memsz = 5 * sizeof(Elf64_Dyn);
  std::memcpy(image.data() + ehdr.e_phoff, &phdr, sizeof(phdr));

  auto *dyn = reinterpret_cast<Elf64_Dyn *>(image.data() + dyn_offset);
  dyn[0].d_tag = DT_SYMTAB;
  dyn[0].d_un.d_val = symtab_offset;
  dyn[1].d_tag = DT_STRTAB;
  dyn[1].d_un.d_val = strtab_offset;
  dyn[2].d_tag = DT_STRSZ;
  dyn[2].d_un.d_val = 1 + symbol_name.size() + (include_symbol_terminator ? 1 : 0);
  dyn[3].d_tag = DT_HASH;
  dyn[3].d_un.d_val = hash_offset;
  dyn[4].d_tag = DT_NULL;

  image[strtab_offset] = '\0';
  std::memcpy(image.data() + strtab_offset + 1, symbol_name.data(), symbol_name.size());

  auto *sym = reinterpret_cast<Elf64_Sym *>(image.data() + symtab_offset);
  sym[1].st_name = 1;
  sym[1].st_value = kernel_descriptor_offset;

  auto *hash = reinterpret_cast<uint32_t *>(image.data() + hash_offset);
  hash[1] = 2; // nchain: null symbol + kernel descriptor symbol.

  return image;
}

TEST(KernelSymbolTest, RejectsDynstrSymbolWithoutTerminator) {
  constexpr uint64_t kernel_descriptor_offset = 0x800;
  constexpr std::string_view symbol_name = "unterminated_kernel.kd";
  auto image = make_loaded_kernel_symbol_elf(kernel_descriptor_offset, symbol_name,
                                             /*include_symbol_terminator=*/false);

  EXPECT_TRUE(
      find_kernel_symbol(image.data() + kernel_descriptor_offset, image.data(), image.size())
          .empty());
}

TEST(GpuMemoryTest, ReadWriteRoundTrip) {
  VmFixture f;
  auto *mem = f.mem();

  mem->write32(0x1000, 0xDEADBEEF);
  EXPECT_EQ(mem->read32(0x1000), 0xDEADBEEF);

  mem->write64(0x2000, 0x0123456789ABCDEFULL);
  EXPECT_EQ(mem->read64(0x2000), 0x0123456789ABCDEFULL);
}

TEST(GpuMemoryTest, LoadImage) {
  VmFixture f;
  auto *mem = f.mem();

  const uint32_t program[] = {SOPP_S_NOP, SOPP_S_ENDPGM};
  mem->load_image(reinterpret_cast<const uint8_t *>(program), sizeof(program), 0x0);

  EXPECT_EQ(mem->read32(0x0), SOPP_S_NOP);
  EXPECT_EQ(mem->read32(0x4), SOPP_S_ENDPGM);
}

TEST(GpuMemoryTest, SparsePages) {
  VmFixture f;
  auto *mem = f.mem();

  mem->write32(0x0, 42);
  mem->write32(0x100000, 99);
  EXPECT_EQ(mem->read32(0x0), 42u);
  EXPECT_EQ(mem->read32(0x100000), 99u);
  EXPECT_EQ(mem->read32(0x50000), 0u);
}

TEST(GpuMemoryTest, FindHostRangeUsesVmidPageTable) {
  amdgpu::GpuMemory mem("vmid_range_mem");
  KfdProcess process(/*process_id=*/123);
  alignas(4096) std::array<uint8_t, 3 * amdgpu::GpuMemory::PAGE_SIZE> backing{};
  constexpr uint64_t kGpuVa = 0x100000;

  mem.register_process(process.process_id(), &process.page_table_, &process.page_table_mutex_);
  process.map_pages(kGpuVa, backing.data(), backing.size());

  auto [host_base, size] =
      mem.find_host_range(kGpuVa + amdgpu::GpuMemory::PAGE_SIZE + 0x123, process.process_id());
  EXPECT_EQ(host_base, reinterpret_cast<uint64_t>(backing.data()));
  EXPECT_EQ(size, backing.size());

  mem.unregister_process(process.process_id());
}

TEST(GpuMemoryTest, FindHostRangeStopsAtNonContiguousVmidHostPages) {
  amdgpu::GpuMemory mem("vmid_noncontiguous_range_mem");
  KfdProcess process(/*process_id=*/124);
  alignas(4096) std::array<uint8_t, 3 * amdgpu::GpuMemory::PAGE_SIZE> backing{};
  constexpr uint64_t kGpuVa = 0x200000;

  mem.register_process(process.process_id(), &process.page_table_, &process.page_table_mutex_);
  process.map_pages(kGpuVa, backing.data(), amdgpu::GpuMemory::PAGE_SIZE);
  process.map_pages(kGpuVa + amdgpu::GpuMemory::PAGE_SIZE,
                    backing.data() + 2 * amdgpu::GpuMemory::PAGE_SIZE,
                    amdgpu::GpuMemory::PAGE_SIZE);

  auto [first_host_base, first_size] = mem.find_host_range(kGpuVa, process.process_id());
  EXPECT_EQ(first_host_base, reinterpret_cast<uint64_t>(backing.data()));
  EXPECT_EQ(first_size, amdgpu::GpuMemory::PAGE_SIZE);

  auto [second_host_base, second_size] =
      mem.find_host_range(kGpuVa + amdgpu::GpuMemory::PAGE_SIZE, process.process_id());
  EXPECT_EQ(second_host_base,
            reinterpret_cast<uint64_t>(backing.data() + 2 * amdgpu::GpuMemory::PAGE_SIZE));
  EXPECT_EQ(second_size, amdgpu::GpuMemory::PAGE_SIZE);

  mem.unregister_process(process.process_id());
}

TEST(RdnaDispatchTest, WgpModeCombinesSiblingCuLdsCapacity) {
  constexpr uint32_t kPerCuLdsBytes = 64 * 1024;
  constexpr uint32_t kWgpLdsBytes = 2 * kPerCuLdsBytes;
  const uint32_t code[] = {SOPP_S_ENDPGM};

  for (const char *arch : {"rdna1", "rdna2", "rdna3", "rdna3_5", "rdna4"}) {
    SCOPED_TRACE(arch);
    VmFixture f(arch, 2, 10, /*lds_size_kb=*/64, /*sgprs_per_wf=*/128);
    auto *snap = f.capture_halts();
    uint64_t ko = f.write_kernel(0x1000, code, sizeof(code), 104, 64, 2, kWgpLdsBytes,
                                 /*wgp_mode=*/true);
    test::AqlQueue queue(f.mem(), f.cp());
    queue.dispatch(ko, 64, 64);

    EXPECT_NO_THROW(f.engine->run());
    EXPECT_EQ(f.cu(0)->lds().size_bytes(), kPerCuLdsBytes);
    EXPECT_EQ(f.se()->spi().max_wgp_lds_bytes(), kWgpLdsBytes);

    // A WGP-mode wave sees the combined sibling-CU LDS capacity (wave32 splits the
    // 64-thread workgroup into two waves; both see the same combined pool).
    ASSERT_GE(snap->snapshots().size(), 1u);
    for (const auto &wf : snap->snapshots())
      EXPECT_EQ(wf.lds_size_bytes, kWgpLdsBytes);
  }
}

TEST(AqlDispatchTest, InitializesModeFromComputePgmRsrc1) {
  using namespace rocr::llvm::amdhsa;

  uint32_t rsrc1 = 0;
  AMDHSA_BITS_SET(rsrc1, COMPUTE_PGM_RSRC1_FLOAT_ROUND_MODE_32, FLOAT_ROUND_MODE_ZERO);
  AMDHSA_BITS_SET(rsrc1, COMPUTE_PGM_RSRC1_FLOAT_ROUND_MODE_16_64, FLOAT_ROUND_MODE_PLUS_INFINITY);
  AMDHSA_BITS_SET(rsrc1, COMPUTE_PGM_RSRC1_FLOAT_DENORM_MODE_32, FLOAT_DENORM_MODE_FLUSH_SRC);
  AMDHSA_BITS_SET(rsrc1, COMPUTE_PGM_RSRC1_FLOAT_DENORM_MODE_16_64, FLOAT_DENORM_MODE_FLUSH_NONE);
  AMDHSA_BITS_SET(rsrc1, COMPUTE_PGM_RSRC1_ENABLE_DX10_CLAMP, 1);
  AMDHSA_BITS_SET(rsrc1, COMPUTE_PGM_RSRC1_ENABLE_IEEE_MODE, 1);
  AMDHSA_BITS_SET(rsrc1, COMPUTE_PGM_RSRC1_DEBUG_MODE, 1);
  AMDHSA_BITS_SET(rsrc1, COMPUTE_PGM_RSRC1_FP16_OVFL, 1);

  const uint32_t common_mode =
      (FLOAT_ROUND_MODE_ZERO << 0) | (FLOAT_ROUND_MODE_PLUS_INFINITY << 2) |
      (FLOAT_DENORM_MODE_FLUSH_SRC << 4) | (FLOAT_DENORM_MODE_FLUSH_NONE << 6) |
      amdgpu::Wavefront::FP16_OVFL_BIT;

  struct ModeInitCase {
    const char *arch;
    uint32_t endpgm;
    uint32_t expected_mode;
  };

  constexpr uint32_t kGfx9Endpgm = SOPP_S_ENDPGM;
  constexpr uint32_t kGfx11Endpgm = 0xBFB00000u;
  const ModeInitCase cases[] = {
      {"cdna1", kGfx9Endpgm, common_mode | (1u << 8) | (1u << 9) | (1u << 11)},
      {"cdna2", kGfx9Endpgm, common_mode | (1u << 8) | (1u << 9) | (1u << 11)},
      {"cdna3", kGfx9Endpgm, common_mode | (1u << 8) | (1u << 9) | (1u << 11)},
      {"cdna4", kGfx9Endpgm, common_mode | (1u << 8) | (1u << 9) | (1u << 11)},
      {"rdna1", kGfx9Endpgm, common_mode | (1u << 8) | (1u << 9) | (1u << 11)},
      {"rdna2", kGfx9Endpgm, common_mode | (1u << 8) | (1u << 9) | (1u << 11)},
      {"rdna3", kGfx11Endpgm, common_mode | (1u << 8) | (1u << 9) | (1u << 11)},
      {"rdna3_5", kGfx11Endpgm, common_mode | (1u << 8) | (1u << 9) | (1u << 11)},
      {"rdna4", kGfx11Endpgm, common_mode},
      {"gfx1250", kGfx11Endpgm, common_mode},
  };

  for (const ModeInitCase &mode_case : cases) {
    SCOPED_TRACE(mode_case.arch);
    const uint32_t code[] = {mode_case.endpgm};
    VmFixture f(mode_case.arch, 1, 10, /*lds_size_kb=*/64, /*sgprs_per_wf=*/128);
    auto *snap = f.capture_halts();
    uint64_t ko = f.write_kernel(0x1000, code, sizeof(code), 104, 64, 2, 0,
                                 /*wgp_mode=*/false, /*enable_vgpr_workitem_id=*/0,
                                 /*extra_compute_pgm_rsrc1=*/rsrc1);
    test::AqlQueue queue(f.mem(), f.cp());
    queue.dispatch(ko, 64, 64);

    EXPECT_NO_THROW(f.engine->run());
    ASSERT_GE(snap->snapshots().size(), 1u);
    for (const auto &wf : snap->snapshots())
      EXPECT_EQ(wf.mode_raw, mode_case.expected_mode);
  }
}

TEST(RdnaDispatchTest, Gfx1250DoesNotEnableWgpMode) {
  const uint32_t code[] = {SOPP_S_ENDPGM};
  VmFixture f("gfx1250", 2, 10, /*lds_size_kb=*/64, /*sgprs_per_wf=*/128);
  uint64_t ko = f.write_kernel(0x1000, code, sizeof(code), 104, 64, 2, 128 * 1024,
                               /*wgp_mode=*/true);
  test::AqlQueue queue(f.mem(), f.cp());
  queue.dispatch(ko, 64, 64);

  try {
    (void)f.engine->step();
    FAIL() << "gfx1250 must not enable WGP mode from the descriptor bit";
  } catch (const std::runtime_error &error) {
    EXPECT_NE(std::string(error.what()).find("in CU mode"), std::string::npos);
  }
}

TEST(RdnaDispatchTest, LegacySpiQueueRejectsWgpMode) {
  VmFixture f("rdna4", 2);
  amdgpu::DispatchEntry entry{};
  entry.wgp_mode = true;

  EXPECT_THROW(f.se()->spi().enqueue_wg(0, 0, &entry), std::invalid_argument);
}

TEST(RdnaDispatchTest, CuModeRejectsLdsRequestAboveOneCu) {
  const uint32_t code[] = {SOPP_S_ENDPGM};
  VmFixture f("rdna4", 2, 10, /*lds_size_kb=*/64, /*sgprs_per_wf=*/128);
  uint64_t ko = f.write_kernel(0x1000, code, sizeof(code), 104, 64, 2, 64 * 1024 + 1,
                               /*wgp_mode=*/false);
  test::AqlQueue queue(f.mem(), f.cp());
  queue.dispatch(ko, 64, 64);

  try {
    (void)f.engine->step();
    FAIL() << "oversized CU-mode LDS request should fail";
  } catch (const std::runtime_error &error) {
    EXPECT_NE(std::string(error.what())
                  .find("requests 65537 bytes of LDS (65792 bytes after alignment) in CU mode"),
              std::string::npos);
    EXPECT_NE(std::string(error.what()).find("at most 65536 bytes"), std::string::npos);
  }
}

TEST(RdnaDispatchTest, WgpModeRejectsLdsRequestAboveSiblingPair) {
  const uint32_t code[] = {SOPP_S_ENDPGM};
  VmFixture f("rdna4", 2, 10, /*lds_size_kb=*/64, /*sgprs_per_wf=*/128);
  uint64_t ko = f.write_kernel(0x1000, code, sizeof(code), 104, 64, 2, 128 * 1024 + 1,
                               /*wgp_mode=*/true);
  test::AqlQueue queue(f.mem(), f.cp());
  queue.dispatch(ko, 64, 64);

  try {
    (void)f.engine->step();
    FAIL() << "oversized WGP-mode LDS request should fail";
  } catch (const std::runtime_error &error) {
    EXPECT_NE(std::string(error.what())
                  .find("requests 131073 bytes of LDS (131328 bytes after alignment) in WGP mode"),
              std::string::npos);
    EXPECT_NE(std::string(error.what()).find("at most 131072 bytes"), std::string::npos);
  }
}

TEST(RdnaDispatchTest, WgpModeRequiresConfiguredSiblingCuPair) {
  const uint32_t code[] = {SOPP_S_ENDPGM};
  VmFixture f("rdna4", 1, 10, /*lds_size_kb=*/64, /*sgprs_per_wf=*/128);
  uint64_t ko = f.write_kernel(0x1000, code, sizeof(code), 104, 64, 2, 0,
                               /*wgp_mode=*/true);
  test::AqlQueue queue(f.mem(), f.cp());
  queue.dispatch(ko, 64, 64);

  EXPECT_THROW((void)f.engine->step(), std::runtime_error);
}

TEST(GpuMemoryTest, VmidMappedKernelSymbolUsesTranslatedHostPointer) {
  amdgpu::GpuMemory mem("vmid_kernel_symbol_mem");
  KfdProcess process(/*process_id=*/125);
  constexpr uint64_t gpu_va = 0x5400200000;
  constexpr uint64_t kernel_descriptor_offset = 0x800;
  auto image = make_loaded_kernel_symbol_elf(kernel_descriptor_offset, "vmid_kernel.kd");

  mem.register_process(process.process_id(), &process.page_table_, &process.page_table_mutex_);
  process.map_pages(gpu_va, image.data(), image.size());

  uint64_t kernel_object = gpu_va + kernel_descriptor_offset;
  auto [host_base, host_size] = mem.find_host_range(kernel_object, process.process_id());
  ASSERT_NE(host_base, 0u);

  // Kernel descriptors arrive as GPU VAs. Symbol lookup needs the translated
  // host pointer so pointer subtraction against the loaded ELF image produces
  // the descriptor offset recorded in .dynsym.
  auto *mapped_page = mem.translate_debug(kernel_object, process.process_id());
  ASSERT_NE(mapped_page, nullptr);
  auto *kernel_object_host = mapped_page + (kernel_object & amdgpu::GpuMemory::PAGE_MASK);

  EXPECT_NE(reinterpret_cast<uint64_t>(kernel_object_host), kernel_object);
  EXPECT_EQ(find_kernel_symbol(kernel_object_host, reinterpret_cast<const uint8_t *>(host_base),
                               host_size),
            "vmid_kernel");

  mem.unregister_process(process.process_id());
}

TEST(GpuMemoryTest, UnregisterInvalidatesThreadLocalTranslationCaches) {
  amdgpu::GpuMemory memory("memory");
  constexpr uint32_t kPid = 7;
  constexpr uint64_t kBaseVa = 0x40000000;
  constexpr uint64_t kOffset = 0x123;
  constexpr uint64_t kAddr = kBaseVa + kOffset;

  KfdProcess process(kPid);
  std::array<uint8_t, KfdProcess::kPageSize> page{};
  page[kOffset] = 0x5a;
  process.map_pages(kBaseVa, page.data(), page.size(), amdgpu::Mtype::UC);
  memory.register_process(kPid, &process.page_table_, &process.page_table_mutex_,
                          process.page_table_generation());

  EXPECT_EQ(memory.resolve_host_ptr(kAddr, kPid), page.data());
  EXPECT_EQ(memory.read8(kAddr, kPid), page[kOffset]);
  EXPECT_EQ(memory.pte_mtype(kAddr, kPid), amdgpu::Mtype::UC);

  memory.unregister_process(kPid);

  EXPECT_EQ(memory.resolve_host_ptr(kAddr, kPid), nullptr);
  EXPECT_EQ(memory.pte_mtype(kAddr, kPid), amdgpu::Mtype::RW);
}

TEST(GpuMemoryTest, ReregisterProcessInvalidatesThreadLocalTranslationCaches) {
  amdgpu::GpuMemory memory("memory");
  constexpr uint32_t kPid = 7;
  constexpr uint64_t kBaseVa = 0x40000000;
  constexpr uint64_t kOffset = 0x123;
  constexpr uint64_t kAddr = kBaseVa + kOffset;

  KfdProcess first_process(kPid);
  KfdProcess second_process(kPid);
  std::array<uint8_t, KfdProcess::kPageSize> first_page{};
  std::array<uint8_t, KfdProcess::kPageSize> second_page{};
  first_page[kOffset] = 0x11;
  second_page[kOffset] = 0x22;
  first_process.map_pages(kBaseVa, first_page.data(), first_page.size(), amdgpu::Mtype::UC);
  second_process.map_pages(kBaseVa, second_page.data(), second_page.size(), amdgpu::Mtype::CC);

  memory.register_process(kPid, &first_process.page_table_, &first_process.page_table_mutex_,
                          first_process.page_table_generation());
  ASSERT_EQ(memory.resolve_host_ptr(kAddr, kPid), first_page.data());
  ASSERT_EQ(memory.read8(kAddr, kPid), first_page[kOffset]);
  ASSERT_EQ(memory.pte_mtype(kAddr, kPid), amdgpu::Mtype::UC);

  memory.register_process(kPid, &second_process.page_table_, &second_process.page_table_mutex_,
                          second_process.page_table_generation());
  EXPECT_EQ(memory.resolve_host_ptr(kAddr, kPid), second_page.data());
  EXPECT_EQ(memory.read8(kAddr, kPid), second_page[kOffset]);
  EXPECT_EQ(memory.pte_mtype(kAddr, kPid), amdgpu::Mtype::CC);
}

TEST(GpuMemoryTest, PageTableEntryMutationsInvalidateCachedPtes) {
  amdgpu::GpuMemory memory("memory");
  constexpr uint32_t kPid = 7;
  constexpr uint64_t kBaseVa = 0x40000000;
  constexpr uint64_t kOffset = 0x123;
  constexpr uint64_t kAddr = kBaseVa + kOffset;

  KfdProcess process(kPid);
  std::array<uint8_t, KfdProcess::kPageSize> old_page{};
  std::array<uint8_t, KfdProcess::kPageSize> new_page{};
  old_page[kOffset] = 0x11;
  new_page[kOffset] = 0x22;
  process.map_pages(kBaseVa, old_page.data(), old_page.size(), amdgpu::Mtype::UC);
  memory.register_process(kPid, &process.page_table_, &process.page_table_mutex_,
                          process.page_table_generation());

  ASSERT_EQ(memory.read8(kAddr, kPid), old_page[kOffset]);
  ASSERT_EQ(memory.pte_mtype(kAddr, kPid), amdgpu::Mtype::UC);

  process.remap_page_host_ptrs(kBaseVa, old_page.data(), new_page.data(), new_page.size());
  EXPECT_EQ(memory.read8(kAddr, kPid), new_page[kOffset]);

  process.set_page_mtype(kBaseVa, new_page.size(), amdgpu::Mtype::CC);
  EXPECT_EQ(memory.pte_mtype(kAddr, kPid), amdgpu::Mtype::CC);
}

TEST(GpuMemoryTest, ReusedMemoryInstanceInvalidatesThreadLocalTranslationCaches) {
  constexpr uint32_t kPid = 7;
  constexpr uint64_t kBaseVa = 0x40000000;
  constexpr uint64_t kOffset = 0x123;
  constexpr uint64_t kAddr = kBaseVa + kOffset;
  alignas(amdgpu::GpuMemory) unsigned char storage[sizeof(amdgpu::GpuMemory)];

  KfdProcess first_process(kPid);
  std::array<uint8_t, KfdProcess::kPageSize> first_page{};
  first_page[kOffset] = 0x11;
  first_process.map_pages(kBaseVa, first_page.data(), first_page.size(), amdgpu::Mtype::UC);

  auto *first_memory = new (storage) amdgpu::GpuMemory("memory");
  first_memory->register_process(kPid, &first_process.page_table_, &first_process.page_table_mutex_,
                                 first_process.page_table_generation());
  EXPECT_EQ(first_memory->read8(kAddr, kPid), first_page[kOffset]);
  EXPECT_EQ(first_memory->pte_mtype(kAddr, kPid), amdgpu::Mtype::UC);
  first_memory->~GpuMemory();

  KfdProcess second_process(kPid);
  std::array<uint8_t, KfdProcess::kPageSize> second_page{};
  second_page[kOffset] = 0x22;
  second_process.map_pages(kBaseVa, second_page.data(), second_page.size(), amdgpu::Mtype::CC);

  auto *second_memory = new (storage) amdgpu::GpuMemory("memory");
  second_memory->register_process(kPid, &second_process.page_table_,
                                  &second_process.page_table_mutex_,
                                  second_process.page_table_generation());
  EXPECT_EQ(second_memory->read8(kAddr, kPid), second_page[kOffset]);
  EXPECT_EQ(second_memory->pte_mtype(kAddr, kPid), amdgpu::Mtype::CC);
  second_memory->~GpuMemory();
}

TEST(GpuMemoryThreadingTest, AtomicRmwKeepsStorageIdentityAcrossConcurrentMap) {
  amdgpu::GpuMemory memory("memory");
  constexpr uint32_t kVmid = 17;

  void *raw_mapping = mmap(nullptr, KfdProcess::kPageSize, PROT_READ | PROT_WRITE,
                           MAP_PRIVATE | MAP_ANONYMOUS, -1, 0);
  ASSERT_NE(raw_mapping, MAP_FAILED);
  struct Mapping {
    uint8_t *data;
    ~Mapping() { munmap(data, KfdProcess::kPageSize); }
  } mapping{static_cast<uint8_t *>(raw_mapping)};

  KfdProcess process(kVmid);
  memory.register_process(kVmid, &process.page_table_, &process.page_table_mutex_,
                          process.page_table_generation());
  memory.set_process_client_pid(kVmid, getpid());

  auto *target = reinterpret_cast<uint32_t *>(mapping.data + 64);
  *target = 0;

  std::barrier first_atomic_has_read(2);
  std::barrier allow_first_atomic_to_write(2);
  std::thread first_atomic([&] {
    memory.atomic_rmw(
        reinterpret_cast<uint64_t>(target), sizeof(*target),
        [&](uint8_t *storage) {
          uint32_t value = 0;
          std::memcpy(&value, storage, sizeof(value));
          ++value;
          std::memcpy(storage, &value, sizeof(value));
          first_atomic_has_read.arrive_and_wait();
          allow_first_atomic_to_write.arrive_and_wait();
        },
        kVmid);
  });

  first_atomic_has_read.arrive_and_wait();

  const bool acquired_page_table_exclusively = process.page_table_mutex_.try_lock();
  if (acquired_page_table_exclusively)
    process.page_table_mutex_.unlock();
  EXPECT_FALSE(acquired_page_table_exclusively)
      << "fallback atomic must retain the page-table shared lock during its callback";

  allow_first_atomic_to_write.arrive_and_wait();
  first_atomic.join();

  process.map_pages(reinterpret_cast<uint64_t>(mapping.data), mapping.data, KfdProcess::kPageSize);
  memory.atomic_rmw(
      reinterpret_cast<uint64_t>(target), sizeof(*target),
      [](uint8_t *storage) {
        uint32_t value = 0;
        std::memcpy(&value, storage, sizeof(value));
        ++value;
        std::memcpy(storage, &value, sizeof(value));
      },
      kVmid);

  EXPECT_EQ(*target, 2u);
}

TEST(GpuMemoryThreadingTest, ReadersRemainSafeWhilePagesAreRemapped) {
  amdgpu::GpuMemory memory("memory");
  constexpr uint32_t kPid = 7;
  constexpr uint64_t kAddr = 0x40000000;
  constexpr uint32_t kValuePrefix = 0xa5a50000;
  constexpr uint32_t kReaders = 4;
  constexpr uint32_t kRemaps = 2000;

  KfdProcess process(kPid);
  memory.register_process(kPid, &process.page_table_, &process.page_table_mutex_,
                          process.page_table_generation());

  auto initial_page = std::make_unique<uint32_t>(kValuePrefix);
  process.map_pages(kAddr, initial_page.get(), sizeof(*initial_page));
  std::barrier start(kReaders + 1);
  std::barrier initial_read(kReaders + 1);
  std::atomic<bool> stop{false};
  std::atomic<uint64_t> mapped_reads{0};
  std::atomic<uint64_t> invalid_reads{0};
  std::vector<std::thread> readers;
  readers.reserve(kReaders);
  for (uint32_t i = 0; i < kReaders; ++i) {
    readers.emplace_back([&] {
      start.arrive_and_wait();
      const uint32_t initial_value = memory.read32(kAddr, kPid);
      if ((initial_value & 0xffff0000) == kValuePrefix)
        mapped_reads.fetch_add(1, std::memory_order_relaxed);
      else
        invalid_reads.fetch_add(1, std::memory_order_relaxed);
      initial_read.arrive_and_wait();
      while (!stop.load(std::memory_order_acquire)) {
        const uint32_t value = memory.read32(kAddr, kPid);
        if (value == 0)
          continue;
        if ((value & 0xffff0000) == kValuePrefix)
          mapped_reads.fetch_add(1, std::memory_order_relaxed);
        else
          invalid_reads.fetch_add(1, std::memory_order_relaxed);
      }
    });
  }

  start.arrive_and_wait();
  initial_read.arrive_and_wait();
  process.unmap_pages(kAddr, sizeof(*initial_page));
  initial_page.reset();
  for (uint32_t i = 0; i < kRemaps; ++i) {
    auto page = std::make_unique<uint32_t>(kValuePrefix | (i & 0xffff));
    process.map_pages(kAddr, page.get(), sizeof(*page));
    std::this_thread::yield();
    process.unmap_pages(kAddr, sizeof(*page));
    page.reset();
  }
  stop.store(true, std::memory_order_release);

  for (auto &reader : readers)
    reader.join();

  EXPECT_GT(mapped_reads.load(std::memory_order_relaxed), 0u);
  EXPECT_EQ(invalid_reads.load(std::memory_order_relaxed), 0u);
}

TEST(GpuMemoryTest, RegisteredVmidPassthroughMissRespectsUserSpaceLimit) {
  amdgpu::GpuMemory memory("memory");
  memory.set_passthrough(true);
  constexpr uint32_t kPid = 7;
  constexpr uint64_t kUserSpaceLimit = 0x800000000000ULL;

  KfdProcess process(kPid);
  memory.register_process(kPid, &process.page_table_, &process.page_table_mutex_,
                          process.page_table_generation());

  EXPECT_EQ(memory.resolve_host_ptr(kUserSpaceLimit + 0x123, kPid), nullptr);
  EXPECT_EQ(memory.resolve_host_ptr(kUserSpaceLimit + KfdProcess::kPageSize + 0x123, kPid),
            nullptr);
  EXPECT_EQ(memory.resolve_host_ptr(0x4000, kPid), reinterpret_cast<uint8_t *>(0x4000));
}

TEST(GpuMemoryTest, UnregisteredVmidPassthroughRespectsUserSpaceLimit) {
  amdgpu::GpuMemory memory("memory");
  memory.set_passthrough(true);
  constexpr uint64_t kUserSpaceLimit = 0x800000000000ULL;

  EXPECT_EQ(memory.resolve_host_ptr(kUserSpaceLimit), nullptr);
  EXPECT_EQ(memory.resolve_host_ptr(kUserSpaceLimit + 0x123), nullptr);
  EXPECT_EQ(memory.resolve_host_ptr(0x4000), reinterpret_cast<uint8_t *>(0x4000));
}

TEST(GpuMemoryTest, ZeroPassthroughAddressUsesFallbackStorage) {
  amdgpu::GpuMemory memory("memory");
  memory.set_passthrough(true);

  constexpr uint32_t kValue = 0xc001d00d;
  memory.write32(0, kValue);

  EXPECT_EQ(memory.read32(0), kValue);
}

TEST(GpuMemoryTest, NullBackedPteDoesNotInvokeMappedCallback) {
  amdgpu::GpuMemory memory("memory");
  memory.set_passthrough(true);
  constexpr uint32_t kPid = 7;
  constexpr uint64_t kAddr = 0x4000;
  constexpr uint32_t kValue = 0x12345678;

  KfdProcess process(kPid);
  process.page_table_[kAddr >> KfdProcess::kPageShift] = {};
  memory.register_process(kPid, &process.page_table_, &process.page_table_mutex_,
                          process.page_table_generation());

  ASSERT_EQ(memory.resolve_host_ptr(kAddr, kPid), nullptr);
  memory.write32(kAddr, kValue, kPid);
  EXPECT_EQ(memory.read32(kAddr, kPid), kValue);
}

TEST(GpuMemoryTest, RegisteredVmidBlockMissUsesClientMemory) {
  amdgpu::GpuMemory memory("memory");
  constexpr uint32_t kPid = 7;
  constexpr size_t kMappingSize = KfdProcess::kPageSize * 2;

  void *raw_mapping =
      mmap(nullptr, kMappingSize, PROT_READ | PROT_WRITE, MAP_PRIVATE | MAP_ANONYMOUS, -1, 0);
  ASSERT_NE(raw_mapping, MAP_FAILED);
  struct Mapping {
    uint8_t *data = nullptr;
    size_t size = 0;
    ~Mapping() {
      if (data)
        munmap(data, size);
    }
  } mapping{static_cast<uint8_t *>(raw_mapping), kMappingSize};

  for (size_t i = 0; i < kMappingSize; ++i)
    mapping.data[i] = static_cast<uint8_t>((i * 17 + 3) & 0xff);

  KfdProcess process(kPid);
  memory.register_process(kPid, &process.page_table_, &process.page_table_mutex_,
                          process.page_table_generation());
  memory.set_process_client_pid(kPid, getpid());

  constexpr size_t kAccessOffset = KfdProcess::kPageSize - 8;
  constexpr size_t kAccessSize = 32;
  const uint64_t addr = reinterpret_cast<uint64_t>(mapping.data + kAccessOffset);

  std::array<uint8_t, kAccessSize> actual{};
  memory.read_block(addr, std::span<uint8_t>(actual), kPid);
  EXPECT_TRUE(std::equal(actual.begin(), actual.end(), mapping.data + kAccessOffset));

  std::array<uint8_t, kAccessSize> replacement{};
  for (size_t i = 0; i < replacement.size(); ++i)
    replacement[i] = static_cast<uint8_t>(0xa0 + i);

  memory.write_block(addr, std::span<const uint8_t>(replacement), kPid);
  EXPECT_TRUE(std::equal(replacement.begin(), replacement.end(), mapping.data + kAccessOffset));
}

TEST(VmLifecycleTest, CreateAndDestroy) {
  std::string json = R"({"max_ticks":10000,"num_threads":1,
    "vm":{"arch":"cdna3"},
    "topology":{
      "root":{
        "name":"soc","type":"soc",
        "children":[
          {"name":"vram","type":"gpu_memory"},
          {"name":"xcd0","type":"xcd","children":[
            {"name":"l2","type":"l2_cache"},
            {"name":"cp","type":"command_processor"},
            {"name":"se0","type":"shader_engine","children":[
              {"name":"cu[0:3]","type":"compute_unit","config":[
                {"key":"num_wf_slots","value":"10"},
                {"key":"sgprs_per_wf","value":"104"},
                {"key":"vgprs_per_wf","value":"256"},
                {"key":"lds_size_kb","value":"64"}
              ]}
            ]},
            {"name":"se1","type":"shader_engine","children":[
              {"name":"cu[0:3]","type":"compute_unit","config":[
                {"key":"num_wf_slots","value":"10"},
                {"key":"sgprs_per_wf","value":"104"},
                {"key":"vgprs_per_wf","value":"256"},
                {"key":"lds_size_kb","value":"64"}
              ]}
            ]}
          ]}
        ]
      },
      "links":[
        {"src":"xcd0.cp.req_0","dst":"xcd0.se0.cu0.cpl","latency":1,"weight":2},
        {"src":"xcd0.cp.req_1","dst":"xcd0.se0.cu1.cpl","latency":1,"weight":2},
        {"src":"xcd0.cp.req_2","dst":"xcd0.se0.cu2.cpl","latency":1,"weight":2},
        {"src":"xcd0.cp.req_3","dst":"xcd0.se1.cu0.cpl","latency":1,"weight":2},
        {"src":"xcd0.cp.req_4","dst":"xcd0.se1.cu1.cpl","latency":1,"weight":2},
        {"src":"xcd0.cp.req_5","dst":"xcd0.se1.cu2.cpl","latency":1,"weight":2},
        {"src":"xcd0.se0.cu0.req","dst":"xcd0.l2.cpl_0","latency":1,"weight":10},
        {"src":"xcd0.se0.cu1.req","dst":"xcd0.l2.cpl_1","latency":1,"weight":10},
        {"src":"xcd0.se0.cu2.req","dst":"xcd0.l2.cpl_2","latency":1,"weight":10},
        {"src":"xcd0.se1.cu0.req","dst":"xcd0.l2.cpl_3","latency":1,"weight":10},
        {"src":"xcd0.se1.cu1.req","dst":"xcd0.l2.cpl_4","latency":1,"weight":10},
        {"src":"xcd0.se1.cu2.req","dst":"xcd0.l2.cpl_5","latency":1,"weight":10}
      ]
    }
  })";
  auto loaded = config::load_config_from_string(json, rocjitsu::kEmbeddedSchema);
  auto *soc = loaded.soc();

  auto *xcd = soc->xcd(0);
  EXPECT_EQ(xcd->num_shader_engines(), 2u);
  EXPECT_EQ(xcd->shader_engine(0)->num_compute_units(), 3u);
  EXPECT_EQ(xcd->shader_engine(1)->num_compute_units(), 3u);
}

TEST(VmLifecycleTest, MissingArchFails) {
  const char *json = R"({"vm":{"gpu":{"num_shader_engines":1}}})";
  rj_vm_t *handle = nullptr;
  EXPECT_NE(rj_vm_create_from_string(json, RJ_VM_MODE_DEFAULT, &handle), ROCJITSU_STATUS_SUCCESS);
}

TEST(VmLifecycleTest, InvalidArchFails) {
  const char *json = R"({"vm":{"arch":"bogus"}})";
  rj_vm_t *handle = nullptr;
  EXPECT_NE(rj_vm_create_from_string(json, RJ_VM_MODE_DEFAULT, &handle), ROCJITSU_STATUS_SUCCESS);
}

class IsaTest : public ::testing::TestWithParam<std::string> {
protected:
  std::string arch() const { return GetParam(); }
};

TEST_P(IsaTest, RegisterAccess) {
  VmFixture f(arch());

  // Exercise the CU register-file read/write API on a live, resident wavefront.
  // dispatch_wf() allocates a slot without running the kernel to s_endpgm (which
  // would free the slot), so the registers stay readable/writable here.
  auto *cu = f.cu();
  auto *w = cu->dispatch_wf(/*wg_id=*/0, /*pc=*/0x1040, /*num_sgprs=*/104, /*num_vgprs=*/256);
  ASSERT_NE(w, nullptr);
  ASSERT_GE(cu->num_wfs(), 1u);

  EXPECT_EQ(w->wf_size(), 64u);
  EXPECT_EQ(w->num_sgprs(), 104u);
  EXPECT_EQ(w->num_vgprs(), 256u);

  uint32_t sb = w->sgpr_alloc().base;
  uint32_t vb = w->vgpr_alloc().base;
  cu->write_sgpr(sb + 2, 42);
  cu->write_sgpr(sb + 103, 0xFFFFFFFF);
  EXPECT_EQ(cu->read_sgpr(sb + 2), 42u);
  EXPECT_EQ(cu->read_sgpr(sb + 103), 0xFFFFFFFF);
  EXPECT_EQ(cu->read_sgpr(sb + 50), 0u);

  cu->write_vgpr(vb + 1, 0, 100);
  cu->write_vgpr(vb + 1, 63, 200);
  EXPECT_EQ(cu->read_vgpr(vb + 1, 0), 100u);
  EXPECT_EQ(cu->read_vgpr(vb + 1, 63), 200u);
  EXPECT_EQ(cu->read_vgpr(vb + 1, 1), 0u);
}

TEST(RdnaDispatchTest, PackedTidHonorsRequestedComponents) {
  const uint32_t code[] = {SOPP_S_ENDPGM};

  for (const std::string &arch :
       {std::string("rdna3"), std::string("rdna3_5"), std::string("rdna4")}) {
    for (uint32_t component_count = 0; component_count <= 2; ++component_count) {
      SCOPED_TRACE(arch + " component_count=" + std::to_string(component_count));
      VmFixture f(arch, 1, 10, /*lds_size_kb=*/64, /*sgprs_per_wf=*/128);
      auto *snap = f.capture_halts();
      uint64_t ko =
          f.write_kernel(0x1000, code, sizeof(code), 104, 32, 2, 0, false, component_count);

      test::AqlQueue queue(f.mem(), f.cp());
      hsa_kernel_dispatch_packet_t pkt{};
      pkt.header = HSA_PACKET_TYPE_KERNEL_DISPATCH;
      pkt.setup = 3;
      pkt.workgroup_size_x = 8;
      pkt.workgroup_size_y = 4;
      pkt.workgroup_size_z = 2;
      pkt.grid_size_x = 8;
      pkt.grid_size_y = 4;
      pkt.grid_size_z = 2;
      pkt.kernel_object = ko;
      queue.submit(pkt);
      step_until_halted(*f.engine, {f.cu()});

      ASSERT_EQ(snap->snapshots().size(), 2u);
      const auto *wf0 = snap->by_wf_id(0);
      const auto *wf1 = snap->by_wf_id(1);
      ASSERT_NE(wf0, nullptr);
      ASSERT_NE(wf1, nullptr);
      EXPECT_EQ(wf0->vgpr(0, 0), amdgpu::pack_workitem_id({0, 0, 0}, component_count));
      EXPECT_EQ(wf0->vgpr(0, 9), amdgpu::pack_workitem_id({1, 1, 0}, component_count));
      EXPECT_EQ(wf1->vgpr(0, 8), amdgpu::pack_workitem_id({0, 1, 1}, component_count));
      EXPECT_EQ(wf1->vgpr(0, 31), amdgpu::pack_workitem_id({7, 3, 1}, component_count));
    }
  }
}

TEST(CdnaDispatchTest, Wave64PackedTidHonorsRequestedComponents) {
  const uint32_t code[] = {SOPP_S_ENDPGM};

  for (std::string_view arch : {"cdna2", "cdna3", "cdna4"}) {
    for (uint32_t component_count = 0; component_count <= 2; ++component_count) {
      SCOPED_TRACE(::testing::Message() << arch << " component_count=" << component_count);
      VmFixture f(arch);
      auto *snap = f.capture_halts();
      ASSERT_TRUE(f.cp()->packed_tid());
      uint64_t ko =
          f.write_kernel(0x1000, code, sizeof(code), 104, 256, 2, 0, false, component_count);

      test::AqlQueue queue(f.mem(), f.cp());
      hsa_kernel_dispatch_packet_t pkt{};
      pkt.header = HSA_PACKET_TYPE_KERNEL_DISPATCH;
      pkt.setup = 3;
      pkt.workgroup_size_x = 8;
      pkt.workgroup_size_y = 4;
      pkt.workgroup_size_z = 2;
      pkt.grid_size_x = 8;
      pkt.grid_size_y = 4;
      pkt.grid_size_z = 2;
      pkt.kernel_object = ko;
      queue.submit(pkt);
      step_until_halted(*f.engine, {f.cu()});

      ASSERT_EQ(snap->snapshots().size(), 1u);
      const auto *wf = &snap->snapshots().front();
      EXPECT_EQ(wf->wf_size, 64u);
      EXPECT_EQ(wf->vgpr(0, 40), amdgpu::pack_workitem_id({0, 1, 1}, component_count));
      EXPECT_EQ(wf->vgpr(0, 63), amdgpu::pack_workitem_id({7, 3, 1}, component_count));
    }
  }
}

TEST(CdnaDispatchTest, Wave64MasksMultidimensionalGridTailAcrossLane32) {
  VmFixture f("cdna3");
  auto *snap = f.capture_halts();
  const uint32_t code[] = {SOPP_S_ENDPGM};
  uint64_t ko = f.write_kernel(0x1000, code, sizeof(code));

  test::AqlQueue queue(f.mem(), f.cp());
  hsa_kernel_dispatch_packet_t pkt{};
  pkt.header = HSA_PACKET_TYPE_KERNEL_DISPATCH;
  pkt.setup = 2;
  pkt.workgroup_size_x = 8;
  pkt.workgroup_size_y = 6;
  pkt.workgroup_size_z = 1;
  pkt.grid_size_x = 5;
  pkt.grid_size_y = 5;
  pkt.grid_size_z = 1;
  pkt.kernel_object = ko;
  queue.submit(pkt);
  step_until_halted(*f.engine, {f.cu()});

  ASSERT_EQ(snap->snapshots().size(), 1u);
  const auto *wf = &snap->snapshots().front();
  EXPECT_EQ(wf->wf_size, 64u);
  // Five active X lanes in each of five Y rows. The final active row starts
  // at lane 32; lanes 40-47 are past grid_size_y, and lanes 48-63 map beyond
  // workgroup_size_z.
  EXPECT_EQ(wf->exec, 0x0000'001F'1F1F'1F1FULL);
}

TEST(CdnaDispatchTest, Cdna1UnpackedTidHonorsRequestedComponents) {
  const uint32_t code[] = {SOPP_S_ENDPGM};

  for (uint32_t component_count = 0; component_count <= 2; ++component_count) {
    SCOPED_TRACE("component_count=" + std::to_string(component_count));
    VmFixture f("cdna1");
    auto *snap = f.capture_halts();
    ASSERT_FALSE(f.cp()->packed_tid());
    uint64_t ko =
        f.write_kernel(0x1000, code, sizeof(code), 104, 256, 2, 0, false, component_count);

    test::AqlQueue queue(f.mem(), f.cp());
    hsa_kernel_dispatch_packet_t pkt{};
    pkt.header = HSA_PACKET_TYPE_KERNEL_DISPATCH;
    pkt.setup = 3;
    pkt.workgroup_size_x = 8;
    pkt.workgroup_size_y = 4;
    pkt.workgroup_size_z = 2;
    pkt.grid_size_x = 8;
    pkt.grid_size_y = 4;
    pkt.grid_size_z = 2;
    pkt.kernel_object = ko;
    queue.submit(pkt);
    step_until_halted(*f.engine, {f.cu()});

    ASSERT_EQ(snap->snapshots().size(), 1u);
    const auto *wf = &snap->snapshots().front();
    EXPECT_EQ(wf->wf_size, 64u);
    EXPECT_EQ(wf->vgpr(0, 63), 7u);
    EXPECT_EQ(wf->vgpr(1, 63), component_count >= 1 ? 3u : 0u);
    EXPECT_EQ(wf->vgpr(2, 63), component_count >= 2 ? 1u : 0u);
  }
}

TEST(DispatchEntryTest, InitialExecMaskSupportsWave64GridTail) {
  amdgpu::DispatchEntry entry{};
  entry.grid_size_x = 65;
  entry.grid_wgs_x = 2;
  entry.workgroup_size_x = 64;

  EXPECT_EQ(amdgpu::initial_exec_mask_for_wave(entry, 0, 0, 64), ~0ULL);
  EXPECT_EQ(amdgpu::initial_exec_mask_for_wave(entry, 1, 0, 64), 1ULL);
}

TEST(DispatchEntryTest, InitialExecMaskHandles3DTailWithWorkgroupOffset) {
  amdgpu::DispatchEntry entry{};
  entry.workgroup_id_offset = 100;
  entry.grid_size_x = 4;
  entry.grid_size_y = 3;
  entry.grid_size_z = 3;
  entry.grid_wgs_x = 1;
  entry.grid_wgs_y = 2;
  entry.grid_wgs_z = 2;
  entry.workgroup_size_x = 4;
  entry.workgroup_size_y = 2;
  entry.workgroup_size_z = 2;

  EXPECT_EQ(amdgpu::initial_exec_mask_for_wave(entry, 103, 0, 64), 0xFULL);
}

TEST_P(IsaTest, RegisterFileIsolation) {
  VmFixture f(arch(), 1, 2);

  // Two concurrently-resident wavefronts must own disjoint register blocks.
  // dispatch_wf() places both without running to s_endpgm (which would free them),
  // so their register files stay live for the isolation checks below.
  auto *cu = f.cu();
  auto *w0 = cu->dispatch_wf(/*wg_id=*/0, /*pc=*/0x1040, /*num_sgprs=*/104, /*num_vgprs=*/256);
  auto *w1 = cu->dispatch_wf(/*wg_id=*/0, /*pc=*/0x1040, /*num_sgprs=*/104, /*num_vgprs=*/256);
  ASSERT_NE(w0, nullptr);
  ASSERT_NE(w1, nullptr);
  ASSERT_EQ(cu->num_wfs(), 2u);
  ASSERT_NE(w0->sgpr_alloc().base, w1->sgpr_alloc().base);
  ASSERT_NE(w0->vgpr_alloc().base, w1->vgpr_alloc().base);

  cu->write_sgpr(w0->sgpr_alloc().base + 0, 42);
  cu->write_sgpr(w1->sgpr_alloc().base + 0, 99);
  EXPECT_EQ(cu->read_sgpr(w0->sgpr_alloc().base + 0), 42u);
  EXPECT_EQ(cu->read_sgpr(w1->sgpr_alloc().base + 0), 99u);

  cu->write_vgpr(w0->vgpr_alloc().base + 0, 0, 100);
  cu->write_vgpr(w1->vgpr_alloc().base + 0, 0, 200);
  EXPECT_EQ(cu->read_vgpr(w0->vgpr_alloc().base + 0, 0), 100u);
  EXPECT_EQ(cu->read_vgpr(w1->vgpr_alloc().base + 0, 0), 200u);
}

TEST_P(IsaTest, DispatchAndCapacity) {
  VmFixture f(arch(), 1, 2);

  // 2 workgroups of 64 (= 2 wavefronts), CU has 2 slots — fills exactly.
  const uint32_t code[] = {SOPP_S_NOP, SOPP_S_ENDPGM};
  uint64_t ko = f.write_kernel(0x1000, code, sizeof(code));
  test::AqlQueue queue(f.mem(), f.cp());
  queue.dispatch(ko, 128, 64); // grid=128, wg=64 → 2 workgroups
  f.engine->run();

  // Both workgroups ran to completion; their waves freed themselves at s_endpgm.
  EXPECT_EQ(f.cp()->dispatched_count(), 1u);
}

TEST_P(IsaTest, DispatchWfReturnsNullWhenSlotsExhausted) {
  // dispatch_wf() promises nullptr (not an out-of-bounds slot) when the CU is full.
  // The CP relies on can_accept_workgroup() gating, but the API contract must hold
  // independently so it cannot silently drift into OOB access in release builds.
  constexpr uint32_t kSlots = 4;
  VmFixture f(arch(), 1, kSlots);
  auto *cu = f.cu();
  for (uint32_t i = 0; i < kSlots; ++i)
    ASSERT_NE(cu->dispatch_wf(0, 0x1040, 104, 256), nullptr) << "slot " << i;
  EXPECT_EQ(cu->num_wfs(), kSlots);
  // All slots are occupied by resident (non-halted) waves — the next dispatch fails.
  EXPECT_EQ(cu->dispatch_wf(0, 0x1040, 104, 256), nullptr);
}

TEST_P(IsaTest, VendorSpecificExtKernelDispatch) {
  VmFixture f(arch(), 1, 8);
  auto *snap = f.capture_halts();

  const uint32_t code[] = {SOPP_S_NOP, SOPP_S_ENDPGM};
  uint64_t ko = f.write_kernel(0x1000, code, sizeof(code));

  amdgpu::AmdExtKernelDispatchPacket ext{};
  ext.header = HSA_PACKET_TYPE_VENDOR_SPECIFIC;
  ext.amd_format = amdgpu::kHsaAmdPacketTypeExtKernelDispatch;
  ext.setup = 1;
  ext.workgroup_size_x = 64;
  ext.workgroup_size_y = 1;
  ext.workgroup_size_z = 1;
  ext.cluster_count_x = 2;
  ext.cluster_count_y = 1;
  ext.cluster_count_z = 1;
  ext.cluster_size_x = 2;
  ext.cluster_size_y = 1;
  ext.cluster_size_z = 1;
  ext.kernel_object = ko;

  hsa_kernel_dispatch_packet_t raw{};
  std::memcpy(&raw, &ext, sizeof(ext));
  test::AqlQueue queue(f.mem(), f.cp());
  queue.submit(raw);
  f.engine->run();

  // The clustered dispatch produced one consolidated dispatch, and its cluster
  // wavefronts ran and halted (freeing themselves at s_endpgm). Observe the waves
  // at halt rather than as post-run residents.
  EXPECT_EQ(f.cp()->dispatched_count(), 1u);
  EXPECT_GE(snap->snapshots().size(), 1u);
}

TEST_P(IsaTest, VendorSpecificExtKernelDispatchReadsDependencySignalFromGpuMemory) {
  VmFixture f(arch(), 1, 8);

  const uint32_t code[] = {SOPP_S_NOP, SOPP_S_ENDPGM};
  uint64_t ko = f.write_kernel(0x1000, code, sizeof(code));
  constexpr uint64_t kDepSignal = 0x7000;
  constexpr uint32_t kSignalValueOffset = 8;
  f.mem()->write64(kDepSignal + kSignalValueOffset, 1);

  amdgpu::AmdExtKernelDispatchPacket ext{};
  ext.header = HSA_PACKET_TYPE_VENDOR_SPECIFIC;
  ext.amd_format = amdgpu::kHsaAmdPacketTypeExtKernelDispatch;
  ext.setup = 1;
  ext.workgroup_size_x = 64;
  ext.workgroup_size_y = 1;
  ext.workgroup_size_z = 1;
  ext.cluster_count_x = 1;
  ext.cluster_count_y = 1;
  ext.cluster_count_z = 1;
  ext.cluster_size_x = 1;
  ext.cluster_size_y = 1;
  ext.cluster_size_z = 1;
  ext.dep_signal.handle = kDepSignal;
  ext.kernel_object = ko;

  test::AqlQueue queue(f.mem(), f.cp());
  queue.submit(ext);
  (void)f.engine->step();
  EXPECT_EQ(f.cp()->dispatched_count(), 0u);

  f.mem()->write64(kDepSignal + kSignalValueOffset, 0);
  f.engine->run();
  EXPECT_EQ(f.cp()->dispatched_count(), 1u);
}

TEST_P(IsaTest, VendorSpecificBarrierValueConditionsWaitAndResume) {
  struct ConditionCase {
    uint32_t condition;
    int64_t initial_signal_value;
    int64_t ready_signal_value;
    int64_t value;
    int64_t mask;
  };
  constexpr std::array cases{
      ConditionCase{HSA_SIGNAL_CONDITION_EQ, 0x106, 0x105, 0x5, 0xff},
      ConditionCase{HSA_SIGNAL_CONDITION_EQ, 0, std::numeric_limits<int64_t>::min(),
                    std::numeric_limits<int64_t>::min(), std::numeric_limits<int64_t>::min()},
      ConditionCase{HSA_SIGNAL_CONDITION_NE, 4, 5, 4, std::numeric_limits<int64_t>::max()},
      ConditionCase{HSA_SIGNAL_CONDITION_LT, 1, 0, 1, std::numeric_limits<int64_t>::max()},
      ConditionCase{HSA_SIGNAL_CONDITION_GTE, 4, 5, 5, std::numeric_limits<int64_t>::max()},
  };

  for (const auto &test_case : cases) {
    SCOPED_TRACE(test_case.condition);
    VmFixture f(arch(), 1, 8);

    constexpr uint64_t kDepSignal = 0x7000;
    constexpr uint64_t kCompletionSignal = 0x7100;
    constexpr uint32_t kSignalValueOffset = 8;
    f.mem()->write64(kDepSignal + kSignalValueOffset, test_case.initial_signal_value);
    f.mem()->write64(kCompletionSignal + kSignalValueOffset, 1);

    amdgpu::AmdBarrierValuePacket barrier{};
    barrier.header = HSA_PACKET_TYPE_VENDOR_SPECIFIC | (1 << HSA_PACKET_HEADER_BARRIER);
    barrier.amd_format = amdgpu::kHsaAmdPacketTypeBarrierValue;
    barrier.signal.handle = kDepSignal;
    barrier.value = test_case.value;
    barrier.mask = test_case.mask;
    barrier.condition = test_case.condition;
    barrier.completion_signal.handle = kCompletionSignal;

    test::AqlQueue queue(f.mem(), f.cp());
    queue.submit(barrier);
    (void)f.engine->step();

    EXPECT_EQ(f.mem()->read64(test::AqlQueue::DEFAULT_READ_PTR_ADDR), 0u);
    EXPECT_EQ(f.mem()->read64(kCompletionSignal + kSignalValueOffset), 1u);

    f.mem()->write64(kDepSignal + kSignalValueOffset, test_case.ready_signal_value);
    f.engine->run();

    EXPECT_EQ(f.mem()->read64(test::AqlQueue::DEFAULT_READ_PTR_ADDR), 1u);
    EXPECT_EQ(f.mem()->read64(kCompletionSignal + kSignalValueOffset), 0u);
  }
}

TEST_P(IsaTest, VendorSpecificBarrierValueAllowsNullSignals) {
  VmFixture f(arch(), 1, 8);

  constexpr uint64_t kCompletionSignal = 0x7100;
  constexpr uint32_t kSignalValueOffset = 8;
  f.mem()->write64(kCompletionSignal + kSignalValueOffset, 1);

  amdgpu::AmdBarrierValuePacket barrier{};
  barrier.header = HSA_PACKET_TYPE_VENDOR_SPECIFIC | (1 << HSA_PACKET_HEADER_BARRIER);
  barrier.amd_format = amdgpu::kHsaAmdPacketTypeBarrierValue;
  barrier.condition = HSA_SIGNAL_CONDITION_EQ;
  barrier.completion_signal.handle = kCompletionSignal;

  test::AqlQueue queue(f.mem(), f.cp());
  queue.submit(barrier);
  barrier.completion_signal.handle = 0;
  queue.submit(barrier);
  f.engine->run();

  EXPECT_EQ(f.mem()->read64(test::AqlQueue::DEFAULT_READ_PTR_ADDR), 2u);
  EXPECT_EQ(f.mem()->read64(kCompletionSignal + kSignalValueOffset), 0u);
}

TEST_P(IsaTest, VendorSpecificBarrierValueRejectsInvalidCondition) {
  VmFixture f(arch(), 1, 8);

  constexpr uint64_t kDepSignal = 0x7000;
  constexpr uint32_t kSignalValueOffset = 8;
  f.mem()->write64(kDepSignal + kSignalValueOffset, 1);

  amdgpu::AmdBarrierValuePacket barrier{};
  barrier.header = HSA_PACKET_TYPE_VENDOR_SPECIFIC | (1 << HSA_PACKET_HEADER_BARRIER);
  barrier.amd_format = amdgpu::kHsaAmdPacketTypeBarrierValue;
  barrier.signal.handle = kDepSignal;
  barrier.mask = std::numeric_limits<int64_t>::max();
  barrier.condition = 99;

  test::AqlQueue queue(f.mem(), f.cp());
  queue.submit(barrier);

  EXPECT_THROW((void)f.engine->step(), std::runtime_error);
  EXPECT_EQ(f.mem()->read64(test::AqlQueue::DEFAULT_READ_PTR_ADDR), 0u);
}

TEST_P(IsaTest, VendorSpecificBarrierValueOrdersQueueEntries) {
  VmFixture f(arch(), 1, 8);

  const uint32_t code[] = {SOPP_S_NOP, SOPP_S_ENDPGM};
  uint64_t ko = f.write_kernel(0x1000, code, sizeof(code));
  constexpr uint64_t kBarrierCompletionSignal = 0x7000;
  constexpr uint64_t kLaterCompletionSignal = 0x7100;
  constexpr uint32_t kSignalValueOffset = 8;
  f.mem()->write64(kBarrierCompletionSignal + kSignalValueOffset, 1);
  f.mem()->write64(kLaterCompletionSignal + kSignalValueOffset, 1);

  hsa_kernel_dispatch_packet_t dispatch{};
  dispatch.header = HSA_PACKET_TYPE_KERNEL_DISPATCH;
  dispatch.setup = 1;
  dispatch.workgroup_size_x = 64;
  dispatch.workgroup_size_y = 1;
  dispatch.workgroup_size_z = 1;
  dispatch.grid_size_x = 64;
  dispatch.grid_size_y = 1;
  dispatch.grid_size_z = 1;
  dispatch.kernel_object = ko;

  amdgpu::AmdBarrierValuePacket barrier{};
  barrier.header = HSA_PACKET_TYPE_VENDOR_SPECIFIC | (1 << HSA_PACKET_HEADER_BARRIER);
  barrier.amd_format = amdgpu::kHsaAmdPacketTypeBarrierValue;
  barrier.completion_signal.handle = kBarrierCompletionSignal;

  test::AqlQueue queue(f.mem(), f.cp());
  queue.submit(dispatch);
  queue.submit(barrier);
  dispatch.completion_signal.handle = kLaterCompletionSignal;
  queue.submit(dispatch);
  (void)f.engine->step();

  EXPECT_EQ(f.cu()->num_wfs(), 1u);
  EXPECT_EQ(f.mem()->read64(kBarrierCompletionSignal + kSignalValueOffset), 1u);
  EXPECT_EQ(f.mem()->read64(kLaterCompletionSignal + kSignalValueOffset), 1u);

  f.engine->run();

  EXPECT_EQ(f.mem()->read64(kBarrierCompletionSignal + kSignalValueOffset), 0u);
  EXPECT_EQ(f.mem()->read64(kLaterCompletionSignal + kSignalValueOffset), 0u);
}

TEST_P(IsaTest, NonKernelBarrierPacketsOrderQueueEntries) {
  constexpr std::array packet_types{
      HSA_PACKET_TYPE_BARRIER_AND,
      HSA_PACKET_TYPE_BARRIER_OR,
      HSA_PACKET_TYPE_VENDOR_SPECIFIC,
  };

  for (const auto packet_type : packet_types) {
    SCOPED_TRACE(packet_type);
    VmFixture f(arch(), 1, 8);

    const uint32_t code[] = {SOPP_S_NOP, SOPP_S_ENDPGM};
    uint64_t ko = f.write_kernel(0x1000, code, sizeof(code));
    constexpr uint64_t kBarrierCompletionSignal = 0x7000;
    constexpr uint64_t kLaterCompletionSignal = 0x7100;
    constexpr uint32_t kSignalValueOffset = 8;
    f.mem()->write64(kBarrierCompletionSignal + kSignalValueOffset, 1);
    f.mem()->write64(kLaterCompletionSignal + kSignalValueOffset, 1);

    hsa_kernel_dispatch_packet_t dispatch{};
    dispatch.header = HSA_PACKET_TYPE_KERNEL_DISPATCH;
    dispatch.setup = 1;
    dispatch.workgroup_size_x = 64;
    dispatch.workgroup_size_y = 1;
    dispatch.workgroup_size_z = 1;
    dispatch.grid_size_x = 64;
    dispatch.grid_size_y = 1;
    dispatch.grid_size_z = 1;
    dispatch.kernel_object = ko;

    hsa_kernel_dispatch_packet_t barrier{};
    barrier.header = packet_type | (1 << HSA_PACKET_HEADER_BARRIER);
    if (packet_type == HSA_PACKET_TYPE_VENDOR_SPECIFIC)
      barrier.setup = amdgpu::kAmdAqlFormatPm4Ib;
    barrier.completion_signal.handle = kBarrierCompletionSignal;

    test::AqlQueue queue(f.mem(), f.cp());
    queue.submit(dispatch);
    queue.submit(barrier);
    dispatch.completion_signal.handle = kLaterCompletionSignal;
    queue.submit(dispatch);
    (void)f.engine->step();

    EXPECT_EQ(f.cu()->num_wfs(), 1u);
    EXPECT_EQ(f.mem()->read64(kBarrierCompletionSignal + kSignalValueOffset), 1u);
    EXPECT_EQ(f.mem()->read64(kLaterCompletionSignal + kSignalValueOffset), 1u);

    f.engine->run();

    EXPECT_EQ(f.mem()->read64(kBarrierCompletionSignal + kSignalValueOffset), 0u);
    EXPECT_EQ(f.mem()->read64(kLaterCompletionSignal + kSignalValueOffset), 0u);
  }
}

TEST_P(IsaTest, VendorSpecificRejectsUnsupportedFormats) {
  constexpr std::array<uint8_t, 2> unsupported_formats{0, amdgpu::kHsaAmdPacketTypeReserved200};

  for (const auto amd_format : unsupported_formats) {
    SCOPED_TRACE(static_cast<unsigned>(amd_format));
    VmFixture f(arch(), 1, 8);

    amdgpu::AmdExtKernelDispatchPacket packet{};
    packet.header = HSA_PACKET_TYPE_VENDOR_SPECIFIC;
    packet.amd_format = amd_format;

    test::AqlQueue queue(f.mem(), f.cp());
    queue.submit(packet);

    EXPECT_THROW((void)f.engine->step(), std::runtime_error);
    EXPECT_EQ(f.mem()->read64(test::AqlQueue::DEFAULT_READ_PTR_ADDR), 0u);
  }
}

TEST(ClusterDispatchTest, RejectsClusterThatCannotFitWithoutSpinning) {
  VmFixture f("gfx1250", 1, 1);

  const uint32_t code[] = {0xBFB00000u}; // s_endpgm
  uint64_t ko = f.write_kernel(0x1000, code, sizeof(code));
  test::AqlQueue queue(f.mem(), f.cp());
  queue.dispatch_clustered(ko, /*cluster_count_x=*/1, /*cluster_size_x=*/2,
                           /*workgroup_size_x=*/32);

  EXPECT_THROW((void)f.engine->step(), std::runtime_error);
  EXPECT_FALSE(f.cu()->has_active_wfs());
  EXPECT_TRUE(f.cp()
                  ->cluster_lds_targets(/*dispatch_id=*/1, /*wg_id=*/0,
                                        /*mcast_mask=*/0x3)
                  .empty());
}

TEST(ClusterDispatchTest, AccountsForPerWorkgroupLdsAlignmentWhenPlanningCluster) {
  VmFixture f("gfx1250", 1, 3, /*lds_size_kb=*/1);

  const uint32_t code[] = {0xBFB00000u}; // s_endpgm
  uint64_t ko = f.write_kernel(0x1000, code, sizeof(code));
  test::AqlQueue queue(f.mem(), f.cp());
  queue.dispatch_clustered(ko, /*cluster_count_x=*/1, /*cluster_size_x=*/3,
                           /*workgroup_size_x=*/32, /*kernarg_addr=*/0,
                           /*group_segment_size=*/257);

  EXPECT_THROW((void)f.engine->step(), std::runtime_error);
  EXPECT_FALSE(f.cu()->has_active_wfs());
}

TEST(ClusterDispatchTest, ReclaimsLdsBetweenClusterWaves) {
  VmFixture f("cdna3", 2, 1, /*lds_size_kb=*/1);

  const uint32_t code[] = {SOPP_S_ENDPGM};
  uint64_t ko = f.write_kernel(0x1000, code, sizeof(code));
  constexpr uint64_t kSignal = 0x7000;
  constexpr uint32_t kSignalValueOffset = 8;
  f.mem()->write64(kSignal + kSignalValueOffset, 1);

  amdgpu::AmdExtKernelDispatchPacket ext{};
  ext.header = HSA_PACKET_TYPE_VENDOR_SPECIFIC;
  ext.amd_format = amdgpu::kHsaAmdPacketTypeExtKernelDispatch;
  ext.setup = 1;
  ext.workgroup_size_x = 32;
  ext.workgroup_size_y = 1;
  ext.workgroup_size_z = 1;
  ext.cluster_count_x = 2;
  ext.cluster_count_y = 1;
  ext.cluster_count_z = 1;
  ext.cluster_size_x = 2;
  ext.cluster_size_y = 1;
  ext.cluster_size_z = 1;
  ext.group_segment_size = 769;
  ext.kernel_object = ko;
  ext.completion_signal.handle = kSignal;

  test::AqlQueue queue(f.mem(), f.cp());
  queue.submit(ext);

  EXPECT_NO_THROW(f.engine->run());
  EXPECT_EQ(f.mem()->read64(kSignal + kSignalValueOffset), 0u);
  EXPECT_FALSE(f.cu(0)->has_active_wfs());
  EXPECT_FALSE(f.cu(1)->has_active_wfs());
}

TEST(ClusterDispatchTest, Rdna4ExtendedDispatchKeepsOrdinaryTtmpWorkgroupIds) {
  VmFixture f("rdna4", 1, 8, /*lds_size_kb=*/64, /*sgprs_per_wf=*/128);
  auto *snap = f.capture_halts();

  const uint32_t code[] = {SOPP_S_NOP, SOPP_S_ENDPGM};
  uint64_t ko = f.write_kernel(0x1000, code, sizeof(code), /*sgprs=*/128);

  amdgpu::AmdExtKernelDispatchPacket ext{};
  ext.header = HSA_PACKET_TYPE_VENDOR_SPECIFIC;
  ext.amd_format = amdgpu::kHsaAmdPacketTypeExtKernelDispatch;
  ext.setup = 3;
  ext.workgroup_size_x = 32;
  ext.workgroup_size_y = 1;
  ext.workgroup_size_z = 1;
  ext.cluster_count_x = 1;
  ext.cluster_count_y = 1;
  ext.cluster_count_z = 1;
  ext.cluster_size_x = 2;
  ext.cluster_size_y = 2;
  ext.cluster_size_z = 2;
  ext.kernel_object = ko;

  test::AqlQueue queue(f.mem(), f.cp());
  queue.submit(ext);
  step_until_halted(*f.engine, {f.cu()});

  ASSERT_EQ(snap->snapshots().size(), 8u);
  std::array<bool, 8> seen{};
  for (const auto &wf : snap->snapshots()) {
    const uint32_t workgroup_id = wf.wg_id;
    ASSERT_LT(workgroup_id, seen.size());
    seen[workgroup_id] = true;

    const uint32_t workgroup_x = workgroup_id % 2;
    const uint32_t workgroup_y = (workgroup_id / 2) % 2;
    const uint32_t workgroup_z = workgroup_id / 4;
    EXPECT_EQ(wf.sgpr(114), 0u);
    EXPECT_EQ(wf.sgpr(115), (workgroup_z << 16) | workgroup_y);
    EXPECT_EQ(wf.sgpr(117), workgroup_x);
  }
  for (bool was_seen : seen)
    EXPECT_TRUE(was_seen);
}

TEST(ClusterDispatchTest, RejectsExtKernelDispatchWithZeroClusterShape) {
  VmFixture f("cdna3", 1, 8);

  const uint32_t code[] = {SOPP_S_ENDPGM};
  uint64_t ko = f.write_kernel(0x1000, code, sizeof(code));

  amdgpu::AmdExtKernelDispatchPacket ext{};
  ext.header = HSA_PACKET_TYPE_VENDOR_SPECIFIC;
  ext.amd_format = amdgpu::kHsaAmdPacketTypeExtKernelDispatch;
  ext.setup = 1;
  ext.workgroup_size_x = 32;
  ext.workgroup_size_y = 1;
  ext.workgroup_size_z = 1;
  ext.cluster_count_x = 0;
  ext.cluster_count_y = 1;
  ext.cluster_count_z = 1;
  ext.cluster_size_x = 2;
  ext.cluster_size_y = 1;
  ext.cluster_size_z = 1;
  ext.kernel_object = ko;

  test::AqlQueue queue(f.mem(), f.cp());
  queue.submit(ext);

  EXPECT_THROW((void)f.engine->step(), std::runtime_error);
}

TEST(ClusterDispatchTest, RejectsExtKernelDispatchGridOverflow) {
  VmFixture f("cdna3", 1, 8);

  const uint32_t code[] = {SOPP_S_ENDPGM};
  uint64_t ko = f.write_kernel(0x1000, code, sizeof(code));

  amdgpu::AmdExtKernelDispatchPacket ext{};
  ext.header = HSA_PACKET_TYPE_VENDOR_SPECIFIC;
  ext.amd_format = amdgpu::kHsaAmdPacketTypeExtKernelDispatch;
  ext.setup = 1;
  ext.workgroup_size_x = 64;
  ext.workgroup_size_y = 1;
  ext.workgroup_size_z = 1;
  ext.cluster_count_x = std::numeric_limits<uint32_t>::max();
  ext.cluster_count_y = 1;
  ext.cluster_count_z = 1;
  ext.cluster_size_x = 2;
  ext.cluster_size_y = 1;
  ext.cluster_size_z = 1;
  ext.kernel_object = ko;

  test::AqlQueue queue(f.mem(), f.cp());
  queue.submit(ext);

  EXPECT_THROW((void)f.engine->step(), std::runtime_error);
}

TEST_P(IsaTest, DispatchCreatesWavefronts) {
  VmFixture f(arch(), 2);
  auto *snap = f.capture_halts();

  const uint32_t code[] = {SOPP_S_NOP, SOPP_S_ENDPGM};
  uint64_t ko = f.write_kernel(0x1000, code, sizeof(code));
  test::AqlQueue queue(f.mem(), f.cp());
  queue.dispatch(ko, 128); // 2 workgroups of 64
  step_until_halted(*f.engine, {f.se()->compute_unit(0), f.se()->compute_unit(1)});

  // Round-robin placement puts one workgroup (one wave) on each CU.
  EXPECT_EQ(snap->for_cu(f.se()->compute_unit(0)).size(), 1u);
  EXPECT_EQ(snap->for_cu(f.se()->compute_unit(1)).size(), 1u);
}

TEST_P(IsaTest, MultipleWavesPerWorkgroup) {
  VmFixture f(arch());
  auto *snap = f.capture_halts();

  const uint32_t code[] = {SOPP_S_NOP, SOPP_S_ENDPGM};
  uint64_t ko = f.write_kernel(0x1000, code, sizeof(code));
  {
    test::AqlQueue queue(f.mem(), f.cp());
    hsa_kernel_dispatch_packet_t pkt{};
    pkt.header = HSA_PACKET_TYPE_KERNEL_DISPATCH;
    pkt.setup = 1;
    pkt.workgroup_size_x = 192; // 3 wavefronts per workgroup
    pkt.workgroup_size_y = 1;
    pkt.workgroup_size_z = 1;
    pkt.grid_size_x = 192; // 1 workgroup
    pkt.grid_size_y = 1;
    pkt.grid_size_z = 1;
    pkt.kernel_object = ko;
    pkt.kernarg_address = nullptr;
    queue.submit(pkt);
  }
  step_until_halted(*f.engine, {f.cu()});

  // One workgroup of 192 work-items = 3 wavefronts, all on the single CU.
  EXPECT_EQ(snap->for_cu(f.cu()).size(), 3u);
}

TEST_P(IsaTest, RunToCompletion) {
  std::string json = R"({"max_ticks":10000,"num_threads":1,"vm":{"arch":")" + arch() +
                     R"("},)"
                     R"("topology":{"root":{"name":"soc","type":"soc","children":[)"
                     R"({"name":"vram","type":"gpu_memory"},)"
                     R"({"name":"xcd0","type":"xcd","children":[)"
                     R"({"name":"l2","type":"l2_cache"},)"
                     R"({"name":"cp","type":"command_processor"},)"
                     R"({"name":"se0","type":"shader_engine","children":[)"
                     R"({"name":"cu[0:1]","type":"compute_unit","config":[)"
                     R"({"key":"num_wf_slots","value":"10"},)"
                     R"({"key":"sgprs_per_wf","value":"104"},)"
                     R"({"key":"vgprs_per_wf","value":"256"},)"
                     R"({"key":"lds_size_kb","value":"64"})"
                     R"(]}]}]}]},"links":[)"
                     R"({"src":"xcd0.cp.req_0","dst":"xcd0.se0.cu0.cpl","latency":1,"weight":2},)"
                     R"({"src":"xcd0.se0.cu0.req","dst":"xcd0.l2.cpl_0","latency":1,"weight":10})"
                     R"(]}})";

  rj_vm_t *handle = nullptr;
  ASSERT_EQ(rj_vm_create_from_string(json.c_str(), RJ_VM_MODE_DEFAULT, &handle),
            ROCJITSU_STATUS_SUCCESS);

  uint64_t ticks = 0;
  EXPECT_EQ(rj_vm_run(handle, &ticks), ROCJITSU_STATUS_SUCCESS);

  rj_vm_destroy(handle);
}

namespace enc {

constexpr uint32_t SGPR(uint32_t idx) { return idx; }
constexpr uint32_t VGPR_SRC(uint32_t idx) { return 256 + idx; }
constexpr uint32_t INLINE_CONST(uint32_t val) { return 128 + val; }

// SOPP: encoding[31:23]=0x17F, op[22:16], simm16[15:0]
constexpr uint32_t sopp(uint32_t op, uint16_t simm16 = 0) {
  return (0x17Fu << 23) | (op << 16) | simm16;
}

constexpr uint32_t s_branch(int16_t off) { return sopp(2, static_cast<uint16_t>(off)); }
constexpr uint32_t s_cbranch_scc0(int16_t off) { return sopp(4, static_cast<uint16_t>(off)); }
constexpr uint32_t s_cbranch_scc1(int16_t off) { return sopp(5, static_cast<uint16_t>(off)); }

// SOP1: encoding[31:23]=0x17D, sdst[22:16], op[15:8], ssrc0[7:0]
constexpr uint32_t sop1(uint32_t op, uint32_t sdst, uint32_t ssrc0) {
  return (0x17Du << 23) | (sdst << 16) | (op << 8) | ssrc0;
}
constexpr uint32_t s_mov_b32(uint32_t sdst, uint32_t ssrc0) { return sop1(0, sdst, ssrc0); }

// SOP2: encoding[31:30]=0x2, op[29:23], sdst[22:16], ssrc1[15:8], ssrc0[7:0]
constexpr uint32_t sop2(uint32_t op, uint32_t sdst, uint32_t ssrc0, uint32_t ssrc1) {
  return (0x2u << 30) | (op << 23) | (sdst << 16) | (ssrc1 << 8) | ssrc0;
}
constexpr uint32_t s_add_u32(uint32_t sdst, uint32_t s0, uint32_t s1) {
  return sop2(0, sdst, s0, s1);
}
constexpr uint32_t s_add_i32(uint32_t sdst, uint32_t s0, uint32_t s1) {
  return sop2(2, sdst, s0, s1);
}
// SOPC: encoding[31:23]=0x17E, op[22:16], ssrc1[15:8], ssrc0[7:0]
constexpr uint32_t sopc(uint32_t op, uint32_t ssrc0, uint32_t ssrc1) {
  return (0x17Eu << 23) | (op << 16) | (ssrc1 << 8) | ssrc0;
}
constexpr uint32_t s_cmp_eq_i32(uint32_t s0, uint32_t s1) { return sopc(0, s0, s1); }
constexpr uint32_t s_cmp_gt_i32(uint32_t s0, uint32_t s1) { return sopc(2, s0, s1); }
// VOP1: encoding[31:25]=0x3F, vdst[24:17], op[16:9], src0[8:0]
constexpr uint32_t vop1(uint32_t op, uint32_t vdst, uint32_t src0) {
  return (0x3Fu << 25) | (vdst << 17) | (op << 9) | src0;
}
constexpr uint32_t v_mov_b32(uint32_t vdst, uint32_t src0) { return vop1(1, vdst, src0); }

constexpr std::array<uint32_t, 2> vop3_cdna(uint32_t op, uint32_t vdst, uint32_t src0,
                                            uint32_t src1, uint32_t src2 = 0, uint32_t opsel = 0) {
  return {vdst | ((opsel & 0xFu) << 11) | ((op & 0x3FFu) << 16) | (0x34u << 26),
          (src0 & 0x1FFu) | ((src1 & 0x1FFu) << 9) | ((src2 & 0x1FFu) << 18)};
}

// VOP2: encoding[31]=0, op[30:25], vdst[24:17], vsrc1[16:9], src0[8:0]
constexpr uint32_t vop2(uint32_t op, uint32_t vdst, uint32_t src0, uint32_t vsrc1) {
  return (op << 25) | (vdst << 17) | (vsrc1 << 9) | src0;
}
constexpr uint32_t v_add_f32(uint32_t vdst, uint32_t s0, uint32_t vs1) {
  return vop2(1, vdst, s0, vs1);
}
constexpr uint32_t v_mul_f32(uint32_t vdst, uint32_t s0, uint32_t vs1) {
  return vop2(5, vdst, s0, vs1);
}
constexpr uint32_t v_add_u32(uint32_t vdst, uint32_t s0, uint32_t vs1) {
  return vop2(52, vdst, s0, vs1);
}
constexpr uint32_t v_cndmask_b32(uint32_t vdst, uint32_t s0, uint32_t vs1) {
  return vop2(0, vdst, s0, vs1);
}

// VOPC: encoding[31:25]=0x3E, op[24:17], vsrc1[16:9], src0[8:0]
constexpr uint32_t vopc(uint32_t op, uint32_t src0, uint32_t vsrc1) {
  return (0x3Eu << 25) | (op << 17) | (vsrc1 << 9) | src0;
}
constexpr uint32_t v_cmp_eq_f32(uint32_t s0, uint32_t vs1) { return vopc(66, s0, vs1); }

// DS: 64-bit instruction.
// dword0: offset0[7:0], offset1[15:8], gds[16], op[24:17], acc[25], encoding[31:26]=0x36
// dword1: addr[7:0], data0[15:8], data1[23:16], vdst[31:24]
constexpr uint32_t ds_lo(uint32_t op, uint8_t offset0 = 0, uint8_t offset1 = 0, uint8_t acc = 0) {
  return (0x36u << 26) | (static_cast<uint32_t>(acc) << 25) | (op << 17) |
         (static_cast<uint32_t>(offset1) << 8) | offset0;
}
constexpr uint32_t ds_hi(uint32_t vdst, uint32_t data0, uint32_t addr, uint32_t data1 = 0) {
  return (vdst << 24) | (data1 << 16) | (data0 << 8) | addr;
}

// FLAT (64-bit): CDNA3/4 layout.
// dword0: offset[11:0], pad_12[12], lds[13], seg[15:14], sc0[16], nt[17],
//         op[24:18], sc1[25], encoding[31:26]=0x37
// dword1: addr[7:0], data[15:8], saddr[22:16], acc[23], vdst[31:24]
constexpr uint32_t flat_lo(uint32_t op, uint32_t seg = 0, uint32_t sc0 = 0) {
  return (0x37u << 26) | (op << 18) | (sc0 << 16) | (seg << 14);
}
constexpr uint32_t flat_hi(uint32_t vdst, uint32_t data, uint32_t addr, uint32_t saddr = 0x7F) {
  return (vdst << 24) | (saddr << 16) | (data << 8) | addr;
}

constexpr uint32_t S_WAITCNT_0 = sopp(12, 0);
constexpr uint32_t S_ENDPGM = sopp(1, 0);

} // namespace enc

TEST(AqlDispatchTest, Fp16OvflDescriptorControlsFp8ConversionResult) {
  using namespace rocr::llvm::amdhsa;
  using namespace enc;

  auto run_case = [](bool fp16_ovfl, uint32_t expected_v2) {
    uint32_t rsrc1 = 0;
    const uint32_t fp16_ovfl_bit = fp16_ovfl ? 1u : 0u;
    AMDHSA_BITS_SET(rsrc1, COMPUTE_PGM_RSRC1_FP16_OVFL, fp16_ovfl_bit);

    const auto cvt_pk_fp8 = vop3_cdna(cdna4::kVCvtPkFp8F32Vop3, 2, VGPR_SRC(5), VGPR_SRC(6));
    const uint32_t code[] = {
        s_mov_b32(SGPR(5), 255), std::bit_cast<uint32_t>(500.0f),
        s_mov_b32(SGPR(6), 255), std::bit_cast<uint32_t>(500.0f),
        v_mov_b32(5, SGPR(5)),   v_mov_b32(6, SGPR(6)),
        cvt_pk_fp8[0], // v_cvt_pk_fp8_f32 v2.l, v5, v6
        cvt_pk_fp8[1],           S_ENDPGM,
    };

    VmFixture f("cdna4");
    auto *snap = f.capture_halts();
    uint64_t ko = f.write_kernel(0x1000, code, sizeof(code), 104, 64, 2, 0,
                                 /*wgp_mode=*/false, /*enable_vgpr_workitem_id=*/0,
                                 /*extra_compute_pgm_rsrc1=*/rsrc1);
    test::AqlQueue queue(f.mem(), f.cp());
    queue.dispatch(ko, 64, 64);

    ASSERT_NO_THROW(f.engine->run());
    ASSERT_EQ(snap->snapshots().size(), 1u);
    EXPECT_EQ((snap->snapshots().front().mode_raw & amdgpu::Wavefront::FP16_OVFL_BIT) != 0,
              fp16_ovfl);
    EXPECT_EQ(snap->snapshots().front().vgpr(2, 0), expected_v2);
  };

  run_case(false, 0x00007F7Fu);
  run_case(true, 0x00007E7Eu);
}

TEST(RdnaDispatchTest, WgpModeRoutesDsWritesThroughSiblingLdsPool) {
  constexpr uint32_t kPerCuLdsBytes = 64 * 1024;
  constexpr uint32_t kWgpLdsBytes = 2 * kPerCuLdsBytes;
  constexpr uint32_t kAddress = kPerCuLdsBytes + 4;
  constexpr uint32_t kValue = 0xC001D00D;
  constexpr uint32_t kDsStoreB32 = 26;
  constexpr uint32_t kDsLoadB32 = 108;
  constexpr uint32_t kRdna4WaitcntLgkm0 = 0xBF89FC07;
  constexpr uint32_t kRdna4Endpgm = 0xBFB00000;

  using namespace enc;
  const uint32_t code[] = {
      s_mov_b32(SGPR(4), 255),
      kAddress,
      s_mov_b32(SGPR(5), 255),
      kValue,
      v_mov_b32(0, SGPR(4)),
      v_mov_b32(1, SGPR(5)),
      ds_lo(kDsStoreB32),
      ds_hi(/*vdst=*/0, /*data0=*/1, /*addr=*/0),
      kRdna4WaitcntLgkm0,
      ds_lo(kDsLoadB32),
      ds_hi(/*vdst=*/2, /*data0=*/0, /*addr=*/0),
      kRdna4WaitcntLgkm0,
      kRdna4Endpgm,
  };

  VmFixture f("rdna4", 2, 10, /*lds_size_kb=*/64, /*sgprs_per_wf=*/128);
  auto *snap = f.capture_halts();
  uint64_t ko = f.write_kernel(0x1000, code, sizeof(code), 128, 64, 2, kWgpLdsBytes,
                               /*wgp_mode=*/true);
  test::AqlQueue queue(f.mem(), f.cp());
  queue.dispatch(ko, 32, 32);

  ASSERT_NO_THROW(f.engine->run());

  // The DS store wrote an offset in the sibling CU's half of the combined WGP LDS
  // pool, then the DS load read it back into v2. A correct v2 proves the store and
  // load both resolved through the combined pool (kAddress > per-CU capacity, so it
  // only succeeds if the sibling half is addressable).
  ASSERT_EQ(snap->snapshots().size(), 1u);
  const auto &wf = snap->snapshots().front();
  EXPECT_EQ(wf.lds_size_bytes, kWgpLdsBytes);
  EXPECT_EQ(wf.vgpr(2, 0), kValue);
}

// Runs a self-contained scalar/vector program to s_endpgm and exposes the
// wavefront's final architectural state. Because a wave frees its registers the
// instant it halts, the final state is captured at halt time via HaltSnapshotPlugin
// and served from that snapshot rather than from the (now-freed) slot.
struct ExecFixture {
  VmFixture f;
  std::string arch_;
  test::HaltSnapshotPlugin *snap_ = nullptr;

  explicit ExecFixture(const std::string &arch) : f(arch), arch_(arch) {
    snap_ = f.capture_halts();
  }

  bool is_cdna4() const { return arch_ == "cdna4"; }
  uint32_t sopp_bytes() const { return 4u; }

  std::vector<uint32_t> sopp(uint32_t word) const { return {word}; }

  static std::vector<uint32_t> cat(std::initializer_list<std::vector<uint32_t>> parts) {
    std::vector<uint32_t> result;
    for (const auto &p : parts)
      result.insert(result.end(), p.begin(), p.end());
    return result;
  }

  void load_program(const std::vector<uint32_t> &words, uint64_t base = 0x1000) {
    uint64_t ko = f.write_kernel(base, words.data(), words.size() * sizeof(uint32_t));
    test::AqlQueue queue(f.mem(), f.cp());
    queue.dispatch(ko, 64);
    step_until_halted(*f.engine, {f.cu()});
  }

  amdgpu::ComputeUnitCore *cu() { return f.cu(); }

  /// Snapshot of the (single) halted wavefront's final state.
  const test::WavefrontSnapshot &snap() const {
    EXPECT_EQ(snap_->snapshots().size(), 1u);
    return snap_->snapshots().front();
  }

  bool halted() const { return !snap_->snapshots().empty(); }
  uint32_t status_raw() const { return snap().status; }
  uint64_t vcc() const { return snap().vcc; }
  uint32_t read_sgpr(uint32_t idx) { return snap().sgpr(idx); }
  uint32_t read_vgpr(uint32_t reg, uint32_t lane) { return snap().vgpr(reg, lane); }
};

TEST_P(IsaTest, StepExecutesAndHalts) {
  VmFixture f(arch());
  auto *snap = f.capture_halts();

  auto prog = ExecFixture::cat({{SOPP_S_NOP}, {SOPP_S_NOP}, {SOPP_S_ENDPGM}});
  uint64_t ko = f.write_kernel(0x0, prog.data(), prog.size() * sizeof(uint32_t));
  test::AqlQueue queue(f.mem(), f.cp());
  queue.dispatch(ko, 64);
  step_until_halted(*f.engine, {f.cu()});

  // The wave executed and halted (freeing itself at s_endpgm).
  EXPECT_EQ(snap->snapshots().size(), 1u);
  EXPECT_EQ(f.cu()->num_wfs(), 0u);
}

TEST_P(IsaTest, RoundRobinScheduling) {
  VmFixture f(arch());

  auto prog_a = ExecFixture::cat({{SOPP_S_NOP}, {SOPP_S_NOP}, {SOPP_S_ENDPGM}});
  auto prog_b = ExecFixture::cat({{SOPP_S_NOP}, {SOPP_S_ENDPGM}});
  uint64_t ko_a = f.write_kernel(0x0, prog_a.data(), prog_a.size() * sizeof(uint32_t));
  uint64_t ko_b = f.write_kernel(0x2000, prog_b.data(), prog_b.size() * sizeof(uint32_t));
  test::AqlQueue queue(f.mem(), f.cp());

  // Verify each dispatch executes and the CP tracks them.
  queue.dispatch(ko_a, 64);
  step_until_halted(*f.engine, {f.cu()});
  EXPECT_EQ(f.cp()->dispatched_count(), 1u);

  queue.dispatch(ko_b, 64);
  step_until_halted(*f.engine, {f.cu()});
  EXPECT_EQ(f.cp()->dispatched_count(), 2u);
}

TEST_P(IsaTest, EngineRunsToCompletion) {
  VmFixture f(arch());
  auto *snap = f.capture_halts();

  auto prog = ExecFixture::cat({{SOPP_S_NOP}, {SOPP_S_ENDPGM}});
  uint64_t ko = f.write_kernel(0x0, prog.data(), prog.size() * sizeof(uint32_t));
  test::AqlQueue queue(f.mem(), f.cp());
  queue.dispatch(ko, 64);
  step_until_halted(*f.engine, {f.cu()});

  EXPECT_EQ(snap->snapshots().size(), 1u);
  EXPECT_EQ(f.cu()->num_wfs(), 0u);
}

TEST_P(IsaTest, SMovB32_InlineConst) {
  ExecFixture fx(arch());
  fx.load_program({enc::s_mov_b32(0, enc::INLINE_CONST(42)), SOPP_S_ENDPGM});
  // After engine.run(), the wavefront has executed all instructions and halted.
  // The user_sgprs (s0,s1) are set by init_wavefront_regs (kernarg ptr = 0),
  // and s2 is workgroup_id. Our instruction writes s0. Check final state.
  EXPECT_EQ(fx.read_sgpr(0), 42u);
}

TEST_P(IsaTest, SMovB32_SgprToSgpr) {
  ExecFixture fx(arch());
  fx.load_program({enc::s_mov_b32(1, enc::SGPR(0)), SOPP_S_ENDPGM});
  // s0 was set to 0 by init_wavefront_regs (kernarg low word = 0).
  // s_mov_b32 s1, s0 -> s1 = 0.
  EXPECT_EQ(fx.read_sgpr(1), 0u);
}

TEST_P(IsaTest, SAddI32_NoOverflow) {
  ExecFixture fx(arch());
  // Use inline constants to avoid relying on pre-set register values.
  fx.load_program({enc::s_mov_b32(0, enc::INLINE_CONST(10)),
                   enc::s_mov_b32(1, enc::INLINE_CONST(20)),
                   enc::s_add_i32(2, enc::SGPR(0), enc::SGPR(1)), SOPP_S_ENDPGM});
  EXPECT_EQ(fx.read_sgpr(2), 30u);
}

TEST_P(IsaTest, SAddI32_Overflow) {
  ExecFixture fx(arch());
  // Load INT32_MAX into s0 and 1 into s1, then add.
  // We need a literal constant for 0x7FFFFFFF. Use s_mov + literal.
  // Actually, inline const only goes to 64. We'll verify with a simpler approach:
  // after run, check final register state.
  fx.load_program({enc::s_mov_b32(0, enc::INLINE_CONST(64)), // s0 = 64
                   enc::s_mov_b32(1, enc::INLINE_CONST(64)), // s1 = 64
                   enc::s_add_i32(2, enc::SGPR(0), enc::SGPR(1)), SOPP_S_ENDPGM});
  EXPECT_EQ(fx.read_sgpr(2), 128u);
}

TEST_P(IsaTest, SAddU32_Carry) {
  ExecFixture fx(arch());
  // Use inline const -1 (= 0xFFFFFFFF) and 1.
  constexpr uint32_t NEG_1 = 193;                           // inline constant for -1
  fx.load_program({enc::s_mov_b32(0, NEG_1),                // s0 = 0xFFFFFFFF
                   enc::s_mov_b32(1, enc::INLINE_CONST(1)), // s1 = 1
                   enc::s_add_u32(2, enc::SGPR(0), enc::SGPR(1)), SOPP_S_ENDPGM});
  EXPECT_EQ(fx.read_sgpr(2), 0u); // wraps
}

TEST_P(IsaTest, SAddU32_NoCarry) {
  ExecFixture fx(arch());
  fx.load_program({enc::s_mov_b32(0, enc::INLINE_CONST(10)),
                   enc::s_mov_b32(1, enc::INLINE_CONST(20)),
                   enc::s_add_u32(2, enc::SGPR(0), enc::SGPR(1)), SOPP_S_ENDPGM});
  EXPECT_EQ(fx.read_sgpr(2), 30u);
}

TEST_P(IsaTest, SCmpEqI32_Equal) {
  ExecFixture fx(arch());
  fx.load_program({enc::s_mov_b32(0, enc::INLINE_CONST(42)),
                   enc::s_mov_b32(1, enc::INLINE_CONST(42)),
                   enc::s_cmp_eq_i32(enc::SGPR(0), enc::SGPR(1)), SOPP_S_ENDPGM});
  EXPECT_EQ(fx.status_raw() & 1u, 1u); // SCC=1
}

TEST_P(IsaTest, SCmpEqI32_NotEqual) {
  ExecFixture fx(arch());
  fx.load_program({enc::s_mov_b32(0, enc::INLINE_CONST(42)),
                   enc::s_mov_b32(1, enc::INLINE_CONST(43)),
                   enc::s_cmp_eq_i32(enc::SGPR(0), enc::SGPR(1)), SOPP_S_ENDPGM});
  EXPECT_EQ(fx.status_raw() & 1u, 0u); // SCC=0
}

TEST_P(IsaTest, SCmpGtI32) {
  ExecFixture fx(arch());
  constexpr uint32_t NEG_5 = 128 + 5 + 64;    // inline constant -5 = 197
  constexpr uint32_t NEG_10 = 128 + 10 + 64;  // inline constant -10 = 202
  fx.load_program({enc::s_mov_b32(0, NEG_5),  // s0 = -5
                   enc::s_mov_b32(1, NEG_10), // s1 = -10
                   enc::s_cmp_gt_i32(enc::SGPR(0), enc::SGPR(1)), SOPP_S_ENDPGM});
  EXPECT_EQ(fx.status_raw() & 1u, 1u); // -5 > -10, SCC=1
}

TEST_P(IsaTest, SBranch_Forward) {
  ExecFixture fx(arch());
  uint32_t ss = fx.sopp_bytes();
  int16_t off = static_cast<int16_t>((2 * ss - 4) / 4);
  auto prog =
      ExecFixture::cat({fx.sopp(enc::s_branch(off)), fx.sopp(SOPP_S_NOP), fx.sopp(SOPP_S_ENDPGM)});
  uint64_t ko = fx.f.write_kernel(0x1000, prog.data(), prog.size() * sizeof(uint32_t));
  test::AqlQueue queue(fx.f.mem(), fx.f.cp());
  queue.dispatch(ko, 64);
  step_until_halted(*fx.f.engine, {fx.cu()});
  EXPECT_TRUE(fx.halted());
}

TEST_P(IsaTest, SCbranchScc0_Taken) {
  ExecFixture fx(arch());
  // s_cmp_eq_i32 s0, s1 -> SCC=0 (they differ after init: s0=kernarg_lo, s1=kernarg_hi).
  // Then s_cbranch_scc0 skips to s_endpgm.
  uint32_t ss = fx.sopp_bytes();
  int16_t off = static_cast<int16_t>((2 * ss - 4) / 4);
  // Ensure SCC=0: compare two different values.
  auto prog = ExecFixture::cat({{enc::s_mov_b32(3, enc::INLINE_CONST(0))},
                                {enc::s_mov_b32(4, enc::INLINE_CONST(1))},
                                {enc::s_cmp_eq_i32(enc::SGPR(3), enc::SGPR(4))},
                                fx.sopp(enc::s_cbranch_scc0(off)),
                                fx.sopp(SOPP_S_NOP),
                                fx.sopp(SOPP_S_ENDPGM)});
  uint64_t ko = fx.f.write_kernel(0x1000, prog.data(), prog.size() * sizeof(uint32_t));
  test::AqlQueue queue(fx.f.mem(), fx.f.cp());
  queue.dispatch(ko, 64);
  step_until_halted(*fx.f.engine, {fx.cu()});
  EXPECT_TRUE(fx.halted());
}

TEST_P(IsaTest, SCbranchScc0_NotTaken) {
  ExecFixture fx(arch());
  // Ensure SCC=1: compare two equal values, then s_cbranch_scc0 should not branch.
  uint32_t ss = fx.sopp_bytes();
  int16_t off = static_cast<int16_t>((2 * ss - 4) / 4);
  auto prog = ExecFixture::cat({{enc::s_mov_b32(3, enc::INLINE_CONST(5))},
                                {enc::s_mov_b32(4, enc::INLINE_CONST(5))},
                                {enc::s_cmp_eq_i32(enc::SGPR(3), enc::SGPR(4))},
                                fx.sopp(enc::s_cbranch_scc0(off)),
                                fx.sopp(SOPP_S_NOP),
                                fx.sopp(SOPP_S_ENDPGM)});
  uint64_t ko = fx.f.write_kernel(0x1000, prog.data(), prog.size() * sizeof(uint32_t));
  test::AqlQueue queue(fx.f.mem(), fx.f.cp());
  queue.dispatch(ko, 64);
  step_until_halted(*fx.f.engine, {fx.cu()});
  EXPECT_TRUE(fx.halted());
}

TEST_P(IsaTest, SCbranchScc1_Taken) {
  ExecFixture fx(arch());
  // Ensure SCC=1: compare two equal values, then s_cbranch_scc1 should branch.
  uint32_t ss = fx.sopp_bytes();
  int16_t off = static_cast<int16_t>((2 * ss - 4) / 4);
  auto prog = ExecFixture::cat({{enc::s_mov_b32(3, enc::INLINE_CONST(7))},
                                {enc::s_mov_b32(4, enc::INLINE_CONST(7))},
                                {enc::s_cmp_eq_i32(enc::SGPR(3), enc::SGPR(4))},
                                fx.sopp(enc::s_cbranch_scc1(off)),
                                fx.sopp(SOPP_S_NOP),
                                fx.sopp(SOPP_S_ENDPGM)});
  uint64_t ko = fx.f.write_kernel(0x1000, prog.data(), prog.size() * sizeof(uint32_t));
  test::AqlQueue queue(fx.f.mem(), fx.f.cp());
  queue.dispatch(ko, 64);
  step_until_halted(*fx.f.engine, {fx.cu()});
  EXPECT_TRUE(fx.halted());
}

TEST_P(IsaTest, SEndpgm_Halts) {
  ExecFixture fx(arch());
  fx.load_program({SOPP_S_ENDPGM});
  EXPECT_TRUE(fx.halted());
}

TEST_P(IsaTest, VMovB32_PerLane) {
  ExecFixture fx(arch());
  // V_MOV_B32 v2, v1 -- use v2 as dest to avoid clobbering v0 (lane id)
  // Then check v2 after completion.
  fx.load_program({enc::v_mov_b32(2, enc::VGPR_SRC(0)), SOPP_S_ENDPGM});
  // After run: v0 was set to lane index by init_wavefront_regs.
  // v_mov_b32 v2, v0 copies lane index to v2.
  EXPECT_EQ(fx.read_vgpr(2, 0), 0u);
  EXPECT_EQ(fx.read_vgpr(2, 1), 1u);
  EXPECT_EQ(fx.read_vgpr(2, 63), 63u);
}

TEST_P(IsaTest, VAddF32_PerLane) {
  ExecFixture fx(arch());
  // We need to set up v registers before execution. But with AQL dispatch,
  // the engine runs to completion. So we encode a self-contained program:
  // v_mov_b32 v3, inline_1.5f -- but inline floats are limited.
  // Instead, test that v_add_f32 of v0 (lane index) + v0 = 2*lane_index as float.
  // Actually this won't work since v0 contains integer lane indices, not floats.
  // Let's just verify the instruction halts correctly and check the result.
  // Use v_add_f32 with inline constant 1.0 (0x3F800000 = inline 242).
  // Inline float 1.0 = src code 242.
  fx.load_program({enc::v_add_f32(2, 242, 0), // v2 = 1.0 + v0_as_float
                   SOPP_S_ENDPGM});
  // v0[lane 0] = 0 (int), as float = 0.0. 1.0 + 0.0 = 1.0
  EXPECT_EQ(std::bit_cast<float>(fx.read_vgpr(2, 0)), 1.0f);
}

TEST_P(IsaTest, VMulF32_PerLane) {
  ExecFixture fx(arch());
  // v_mul_f32 v2, 1.0, v0 -> v2 = 1.0 * v0_as_float
  fx.load_program({enc::v_mul_f32(2, 242, 0), SOPP_S_ENDPGM}); // 242 = inline 1.0f
  // v0[lane 0] = 0 (int) = 0.0 as float. 1.0 * 0.0 = 0.0
  EXPECT_EQ(std::bit_cast<float>(fx.read_vgpr(2, 0)), 0.0f);
}

TEST_P(IsaTest, VAddU32_PerLane) {
  ExecFixture fx(arch());
  // v_add_u32 v2, v0, v0 -> v2 = 2 * lane_index
  fx.load_program({enc::v_add_u32(2, enc::VGPR_SRC(0), 0), SOPP_S_ENDPGM});
  EXPECT_EQ(fx.read_vgpr(2, 0), 0u);
  EXPECT_EQ(fx.read_vgpr(2, 1), 2u);
  EXPECT_EQ(fx.read_vgpr(2, 3), 6u);
}

TEST_P(IsaTest, VCmpEqF32_SetsVCC) {
  ExecFixture fx(arch());
  // Compare v0 (lane index as float-bits) with inline 0 (integer 0).
  // Lane 0: v0=0, compared with 0 -> equal -> VCC[0]=1.
  // Lane 1: v0=1, compared with 0 -> not equal -> VCC[1]=0.
  fx.load_program({enc::v_cmp_eq_f32(enc::INLINE_CONST(0), 0), SOPP_S_ENDPGM});
  uint64_t vcc = fx.vcc();
  EXPECT_TRUE(vcc & (1ULL << 0));  // lane 0: 0.0 == 0.0
  EXPECT_FALSE(vcc & (1ULL << 1)); // lane 1: int 1 as float != 0.0
}

TEST_P(IsaTest, VCndmaskB32) {
  ExecFixture fx(arch());
  // v_cndmask_b32 v2, v0, v1 -- selects v1 where VCC set, v0 otherwise.
  // After init: v0 = lane_index. v1 = 0. We can't set VCC before run.
  // Instead, set VCC via v_cmp first, then use v_cndmask.
  // v_cmp_eq_f32 v0, 0 -> VCC[0]=1 (lane 0 = 0 == 0), VCC[1]=0 (1 != 0)
  // v_mov_b32 v1, inline 99
  // v_cndmask_b32 v2, v0, v1 -> lane 0: VCC=1 -> v1=99; lane 1: VCC=0 -> v0=1
  fx.load_program({enc::v_cmp_eq_f32(enc::INLINE_CONST(0), 0), // VCC from v0 == 0
                   enc::v_mov_b32(1, enc::INLINE_CONST(42)),   // v1 = 42 (all lanes)
                   enc::v_cndmask_b32(2, enc::VGPR_SRC(0), 1), // v2 = VCC ? v1 : v0
                   SOPP_S_ENDPGM});
  EXPECT_EQ(fx.read_vgpr(2, 0), 42u); // VCC[0]=1 -> v1=42
  EXPECT_EQ(fx.read_vgpr(2, 1), 1u);  // VCC[1]=0 -> v0=1 (lane index)
}

TEST_P(IsaTest, ExecMask_PreservesInactiveLanes) {
  ExecFixture fx(arch());
  // We can't set EXEC before run. Instead, verify that the engine runs to completion.
  // This test is simplified to just verify halting behavior.
  fx.load_program({enc::v_mov_b32(2, enc::VGPR_SRC(0)), SOPP_S_ENDPGM});
  EXPECT_TRUE(fx.halted());
  EXPECT_EQ(fx.read_vgpr(2, 0), 0u);
  EXPECT_EQ(fx.read_vgpr(2, 1), 1u);
}

TEST_P(IsaTest, MultiInstructionProgram) {
  ExecFixture fx(arch());
  fx.load_program({
      enc::s_mov_b32(3, enc::INLINE_CONST(10)),
      enc::s_mov_b32(4, enc::INLINE_CONST(20)),
      enc::s_add_i32(5, enc::SGPR(3), enc::SGPR(4)),
      SOPP_S_ENDPGM,
  });
  EXPECT_EQ(fx.read_sgpr(3), 10u);
  EXPECT_EQ(fx.read_sgpr(4), 20u);
  EXPECT_EQ(fx.read_sgpr(5), 30u);
  EXPECT_TRUE(fx.halted());
}

TEST_P(IsaTest, BranchLoop) {
  ExecFixture fx(arch());
  // Scalar loop: s3 starts at 3, each iteration subtracts 1, loop back if s3 > 0.
  constexpr uint32_t NEG_1 = 193; // inline constant -1
  auto prog = ExecFixture::cat({
      {enc::s_mov_b32(3, enc::INLINE_CONST(3))}, // s3 = 3
      // loop:
      {enc::s_add_i32(3, enc::SGPR(3), NEG_1)},                // s3 -= 1
      {enc::s_cmp_gt_i32(enc::SGPR(3), enc::INLINE_CONST(0))}, // s3 > 0?
      fx.sopp(enc::s_cbranch_scc1(-3)),                        // if SCC=1 goto loop
      fx.sopp(SOPP_S_ENDPGM),
  });
  uint64_t ko = fx.f.write_kernel(0x1000, prog.data(), prog.size() * sizeof(uint32_t));
  test::AqlQueue queue(fx.f.mem(), fx.f.cp());
  queue.dispatch(ko, 64);
  step_until_halted(*fx.f.engine, {fx.cu()});

  EXPECT_EQ(fx.read_sgpr(3), 0u);
  EXPECT_TRUE(fx.halted());
}

INSTANTIATE_TEST_SUITE_P(Cdna, IsaTest, ::testing::Values("cdna3", "cdna4"),
                         [](const auto &info) { return info.param; });

// ---------------------------------------------------------------------------
// MFMA accumulation unit tests
// ---------------------------------------------------------------------------

TEST_P(IsaTest, MfmaF16Accumulation) {
  VmFixture f(arch());

  // Resident scratch wave: this test drives MFMA directly and reads the VGPR file,
  // so the wave must stay allocated (not run to s_endpgm, which would free it).
  auto *cu = f.cu();
  auto *wf = f.dispatch_scratch_wf();
  ASSERT_NE(wf, nullptr);
  uint32_t vb = wf->vgpr_alloc().base;

  // Test v_mfma_f32_16x16x32_f16: M=16, N=16, K=32, B=1, in_bits=16
  // Set up src0 (A matrix) at vb+10 (8 VGPRs for 32 FP16 packed as 16 dwords)
  // Set up src1 (B matrix) at vb+18 (8 VGPRs)
  // Set up dst/src2 (accumulator) at vb+256 (4 VGPRs, AccVGPR bank)

  // Fill A and B with known FP16 values (all 1.0h = 0x3C00)
  uint32_t packed_ones = 0x3C003C00; // two FP16 1.0 values
  for (uint32_t r = 0; r < 8; r++)
    for (uint32_t lane = 0; lane < 64; lane++) {
      cu->write_vgpr(vb + 10 + r, lane, packed_ones); // A
      cu->write_vgpr(vb + 18 + r, lane, packed_ones); // B
    }

  // Zero the accumulator (AccVGPR bank at +256)
  for (uint32_t r = 0; r < 4; r++)
    for (uint32_t lane = 0; lane < 64; lane++)
      cu->write_vgpr(vb + 256 + r, lane, 0);

  // Execute MFMA: D[16x16] = 0 + A[16x32] * B[32x16]
  // With all-ones inputs, each output element = sum of K=32 products of 1.0*1.0 = 32.0
  uint32_t dst = vb + 256;
  uint32_t s0 = vb + 10;
  uint32_t s1 = vb + 18;
  uint32_t s2 = vb + 256;
  uint32_t const_acc = amdgpu::ACC_FROM_VGPR;
  amdgpu::exec_f32(*cu, 16, 16, 32, 1, 16, dst, s0, s1, s2, amdgpu::extract_f16,
                   amdgpu::extract_f16, const_acc);

  // Verify: every output element should be exactly 32.0f
  float expected = 32.0f;
  uint32_t expected_bits = std::bit_cast<uint32_t>(expected);
  uint32_t mismatches = 0;
  for (uint32_t row = 0; row < 16; row++) {
    for (uint32_t col = 0; col < 16; col++) {
      auto out = amdgpu::output_loc_32(16, 16, row, col, 0);
      uint32_t got = cu->read_vgpr(dst + out.reg, out.lane);
      if (got != expected_bits) {
        if (mismatches < 5)
          ADD_FAILURE() << "MFMA ones: C[" << row << "][" << col
                        << "] = " << std::bit_cast<float>(got) << " (expected " << expected << ")"
                        << " reg=" << out.reg << " lane=" << out.lane;
        mismatches++;
      }
    }
  }
  EXPECT_EQ(mismatches, 0u) << mismatches << "/256 elements differ";
}

TEST_P(IsaTest, MfmaF16AccumulationPatterned) {
  VmFixture f(arch());

  // Resident scratch wave (see MfmaF16Accumulation): driven directly, not to endpgm.
  auto *cu = f.cu();
  auto *wf = f.dispatch_scratch_wf();
  ASSERT_NE(wf, nullptr);
  uint32_t vb = wf->vgpr_alloc().base;

  // Test with patterned data: A[i][k] = (i+1) as FP16, B[k][j] = 1.0h
  // Expected: C[i][j] = (i+1) * K = (i+1) * 32
  for (uint32_t r = 0; r < 8; r++)
    for (uint32_t lane = 0; lane < 64; lane++) {
      // B = all ones
      cu->write_vgpr(vb + 18 + r, lane, 0x3C003C00);
      // A = row-dependent value: use lane to determine row
      // input_loc maps (row, k) to (vgpr_offset, lane, sub_element)
      // For simplicity, just use uniform values per lane
      cu->write_vgpr(vb + 10 + r, lane, 0x3C003C00); // 1.0
    }

  // Zero accumulator
  for (uint32_t r = 0; r < 4; r++)
    for (uint32_t lane = 0; lane < 64; lane++)
      cu->write_vgpr(vb + 256 + r, lane, 0);

  // Execute MFMA with zero accumulator (const_acc = 0.0f)
  uint32_t dst = vb + 256;
  amdgpu::exec_f32(*cu, 16, 16, 32, 1, 16, dst, vb + 10, vb + 18, dst, amdgpu::extract_f16,
                   amdgpu::extract_f16, std::bit_cast<uint32_t>(0.0f));

  // With const_acc=0.0, all outputs should be 32.0 (sum of 32 ones*ones)
  float expected = 32.0f;
  uint32_t mismatches = 0;
  for (uint32_t row = 0; row < 16; row++) {
    for (uint32_t col = 0; col < 16; col++) {
      auto out = amdgpu::output_loc_32(16, 16, row, col, 0);
      float got = std::bit_cast<float>(cu->read_vgpr(dst + out.reg, out.lane));
      if (got != expected) {
        if (mismatches < 5)
          ADD_FAILURE() << "MFMA patterned: C[" << row << "][" << col << "] = " << got
                        << " (expected " << expected << ")";
        mismatches++;
      }
    }
  }
  EXPECT_EQ(mismatches, 0u);
}

void init_mfma_f64_neg_inputs(amdgpu::ComputeUnitCore *cu, uint32_t s0, uint32_t s1, uint32_t s2,
                              double a = 1.0, double b = 1.0, double c = 1.0) {
  uint64_t a_bits = std::bit_cast<uint64_t>(a);
  uint64_t b_bits = std::bit_cast<uint64_t>(b);
  uint64_t c_bits = std::bit_cast<uint64_t>(c);
  for (uint32_t lane = 0; lane < 64; ++lane) {
    cu->write_vgpr(s0, lane, static_cast<uint32_t>(a_bits));
    cu->write_vgpr(s0 + 1, lane, static_cast<uint32_t>(a_bits >> 32));
    cu->write_vgpr(s1, lane, static_cast<uint32_t>(b_bits));
    cu->write_vgpr(s1 + 1, lane, static_cast<uint32_t>(b_bits >> 32));
    cu->write_vgpr(s2, lane, static_cast<uint32_t>(c_bits));
    cu->write_vgpr(s2 + 1, lane, static_cast<uint32_t>(c_bits >> 32));
  }
}

void expect_mfma_f64_outputs(amdgpu::ComputeUnitCore *cu, uint32_t dst, double expected) {
  uint64_t expected_bits = std::bit_cast<uint64_t>(expected);
  uint32_t mismatches = 0;
  for (uint32_t b = 0; b < 4; ++b) {
    for (uint32_t row = 0; row < 4; ++row) {
      for (uint32_t col = 0; col < 4; ++col) {
        auto out = amdgpu::output_loc_64(4, 4, row, col, b);
        uint32_t lo = cu->read_vgpr(dst + out.reg, out.lane);
        uint32_t hi = cu->read_vgpr(dst + out.reg + 1, out.lane);
        uint64_t got_bits = static_cast<uint64_t>(hi) << 32 | lo;
        if (got_bits != expected_bits) {
          if (mismatches < 5)
            ADD_FAILURE() << "F64 output mismatch b=" << b << " row=" << row << " col=" << col
                          << " expected=" << expected << " got=" << std::bit_cast<double>(got_bits);
          ++mismatches;
        }
      }
    }
  }
  EXPECT_EQ(mismatches, 0u);
}

void expect_mfma_f64_neg_modifier(const std::string &arch) {
  VmFixture f(arch);

  // Resident scratch wave (see MfmaF16Accumulation): driven directly, not to endpgm.
  auto *cu = f.cu();
  auto *wf = f.dispatch_scratch_wf();
  ASSERT_NE(wf, nullptr);
  uint32_t vb = wf->vgpr_alloc().base;
  uint32_t dst = vb + amdgpu::ACC_VGPR_OFFSET;
  uint32_t s0 = vb + 10;
  uint32_t s1 = vb + 20;
  uint32_t s2 = dst;

  init_mfma_f64_neg_inputs(cu, s0, s1, s2);

  // CDNA f64 MFMA uses the BLGP bit range as NEG[2:0]. NEG=5 negates A and C:
  // D = -C + (-A * B) * K = -1 + (-1 * 1) * 4 = -5.
  amdgpu::exec_f64(*cu, 4, 4, 4, 4, dst, s0, s1, s2, amdgpu::ACC_FROM_VGPR, 5);
  expect_mfma_f64_outputs(cu, dst, -5.0);

  init_mfma_f64_neg_inputs(cu, s0, s1, s2);

  // NEG=2 isolates the B operand negate bit:
  // D = C + (A * -B) * K = 1 + (1 * -1) * 4 = -3.
  amdgpu::exec_f64(*cu, 4, 4, 4, 4, dst, s0, s1, s2, amdgpu::ACC_FROM_VGPR, 2);
  expect_mfma_f64_outputs(cu, dst, -3.0);
}

TEST(MfmaF64Cdna3Test, NegModifier) { expect_mfma_f64_neg_modifier("cdna3"); }

TEST(MfmaF64Cdna4Test, NegModifier) { expect_mfma_f64_neg_modifier("cdna4"); }

TEST(MfmaF64Cdna4Test, GeneratedInstructionUsesBlgpNegModifier) {
  VmFixture f("cdna4");

  // Resident scratch wave (see MfmaF16Accumulation): driven directly, not to endpgm.
  auto *cu = f.cu();
  auto *wf = f.dispatch_scratch_wf();
  ASSERT_NE(wf, nullptr);
  uint32_t vb = wf->vgpr_alloc().base;
  constexpr uint32_t kSrc0 = 10;
  constexpr uint32_t kSrc1 = 20;
  constexpr uint32_t kDst = 0;
  uint32_t dst = vb + amdgpu::ACC_VGPR_OFFSET + kDst;

  cdna4::Vop3pMfmaMachineInst raw{};
  raw.vdst = kDst;
  raw.acc_cd = 1;
  raw.src0 = 256 + kSrc0;
  raw.src1 = 256 + kSrc1;
  raw.src2 = 256 + kDst;

  const struct {
    uint32_t blgp;
    double expected;
  } cases[] = {
      {0, 29.0},
      {2, -19.0},
      {5, -29.0},
  };

  for (const auto &test : cases) {
    SCOPED_TRACE(test.blgp);
    init_mfma_f64_neg_inputs(cu, vb + kSrc0, vb + kSrc1, dst, 2.0, 3.0, 5.0);
    raw.blgp = test.blgp;
    cdna4::VMfmaF644x4x44bF64Vop3pMfma inst(reinterpret_cast<const cdna4::MachineInst *>(&raw));
    inst.execute_impl(*wf);

    expect_mfma_f64_outputs(cu, dst, test.expected);
  }
}

// ---------------------------------------------------------------------------
// Atomic stress tests
// ---------------------------------------------------------------------------

// Dispatches multiple wavefronts that all atomically add 1 to LDS[0].
// If atomics are truly atomic, the final value must equal the total number
// of active lanes across all wavefronts.
TEST(AtomicStressTest, DsAddRtnU32_MultiWavefront) {
  // 3 wavefronts × 64 lanes = 192 total atomic adds.
  VmFixture f("cdna4", 1, 10);

  // Kernel:
  //   v_mov_b32 v1, 1           // data0 = 1
  //   v_mov_b32 v3, 0           // addr = LDS offset 0
  //   ds_add_rtn_u32 v2, v3, v1 // V2 = old LDS[0]; LDS[0] += 1
  //   s_waitcnt lgkm:0
  //   s_endpgm
  using namespace enc;
  const uint32_t code[] = {
      v_mov_b32(1, INLINE_CONST(1)),              // v1 = 1
      v_mov_b32(3, INLINE_CONST(0)),              // v3 = 0 (LDS addr)
      ds_lo(32),                                  // ds_add_rtn_u32 (op=32), offset0=0
      ds_hi(/*vdst=*/2, /*data0=*/1, /*addr=*/3), // v2=result, v1=data, v3=addr
      S_WAITCNT_0,
      S_ENDPGM,
  };
  uint64_t ko = f.write_kernel(0x1000, code, sizeof(code));

  // Initialize LDS[0] = 0.
  f.cu()->lds().write32(0, 0);

  // Dispatch 192 workitems = 3 wavefronts of 64 lanes.
  test::AqlQueue queue(f.mem(), f.cp());
  queue.dispatch(ko, 192, 192); // 1 workgroup of 192 threads
  f.engine->run();

  // All 192 lanes should have atomically added 1.
  uint32_t final_val = f.cu()->lds().read32(0);
  EXPECT_EQ(final_val, 192u) << "LDS atomic add result should be 192 (3 waves × 64 lanes)";
}

// Same test but with 4 workgroups dispatched independently.
// All share the same CU (and thus same LDS), so atomics must be correct
// across workgroup boundaries within one CU.
TEST(AtomicStressTest, DsAddRtnU32_MultiWorkgroup) {
  VmFixture f("cdna4", 1, 10);

  using namespace enc;
  const uint32_t code[] = {
      v_mov_b32(1, INLINE_CONST(1)),
      v_mov_b32(3, INLINE_CONST(0)),
      ds_lo(32),
      ds_hi(2, 1, 3),
      S_WAITCNT_0,
      S_ENDPGM,
  };
  uint64_t ko = f.write_kernel(0x1000, code, sizeof(code));

  f.cu()->lds().write32(0, 0);

  // 4 workgroups × 64 threads each = 4 wavefronts = 256 atomic adds.
  test::AqlQueue queue(f.mem(), f.cp());
  queue.dispatch(ko, 256, 64);
  f.engine->run();

  uint32_t final_val = f.cu()->lds().read32(0);
  EXPECT_EQ(final_val, 256u) << "LDS atomic add across 4 workgroups should be 256";
}

// Non-RTN DS atomic add: verify LDS gets the correct sum even without
// returning the old value.
TEST(AtomicStressTest, DsAddU32_NoReturn) {
  VmFixture f("cdna4", 1, 10);

  using namespace enc;
  const uint32_t code[] = {
      v_mov_b32(1, INLINE_CONST(1)),
      v_mov_b32(3, INLINE_CONST(0)),
      ds_lo(0),       // ds_add_u32 (op=0), no return
      ds_hi(0, 1, 3), // vdst unused, data0=v1, addr=v3
      S_WAITCNT_0,
      S_ENDPGM,
  };
  uint64_t ko = f.write_kernel(0x1000, code, sizeof(code));

  f.cu()->lds().write32(0, 100); // Start at 100.

  test::AqlQueue queue(f.mem(), f.cp());
  queue.dispatch(ko, 128, 128); // 2 wavefronts = 128 lanes
  f.engine->run();

  uint32_t final_val = f.cu()->lds().read32(0);
  EXPECT_EQ(final_val, 228u) << "100 + 128 atomic adds = 228";
}

// Global (L2) atomic add: multiple wavefronts atomically increment a global
// memory location. Exercises the L2 cache's striped-mutex atomic_rmw path.
// Global (L2) atomic add: multiple wavefronts atomically increment a global
// memory location. Exercises the L2 cache's striped-mutex atomic_rmw path.
// Uses SGPR pair s4:s5 as the base address (saddr) with VGPR v4=0 (offset).
TEST(AtomicStressTest, GlobalAtomicAdd_L2) {
  VmFixture f("cdna4", 1, 10);

  constexpr uint64_t TARGET_ADDR = 0x2000ULL;

  // Kernel uses literal constants (ssrc0=255 + next dword) to load the
  // target address into s4:s5, since the address doesn't fit in an inline
  // constant (0-64 range).
  using namespace enc;
  const uint32_t code[] = {
      s_mov_b32(SGPR(4), 255),             // s4 = literal (next dword)
      static_cast<uint32_t>(TARGET_ADDR),  // literal: 0x2000
      s_mov_b32(SGPR(5), INLINE_CONST(0)), // s5 = 0 (high 32 bits)
      v_mov_b32(1, INLINE_CONST(1)),       // v1 = 1
      v_mov_b32(4, INLINE_CONST(0)),       // v4 = 0 (offset)
      flat_lo(66, /*seg=*/2, /*sc0=*/1),   // flat_atomic_add, GLOBAL, return
      flat_hi(/*vdst=*/2, /*data=*/1, /*addr=*/4, /*saddr=*/4),
      S_WAITCNT_0,
      S_ENDPGM,
  };
  uint64_t ko = f.write_kernel(0x1000, code, sizeof(code));

  f.mem()->write32(TARGET_ADDR, 0);

  // 3 wavefronts × 64 lanes = 192 global atomic adds through L2.
  test::AqlQueue queue(f.mem(), f.cp());
  queue.dispatch(ko, 192, 192);
  f.engine->run();
  f.cu()->flush_all();

  uint32_t final_val = f.mem()->read32(TARGET_ADDR);
  EXPECT_EQ(final_val, 192u) << "Global atomic add through L2 should be 192 (3 waves × 64 lanes)";
}

// Multiple workgroups all atomically add to the same global address.
TEST(AtomicStressTest, GlobalAtomicAdd_MultiWorkgroup) {
  VmFixture f("cdna4", 1, 10);

  constexpr uint64_t TARGET_ADDR = 0x3000ULL;

  using namespace enc;
  const uint32_t code[] = {
      s_mov_b32(SGPR(4), 255),
      static_cast<uint32_t>(TARGET_ADDR),
      s_mov_b32(SGPR(5), INLINE_CONST(0)),
      v_mov_b32(1, INLINE_CONST(1)),
      v_mov_b32(4, INLINE_CONST(0)),
      flat_lo(66, 2, 1),
      flat_hi(2, 1, 4, 4),
      S_WAITCNT_0,
      S_ENDPGM,
  };
  uint64_t ko = f.write_kernel(0x1000, code, sizeof(code));

  f.mem()->write32(TARGET_ADDR, 1000);

  // 4 workgroups × 64 threads = 4 wavefronts = 256 atomic adds.
  test::AqlQueue queue(f.mem(), f.cp());
  queue.dispatch(ko, 256, 64);
  f.engine->run();
  f.cu()->flush_all();

  uint32_t final_val = f.mem()->read32(TARGET_ADDR);
  EXPECT_EQ(final_val, 1256u) << "1000 + 256 global atomic adds = 1256";
}

// Verify that ds_read_b64_tr_b16 with acc=1 writes to AccVGPR (vb+256+vdst),
// not to VGPR (vb+vdst).
TEST(DsTransposeTest, ReadB64TrB16_AccBit) {
  VmFixture f("cdna4", 1, 10);
  auto *snap = f.capture_halts();

  constexpr uint32_t VDST = 4;
  constexpr uint32_t ADDR_REG = 0;
  constexpr uint32_t DS_OP = 227; // ds_read_b64_tr_b16

  // Kernel:
  //   v_mov_b32 v0, 0          ; addr = LDS offset 0
  //   v_mov_b32 v4, 0x42       ; sentinel in VGPR v4
  //   v_mov_b32 v5, 0x42       ; sentinel in VGPR v5
  //   ds_read_b64_tr_b16 a[4:5], v0  ; acc=1: write to AccVGPR
  //   s_waitcnt lgkmcnt(0)
  //   s_endpgm
  using namespace enc;
  const uint32_t code[] = {
      v_mov_b32(ADDR_REG, INLINE_CONST(0)),
      v_mov_b32(VDST, INLINE_CONST(42)),
      v_mov_b32(VDST + 1, INLINE_CONST(42)),
      ds_lo(DS_OP, /*offset0=*/0, /*offset1=*/0, /*acc=*/1),
      ds_hi(VDST, /*data0=*/0, ADDR_REG),
      S_WAITCNT_0,
      S_ENDPGM,
  };
  uint64_t ko = f.write_kernel(0x1000, code, sizeof(code));

  auto *cu = f.cu();

  // Write a known non-zero pattern to LDS.
  for (uint32_t i = 0; i < 256; ++i)
    cu->lds().write32(i * 4, 0xDEADBEEF);

  test::AqlQueue queue(f.mem(), f.cp());
  queue.dispatch(ko, 64);
  f.engine->run();

  ASSERT_EQ(snap->snapshots().size(), 1u);
  const auto &wf = snap->snapshots().front();

  // VGPR v4 should still hold the sentinel (42), not overwritten by ds_read.
  uint32_t vgpr_val = wf.vgpr(VDST, 0);
  EXPECT_EQ(vgpr_val, 42u) << "VGPR v" << VDST << " should NOT have been written when acc=1";

  // AccVGPR a4 should have been written with LDS data (not 42, not 0).
  uint32_t acc_val = wf.vgpr(256 + VDST, 0);
  EXPECT_NE(acc_val, 0u) << "AccVGPR a" << VDST << " should have been written by ds_read";
  EXPECT_NE(acc_val, 42u) << "AccVGPR a" << VDST
                          << " should contain LDS data, not the VGPR sentinel";
}

// Write-through scalar stores must use the store's VMID. Two page tables map
// the same GPU VA to different host pages; a store under VMID 7 followed by the
// no-op writeback_all(8) must land only in VMID 7's backing.
TEST(L1ScalarCacheVmidTest, WriteThroughStoreUsesStoreVmid) {
  constexpr uint32_t kVmidA = 7;
  constexpr uint32_t kVmidB = 8;
  constexpr uint64_t kSharedVa = 0x40000; // page-aligned, aliased across procs.
  constexpr uint32_t kStoreWord = 0xA5A5A5A5u;

  amdgpu::GpuMemory mem("test.vram");

  // Two processes whose page tables map the same VA to different host buffers.
  KfdProcess proc_a(kVmidA);
  KfdProcess proc_b(kVmidB);
  alignas(4096) std::array<uint8_t, 4096> backing_a{};
  alignas(4096) std::array<uint8_t, 4096> backing_b{};
  proc_a.map_pages(kSharedVa, backing_a.data(), backing_a.size());
  proc_b.map_pages(kSharedVa, backing_b.data(), backing_b.size());
  mem.register_process(kVmidA, &proc_a.page_table_, &proc_a.page_table_mutex_,
                       proc_a.page_table_generation());
  mem.register_process(kVmidB, &proc_b.page_table_, &proc_b.page_table_mutex_,
                       proc_b.page_table_generation());

  amdgpu::L2Cache l2("test.l2");
  l2.set_backing_memory(&mem);

  amdgpu::L1ScalarCache k_cache(&l2);
  k_cache.set_memory(&mem);

  // Store a dword under VMID A. Write-through must publish it immediately
  // through VMID A's page table.
  k_cache.store(kSharedVa, /*num_dwords=*/1, &kStoreWord, kVmidA);
  EXPECT_EQ(mem.read32(kSharedVa, kVmidA), kStoreWord);
  EXPECT_NE(mem.read32(kSharedVa, kVmidB), kStoreWord);

  // A writeback request from another process is harmless because K$ has no
  // dirty data to publish.
  k_cache.writeback_all(kVmidB);
  l2.flush_all();

  EXPECT_EQ(mem.read32(kSharedVa, kVmidA), kStoreWord)
      << "write-through store must remain in VMID A";
  EXPECT_NE(mem.read32(kSharedVa, kVmidB), kStoreWord)
      << "store must NOT leak into VMID B's address space";

  mem.unregister_process(kVmidA);
  mem.unregister_process(kVmidB);
}

TEST(DoorbellMonitorLifecycle, RetiresAfterLastQueueAndRestartsOnNewQueue) {
  // Regression for the leaked idle CP doorbell poller: the monitor must self-exit
  // once the last host-accessible (KFD) queue is destroyed, and a later queue
  // registration must start a fresh monitor. Uses only the queue-registration
  // lifecycle (no dispatch), so the monitor thread runs on its own and its state
  // is observed through the test-only accessor.
  VmFixture f("cdna4", /*num_cus=*/1);
  amdgpu::CommandProcessor *cp = f.cp();

  auto wait_for_monitor = [&](bool expected) {
    // The loop retires on a 100us cadence; allow generous slack for CI load.
    for (int i = 0; i < 2000 && cp->doorbell_monitor_running_for_test() != expected; ++i)
      std::this_thread::sleep_for(std::chrono::milliseconds(1));
    return cp->doorbell_monitor_running_for_test();
  };

  EXPECT_FALSE(cp->doorbell_monitor_running_for_test())
      << "no monitor should run before any host-accessible queue is registered";

  amdgpu::HwQueue queue{};
  queue.process_id = 1;
  queue.queue_id = 7;
  queue.host_accessible = true;
  cp->register_queue(queue);
  EXPECT_TRUE(wait_for_monitor(true)) << "registering a KFD queue must start the monitor";

  cp->unregister_queue(queue.queue_id, queue.process_id);
  EXPECT_FALSE(wait_for_monitor(false))
      << "monitor must self-exit after the last host-accessible queue is destroyed";

  // A new queue landing on a CP whose monitor retired must get polling back.
  amdgpu::HwQueue queue2{};
  queue2.process_id = 1;
  queue2.queue_id = 8;
  queue2.host_accessible = true;
  cp->register_queue(queue2);
  EXPECT_TRUE(wait_for_monitor(true)) << "a new KFD queue must restart a retired monitor";

  cp->unregister_queue(queue2.queue_id, queue2.process_id);
  EXPECT_FALSE(wait_for_monitor(false)) << "monitor must retire again after the last queue";
}

} // namespace