Skip to content

fix(rocjitsu): disable zero-count TDM descriptors #9941

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
newling merged 2 commits into from
Aug 17, 2026
Merged

fix(rocjitsu): disable zero-count TDM descriptors #9941

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
9352152
fix(rocjitsu): disable zero-count TDM descriptors
newling Aug 10, 2026
abdd743
fix(rocjitsu): reject unsupported TDM count encodings
newling Aug 11, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
31 changes: 21 additions & 10 deletions emulation/rocjitsu/lib/rocjitsu/src/rocjitsu/isa/arch/amdgpu/shared/tensor_dma.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -58,6 +58,7 @@ struct TensorDmaDescriptor {
std::array<uint32_t, 8> d1{};
std::array<uint32_t, 4> d2{};
std::array<uint32_t, 4> d3{};
uint32_t count = 0;
uint64_t global_base = 0;
uint32_t lds_base = 0;
uint32_t elem_size = 0;
Expand All @@ -78,6 +79,8 @@ struct TensorDmaDescriptor {
bool iterate = false;
bool pad = false;

bool active() const { return count != 0; }

uint32_t rank() const {
if (gather) {
if (tensor_dims[1] != 0)
Expand Down Expand Up @@ -112,6 +115,7 @@ inline TensorDmaDescriptor parse_descriptor(std::array<uint32_t, 4> d0, std::arr

// Descriptor bit layout follows LLVM MLIR's gfx1250 TDM lowering in
// mlir/lib/Conversion/AMDGPUToROCDL/AMDGPUToROCDL.cpp.
desc.count = d0[0] & 0x3u;
Comment thread
newling marked this conversation as resolved.
Resolved
desc.gather_indices_32bit = (d0[0] & (1u << 30)) != 0;
desc.gather = (d0[0] & (1u << 31)) != 0;
desc.lds_base = d0[1];
Expand Down Expand Up @@ -180,6 +184,10 @@ inline TensorDmaDescriptor parse_descriptor(std::array<uint32_t, 4> d0, std::arr
}

inline void validate_supported_descriptor(const TensorDmaDescriptor &desc, bool store_from_lds) {
// The in-tree HIP descriptor API and LLVM lowering only produce the boolean
// count encodings 0 (disabled) and 1 (active).
if (desc.count > 1)
throw util::UnimplementedInst("tensor DMA count encoding");
// LLVM MLIR's AMDGPU dialect documents padding only for memory-to-LDS copies.
if (desc.pad && store_from_lds)
throw util::UnimplementedInst("tensor DMA padded store descriptor");
Expand Down Expand Up @@ -340,22 +348,25 @@ TensorDmaDescriptor read_descriptor(const Inst &inst, const Wavefront &wf) {
read_sgpr_group<4>(wf, inst.vaddr3.encoding_value(), true));
}

template <typename Inst>
void execute_tensor_dma(const Inst &inst, Wavefront &wf, bool store_from_lds) {
ScopedWaitCounter counter(wf, WaitCounterType::TENSORCNT);
const auto desc = read_descriptor(inst, wf);
if (!desc.active())
return;
copy_tensor(desc, wf, store_from_lds);
if (desc.atomic_barrier)
arrive_atomic_barrier(desc, wf);
}

} // namespace tensor_dma_detail

template <typename Inst> void execute_tensor_load_to_lds(const Inst &inst, Wavefront &wf) {
tensor_dma_detail::ScopedWaitCounter counter(wf, WaitCounterType::TENSORCNT);
const auto desc = tensor_dma_detail::read_descriptor(inst, wf);
tensor_dma_detail::copy_tensor(desc, wf, false);
if (desc.atomic_barrier)
tensor_dma_detail::arrive_atomic_barrier(desc, wf);
tensor_dma_detail::execute_tensor_dma(inst, wf, false);
}

template <typename Inst> void execute_tensor_store_from_lds(const Inst &inst, Wavefront &wf) {
tensor_dma_detail::ScopedWaitCounter counter(wf, WaitCounterType::TENSORCNT);
const auto desc = tensor_dma_detail::read_descriptor(inst, wf);
tensor_dma_detail::copy_tensor(desc, wf, true);
if (desc.atomic_barrier)
tensor_dma_detail::arrive_atomic_barrier(desc, wf);
tensor_dma_detail::execute_tensor_dma(inst, wf, true);
}

} // namespace amdgpu
Expand Down
87 changes: 87 additions & 0 deletions emulation/rocjitsu/tests/cdna5_tensor_dma_test.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -352,6 +352,93 @@ TEST(Gfx1250ExecutionTest, TensorDmaIterateCopiesMultipleTiles) {
EXPECT_EQ(cu->lds().read32(wf->lds_base() + 7 * 4), kSentinel);
}

TEST(Gfx1250ExecutionTest, TensorDmaZeroCountDisablesLoadStoreAndBarrier) {
Gfx1250Sim sim;
auto *cu = sim.cu();
auto *wf = cu->dispatch_wf(0, 0, kGfx1250ScalarSlots, 32);
ASSERT_NE(wf, nullptr);
wf->set_lds_base(cu->allocate_lds(256));

constexpr uint64_t kLoadGlobal = 0x133000;
constexpr uint64_t kStoreGlobal = 0x134000;
constexpr uint32_t kElements = 4;
constexpr uint32_t kBarrierLdsAddr = 64;
constexpr uint32_t kLdsSentinel = 0xC0DEC0DEu;
constexpr uint32_t kGlobalSentinel = 0xA11CA11Cu;

write_tensor_dma_d0(*cu, *wf, 0, kLoadGlobal);
write_wave_sgpr(*cu, *wf, 0, 0); // Count zero disables this descriptor.
write_wave_sgpr(*cu, *wf, 12, (2u << 16) | (1u << 18)); // i32, atomic barrier enabled.
write_wave_sgpr(*cu, *wf, 13, (kElements << 16) | (kBarrierLdsAddr >> 3));
write_wave_sgpr(*cu, *wf, 14, 0);
write_wave_sgpr(*cu, *wf, 15, kElements << 16);
write_wave_sgpr(*cu, *wf, 16, 0);
write_wave_sgpr(*cu, *wf, 17, 0);
write_wave_sgpr(*cu, *wf, 18, 0);
write_wave_sgpr(*cu, *wf, 19, 0);

for (uint32_t i = 0; i < kElements; ++i) {
write_global_u32(*sim.memory, kLoadGlobal + i * sizeof(uint32_t), 0x47000000u + i);
write_global_u32(*sim.memory, kStoreGlobal + i * sizeof(uint32_t), kGlobalSentinel);
cu->lds().write32(wf->lds_base() + i * sizeof(uint32_t), kLdsSentinel);
}
cu->lds().write64(wf->lds_base() + kBarrierLdsAddr, 0);

const std::array<uint32_t, 3> load_words = {0xd0710001u, 0x7c000000u, 0x7c7c0c00u};
auto load = decode_gfx1250(load_words, "tensor_load_to_lds");
ASSERT_NE(load, nullptr);
load->execute(*load, wf);

for (uint32_t i = 0; i < kElements; ++i)
EXPECT_EQ(cu->lds().read32(wf->lds_base() + i * sizeof(uint32_t)), kLdsSentinel);
EXPECT_EQ(cu->lds().read64(wf->lds_base() + kBarrierLdsAddr), 0u);
EXPECT_TRUE(wf->wait_counters().empty());

for (uint32_t i = 0; i < kElements; ++i)
cu->lds().write32(wf->lds_base() + i * sizeof(uint32_t), 0x48000000u + i);
write_tensor_dma_d0(*cu, *wf, 0, kStoreGlobal);
write_wave_sgpr(*cu, *wf, 0, 0); // Keep the store descriptor disabled.

const std::array<uint32_t, 3> store_words = {0xd0714001u, 0x7c000000u, 0x7c7c0c00u};
auto store = decode_gfx1250(store_words, "tensor_store_from_lds");
ASSERT_NE(store, nullptr);
store->execute(*store, wf);

for (uint32_t i = 0; i < kElements; ++i)
EXPECT_EQ(read_global_u32(*sim.memory, kStoreGlobal + i * sizeof(uint32_t)), kGlobalSentinel);
EXPECT_EQ(cu->lds().read64(wf->lds_base() + kBarrierLdsAddr), 0u);
EXPECT_TRUE(wf->wait_counters().empty());
}

TEST(Gfx1250ExecutionTest, TensorDmaUnsupportedCountEncodingsAreRejected) {
Gfx1250Sim sim;
auto *cu = sim.cu();
auto *wf = cu->dispatch_wf(0, 0, kGfx1250ScalarSlots, 32);
ASSERT_NE(wf, nullptr);
wf->set_lds_base(cu->allocate_lds(256));

write_tensor_dma_d0(*cu, *wf, 0, 0x138000);
write_wave_sgpr(*cu, *wf, 12, 2u << 16); // i32.
write_wave_sgpr(*cu, *wf, 13, 1u << 16);
write_wave_sgpr(*cu, *wf, 14, 0);
write_wave_sgpr(*cu, *wf, 15, 1u << 16);
write_wave_sgpr(*cu, *wf, 16, 0);
write_wave_sgpr(*cu, *wf, 17, 0);
write_wave_sgpr(*cu, *wf, 18, 0);
write_wave_sgpr(*cu, *wf, 19, 0);

const std::array<uint32_t, 3> load_words = {0xd0710001u, 0x7c000000u, 0x7c7c0c00u};
auto load = decode_gfx1250(load_words, "tensor_load_to_lds");
ASSERT_NE(load, nullptr);

for (uint32_t count : {2u, 3u}) {
SCOPED_TRACE("count=" + std::to_string(count));
write_wave_sgpr(*cu, *wf, 0, count);
EXPECT_THROW(load->execute(*load, wf), util::UnimplementedInst);
EXPECT_TRUE(wf->wait_counters().empty());
}
}

TEST(Gfx1250ExecutionTest, TensorDmaAtomicBarrierArrivesAfterCopy) {
Gfx1250Sim sim;
auto *cu = sim.cu();
Expand Down
Loading