derek/TheRock
1
0 Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity
Files
main
TheRock/rocm-libraries/shared/rocroller/test/catch/TensorDataMoverTest.cpp
T
derek.holloway 7502bf64a4
CodeQL / Analyze (actions) (push) Canceled after 0s
Details
CodeQL / Analyze (python) (push) Canceled after 0s
Details
pre-commit / pre-commit (push) Canceled after 0s
Details
Multi-Arch CI / setup (push) Canceled after 0s
Details
Publish build_manylinux_x86_64 images / publish_build_manylinux_x86_64 (push) Canceled after 0s
Details
Setup (Multi-Arch) / setup (push) Canceled after 0s
Details
Multi-Arch CI / Linux::${{ fromJSON(needs.setup.outputs.linux_build_config || '{}').build_variant_label || 'skip' }} (push) Canceled after 0s
Details
Publish TheRock Docker image / build-and-push-image (push) Canceled after 0s
Details
Publish no_rocm_image_ubi10 image / publish_no_rocm_image_ubi10 (push) Canceled after 0s
Details
Multi-Arch CI / Windows::${{ fromJSON(needs.setup.outputs.windows_build_config || '{}').build_variant_label || 'skip' }} (push) Canceled after 0s
Details
Publish TheRock Docker image / build-and-push-image (push) Canceled after 0s
Details
Multi-Arch CI / Manifest Diff (push) Canceled after 0s
Details
TheRock Manifest Diff Report / Generate Manifest Diff Report (push) Canceled after 0s
Details
Multi-Arch CI / CI Summary (push) Canceled after 0s
Details
Publish no_rocm_image_ubuntu24_04 images / publish_no_rocm_image_ubuntu24_04 (push) Canceled after 0s
Details
Publish no_rocm_image_ubuntu24_04_ocl_rt images / publish_no_rocm_image_ubuntu24_04_ocl_rt (push) Canceled after 0s
Details
Publish TheRock Docker image / build-and-push-image (push) Canceled after 0s
Details
Publish TheRock Docker image / build-and-push-image (push) Canceled after 0s
Details
Publish no_rocm_image_ubuntu24_04 images / publish_no_rocm_image_ubuntu24_04_rocgdb (push) Canceled after 0s
Details
Publish no_rocm_image_ubuntu24_04_rocgdb images / publish_no_rocm_image_ubuntu24_04_rocgdb (push) Canceled after 0s
Details
Publish no_rocm_image_ubuntu24_04 images / publish_no_rocm_image_ubuntu24_04_ocl_rt (push) Canceled after 0s
Details
Publish TheRock Docker image / build-and-push-image (push) Canceled after 0s
Details
Unit Tests / Unit Tests :: ubuntu-24.04 (push) Canceled after 0s
Details
Unit Tests / Unit Tests :: windows-2022 (push) Canceled after 0s
Details
Publish no_rocm_image_ubuntu24_04 images / publish_no_rocm_image_ubuntu24_04_media (push) Canceled after 0s
Details
Unit Tests / Unit Tests Summary (push) Canceled after 0s
Details
Publish no_rocm_image_ubuntu24_04 images / publish_no_rocm_image_ubuntu24_04_openmpi (push) Canceled after 0s
Details
Flatten subprojects into a single base monorepo
2026-09-02 16:34:47 -07:00

553 lines
21 KiB
C++
Raw Permalink Blame History

// Copyright Advanced Micro Devices, Inc., or its affiliates.
// SPDX-License-Identifier: MIT
#include <rocRoller/CodeGen/ArgumentLoader.hpp>
#include <rocRoller/CodeGen/MemoryInstructions.hpp>
#include <rocRoller/CodeGen/TensorDataMover.hpp>
#include <rocRoller/Utilities/Error.hpp>
#include <catch2/catch_template_test_macros.hpp>
#include <catch2/catch_test_macros.hpp>
#include <catch2/generators/catch_generators.hpp>
#include <common/SourceMatcher.hpp>
#include "CustomMatchers.hpp"
#include "TestContext.hpp"
#include "TestKernels.hpp"
#include "rocRoller/Expression.hpp"
using namespace rocRoller;
// using namespace TensorDataMover;
namespace TensorDataMoverTest
{
static auto L(const auto& x)
{
return Expression::literal(x);
}
static auto L(const auto& x, DataType dt)
{
return Expression::literal(x, dt);
};
TEST_CASE("Packing and codegen of TensorDataMover descriptors", "[codegen][utility]")
{
SECTION("A descriptor with all bitfields as literals should be fully packed")
{
auto context = TestContext::ForDefaultTarget();
auto ctx = context.get();
auto kb = [&]() -> Generator<Instruction> {
auto tdmRegs = Register::Value::Placeholder(
ctx, Register::Type::Scalar, {DataType::None, PointerType::TDM}, 1);
tdmRegs->allocateNow();
auto tdmExpr = L(TDM{});
tdmExpr = TDMDescriptor::SetDefaults(tdmExpr, ctx);
tdmExpr = TDMDescriptor::SetLDSAddress(tdmExpr, L(0x8BADF00Du, DataType::UInt32));
tdmExpr = TDMDescriptor::SetGlobalAddress(
tdmExpr, L(0x01600DBEEFDEADBDull, DataType::UInt64));
co_yield Expression::generate(tdmRegs, tdmExpr, ctx);
};
ctx->schedule(kb());
auto expectedCode = R"(
s_mov_b32 s0, 1
s_mov_b32 s1, 2343432205
s_mov_b32 s2, 4024348093
s_mov_b32 s3, 2170555838
s_mov_b32 s4, 0
s_mov_b32 s5, 0
s_mov_b32 s6, 0
s_mov_b32 s7, 0
s_mov_b32 s8, 0
s_mov_b32 s9, 0
s_mov_b32 s10, 0
s_mov_b32 s11, 0
)";
CHECK(NormalizedSource(context.output()) == NormalizedSource(expectedCode));
}
SECTION("Can set fields that cross sgpr boundaries with literals")
{
auto context = TestContext::ForDefaultTarget();
auto ctx = context.get();
auto kb = [&]() -> Generator<Instruction> {
auto tdmRegs = Register::Value::Placeholder(
ctx, Register::Type::Scalar, {DataType::None, PointerType::TDM}, 1);
tdmRegs->allocateNow();
auto tdmExpr = L(TDM{});
tdmExpr = TDMDescriptor::SetDefaults(tdmExpr, ctx);
tdmExpr = TDMDescriptor::SetTensorDims(
tdmExpr, L(0x8BADF00D, DataType::UInt32), L(0x8BADBEEF, DataType::UInt32));
tdmExpr = TDMDescriptor::SetTileDims(tdmExpr,
L(0x8BAD, DataType::UInt32),
L(0xDECA, DataType::UInt32),
L(0xBEEF, DataType::UInt32));
co_yield Expression::generate(tdmRegs, tdmExpr, ctx);
};
ctx->schedule(kb());
auto expectedCode = R"(
s_mov_b32 s0, 1
s_mov_b32 s1, 0
s_mov_b32 s2, 0
s_mov_b32 s3, 2147483648
s_mov_b32 s4, 0
s_mov_b32 s5, 4027383808
s_mov_b32 s6, 3203369901
s_mov_b32 s7, 2343406509
s_mov_b32 s8, 3203391178
s_mov_b32 s9, 0
s_mov_b32 s10, 0
s_mov_b32 s11, 0
)";
CHECK(NormalizedSource(context.output()) == NormalizedSource(expectedCode));
}
SECTION("Update ldsAddress should emit only one additional instruction")
{
auto context = TestContext::ForDefaultTarget();
auto ctx = context.get();
auto kb = [&]() -> Generator<Instruction> {
auto tdmRegs = Register::Value::Placeholder(
ctx, Register::Type::Scalar, {DataType::None, PointerType::TDM}, 1);
auto tdmExpr = L(TDM{});
tdmExpr = TDMDescriptor::SetDefaults(tdmExpr, ctx);
tdmExpr = TDMDescriptor::SetLDSAddress(tdmExpr, L(0x8BADF00D));
co_yield Expression::generate(tdmRegs, tdmExpr, ctx);
tdmExpr = TDMDescriptor::SetLDSAddress(tdmRegs->expression(), L(0xDECAFBAD));
co_yield Expression::generate(tdmRegs, tdmExpr, ctx);
auto ldsAddress = Register::Value::Placeholder(
ctx, Register::Type::Scalar, DataType::UInt32, 1);
ldsAddress->allocateNow();
tdmExpr
= TDMDescriptor::SetLDSAddress(tdmRegs->expression(), ldsAddress->expression());
co_yield Expression::generate(tdmRegs, tdmExpr, ctx);
};
ctx->schedule(kb());
const auto expectedCode = R"(
s_mov_b32 s0, 1
s_mov_b32 s1, 2343432205
s_mov_b32 s2, 0
s_mov_b32 s3, 2147483648
s_mov_b32 s4, 0
s_mov_b32 s5, 0
s_mov_b32 s6, 0
s_mov_b32 s7, 0
s_mov_b32 s8, 0
s_mov_b32 s9, 0
s_mov_b32 s10, 0
s_mov_b32 s11, 0
// update to ldsAddress field with literal
s_mov_b32 s1, 3737844653
// update to ldsAddress field with SGPR
s_mov_b32 s1, s12
)";
CHECK(NormalizedSource(context.output()) == NormalizedSource(expectedCode));
}
SECTION("Update globalAddress (57 bits) without touching reserved bits")
{
auto context = TestContext::ForDefaultTarget();
auto ctx = context.get();
auto kb = [&]() -> Generator<Instruction> {
auto tdmRegs = Register::Value::Placeholder(
ctx, Register::Type::Scalar, {DataType::None, PointerType::TDM}, 1);
auto tdmExpr = L(TDM{});
tdmExpr = TDMDescriptor::SetDefaults(tdmExpr, ctx);
tdmExpr = TDMDescriptor::SetGlobalAddress(tdmExpr, L(0x01600DBEEFDEADBD));
co_yield Expression::generate(tdmRegs, tdmExpr, ctx);
auto globalAddress = Register::Value::Placeholder(
ctx, Register::Type::Scalar, DataType::UInt64, 1);
globalAddress->allocateNow();
tdmExpr = TDMDescriptor::SetGlobalAddress(tdmRegs->expression(),
globalAddress->expression());
co_yield Expression::generate(tdmRegs, tdmExpr, ctx);
};
ctx->schedule(kb());
const auto expectedCode = R"(
s_mov_b32 s0, 1
s_mov_b32 s1, 0
s_mov_b32 s2, 4024348093
s_mov_b32 s3, 2170555838
s_mov_b32 s4, 0
s_mov_b32 s5, 0
s_mov_b32 s6, 0
s_mov_b32 s7, 0
s_mov_b32 s8, 0
s_mov_b32 s9, 0
s_mov_b32 s10, 0
s_mov_b32 s11, 0
// using bfe to avoid touching reserved bits
s_bfe_u32 s14, s13, 1638400
s_and_b32 s15, 33554431, s14
s_and_b32 s14, 4261412864, s3
s_or_b32 s16, s15, s14
s_mov_b32 s2, s12
s_mov_b32 s3, s16
)";
CHECK(NormalizedSource(context.output()) == NormalizedSource(expectedCode));
}
}
class TensorDataMoverKernel : public AssemblyTestKernel
{
public:
TensorDataMoverKernel(ContextPtr context,
DataType datatype,
uint32_t tensorX,
uint32_t tensorY,
uint32_t tensorDim0,
uint32_t tensorDim1,
uint32_t tileDim0,
uint32_t tileDim1)
: AssemblyTestKernel(context)
, m_datatype(datatype)
, m_tensorX(tensorX)
, m_tensorY(tensorY)
, m_tensorDim0(tensorDim0)
, m_tensorDim1(tensorDim1)
, m_tileDim0(tileDim0)
, m_tileDim1(tileDim1)
{
}
void generate() override
{
const uint workitemCountX = 32u;
auto k = m_context->kernel();
auto one = Expression::literal(1);
auto workitemCountXExpr = Expression::literal(workitemCountX);
k->setKernelDimensions(1);
k->setWorkgroupSize({workitemCountX, 1, 1});
k->setWorkitemCount({workitemCountXExpr, one, one});
k->addArgument(
{"outputTile", {m_datatype, PointerType::PointerGlobal}, DataDirection::WriteOnly});
k->addArgument(
{"inputTensor", {m_datatype, PointerType::PointerGlobal}, DataDirection::ReadOnly});
m_context->schedule(k->preamble());
m_context->schedule(k->prolog());
const auto dataSize = [](auto datatype) -> TDMDescriptor::DataSize {
const size_t size = DataTypeInfo::Get(datatype).elementBytes;
switch(size)
{
case 1:
return TDMDescriptor::DataSize::OneByte;
case 2:
return TDMDescriptor::DataSize::TwoBytes;
case 4:
return TDMDescriptor::DataSize::FourBytes;
case 8:
return TDMDescriptor::DataSize::EightBytes;
default:
Throw<FatalError>(fmt::format(
"Invalid datatype {}. TDM does not support moving data with {} bytes.",
toString(datatype),
size));
}
}(m_datatype);
auto kb = [&]() -> Generator<Instruction> {
Register::ValuePtr sOutputTile, sInputTensor;
co_yield m_context->argLoader()->getValue("outputTile", sOutputTile);
co_yield m_context->argLoader()->getValue("inputTensor", sInputTensor);
auto tdmRegs = Register::Value::Placeholder(
m_context, Register::Type::Scalar, {DataType::None, PointerType::TDM}, 1);
// TODO: remove this once assembler bug is fixed.
tdmRegs->allocateNow();
auto sLDSAddress = Register::Value::Placeholder(
m_context, Register::Type::Scalar, DataType::UInt32, 1);
auto lds
= Register::Value::AllocateLDS(m_context, m_datatype, m_tileDim0 * m_tileDim1);
co_yield m_context->copier()->copy(
sLDSAddress, Register::Value::Literal(lds->getLDSAllocation()->offset()));
auto tdmExpr = L(TDM{});
tdmExpr = TDMDescriptor::SetDefaults(tdmExpr, m_context);
tdmExpr = TDMDescriptor::SetLDSAddress(tdmExpr, sLDSAddress->expression());
tdmExpr = TDMDescriptor::SetGlobalAddress(tdmExpr, sInputTensor->expression());
tdmExpr = TDMDescriptor::SetDataSize(tdmExpr, dataSize);
tdmExpr = TDMDescriptor::SetTensorDims(tdmExpr, L(m_tensorDim0), L(m_tensorDim1));
tdmExpr = TDMDescriptor::SetTensorStrides(
tdmExpr, L(m_tensorDim1, DataType::UInt64), L(1, DataType::UInt64));
// tileDim0 is tile's fast-moving dimension
tdmExpr = TDMDescriptor::SetTileDims(tdmExpr, L(m_tileDim1), L(m_tileDim0));
co_yield Expression::generate(tdmRegs, tdmExpr, m_context);
co_yield m_context->mem()->loadTensorToLDS(tdmRegs).map(
MemoryInstructions::addExtraDst(lds));
tdmExpr = TDMDescriptor::SetGlobalAddress(tdmRegs->expression(),
sOutputTile->expression());
// tileDim0 is tile's fast-moving dimension
tdmExpr = TDMDescriptor::SetTensorDims(tdmExpr, L(m_tileDim1), L(m_tileDim0));
tdmExpr = TDMDescriptor::SetTensorStrides(
tdmExpr, L(m_tileDim1, DataType::UInt64), L(1, DataType::UInt64));
co_yield Expression::generate(tdmRegs, tdmExpr, m_context);
co_yield m_context->mem()->storeTensorFromLDS(tdmRegs).map(
MemoryInstructions::addExtraSrc(lds));
};
m_context->schedule(kb());
m_context->schedule(k->postamble());
m_context->schedule(k->amdgpu_metadata());
}
template <typename StorageType>
void run()
{
std::vector<StorageType> inputTensor(m_tensorDim0 * m_tensorDim1);
constexpr StorageType deadValue = []() -> StorageType {
constexpr size_t numBytes = sizeof(StorageType);
switch(numBytes)
{
case 1:
return static_cast<uint8_t>(0xFF);
case 2:
return static_cast<uint16_t>(0xDEAD);
case 4:
return static_cast<uint32_t>(0xDEADBEEF);
case 8:
return static_cast<uint64_t>(0xDEADBEEF8BADF00DULL);
default:
Throw<FatalError>(
fmt::format("TDM does not support moving data with {} bytes.", numBytes));
}
}();
for(size_t i = 0; i < m_tensorDim0; ++i)
{
for(size_t j = 0; j < m_tensorDim1; ++j)
{
if((m_tensorX <= i && i < (m_tensorX + m_tileDim0))
&& (m_tensorY <= j && j < (m_tensorY + m_tileDim1)))
{
inputTensor[i * m_tensorDim1 + j]
= static_cast<StorageType>(i * m_tileDim1 + j);
}
else
{
inputTensor[i * m_tensorDim1 + j] = deadValue;
}
}
}
const auto debugTensor = false;
auto printTensor = [deadValue](auto v, size_t M, size_t N) {
if(!debugTensor)
{
return;
}
for(size_t i = 0; i < M; ++i)
{
for(size_t j = 0; j < N; ++j)
{
const auto x = static_cast<uint32_t>(v[i * N + j]);
if(x == deadValue)
{
std::cout << fmt::format("{:4} ", "xxxx");
}
else
{
std::cout << fmt::format("{:4} ", x);
}
}
std::cout << "\n";
}
};
printTensor(inputTensor, m_tensorDim0, m_tensorDim1);
std::vector<StorageType> outputTile(m_tileDim0 * m_tileDim1);
auto d_outputTile = make_shared_device<StorageType>(outputTile.size());
auto d_inputTensor = make_shared_device(inputTensor);
this->operator()(
KernelInvocation{.workitemCount = {32, 1, 1}, .workgroupSize = {32, 1, 1}},
d_outputTile.get(),
d_inputTensor.get() + (m_tensorX * m_tensorDim1 + m_tensorY));
CHECK_THAT(hipMemcpy(outputTile.data(),
d_outputTile.get(),
sizeof(StorageType) * outputTile.size(),
hipMemcpyDefault),
HasHipSuccess(0));
for(size_t i = 0; i < m_tileDim0; ++i)
{
for(size_t j = 0; j < m_tileDim1; ++j)
{
CHECK(inputTensor[(i + m_tensorX) * m_tensorDim1 + (j + m_tensorY)]
== outputTile[i * m_tileDim1 + j]);
}
}
printTensor(outputTile, m_tileDim0, m_tileDim1);
}
template <typename TestType>
static void testBody(bool run = false);
private:
DataType m_datatype;
uint32_t m_tensorX, m_tensorY;
uint32_t m_tensorDim0, m_tensorDim1;
uint32_t m_tileDim0, m_tileDim1;
};
template <typename TestType>
void TensorDataMoverKernel::testBody(bool runOnGPU)
{
constexpr auto datatype = [numBytes = sizeof(TestType)]() {
switch(numBytes)
{
case 1:
return DataType::UInt8;
case 2:
return DataType::UInt16;
case 4:
return DataType::UInt32;
case 8:
return DataType::UInt64;
default:
Throw<FatalError>(
fmt::format("Test does not support datatypes with {} bytes.", numBytes));
}
}();
// clang-format off
auto tdmParams = GENERATE(
std::tuple( 0, 0, 32, 32, 16, 16),
std::tuple( 0, 16, 32, 32, 16, 16),
std::tuple(16, 0, 32, 32, 16, 16),
std::tuple(16, 16, 32, 32, 16, 16),
std::tuple( 0, 0, 512, 512, 16, 128),
std::tuple(128, 0, 512, 512, 128, 16),
std::tuple( 0, 256, 512, 512, 32, 256),
std::tuple(256, 0, 512, 512, 256, 32),
std::tuple( 0, 0, 1024, 512, 32, 256),
std::tuple( 0, 256, 1024, 512, 256, 32),
std::tuple(512, 256, 1024, 512, 32, 256),
std::tuple(512, 256, 1024, 512, 256, 32)
);
auto kernelTarget = GENERATE(GPUArchTargetGFX1250Rev0, GPUArchTargetGFX1250Rev1);
// clang-format on
DYNAMIC_SECTION(
"Test the following params: " << fmt::format("{} {}", toString(datatype), tdmParams))
{
const auto [tensorX, tensorY, tensorDim0, tensorDim1, tileDim0, tileDim1] = tdmParams;
AssertFatal(tensorX < tensorDim0 && tensorY < tensorDim1 && tileDim0 <= tensorDim0
&& tileDim1 <= tensorDim1,
"Invalid params.",
ShowValue(datatype),
ShowValue(tensorX),
ShowValue(tensorY),
ShowValue(tensorDim0),
ShowValue(tensorDim1),
ShowValue(tileDim0),
ShowValue(tileDim1));
auto context = TestContext::ForTarget(kernelTarget,
KernelOptions{},
datatype,
tensorX,
tensorY,
tensorDim0,
tensorDim1,
tileDim0,
tileDim1);
auto ctx = context.get();
TensorDataMoverKernel tdmKernel(
ctx, datatype, tensorX, tensorY, tensorDim0, tensorDim1, tileDim0, tileDim1);
if(runOnGPU)
{
auto currDeviceCtx = TestContext::ForTestDevice();
const auto& currDeviceArch = currDeviceCtx->targetArchitecture();
if(!currDeviceArch.HasCapability(GPUCapability::HasTDM))
{
SKIP(fmt::format("Target {} does not support TDM",
toString(currDeviceArch.target())));
}
if(kernelTarget != currDeviceCtx->targetArchitecture().target())
{
SKIP(fmt::format("Cannot run kernel for {} on {} device",
toString(kernelTarget),
toString(currDeviceArch.target())));
}
tdmKernel.run<TestType>();
}
else
{
tdmKernel.generate();
}
const auto code = context.output();
CHECK(countSubstring(code, "tensor_load_to_lds ") == 1);
CHECK(countSubstring(code, "tensor_store_from_lds ") == 1);
}
}
TEMPLATE_TEST_CASE("Run TDM kernels that copy tiles of a tensor",
"[memory-instructions][gpu]",
uint64_t,
uint32_t,
uint16_t,
uint8_t)
{
TensorDataMoverKernel::testBody<TestType>(/*runOnGPU*/ true);
}
TEMPLATE_TEST_CASE("Assemble TDM kernels that copy tiles of a tensor",
"[memory-instructions][codegen]",
uint64_t,
uint32_t,
uint16_t,
uint8_t)
{
TensorDataMoverKernel::testBody<TestType>();
}
}
Reference in new issue View Git Blame Copy Permalink