ROCm/composable_kernel
1
0
Fork 0
mirror of https://github.com/ROCm/composable_kernel.git synced 2026-07-14 19:18:35 +00:00
Code Issues Packages Projects Releases Wiki Activity

[rocm-libraries] ROCm/rocm-libraries#7179 (commit 05edc86)

Browse Source 
[CK] Add rocm_ck schema engine: Signature, resolve(),
 ArchProperties (#7179)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

## Summary

A `Signature` is a directed compute graph: tensors are nodes, operators
are edges. Shared names between operator outputs and inputs form the
graph structure. `resolve()` walks this graph at compile time
(`consteval`), inferring dtype, rank, and layout for every tensor —
invalid configs become compiler errors, not runtime crashes.

**Key design decisions:**

- **Operators teach the system about tensors.** `GemmOp` implies rank 2
and Row/Col/Row layout. `AddOp` and `ReluOp` propagate from connected
slots. The dtype cascade fills in the rest: per-tensor → signature-wide
→ error.

- **Adding a new op is zero lines in the resolution engine** if it's
structurally binary (`lhs/rhs/out`) or unary (`in/out`) — C++20 concepts
handle dispatch automatically. Only ops with special semantics need
explicit branches.

- **TargetSet is a compile-time bitset over GPU targets.** The wave tile
validation table is the single source of truth for valid instruction
shapes, traced from CK Tile's WarpGemmDispatcher. FP8 tiles are
available on gfx942+ via IterateK composition, not gfx950-only.

**Reading order:** signature.hpp (data model) → arch_properties.hpp
(TargetSet, wave tiles) → resolve.hpp (resolution engine).

3 new headers, 3 unit tests (including diamond DAG coverage), 3
compile-fail tests. Introduces tests/compile_fail/ infrastructure.

**Stack**: PR 2 of 3 porting the rocm_ck constexpr schema from
experimental to production.
1. Foundation types — DataType, Layout, Args, Ops (#7114)
2. **This PR** — Schema engine (graph resolution)
3. Spec factories — GemmSpec, makeSpec() (#7180 )

Note: We also removed `FmhaBwdOp` for clarity, since that was introduced
early and doesn't have tests set up.

**Depends on**: #7114

## Test plan

- [x] ctest --test-dir build --output-on-failure — unit tests +
compile-fail tests pass
- [x] Compile-fail tests correctly reject: mixed CDNA+RDNA TargetSet,
conflicting layouts, empty quantization scale names
This commit is contained in:
John Shumway
2026-05-22 16:44:41 +00:00
committed by assistant-librarian[bot]
parent 5a9e1d757d
commit a1dcf5443c
22 changed files with 2428 additions and 24 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
rocm_ck/README.md
View File

@@ -91,6 +91,7 @@ rocm_ck/
└── tests/
├── CMakeLists.txt # Test tiers: ROCM_CK_SMOKE, ROCM_CK_KERNEL
├── unit/ # Fast host-only tests (< 1s, no GPU)
├── compile_fail/ # Static assertion tests — verify invalid configs fail at compile time
└── kernel/ # (planned) GPU kernel tests
```

388
rocm_ck/include/rocm_ck/arch_properties.hpp Normal file
View File

@@ -0,0 +1,388 @@
// Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
// SPDX-License-Identifier: MIT
//
// Role: types — GPU target properties and target set (consteval bitset).
//
// UPSTREAM CANDIDATE: This header prototypes functionality that should
// eventually live in ck_tile/core/arch/arch_properties.hpp as the single
// source of truth for GPU target metadata. It is intentionally free of
// device code, HIP runtime, and CK Tile dependencies so it can be used
// in both host-side consteval validation and device-side template selection.
//
// When upstreaming, the valid tile table should be generated from or
// verified against CK Tile's MMA selector specializations
// (core/arch/mma/mfma/mfma_gfx9.hpp, wmma/wmma_gfx11.hpp, etc.).
//
// No runtime code, no CK deps. Pure C++20 constexpr/consteval.
#pragma once
#include <rocm_ck/datatype.hpp>
#include <rocm_ck/gpu_target.hpp>
#include <bit>
#include <cstdint>
namespace rocm_ck {
// ============================================================================
// Per-target properties
// ============================================================================
/// ISA architecture family.
///
/// CDNA = data-center compute (Instinct GPUs), MFMA instructions, wave64.
/// RDNA = gaming/workstation (Radeon GPUs), WMMA instructions, wave32.
enum class ArchFamily
{
CDNA,
RDNA
};
/// Consteval properties for a single GPU target.
struct TargetProperties
{
int wavefront_size;
ArchFamily arch_family;
};
/// Single source of truth for per-target metadata.
constexpr TargetProperties properties(GpuTarget target)
{
switch(target)
{
case GpuTarget::gfx90a: return {.wavefront_size = 64, .arch_family = ArchFamily::CDNA};
case GpuTarget::gfx942: return {.wavefront_size = 64, .arch_family = ArchFamily::CDNA};
case GpuTarget::gfx950: return {.wavefront_size = 64, .arch_family = ArchFamily::CDNA};
case GpuTarget::gfx1100: return {.wavefront_size = 32, .arch_family = ArchFamily::RDNA};
case GpuTarget::gfx1101: return {.wavefront_size = 32, .arch_family = ArchFamily::RDNA};
case GpuTarget::gfx1102: return {.wavefront_size = 32, .arch_family = ArchFamily::RDNA};
case GpuTarget::gfx1150: return {.wavefront_size = 32, .arch_family = ArchFamily::RDNA};
case GpuTarget::gfx1151: return {.wavefront_size = 32, .arch_family = ArchFamily::RDNA};
case GpuTarget::_count:
default: throw "unsupported GpuTarget — add a case to properties() for new targets";
}
}
/// True if the target uses CDNA architecture (MFMA instructions, wave64).
constexpr bool isCDNA(GpuTarget target)
{
return properties(target).arch_family == ArchFamily::CDNA;
}
/// True if the target uses RDNA architecture (WMMA instructions, wave32).
constexpr bool isRDNA(GpuTarget target)
{
return properties(target).arch_family == ArchFamily::RDNA;
}
/// Wavefront size for a given GPU target.
constexpr int wavefrontSize(GpuTarget target) { return properties(target).wavefront_size; }
// ============================================================================
// TargetSet — consteval bitset over GpuTarget values
// ============================================================================
/// Compile-time set of GPU targets. Structural type (usable as NTTP).
///
/// Storage: each GpuTarget maps to a bit in a uint64_t.
///
/// Named constructors express CK Tile's 3-level hierarchy:
/// Architecture: TargetSet::cdna(), TargetSet::rdna(), TargetSet::all()
/// Family: TargetSet::family_gfx9(), family_gfx94(), family_gfx11(), family_gfx115()
/// Specific: TargetSet::only(GpuTarget::gfx942)
///
/// CK Tile mapping:
/// TargetSet::cdna() → enable_if_target_arch_cdna_t
/// TargetSet::rdna() → enable_if_target_arch_rdna_t
/// TargetSet::family_gfx9() → enable_if_target_family_gfx9_t
/// TargetSet::family_gfx11() → __gfx11__ preprocessor grouping
/// TargetSet::family_gfx115() → __gfx115__ preprocessor grouping
/// TargetSet::only(gfx942, gfx950) → enable_if_target_id_t<T, GFX942, GFX950>
/// TargetSet::cdna().excluding(gfx90a) → is_any_value_of(T::TARGET_ID, GFX942, GFX950)
struct TargetSet
{
uint64_t bits = 0;
// ---- Bit index mapping ------------------------------------------------
// Bit positions equal GpuTarget enum values. Adding a target:
// 1. Add enum value to GpuTarget before _count
// 2. Add to architecture/family named constructors
static constexpr int kNumTargets = static_cast<int>(GpuTarget::_count);
static constexpr int bitIndex(GpuTarget target)
{
if(target >= GpuTarget::_count)
throw "GpuTarget out of range — value must be a valid enum member, not _count";
return static_cast<int>(target);
}
static constexpr GpuTarget targetAt(int index)
{
if(index < 0 || index >= kNumTargets)
throw "TargetSet index out of range [0, kNumTargets)";
return static_cast<GpuTarget>(index);
}
// ---- Constructors -----------------------------------------------------
/// Default: empty set.
constexpr TargetSet() = default;
/// Implicit conversion from a single GpuTarget.
constexpr TargetSet(GpuTarget target) : bits(uint64_t{1} << bitIndex(target)) {}
// ---- Named constructors: architecture level ---------------------------
/// All real GPU targets.
static constexpr TargetSet all() { return fromBits((uint64_t{1} << kNumTargets) - 1); }
/// All CDNA targets (MFMA, wave64): gfx90a, gfx942, gfx950.
static constexpr TargetSet cdna()
{
return only(GpuTarget::gfx90a, GpuTarget::gfx942, GpuTarget::gfx950);
}
/// All RDNA targets (WMMA, wave32): gfx1100, gfx1101, gfx1102, gfx1150, gfx1151.
static constexpr TargetSet rdna()
{
return only(GpuTarget::gfx1100, GpuTarget::gfx1101, GpuTarget::gfx1102)
.union_with(only(GpuTarget::gfx1150, GpuTarget::gfx1151));
}
// ---- Named constructors: family level ---------------------------------
// Matches CK Tile's __gfx9__, __gfx94__, __gfx11__ groupings.
/// GFX9 family (CDNA): gfx90a, gfx942, gfx950.
static constexpr TargetSet family_gfx9()
{
return only(GpuTarget::gfx90a, GpuTarget::gfx942, GpuTarget::gfx950);
}
/// GFX94 subfamily (CDNA 3+): gfx942, gfx950.
static constexpr TargetSet family_gfx94() { return only(GpuTarget::gfx942, GpuTarget::gfx950); }
/// GFX11 family (RDNA 3/3.5): gfx1100, gfx1101, gfx1102, gfx1150, gfx1151.
/// Matches CK Tile's __gfx11__ preprocessor grouping.
static constexpr TargetSet family_gfx11() { return rdna(); }
/// GFX115 subfamily (RDNA 3.5): gfx1150, gfx1151.
/// Matches CK Tile's __gfx115__ preprocessor grouping.
static constexpr TargetSet family_gfx115()
{
return only(GpuTarget::gfx1150, GpuTarget::gfx1151);
}
// ---- Named constructors: specific targets -----------------------------
/// Set containing exactly one target.
static constexpr TargetSet only(GpuTarget t) { return fromBits(uint64_t{1} << bitIndex(t)); }
/// Set containing exactly two targets.
static constexpr TargetSet only(GpuTarget a, GpuTarget b)
{
return only(a).union_with(only(b));
}
/// Set containing exactly three targets.
static constexpr TargetSet only(GpuTarget a, GpuTarget b, GpuTarget c)
{
return only(a).union_with(only(b)).union_with(only(c));
}
// ---- Set operations ---------------------------------------------------
/// Remove one target from the set.
constexpr TargetSet excluding(GpuTarget t) const
{
return fromBits(bits & ~(uint64_t{1} << bitIndex(t)));
}
/// Union: targets in either set.
constexpr TargetSet union_with(TargetSet other) const { return fromBits(bits | other.bits); }
/// Intersection: targets in both sets.
constexpr TargetSet intersect_with(TargetSet other) const
{
return fromBits(bits & other.bits);
}
/// Difference: targets in this set but not in other.
constexpr TargetSet minus(TargetSet other) const { return fromBits(bits & ~other.bits); }
// ---- Operators (delegate to named methods) ----------------------------
friend constexpr TargetSet operator|(TargetSet a, TargetSet b) { return a.union_with(b); }
friend constexpr TargetSet operator&(TargetSet a, TargetSet b) { return a.intersect_with(b); }
friend constexpr TargetSet operator-(TargetSet a, TargetSet b) { return a.minus(b); }
friend constexpr bool operator==(TargetSet a, TargetSet b) { return a.bits == b.bits; }
// ---- Queries ----------------------------------------------------------
/// True if target is in the set.
constexpr bool contains(GpuTarget t) const
{
return (bits & (uint64_t{1} << bitIndex(t))) != 0;
}
/// True if the set is empty.
constexpr bool is_empty() const { return bits == 0; }
/// True if the set contains exactly one target.
constexpr bool is_single_target() const { return bits != 0 && (bits & (bits - 1)) == 0; }
/// Number of targets in the set.
constexpr int count() const { return std::popcount(bits); }
/// Wavefront size, if uniform across all targets in the set.
/// Compile error if targets disagree or set is empty.
constexpr int wavefront_size() const
{
if(is_empty())
throw "wavefront_size() called on empty TargetSet";
int wf = -1;
for(int i = 0; i < kNumTargets; ++i)
{
if((bits & (uint64_t{1} << i)) == 0)
continue;
int target_wf = wavefrontSize(targetAt(i));
if(wf == -1)
wf = target_wf;
else if(wf != target_wf)
throw "wavefront_size() requires all targets in the set to have "
"the same wavefront size — this set mixes wave64 (CDNA) and "
"wave32 (RDNA) targets. Split with intersect_with(cdna()) or "
"intersect_with(rdna()).";
}
return wf;
}
/// True if all targets in the set are CDNA.
constexpr bool is_all_cdna() const { return !is_empty() && minus(cdna()).is_empty(); }
/// True if all targets in the set are RDNA.
constexpr bool is_all_rdna() const { return !is_empty() && minus(rdna()).is_empty(); }
/// Get the single target (compile error if set has != 1 target).
constexpr GpuTarget single_target() const
{
if(!is_single_target())
throw "single_target() requires exactly one target in the set";
return targetAt(std::countr_zero(bits));
}
// ---- Iteration --------------------------------------------------------
/// Call fn(GpuTarget) for each target in the set.
template <typename Fn>
constexpr void for_each(Fn fn) const
{
for(int i = 0; i < kNumTargets; ++i)
if(bits & (uint64_t{1} << i))
fn(targetAt(i));
}
private:
explicit constexpr TargetSet(uint64_t b) : bits(b) {}
static constexpr TargetSet fromBits(uint64_t b) { return TargetSet{b}; }
};
// ============================================================================
// Wave tile validation — single source of truth
// ============================================================================
// Based on CK Tile's WarpGemmDispatcher specializations.
// See: ck_tile/core/arch/mma/mfma/mfma_gfx9.hpp (MFMA builtins)
// ck_tile/core/arch/mma/wmma/wmma_gfx11.hpp (WMMA builtins)
// ck_tile/ops/gemm/warp/warp_gemm_dispatcher.hpp (dispatch table)
/// Check if (a_dtype, m, n, k) maps to a valid wave instruction shape
/// on a specific target.
consteval bool isValidWaveTile(DataType a_dtype, int m, int n, int k, GpuTarget target)
{
// RDNA targets: WMMA — fixed 16x16x16 tile shape
if(isRDNA(target))
{
if(m != 16 || n != 16 || k != 16)
return false;
// RDNA (gfx11xx) WMMA: fp16, bf16, int8 — all targets share 16×16×16 tile
return a_dtype == DataType::FP16 || a_dtype == DataType::BF16 || a_dtype == DataType::I8;
}
// CDNA MFMA tiles — common across gfx90a, gfx942, gfx950
if(a_dtype == DataType::FP32)
{
if(m == 16 && n == 16 && (k == 4 || k == 8 || k == 16))
return true;
if(m == 32 && n == 32 && (k == 4 || k == 8))
return true;
}
if(a_dtype == DataType::FP16)
{
if(m == 16 && n == 16 && (k == 16 || k == 32))
return true;
if(m == 32 && n == 32 && (k == 8 || k == 16))
return true;
}
if(a_dtype == DataType::BF16)
{
if(m == 16 && n == 16 && (k == 16 || k == 32))
return true;
if(m == 32 && n == 32 && (k == 8 || k == 16))
return true;
}
// INT8 MFMA — int8x int8→int32 accumulation
if(a_dtype == DataType::I8)
{
if(m == 32 && n == 32 && k == 16)
return true;
if(m == 16 && n == 16 && k == 32)
return true;
}
// FP8/BF8 MFMA — architecture-dependent
if(a_dtype == DataType::FP8_FNUZ || a_dtype == DataType::BF8_FNUZ)
{
// gfx90a: no FP8 MFMA support
if(target == GpuTarget::gfx90a)
return false;
// gfx942+: base FP8 MFMA tiles
if(m == 32 && n == 32 && k == 16)
return true;
if(m == 16 && n == 16 && k == 32)
return true;
// gfx942+: IterateK compositions of base FP8 MFMA
if(m == 32 && n == 32 && (k == 32 || k == 64))
return true;
if(m == 16 && n == 16 && k == 64)
return true;
}
// FP8_OCP/BF8_OCP — not yet supported
if(a_dtype == DataType::FP8_OCP || a_dtype == DataType::BF8_OCP)
throw "FP8_OCP/BF8_OCP not yet supported in GEMM — use FP8_FNUZ/BF8_FNUZ";
return false;
}
/// Check if a tile is valid for ALL targets in a set (intersection semantics).
consteval bool isValidWaveTile(DataType a_dtype, int m, int n, int k, TargetSet targets)
{
if(targets.is_empty())
throw "isValidWaveTile called with empty TargetSet";
for(int i = 0; i < TargetSet::kNumTargets; ++i)
{
if(!(targets.bits & (uint64_t{1} << i)))
continue;
if(!isValidWaveTile(a_dtype, m, n, k, TargetSet::targetAt(i)))
return false;
}
return true;
}
} // namespace rocm_ck

1
rocm_ck/include/rocm_ck/gpu_target.hpp
View File

@@ -19,6 +19,7 @@ enum class GpuTarget : uint8_t
gfx1102, // RDNA 3
gfx1150, // RDNA 3.5
gfx1151, // RDNA 3.5
_count // must be last — new targets go above this line
};
} // namespace rocm_ck

22
rocm_ck/include/rocm_ck/ops.hpp
View File

@@ -103,25 +103,6 @@ struct ScaleOp
std::string_view scale;
};
// Fused multi-head attention backward pass.
// Implemented as a single CK Tile kernel, not a chain of ops.
// Feature flags (mask, dropout, bias, deterministic) belong in the Algorithm.
struct FmhaBwdOp
{
std::string_view q; // query
std::string_view k; // key
std::string_view v; // value
std::string_view lse; // log-sum-exp from forward pass
std::string_view do_; // output gradient
std::string_view d; // dot(output_grad, output)
std::string_view dq; // query gradient
std::string_view dk; // key gradient
std::string_view dv; // value gradient
DataType acc_dtype = DataType::FP32;
};
// The closed set of supported operators. std::monostate marks empty slots.
using Op = std::variant<std::monostate,
GemmOp,
@@ -133,7 +114,6 @@ using Op = std::variant<std::monostate,
SiluOp,
SigmoidOp,
SoftmaxOp,
ScaleOp,
FmhaBwdOp>;
ScaleOp>;
} // namespace rocm_ck

508
rocm_ck/include/rocm_ck/resolve.hpp Normal file
View File

@@ -0,0 +1,508 @@
// Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
// SPDX-License-Identifier: MIT
//
// Role: meta — consteval resolve(), C++20 concepts. No runtime, no CK deps.
//
// Signature resolution: resolves a Signature into concrete tensor descriptors.
//
// resolve() walks the operator graph, collects tensor slots with operator-implied
// defaults, propagates rank/layout through connected tensors, merges explicit
// Tensor entries, and applies the dtype cascade. All at compile time (consteval).
//
// Op dispatch uses C++20 concepts to classify ops by structural shape:
// BinaryOpLike — has {lhs, rhs, out} string_view members
// UnaryOpLike — has {in, out} string_view members
// A single visitOp() function is the only place the Op variant type list
// appears. Adding a new op requires one line in visitOp(); concepts handle
// generic registration and propagation automatically.
//
// This header has NO CK Tile dependency.
#pragma once
#include <rocm_ck/signature.hpp>
#include <rocm_ck/resolved_tensor.hpp>
#include <array>
#include <type_traits>
namespace rocm_ck {
// ============================================================================
// Op structural concepts — classify ops by their tensor slot shape
// ============================================================================
/// Ops with three tensor slots: lhs, rhs, out (e.g., AddOp, MulOp).
/// Excludes ops that also have {in} to prevent overlap with UnaryOpLike.
template <typename T>
concept BinaryOpLike = requires(const T& t) {
t.lhs;
t.rhs;
t.out;
} && !requires(const T& t) { t.in; };
/// Ops with two tensor slots: in, out (e.g., ReluOp, SigmoidOp, ScaleOp).
/// Excludes ops that also have {lhs, rhs} to prevent overlap with BinaryOpLike.
template <typename T>
concept UnaryOpLike = requires(const T& t) {
t.in;
t.out;
} && !requires(const T& t) {
t.lhs;
t.rhs;
};
// ============================================================================
// visitOp — single consteval dispatch point for all Op types
//
// Adding a new Op alternative requires no changes here — std::visit
// enforces exhaustiveness at compile time via the visitor's operator().
// ============================================================================
template <typename F>
consteval auto visitOp(const Op& op, F&& f)
{
return std::visit(std::forward<F>(f), op);
}
// ============================================================================
// ResolvedSignature
// ============================================================================
/// Resolved signature: all tensors and scalars have concrete metadata.
/// Rank and layout are concrete for tensors with operator-implied defaults
/// or explicit Tensor entries; may remain 0/Auto for tensors where the
/// operation type does not specify them (e.g., standalone AddOp).
struct ResolvedSignature
{
int num_tensors = 0;
std::array<ResolvedTensor, kMaxTensors> tensors = {};
int num_scalars = 0;
std::array<Scalar, kMaxScalars> scalars = {};
/// Find a resolved tensor by name. Compile-time error if not found.
consteval ResolvedTensor tensor(std::string_view name) const
{
for(int i = 0; i < num_tensors; ++i)
if(tensors[i].name == name)
return tensors[i];
throw "tensor() — name not found in resolved signature; "
"check that it appears in an op slot or Tensor entry";
}
/// Find a resolved tensor's slot index by name. Compile-time error if not found.
consteval int tensorIndex(std::string_view name) const
{
for(int i = 0; i < num_tensors; ++i)
if(tensors[i].name == name)
return i;
throw "tensorIndex() — name not found in resolved signature; "
"check that it appears in an op slot or Tensor entry";
}
/// Find a resolved scalar by name. Compile-time error if not found.
consteval Scalar scalar(std::string_view name) const
{
for(int i = 0; i < num_scalars; ++i)
if(scalars[i].name == name)
return scalars[i];
throw "scalar() — name not found; "
"add a Scalar entry with this name to the Signature";
}
/// Find a resolved scalar's slot index by name. Compile-time error if not found.
consteval int scalarIndex(std::string_view name) const
{
for(int i = 0; i < num_scalars; ++i)
if(scalars[i].name == name)
return i;
throw "scalarIndex() — name not found; "
"add a Scalar entry with this name to the Signature";
}
/// Find a tensor slot index by name. Returns -1 if not found.
/// Unlike tensorIndex() (consteval), callable at both compile time and runtime.
constexpr int findTensor(std::string_view name) const
{
for(int i = 0; i < num_tensors; ++i)
if(tensors[i].name == name)
return i;
return -1;
}
/// Find a scalar slot index by name. Returns -1 if not found.
/// Unlike scalarIndex() (consteval), callable at both compile time and runtime.
constexpr int findScalar(std::string_view name) const
{
for(int i = 0; i < num_scalars; ++i)
if(scalars[i].name == name)
return i;
return -1;
}
};
// ============================================================================
// Op slot visitors — concept-driven, generic handling for binary/unary ops
// ============================================================================
/// Register all tensor slots of an op. Returns the output tensor name.
///
/// Uses concepts for generic dispatch:
/// GemmOp — special case: sets operator-implied rank/layout defaults
/// ScaleOp — special case: validates scalar reference against sig.scalars[]
/// BinaryOpLike — generic: registers lhs, rhs, out
/// UnaryOpLike — generic: registers in, out
///
/// Adding a new BinaryOpLike or UnaryOpLike op requires no changes here.
consteval std::string_view collectTensorSlotsFromOp(const Op& op,
const Signature& sig,
auto& find_or_add,
auto& set_if_unknown)
{
return visitOp(op, [&](const auto& typed_op) -> std::string_view {
using T = std::remove_cvref_t<decltype(typed_op)>;
if constexpr(std::is_same_v<T, std::monostate>)
{
return {};
}
else if constexpr(std::is_same_v<T, GemmOp>)
{
// GemmOp has operator-implied rank/layout defaults
set_if_unknown(find_or_add(typed_op.lhs), 2, Layout::Row);
set_if_unknown(find_or_add(typed_op.rhs), 2, Layout::Col);
set_if_unknown(find_or_add(typed_op.out), 2, Layout::Row);
return typed_op.out;
}
else if constexpr(std::is_same_v<T, ScaleOp>)
{
// ScaleOp is unary + scalar reference validation
find_or_add(typed_op.in);
find_or_add(typed_op.out);
if(typed_op.scale.empty())
throw "ScaleOp.scale must name a Scalar parameter";
bool found_scalar = false;
for(int si = 0; si < kMaxScalars; ++si)
{
if(!sig.scalars[si].name.empty() && sig.scalars[si].name == typed_op.scale)
{
found_scalar = true;
break;
}
}
if(!found_scalar)
throw "ScaleOp.scale references undeclared Scalar — "
"add a matching Scalar entry to the Signature";
return typed_op.out;
}
else if constexpr(BinaryOpLike<T>)
{
find_or_add(typed_op.lhs);
find_or_add(typed_op.rhs);
find_or_add(typed_op.out);
return typed_op.out;
}
else if constexpr(UnaryOpLike<T>)
{
find_or_add(typed_op.in);
find_or_add(typed_op.out);
return typed_op.out;
}
else
{
throw "unhandled Op type in collectTensorSlotsFromOp — "
"add explicit handling or satisfy BinaryOpLike/UnaryOpLike";
}
});
}
/// Propagate rank/layout through an op's connected tensor slots.
///
/// Binary ops propagate from the first slot with known rank/layout to all others.
/// Unary ops propagate between in and out.
/// GemmOp is skipped (rank/layout already set by collectTensorSlotsFromOp).
///
/// Adding a new BinaryOpLike or UnaryOpLike op requires no changes here.
consteval void propagateRankLayout(const Op& op, auto& propagate_binary, auto& propagate_unary)
{
visitOp(op, [&](const auto& typed_op) {
using T = std::remove_cvref_t<decltype(typed_op)>;
if constexpr(std::is_same_v<T, std::monostate> || std::is_same_v<T, GemmOp>)
{
// monostate: nothing to do
// GemmOp: rank/layout already set in collectTensorSlotsFromOp
}
else if constexpr(BinaryOpLike<T>)
{
propagate_binary(typed_op.lhs, typed_op.rhs, typed_op.out);
}
else if constexpr(UnaryOpLike<T>)
{
propagate_unary(typed_op.in, typed_op.out);
}
else
{
throw "unhandled Op type in propagateRankLayout — "
"add explicit handling or satisfy BinaryOpLike/UnaryOpLike";
}
});
}
// ============================================================================
// resolve()
// ============================================================================
/// Resolve a Signature into concrete tensor and scalar descriptors.
///
/// Phases:
/// 1. Register tensor slots from operators (with op-implied rank/layout)
/// 2. Propagate rank/layout through connected tensors (forward + backward)
/// 3. Merge explicit Tensor entries from sig.tensors (overrides propagation)
/// 4. Apply dtype cascade: explicit tensor -> sig.dtype -> error
/// 5. Collect declared scalars, build result
///
/// GemmOp slots get operator-implied defaults:
/// lhs -> rank 2, Row; rhs -> rank 2, Col; out -> rank 2, Row
///
/// Binary ops (AddOp, MulOp) and unary ops propagate rank/layout from
/// the first connected slot that has known values.
consteval ResolvedSignature resolve(const Signature& sig)
{
// --- Intermediate tracking ---
struct Info
{
std::string_view name;
bool dtype_set = false;
DataType dtype = DataType::FP32;
int rank = 0;
Layout layout = Layout::Auto;
bool has_quantize = false;
Quantization quantize_info{}; // only valid when has_quantize == true
};
Info infos[kMaxTensors] = {};
int num = 0;
// Find tensor by name, return index or -1.
auto find = [&](std::string_view name) -> int {
for(int i = 0; i < num; ++i)
if(infos[i].name == name)
return i;
return -1;
};
// Find tensor by name; add if not present.
auto find_or_add = [&](std::string_view name) -> int {
if(name.empty())
throw "operator slot has empty tensor name";
int idx = find(name);
if(idx >= 0)
return idx;
if(num >= kMaxTensors)
throw "too many unique tensors in Signature (max 16)";
infos[num].name = name;
return num++;
};
// Set rank/layout only if currently unknown. Error on conflicting values.
// Returns true if anything changed.
auto set_if_unknown = [&](int idx, int rank, Layout layout) -> bool {
bool changed = false;
if(rank != 0)
{
if(infos[idx].rank == 0)
{
infos[idx].rank = rank;
changed = true;
}
else if(infos[idx].rank != rank)
throw "conflicting rank for tensor — two operators imply different ranks; "
"check that shared tensor names are intentional";
}
if(layout != Layout::Auto)
{
if(infos[idx].layout == Layout::Auto)
{
infos[idx].layout = layout;
changed = true;
}
else if(infos[idx].layout != layout)
throw "conflicting layout for tensor — two operators imply different layouts; "
"check that shared tensor names are intentional";
}
return changed;
};
// ================================================================
// Phase 1: Register tensor slots from operators
// ================================================================
std::string_view output_names[kMaxOps] = {};
int num_outputs = 0;
for(int i = 0; i < kMaxOps; ++i)
{
std::string_view out_name =
collectTensorSlotsFromOp(sig.ops[i], sig, find_or_add, set_if_unknown);
// SSA uniqueness: each output name may appear at most once
if(!out_name.empty())
{
for(int j = 0; j < num_outputs; ++j)
if(output_names[j] == out_name)
throw "tensor name produced by multiple operators "
"(each output name must be unique)";
output_names[num_outputs++] = out_name;
}
}
// ================================================================
// Phase 2: Propagate rank/layout through connected tensors
// ================================================================
bool changed = false;
auto propagate_binary =
[&](std::string_view lhs_name, std::string_view rhs_name, std::string_view out_name) {
int li = find(lhs_name);
int ri = find(rhs_name);
int oi = find(out_name);
if(li < 0 || ri < 0 || oi < 0)
return;
int src = -1;
if(infos[li].rank != 0)
src = li;
else if(infos[ri].rank != 0)
src = ri;
else if(infos[oi].rank != 0)
src = oi;
if(src >= 0)
{
changed |= set_if_unknown(li, infos[src].rank, infos[src].layout);
changed |= set_if_unknown(ri, infos[src].rank, infos[src].layout);
changed |= set_if_unknown(oi, infos[src].rank, infos[src].layout);
}
};
auto propagate_unary = [&](std::string_view in_name, std::string_view out_name) {
int ii = find(in_name);
int oi = find(out_name);
if(ii < 0 || oi < 0)
return;
if(infos[ii].rank != 0)
changed |= set_if_unknown(oi, infos[ii].rank, infos[ii].layout);
else if(infos[oi].rank != 0)
changed |= set_if_unknown(ii, infos[oi].rank, infos[oi].layout);
};
// Propagate until stable. Chains converge in 1 pass, diamonds in 2-3.
// Cap at 4 to bound compile time. Unreachable with kMaxOps=8 (a graph
// requiring >4 passes would need more ops than fit), but guards against
// future kMaxOps increases.
constexpr int kMaxPropagationPasses = 4;
for(int pass = 0; pass < kMaxPropagationPasses; ++pass)
{
changed = false;
for(int i = 0; i < kMaxOps; ++i)
propagateRankLayout(sig.ops[i], propagate_binary, propagate_unary);
for(int i = kMaxOps - 1; i >= 0; --i)
propagateRankLayout(sig.ops[i], propagate_binary, propagate_unary);
if(!changed)
break;
}
if(changed)
throw "could not infer rank/layout for all tensors — "
"set rank and layout explicitly on Tensor entries, "
"or reduce chained operations";
// ================================================================
// Phase 3: Merge explicit Tensor entries (override propagation)
// ================================================================
for(int i = 0; i < kMaxTensors; ++i)
{
if(sig.tensors[i].name.empty())
{
// Catch entries with metadata but no name (likely a mistake)
if(sig.tensors[i].dtype.has_value() || sig.tensors[i].rank != 0 ||
sig.tensors[i].layout != Layout::Auto || sig.tensors[i].quantize.has_value())
throw "Tensor entry has metadata but no name";
continue;
}
int idx = find_or_add(sig.tensors[i].name);
if(sig.tensors[i].dtype.has_value())
{
infos[idx].dtype_set = true;
infos[idx].dtype = *sig.tensors[i].dtype;
}
if(sig.tensors[i].rank != 0)
infos[idx].rank = sig.tensors[i].rank;
if(sig.tensors[i].layout != Layout::Auto)
infos[idx].layout = sig.tensors[i].layout;
if(sig.tensors[i].quantize.has_value())
{
const auto& q = *sig.tensors[i].quantize;
if(q.scale_name.empty())
throw "Tensor .quantize has empty scale_name";
// Auto-register the scale tensor
int scale_idx = find_or_add(q.scale_name);
if(infos[scale_idx].dtype_set && infos[scale_idx].dtype != q.scale_dtype)
throw "tensor dtype conflicts with quantization scale_dtype "
"for the same name";
infos[scale_idx].dtype_set = true;
infos[scale_idx].dtype = q.scale_dtype;
set_if_unknown(scale_idx, 2, Layout::Row);
// Track quantization on this tensor
infos[idx].has_quantize = true;
infos[idx].quantize_info = q;
}
}
// ================================================================
// Phase 4: dtype cascade
// ================================================================
for(int i = 0; i < num; ++i)
{
if(!infos[i].dtype_set)
{
if(sig.dtype.has_value())
infos[i].dtype = *sig.dtype;
else
throw "tensor dtype unresolvable: set tensor dtype or signature dtype";
}
}
// ================================================================
// Phase 5: Build result
// ================================================================
ResolvedSignature result{};
result.num_tensors = num;
for(int i = 0; i < num; ++i)
{
result.tensors[i] =
ResolvedTensor{infos[i].name, infos[i].dtype, infos[i].rank, infos[i].layout};
if(infos[i].has_quantize)
result.tensors[i].quantize = ResolvedQuantization{infos[i].quantize_info.scale_name,
infos[i].quantize_info.scale_dtype,
infos[i].quantize_info.group_size};
}
// Collect declared scalars (pass-through — no inference needed)
for(int i = 0; i < kMaxScalars; ++i)
{
if(sig.scalars[i].name.empty())
continue;
// Check for duplicate scalar names
for(int j = 0; j < result.num_scalars; ++j)
if(result.scalars[j].name == sig.scalars[i].name)
throw "duplicate scalar name in Signature";
result.scalars[result.num_scalars++] = sig.scalars[i];
}
return result;
}
} // namespace rocm_ck

88
rocm_ck/include/rocm_ck/signature.hpp Normal file
View File

@@ -0,0 +1,88 @@
// Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
// SPDX-License-Identifier: MIT
//
// Role: meta — Signature, Tensor, Scalar. No runtime, no CK deps.
//
// Signature: the complete description of WHAT a kernel computes.
//
// A directed compute graph where tensors are nodes and operators are edges.
// Operators reference tensors by name (SSA-style). Shared names between
// operator outputs and inputs form graph edges.
//
// This header has NO CK Tile dependency. It is included by both host code
// and device code to share the signature definition.
#pragma once
#include <rocm_ck/args.hpp>
#include <rocm_ck/datatype.hpp>
#include <rocm_ck/layout.hpp>
#include <rocm_ck/ops.hpp>
#include <array>
#include <optional>
#include <string_view>
namespace rocm_ck {
/// Quantization metadata for a tensor stored in a packed/quantized format.
///
/// Describes how a quantized tensor (e.g., INT4 weights) is encoded.
/// The scale tensor is a separate physical tensor with per-group factors;
/// group_size is the number of consecutive elements sharing one scale along K.
struct Quantization
{
std::string_view scale_name; // name of the scale tensor (e.g., "scale")
DataType scale_dtype = DataType::FP32; // element type of the scale tensor
int group_size = 128; // elements per quantization group
};
/// A tensor node in the signature's compute graph.
///
/// Default values mean "inherit from the operator slot":
/// dtype = nullopt -> inherit from Signature::dtype cascade
/// rank = 0 -> inherit from operator (e.g., GemmOp implies rank 2)
/// layout = Auto -> inherit from operator (e.g., GemmOp::lhs implies Row)
/// quantize = nullopt -> not quantized
struct Tensor
{
std::string_view name;
std::optional<DataType> dtype = std::nullopt;
int rank = 0;
Layout layout = Layout::Auto;
std::optional<Quantization> quantize = std::nullopt;
};
/// A named scalar parameter (e.g., alpha, beta, scale).
struct Scalar
{
std::string_view name;
DataType dtype = DataType::FP32;
};
/// The complete description of WHAT a kernel computes.
///
/// A directed compute graph where tensors are nodes and operators are edges.
/// Each operator output gets a unique name; shared names form graph edges.
///
/// Example — simple fp16 GEMM:
/// {.dtype = FP16, .ops = {GemmOp{.lhs="A", .rhs="B", .out="C"}}}
///
/// Example — GEMM + bias + ReLU:
/// {.dtype = FP16,
/// .ops = {GemmOp{.lhs="A", .rhs="B", .out="C"},
/// AddOp{.lhs="C", .rhs="bias", .out="D"},
/// ReluOp{.in="D", .out="E"}}}
// kMaxTensors and kMaxScalars are defined in args.hpp (canonical source).
// kMaxOps is Signature-specific (operators, not kernel arguments).
constexpr int kMaxOps = 8;
struct Signature
{
std::optional<DataType> dtype;
std::array<Tensor, kMaxTensors> tensors = {};
std::array<Scalar, kMaxScalars> scalars = {};
std::array<Op, kMaxOps> ops = {};
};
} // namespace rocm_ck

18
rocm_ck/tests/CMakeLists.txt
View File

@@ -4,8 +4,9 @@
# rocm_ck tests
#
# Test tiers:
# ROCM_CK_SMOKE — Fast host-only tests (< 1s total). No GPU required.
# ROCM_CK_KERNEL — GPU kernel tests. Require HIP and a GPU.
# ROCM_CK_SMOKE — Fast host-only tests (< 1s total). No GPU required.
# ROCM_CK_COMPILE_FAIL — Compile-fail tests (~7s). No GPU required.
# ROCM_CK_KERNEL — GPU kernel tests. Require HIP and a GPU.
#
# Usage:
# ninja smoke-rocm-ck # build + run smoke tests
@@ -33,14 +34,20 @@ endfunction()
# ---------------------------------------------------------------------------
# Smoke tests (fast, host-only, no GPU)
# ---------------------------------------------------------------------------
add_rocm_ck_test(rocm_ck_unit
set(ROCM_CK_UNIT_SOURCES
unit/unit_arch_properties.cpp
unit/unit_args.cpp
unit/unit_datatype.cpp
unit/unit_fixed_string.cpp
unit/unit_index_t.cpp
unit/unit_layout.cpp
unit/unit_physical_tensor.cpp
unit/unit_resolve.cpp
unit/unit_signature.cpp
)
set_source_files_properties(${ROCM_CK_UNIT_SOURCES} PROPERTIES LANGUAGE CXX)
add_rocm_ck_test(rocm_ck_unit ${ROCM_CK_UNIT_SOURCES})
target_include_directories(rocm_ck_unit PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/../../include)
add_test(NAME rocm_ck_unit COMMAND rocm_ck_unit)
set_tests_properties(rocm_ck_unit PROPERTIES LABELS "ROCM_CK_SMOKE")
@@ -61,3 +68,8 @@ add_custom_target(check-rocm-ck
DEPENDS build-smoke-rocm-ck
USES_TERMINAL
COMMENT "Running all rocm_ck tests...")
# ---------------------------------------------------------------------------
# Compile-fail tests (expected failures)
# ---------------------------------------------------------------------------
add_subdirectory(compile_fail)

88
rocm_ck/tests/compile_fail/CMakeLists.txt Normal file
View File

@@ -0,0 +1,88 @@
# Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
# SPDX-License-Identifier: MIT
#
# Expected compilation failure tests for rocm_ck schema types.
#
# Each .cpp file must FAIL to compile. The CTest test PASSES when
# compilation FAILS with the expected error message.
#
# The loop auto-discovers .cpp files and sets WILL_FAIL TRUE as a safe
# default. The overrides below upgrade each test to verify the specific
# consteval throw message via PASS_REGULAR_EXPRESSION (which replaces
# WILL_FAIL since the two interact badly — WILL_FAIL inverts the regex
# match result).
#
# Usage:
# ctest --test-dir build -R compile_fail
# ctest --test-dir build -L ROCM_CK_COMPILE_FAIL
file(GLOB COMPILE_FAIL_SOURCES "${CMAKE_CURRENT_SOURCE_DIR}/*.cpp")
foreach(source ${COMPILE_FAIL_SOURCES})
get_filename_component(test_name ${source} NAME_WE)
set(target_name "rocm_ck_compile_fail_${test_name}")
add_library(${target_name} OBJECT EXCLUDE_FROM_ALL ${source})
target_link_libraries(${target_name} PRIVATE rocm_ck)
target_compile_features(${target_name} PRIVATE cxx_std_20)
target_compile_options(${target_name} PRIVATE
-Wno-zero-as-null-pointer-constant # C++20 <=> comparisons to 0
)
add_test(
NAME ${target_name}
COMMAND ${CMAKE_COMMAND} --build ${CMAKE_BINARY_DIR} --target ${target_name}
)
set_tests_properties(${target_name} PROPERTIES
WILL_FAIL TRUE
LABELS "ROCM_CK_COMPILE_FAIL"
)
endforeach()
# ---- Message verification overrides ----
# PASS_REGULAR_EXPRESSION verifies the compiler emits the expected throw
# message. It ignores exit code, so the regex match alone proves compilation
# failed with the correct consteval error. WILL_FAIL must be FALSE here
# because CMake applies WILL_FAIL after PASS_REGULAR_EXPRESSION, inverting
# the result.
# resolve.hpp errors
set_tests_properties(rocm_ck_compile_fail_conflicting_layout PROPERTIES
WILL_FAIL FALSE
PASS_REGULAR_EXPRESSION "conflicting layout for tensor")
set_tests_properties(rocm_ck_compile_fail_quantize_empty_scale_name PROPERTIES
WILL_FAIL FALSE
PASS_REGULAR_EXPRESSION "quantize has empty scale_name")
set_tests_properties(rocm_ck_compile_fail_no_dtype PROPERTIES
WILL_FAIL FALSE
PASS_REGULAR_EXPRESSION "tensor dtype unresolvable")
set_tests_properties(rocm_ck_compile_fail_empty_tensor_name PROPERTIES
WILL_FAIL FALSE
PASS_REGULAR_EXPRESSION "operator slot has empty tensor name")
set_tests_properties(rocm_ck_compile_fail_ssa_violation PROPERTIES
WILL_FAIL FALSE
PASS_REGULAR_EXPRESSION "tensor name produced by multiple operators")
set_tests_properties(rocm_ck_compile_fail_scale_scalar_not_found PROPERTIES
WILL_FAIL FALSE
PASS_REGULAR_EXPRESSION "ScaleOp.scale references undeclared Scalar")
# arch_properties.hpp errors
set_tests_properties(rocm_ck_compile_fail_target_set_mixed_wave PROPERTIES
WILL_FAIL FALSE
PASS_REGULAR_EXPRESSION "wavefront_size.*requires all targets")
set_tests_properties(rocm_ck_compile_fail_fp8_ocp_unsupported PROPERTIES
WILL_FAIL FALSE
PASS_REGULAR_EXPRESSION "FP8_OCP.*not yet supported")
set_tests_properties(rocm_ck_compile_fail_empty_target_set_wavefront PROPERTIES
WILL_FAIL FALSE
PASS_REGULAR_EXPRESSION "wavefront_size.*called on empty")
# fixed_string.hpp errors
set_tests_properties(rocm_ck_compile_fail_tensor_name_too_long PROPERTIES
WILL_FAIL FALSE
PASS_REGULAR_EXPRESSION "FixedString.*exceeds capacity")
# layout.hpp errors
set_tests_properties(rocm_ck_compile_fail_layout_auto_stride PROPERTIES
WILL_FAIL FALSE
PASS_REGULAR_EXPRESSION "leadingDimStride requires Row or Col")

18
rocm_ck/tests/compile_fail/conflicting_layout.cpp Normal file
View File

@@ -0,0 +1,18 @@
// Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
// SPDX-License-Identifier: MIT
//
// Must fail: two GemmOps imply conflicting layouts for the same tensor.
//
// GemmOp1 outputs "C" → implied Row layout.
// GemmOp2 uses "C" as rhs → implied Col layout.
// These conflict: "C" can't be both Row and Col.
//
// Expected error: "conflicting layout for tensor"
#include <rocm_ck/resolve.hpp>
using namespace rocm_ck;
constexpr auto bad = resolve(Signature{.dtype = DataType::FP16,
.ops = {GemmOp{.lhs = "A", .rhs = "B", .out = "C"},
GemmOp{.lhs = "X", .rhs = "C", .out = "Y"}}});

11
rocm_ck/tests/compile_fail/empty_target_set_wavefront.cpp Normal file
View File

@@ -0,0 +1,11 @@
// Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
// SPDX-License-Identifier: MIT
//
// Must fail: wavefront_size() called on empty TargetSet.
// Expected error: "wavefront_size() called on empty TargetSet"
#include <rocm_ck/arch_properties.hpp>
using namespace rocm_ck;
constexpr int bad = TargetSet{}.wavefront_size();

12
rocm_ck/tests/compile_fail/empty_tensor_name.cpp Normal file
View File

@@ -0,0 +1,12 @@
// Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
// SPDX-License-Identifier: MIT
//
// Must fail: GemmOp has empty tensor name in lhs slot.
// Expected error: "operator slot has empty tensor name"
#include <rocm_ck/resolve.hpp>
using namespace rocm_ck;
constexpr auto bad =
resolve(Signature{.dtype = DataType::FP16, .ops = {GemmOp{.lhs = "", .rhs = "B", .out = "C"}}});

11
rocm_ck/tests/compile_fail/fp8_ocp_unsupported.cpp Normal file
View File

@@ -0,0 +1,11 @@
// Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
// SPDX-License-Identifier: MIT
//
// Must fail: FP8_OCP is not yet supported in GEMM wave tiles.
// Expected error: "FP8_OCP/BF8_OCP not yet supported"
#include <rocm_ck/arch_properties.hpp>
using namespace rocm_ck;
constexpr bool bad = isValidWaveTile(DataType::FP8_OCP, 16, 16, 16, GpuTarget::gfx942);

13
rocm_ck/tests/compile_fail/layout_auto_stride.cpp Normal file
View File

@@ -0,0 +1,13 @@
// Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
// SPDX-License-Identifier: MIT
//
// Must fail: Layout::Auto cannot be used for stride computation.
// Expected error: "leadingDimStride requires Row or Col layout"
#include <rocm_ck/layout.hpp>
#include <array>
using namespace rocm_ck;
constexpr auto bad = leadingDimStride(Layout::Auto, std::array{1, 2});

11
rocm_ck/tests/compile_fail/no_dtype.cpp Normal file
View File

@@ -0,0 +1,11 @@
// Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
// SPDX-License-Identifier: MIT
//
// Must fail: no dtype set on tensors or signature.
// Expected error: "tensor dtype unresolvable"
#include <rocm_ck/resolve.hpp>
using namespace rocm_ck;
constexpr auto bad = resolve(Signature{.ops = {GemmOp{.lhs = "A", .rhs = "B", .out = "C"}}});

15
rocm_ck/tests/compile_fail/quantize_empty_scale_name.cpp Normal file
View File

@@ -0,0 +1,15 @@
// Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
// SPDX-License-Identifier: MIT
//
// Must fail: Tensor.quantize has empty scale_name.
// Expected error: "Tensor .quantize has empty scale_name"
#include <rocm_ck/resolve.hpp>
using namespace rocm_ck;
constexpr auto bad = resolve(
Signature{.dtype = DataType::FP16,
.tensors = {Tensor{
.name = "B", .dtype = DataType::I4, .quantize = Quantization{.scale_name = ""}}},
.ops = {GemmOp{.lhs = "A", .rhs = "B", .out = "C"}}});

12
rocm_ck/tests/compile_fail/scale_scalar_not_found.cpp Normal file
View File

@@ -0,0 +1,12 @@
// Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
// SPDX-License-Identifier: MIT
//
// Must fail: ScaleOp references scalar "alpha" but no Scalar entry exists.
// Expected error: "ScaleOp.scale references undeclared Scalar"
#include <rocm_ck/resolve.hpp>
using namespace rocm_ck;
constexpr auto bad = resolve(
Signature{.dtype = DataType::FP16, .ops = {ScaleOp{.in = "A", .out = "B", .scale = "alpha"}}});

13
rocm_ck/tests/compile_fail/ssa_violation.cpp Normal file
View File

@@ -0,0 +1,13 @@
// Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
// SPDX-License-Identifier: MIT
//
// Must fail: two GemmOps produce the same output tensor "C".
// Expected error: "tensor name produced by multiple operators"
#include <rocm_ck/resolve.hpp>
using namespace rocm_ck;
constexpr auto bad = resolve(Signature{.dtype = DataType::FP16,
.ops = {GemmOp{.lhs = "A", .rhs = "B", .out = "C"},
GemmOp{.lhs = "X", .rhs = "Y", .out = "C"}}});

13
rocm_ck/tests/compile_fail/target_set_mixed_wave.cpp Normal file
View File

@@ -0,0 +1,13 @@
// Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
// SPDX-License-Identifier: MIT
//
// Expected failure: wavefront_size() on a mixed CDNA+RDNA TargetSet.
// TargetSet::all() contains both wave64 (CDNA) and wave32 (RDNA) targets.
// Calling wavefront_size() on a mixed set must fail at compile time.
#include <rocm_ck/arch_properties.hpp>
using namespace rocm_ck;
constexpr auto mixed = TargetSet::all();
constexpr int bad_wave = mixed.wavefront_size(); // Must fail: mixed wave64/wave32

11
rocm_ck/tests/compile_fail/tensor_name_too_long.cpp Normal file
View File

@@ -0,0 +1,11 @@
// Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
// SPDX-License-Identifier: MIT
//
// Must fail: tensor name exceeds FixedString<16> capacity (15 chars + null).
// Expected error: "FixedString: input exceeds capacity"
#include <rocm_ck/fixed_string.hpp>
using namespace rocm_ck;
constexpr FixedString<16> bad("ABCDEFGHIJKLMNOP"); // 16 chars, needs 17 with null

487
rocm_ck/tests/unit/unit_arch_properties.cpp Normal file
View File

@@ -0,0 +1,487 @@
// Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
// SPDX-License-Identifier: MIT
#include <rocm_ck/arch_properties.hpp>
#include <gtest/gtest.h>
using ::rocm_ck::ArchFamily;
using ::rocm_ck::DataType;
using ::rocm_ck::GpuTarget;
using ::rocm_ck::isValidWaveTile;
using ::rocm_ck::TargetProperties;
using ::rocm_ck::TargetSet;
// ============================================================================
// TargetProperties
// ============================================================================
TEST(TargetProperties, ReturnsWave64ForCDNA)
{
EXPECT_EQ(properties(GpuTarget::gfx90a).wavefront_size, 64);
EXPECT_EQ(properties(GpuTarget::gfx942).wavefront_size, 64);
EXPECT_EQ(properties(GpuTarget::gfx950).wavefront_size, 64);
}
TEST(TargetProperties, ReturnsWave32ForRDNA)
{
EXPECT_EQ(properties(GpuTarget::gfx1100).wavefront_size, 32);
EXPECT_EQ(properties(GpuTarget::gfx1101).wavefront_size, 32);
EXPECT_EQ(properties(GpuTarget::gfx1102).wavefront_size, 32);
EXPECT_EQ(properties(GpuTarget::gfx1150).wavefront_size, 32);
EXPECT_EQ(properties(GpuTarget::gfx1151).wavefront_size, 32);
}
TEST(TargetProperties, CDNAArchFamily)
{
EXPECT_EQ(properties(GpuTarget::gfx90a).arch_family, ArchFamily::CDNA);
EXPECT_EQ(properties(GpuTarget::gfx942).arch_family, ArchFamily::CDNA);
EXPECT_EQ(properties(GpuTarget::gfx950).arch_family, ArchFamily::CDNA);
}
TEST(TargetProperties, RDNAArchFamily)
{
EXPECT_EQ(properties(GpuTarget::gfx1100).arch_family, ArchFamily::RDNA);
EXPECT_EQ(properties(GpuTarget::gfx1101).arch_family, ArchFamily::RDNA);
EXPECT_EQ(properties(GpuTarget::gfx1102).arch_family, ArchFamily::RDNA);
EXPECT_EQ(properties(GpuTarget::gfx1150).arch_family, ArchFamily::RDNA);
EXPECT_EQ(properties(GpuTarget::gfx1151).arch_family, ArchFamily::RDNA);
}
// ============================================================================
// isCDNA / isRDNA predicates
// ============================================================================
TEST(ArchFamily, IsCDNAForAllCDNATargets)
{
EXPECT_TRUE(isCDNA(GpuTarget::gfx90a));
EXPECT_TRUE(isCDNA(GpuTarget::gfx942));
EXPECT_TRUE(isCDNA(GpuTarget::gfx950));
EXPECT_FALSE(isCDNA(GpuTarget::gfx1100));
EXPECT_FALSE(isCDNA(GpuTarget::gfx1101));
EXPECT_FALSE(isCDNA(GpuTarget::gfx1102));
EXPECT_FALSE(isCDNA(GpuTarget::gfx1150));
EXPECT_FALSE(isCDNA(GpuTarget::gfx1151));
}
TEST(ArchFamily, IsRDNAForAllRDNATargets)
{
EXPECT_TRUE(isRDNA(GpuTarget::gfx1100));
EXPECT_TRUE(isRDNA(GpuTarget::gfx1101));
EXPECT_TRUE(isRDNA(GpuTarget::gfx1102));
EXPECT_TRUE(isRDNA(GpuTarget::gfx1150));
EXPECT_TRUE(isRDNA(GpuTarget::gfx1151));
EXPECT_FALSE(isRDNA(GpuTarget::gfx90a));
EXPECT_FALSE(isRDNA(GpuTarget::gfx942));
EXPECT_FALSE(isRDNA(GpuTarget::gfx950));
}
// ============================================================================
// wavefrontSize free function
// ============================================================================
TEST(WavefrontSize, MatchesTargetProperties)
{
EXPECT_EQ(wavefrontSize(GpuTarget::gfx90a), 64);
EXPECT_EQ(wavefrontSize(GpuTarget::gfx942), 64);
EXPECT_EQ(wavefrontSize(GpuTarget::gfx950), 64);
EXPECT_EQ(wavefrontSize(GpuTarget::gfx1100), 32);
EXPECT_EQ(wavefrontSize(GpuTarget::gfx1101), 32);
EXPECT_EQ(wavefrontSize(GpuTarget::gfx1102), 32);
EXPECT_EQ(wavefrontSize(GpuTarget::gfx1150), 32);
EXPECT_EQ(wavefrontSize(GpuTarget::gfx1151), 32);
}
// ============================================================================
// TargetSet — construction
// ============================================================================
TEST(TargetSet, DefaultConstructsToEmpty)
{
constexpr TargetSet ts{};
EXPECT_TRUE(ts.is_empty());
EXPECT_EQ(ts.count(), 0);
}
TEST(TargetSet, ImplicitConversionFromSingleTarget)
{
constexpr TargetSet ts = GpuTarget::gfx942;
EXPECT_TRUE(ts.is_single_target());
EXPECT_EQ(ts.single_target(), GpuTarget::gfx942);
EXPECT_EQ(ts.count(), 1);
}
// ============================================================================
// TargetSet — named constructors
// ============================================================================
TEST(TargetSet, AllContainsEveryTarget)
{
constexpr auto ts = TargetSet::all();
EXPECT_EQ(ts.count(), 8);
EXPECT_TRUE(ts.contains(GpuTarget::gfx90a));
EXPECT_TRUE(ts.contains(GpuTarget::gfx942));
EXPECT_TRUE(ts.contains(GpuTarget::gfx950));
EXPECT_TRUE(ts.contains(GpuTarget::gfx1100));
EXPECT_TRUE(ts.contains(GpuTarget::gfx1101));
EXPECT_TRUE(ts.contains(GpuTarget::gfx1102));
EXPECT_TRUE(ts.contains(GpuTarget::gfx1150));
EXPECT_TRUE(ts.contains(GpuTarget::gfx1151));
}
TEST(TargetSet, CdnaContainsOnlyCDNATargets)
{
constexpr auto ts = TargetSet::cdna();
EXPECT_EQ(ts.count(), 3);
EXPECT_TRUE(ts.contains(GpuTarget::gfx90a));
EXPECT_TRUE(ts.contains(GpuTarget::gfx942));
EXPECT_TRUE(ts.contains(GpuTarget::gfx950));
EXPECT_FALSE(ts.contains(GpuTarget::gfx1100));
EXPECT_FALSE(ts.contains(GpuTarget::gfx1101));
EXPECT_FALSE(ts.contains(GpuTarget::gfx1102));
EXPECT_FALSE(ts.contains(GpuTarget::gfx1150));
EXPECT_FALSE(ts.contains(GpuTarget::gfx1151));
}
TEST(TargetSet, RdnaContainsOnlyRDNATargets)
{
constexpr auto ts = TargetSet::rdna();
EXPECT_EQ(ts.count(), 5);
EXPECT_TRUE(ts.contains(GpuTarget::gfx1100));
EXPECT_TRUE(ts.contains(GpuTarget::gfx1101));
EXPECT_TRUE(ts.contains(GpuTarget::gfx1102));
EXPECT_TRUE(ts.contains(GpuTarget::gfx1150));
EXPECT_TRUE(ts.contains(GpuTarget::gfx1151));
EXPECT_FALSE(ts.contains(GpuTarget::gfx90a));
}
TEST(TargetSet, FamilyGfx9MatchesCdna) { EXPECT_EQ(TargetSet::family_gfx9(), TargetSet::cdna()); }
TEST(TargetSet, FamilyGfx94ExcludesGfx90a)
{
constexpr auto ts = TargetSet::family_gfx94();
EXPECT_EQ(ts.count(), 2);
EXPECT_TRUE(ts.contains(GpuTarget::gfx942));
EXPECT_TRUE(ts.contains(GpuTarget::gfx950));
EXPECT_FALSE(ts.contains(GpuTarget::gfx90a));
}
TEST(TargetSet, FamilyGfx11MatchesRdna) { EXPECT_EQ(TargetSet::family_gfx11(), TargetSet::rdna()); }
TEST(TargetSet, FamilyGfx115ContainsRDNA35Only)
{
constexpr auto ts = TargetSet::family_gfx115();
EXPECT_EQ(ts.count(), 2);
EXPECT_TRUE(ts.contains(GpuTarget::gfx1150));
EXPECT_TRUE(ts.contains(GpuTarget::gfx1151));
EXPECT_FALSE(ts.contains(GpuTarget::gfx1100));
EXPECT_FALSE(ts.contains(GpuTarget::gfx1101));
EXPECT_FALSE(ts.contains(GpuTarget::gfx1102));
}
TEST(TargetSet, OnlyWithOneTarget)
{
constexpr auto ts = TargetSet::only(GpuTarget::gfx942);
EXPECT_EQ(ts.count(), 1);
EXPECT_TRUE(ts.contains(GpuTarget::gfx942));
}
TEST(TargetSet, OnlyWithTwoTargets)
{
constexpr auto ts = TargetSet::only(GpuTarget::gfx942, GpuTarget::gfx950);
EXPECT_EQ(ts.count(), 2);
EXPECT_TRUE(ts.contains(GpuTarget::gfx942));
EXPECT_TRUE(ts.contains(GpuTarget::gfx950));
}
TEST(TargetSet, OnlyWithThreeTargets)
{
constexpr auto ts = TargetSet::only(GpuTarget::gfx90a, GpuTarget::gfx942, GpuTarget::gfx950);
EXPECT_EQ(ts.count(), 3);
}
// Note: OnlyWithOneTarget, OnlyWithTwoTargets, OnlyWithThreeTargets test the
// variadic arity of TargetSet::only() overloads (1, 2, 3 parameters).
// ============================================================================
// TargetSet — set operations
// ============================================================================
TEST(TargetSet, ExcludingRemovesOneTarget)
{
constexpr auto base = TargetSet::cdna();
constexpr auto without = base.excluding(GpuTarget::gfx90a);
EXPECT_EQ(without.count(), 2);
EXPECT_FALSE(without.contains(GpuTarget::gfx90a));
EXPECT_TRUE(without.contains(GpuTarget::gfx942));
EXPECT_TRUE(without.contains(GpuTarget::gfx950));
}
TEST(TargetSet, UnionCombinesSets)
{
constexpr auto a = TargetSet::only(GpuTarget::gfx90a);
constexpr auto b = TargetSet::only(GpuTarget::gfx1151);
constexpr auto combined = a.union_with(b);
EXPECT_EQ(combined.count(), 2);
EXPECT_TRUE(combined.contains(GpuTarget::gfx90a));
EXPECT_TRUE(combined.contains(GpuTarget::gfx1151));
}
TEST(TargetSet, IntersectReturnsCommonTargets)
{
constexpr auto all = TargetSet::all();
constexpr auto cdna = TargetSet::cdna();
EXPECT_EQ(all.intersect_with(cdna), cdna);
}
TEST(TargetSet, MinusRemovesTargets)
{
constexpr auto all = TargetSet::all();
constexpr auto rdna = TargetSet::rdna();
EXPECT_EQ(all.minus(rdna), TargetSet::cdna());
}
TEST(TargetSet, MinusWithEmptySetIsIdentity)
{
constexpr TargetSet empty{};
constexpr auto cdna = TargetSet::cdna();
EXPECT_EQ(cdna.minus(empty), cdna);
}
TEST(TargetSet, EmptyMinusAnythingIsEmpty)
{
constexpr TargetSet empty{};
constexpr auto cdna = TargetSet::cdna();
EXPECT_TRUE(empty.minus(cdna).is_empty());
}
TEST(TargetSet, EmptyUnionIsIdentity)
{
constexpr TargetSet empty{};
constexpr auto cdna = TargetSet::cdna();
EXPECT_EQ(empty.union_with(cdna), cdna);
EXPECT_EQ(cdna.union_with(empty), cdna);
}
TEST(TargetSet, EmptyIntersectIsEmpty)
{
constexpr TargetSet empty{};
constexpr auto cdna = TargetSet::cdna();
EXPECT_TRUE(empty.intersect_with(cdna).is_empty());
}
// ============================================================================
// TargetSet — operators
// ============================================================================
TEST(TargetSet, OperatorOrDelegatesToUnion)
{
constexpr auto a = TargetSet::only(GpuTarget::gfx942);
constexpr auto b = TargetSet::only(GpuTarget::gfx950);
EXPECT_EQ((a | b).count(), 2);
}
TEST(TargetSet, OperatorAndDelegatesToIntersect)
{
EXPECT_EQ(TargetSet::all() & TargetSet::cdna(), TargetSet::cdna());
}
TEST(TargetSet, OperatorMinusDelegatesToMinus)
{
EXPECT_EQ(TargetSet::all() - TargetSet::rdna(), TargetSet::cdna());
}
TEST(TargetSet, EqualitySameSetReturnsTrue) { EXPECT_EQ(TargetSet::cdna(), TargetSet::cdna()); }
TEST(TargetSet, InequalityDifferentSetsReturnsTrue)
{
EXPECT_NE(TargetSet::cdna(), TargetSet::rdna());
}
// ============================================================================
// TargetSet — queries
// ============================================================================
TEST(TargetSet, ContainsReturnsTrueForMember)
{
EXPECT_TRUE(TargetSet::cdna().contains(GpuTarget::gfx90a));
}
TEST(TargetSet, ContainsReturnsFalseForNonMember)
{
EXPECT_FALSE(TargetSet::cdna().contains(GpuTarget::gfx1151));
}
TEST(TargetSet, IsSingleTargetTrueForOne)
{
EXPECT_TRUE(TargetSet::only(GpuTarget::gfx942).is_single_target());
}
TEST(TargetSet, IsSingleTargetFalseForMultiple)
{
EXPECT_FALSE(TargetSet::cdna().is_single_target());
}
TEST(TargetSet, IsSingleTargetFalseForEmpty) { EXPECT_FALSE(TargetSet{}.is_single_target()); }
TEST(TargetSet, WavefrontSizeUniformForCDNA) { EXPECT_EQ(TargetSet::cdna().wavefront_size(), 64); }
TEST(TargetSet, WavefrontSizeUniformForRDNA) { EXPECT_EQ(TargetSet::rdna().wavefront_size(), 32); }
TEST(TargetSet, WavefrontSizeThrowsOnMixedCDNAAndRDNA)
{
constexpr auto mixed = TargetSet::all();
// wavefront_size() should throw when set mixes wave64 (CDNA) and wave32 (RDNA)
bool caught = false;
try
{
int wf = mixed.wavefront_size();
(void)wf; // suppress unused warning if exception is not thrown
}
catch(...)
{
caught = true;
}
EXPECT_TRUE(caught);
}
TEST(TargetSet, IsAllCdnaPredicateWorks)
{
EXPECT_TRUE(TargetSet::cdna().is_all_cdna());
EXPECT_FALSE(TargetSet::rdna().is_all_cdna());
EXPECT_FALSE(TargetSet::all().is_all_cdna());
EXPECT_FALSE(TargetSet{}.is_all_cdna());
}
TEST(TargetSet, IsAllRdnaPredicateWorks)
{
EXPECT_TRUE(TargetSet::rdna().is_all_rdna());
EXPECT_FALSE(TargetSet::cdna().is_all_rdna());
EXPECT_FALSE(TargetSet::all().is_all_rdna());
EXPECT_FALSE(TargetSet{}.is_all_rdna());
}
// ============================================================================
// TargetSet — iteration
// ============================================================================
TEST(TargetSet, ForEachIteratesAllTargets)
{
int count = 0;
TargetSet::cdna().for_each([&count](GpuTarget) { count++; });
EXPECT_EQ(count, 3);
}
TEST(TargetSet, ForEachOnEmptySetDoesNothing)
{
int count = 0;
TargetSet{}.for_each([&count](GpuTarget) { count++; });
EXPECT_EQ(count, 0);
}
// ============================================================================
// isValidWaveTile — single target
// ============================================================================
TEST(IsValidWaveTile, FP32MFMATilesOnCDNA)
{
EXPECT_TRUE(isValidWaveTile(DataType::FP32, 16, 16, 4, GpuTarget::gfx90a));
EXPECT_TRUE(isValidWaveTile(DataType::FP32, 16, 16, 8, GpuTarget::gfx942));
EXPECT_TRUE(isValidWaveTile(DataType::FP32, 16, 16, 16, GpuTarget::gfx950));
EXPECT_TRUE(isValidWaveTile(DataType::FP32, 32, 32, 4, GpuTarget::gfx90a));
EXPECT_TRUE(isValidWaveTile(DataType::FP32, 32, 32, 8, GpuTarget::gfx942));
}
TEST(IsValidWaveTile, FP32RejectsInvalidTiles)
{
EXPECT_FALSE(isValidWaveTile(DataType::FP32, 32, 32, 16, GpuTarget::gfx90a));
EXPECT_FALSE(isValidWaveTile(DataType::FP32, 8, 8, 4, GpuTarget::gfx90a));
}
TEST(IsValidWaveTile, FP16MFMATilesOnCDNA)
{
EXPECT_TRUE(isValidWaveTile(DataType::FP16, 16, 16, 16, GpuTarget::gfx90a));
EXPECT_TRUE(isValidWaveTile(DataType::FP16, 16, 16, 32, GpuTarget::gfx942));
EXPECT_TRUE(isValidWaveTile(DataType::FP16, 32, 32, 8, GpuTarget::gfx90a));
EXPECT_TRUE(isValidWaveTile(DataType::FP16, 32, 32, 16, GpuTarget::gfx942));
}
TEST(IsValidWaveTile, BF16MatchesFP16Tiles)
{
EXPECT_TRUE(isValidWaveTile(DataType::BF16, 16, 16, 16, GpuTarget::gfx90a));
EXPECT_TRUE(isValidWaveTile(DataType::BF16, 32, 32, 16, GpuTarget::gfx942));
}
TEST(IsValidWaveTile, INT8MFMATilesOnCDNA)
{
EXPECT_TRUE(isValidWaveTile(DataType::I8, 32, 32, 16, GpuTarget::gfx942));
EXPECT_TRUE(isValidWaveTile(DataType::I8, 16, 16, 32, GpuTarget::gfx942));
}
TEST(IsValidWaveTile, FP8FNUZNotSupportedOnGfx90a)
{
EXPECT_FALSE(isValidWaveTile(DataType::FP8_FNUZ, 32, 32, 16, GpuTarget::gfx90a));
}
TEST(IsValidWaveTile, FP8FNUZBaseTilesOnGfx942)
{
EXPECT_TRUE(isValidWaveTile(DataType::FP8_FNUZ, 32, 32, 16, GpuTarget::gfx942));
EXPECT_TRUE(isValidWaveTile(DataType::FP8_FNUZ, 16, 16, 32, GpuTarget::gfx942));
}
TEST(IsValidWaveTile, FP8FNUZIterateKTilesOnGfx942Plus)
{
// IterateK compositions of base FP8 MFMA — available on gfx942+
EXPECT_TRUE(isValidWaveTile(DataType::FP8_FNUZ, 32, 32, 32, GpuTarget::gfx942));
EXPECT_TRUE(isValidWaveTile(DataType::FP8_FNUZ, 32, 32, 64, GpuTarget::gfx942));
EXPECT_TRUE(isValidWaveTile(DataType::FP8_FNUZ, 16, 16, 64, GpuTarget::gfx942));
EXPECT_TRUE(isValidWaveTile(DataType::FP8_FNUZ, 32, 32, 32, GpuTarget::gfx950));
EXPECT_TRUE(isValidWaveTile(DataType::FP8_FNUZ, 32, 32, 64, GpuTarget::gfx950));
EXPECT_TRUE(isValidWaveTile(DataType::FP8_FNUZ, 16, 16, 64, GpuTarget::gfx950));
// Still not on gfx90a
EXPECT_FALSE(isValidWaveTile(DataType::FP8_FNUZ, 32, 32, 32, GpuTarget::gfx90a));
}
TEST(IsValidWaveTile, WMMATilesOnRDNA)
{
// All RDNA targets share identical WMMA: 16×16×16 for FP16, BF16, INT8
EXPECT_TRUE(isValidWaveTile(DataType::FP16, 16, 16, 16, GpuTarget::gfx1100));
EXPECT_TRUE(isValidWaveTile(DataType::BF16, 16, 16, 16, GpuTarget::gfx1100));
EXPECT_TRUE(isValidWaveTile(DataType::I8, 16, 16, 16, GpuTarget::gfx1100));
EXPECT_TRUE(isValidWaveTile(DataType::FP16, 16, 16, 16, GpuTarget::gfx1101));
EXPECT_TRUE(isValidWaveTile(DataType::FP16, 16, 16, 16, GpuTarget::gfx1102));
EXPECT_TRUE(isValidWaveTile(DataType::FP16, 16, 16, 16, GpuTarget::gfx1150));
EXPECT_TRUE(isValidWaveTile(DataType::FP16, 16, 16, 16, GpuTarget::gfx1151));
EXPECT_TRUE(isValidWaveTile(DataType::BF16, 16, 16, 16, GpuTarget::gfx1151));
EXPECT_TRUE(isValidWaveTile(DataType::I8, 16, 16, 16, GpuTarget::gfx1151));
}
TEST(IsValidWaveTile, WMMARejectsNon16x16x16)
{
EXPECT_FALSE(isValidWaveTile(DataType::FP16, 32, 32, 8, GpuTarget::gfx1100));
EXPECT_FALSE(isValidWaveTile(DataType::FP16, 16, 16, 32, GpuTarget::gfx1151));
}
TEST(IsValidWaveTile, WMMARejectsFP32)
{
EXPECT_FALSE(isValidWaveTile(DataType::FP32, 16, 16, 16, GpuTarget::gfx1100));
EXPECT_FALSE(isValidWaveTile(DataType::FP32, 16, 16, 16, GpuTarget::gfx1151));
}
// ============================================================================
// isValidWaveTile — TargetSet (intersection semantics)
// ============================================================================
TEST(IsValidWaveTile, IntersectionAcrossCDNATargets)
{
// FP16 16x16x16 valid on all CDNA
EXPECT_TRUE(isValidWaveTile(DataType::FP16, 16, 16, 16, TargetSet::cdna()));
}
TEST(IsValidWaveTile, IntersectionRejectsWhenAnyTargetFails)
{
// FP8 32x32x16 valid on gfx942 but NOT on gfx90a
EXPECT_FALSE(isValidWaveTile(DataType::FP8_FNUZ, 32, 32, 16, TargetSet::cdna()));
// Valid for gfx94 family only
EXPECT_TRUE(isValidWaveTile(DataType::FP8_FNUZ, 32, 32, 16, TargetSet::family_gfx94()));
}

535
rocm_ck/tests/unit/unit_resolve.cpp Normal file
View File

@@ -0,0 +1,535 @@
// Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
// SPDX-License-Identifier: MIT
#include <rocm_ck/resolve.hpp>
#include <gtest/gtest.h>
using ::rocm_ck::AddOp;
using ::rocm_ck::BinaryOpLike;
using ::rocm_ck::DataType;
using ::rocm_ck::FastGeluOp;
using ::rocm_ck::GeluOp;
using ::rocm_ck::GemmOp;
using ::rocm_ck::kMaxTensors;
using ::rocm_ck::Layout;
using ::rocm_ck::MulOp;
using ::rocm_ck::Quantization;
using ::rocm_ck::ReluOp;
using ::rocm_ck::resolve;
using ::rocm_ck::Scalar;
using ::rocm_ck::ScaleOp;
using ::rocm_ck::SigmoidOp;
using ::rocm_ck::Signature;
using ::rocm_ck::SiluOp;
using ::rocm_ck::SoftmaxOp;
using ::rocm_ck::Tensor;
using ::rocm_ck::UnaryOpLike;
// ============================================================================
// Simple GemmOp resolution
// ============================================================================
TEST(Resolve, ResolvesSimpleGemmToThreeTensors)
{
constexpr auto r = resolve(
Signature{.dtype = DataType::FP16, .ops = {GemmOp{.lhs = "A", .rhs = "B", .out = "C"}}});
EXPECT_EQ(r.num_tensors, 3);
}
TEST(Resolve, CascadesSignatureDtypeToAllGemmTensors)
{
constexpr auto r = resolve(
Signature{.dtype = DataType::FP16, .ops = {GemmOp{.lhs = "A", .rhs = "B", .out = "C"}}});
EXPECT_EQ(r.tensor("A").dtype, DataType::FP16);
EXPECT_EQ(r.tensor("B").dtype, DataType::FP16);
EXPECT_EQ(r.tensor("C").dtype, DataType::FP16);
}
TEST(Resolve, AssignsRank2ToGemmTensors)
{
constexpr auto r = resolve(
Signature{.dtype = DataType::FP16, .ops = {GemmOp{.lhs = "A", .rhs = "B", .out = "C"}}});
EXPECT_EQ(r.tensor("A").rank, 2);
EXPECT_EQ(r.tensor("B").rank, 2);
EXPECT_EQ(r.tensor("C").rank, 2);
}
TEST(Resolve, AssignsRowColRowLayoutToGemmTensors)
{
constexpr auto r = resolve(
Signature{.dtype = DataType::FP16, .ops = {GemmOp{.lhs = "A", .rhs = "B", .out = "C"}}});
EXPECT_EQ(r.tensor("A").layout, Layout::Row);
EXPECT_EQ(r.tensor("B").layout, Layout::Col);
EXPECT_EQ(r.tensor("C").layout, Layout::Row);
}
// ============================================================================
// Custom tensor names
// ============================================================================
TEST(Resolve, AcceptsCustomTensorNames)
{
constexpr auto r = resolve(
Signature{.dtype = DataType::FP16, .ops = {GemmOp{.lhs = "X", .rhs = "Y", .out = "Z"}}});
EXPECT_EQ(r.tensor("X").rank, 2);
EXPECT_EQ(r.tensor("Y").rank, 2);
EXPECT_EQ(r.tensor("Z").rank, 2);
}
// ============================================================================
// dtype cascade
// ============================================================================
TEST(Resolve, CascadesBF16DtypeToAllTensors)
{
constexpr auto r = resolve(
Signature{.dtype = DataType::BF16, .ops = {GemmOp{.lhs = "A", .rhs = "B", .out = "C"}}});
EXPECT_EQ(r.tensor("A").dtype, DataType::BF16);
EXPECT_EQ(r.tensor("C").dtype, DataType::BF16);
}
TEST(Resolve, AllowsPerTensorDtypeOverride)
{
constexpr auto r = resolve( //
Signature{.dtype = DataType::FP16,
.tensors = {Tensor{.name = "C", .dtype = DataType::FP32}},
.ops = {GemmOp{.lhs = "A", .rhs = "B", .out = "C"}}});
EXPECT_EQ(r.tensor("C").dtype, DataType::FP32);
EXPECT_EQ(r.tensor("A").dtype, DataType::FP16); // cascade still applies to A
}
// ============================================================================
// Explicit tensor rank/layout overrides
// ============================================================================
TEST(Resolve, AllowsPerTensorRankOverride)
{
constexpr auto r = resolve( //
Signature{.dtype = DataType::FP16,
.tensors = {Tensor{.name = "A", .rank = 3}},
.ops = {GemmOp{.lhs = "A", .rhs = "B", .out = "C"}}});
EXPECT_EQ(r.tensor("A").rank, 3);
}
TEST(Resolve, AllowsPerTensorLayoutOverride)
{
// Override B from default Col to Row (R×R layout)
constexpr auto r = resolve( //
Signature{.dtype = DataType::FP16,
.tensors = {Tensor{.name = "B", .layout = Layout::Row}},
.ops = {GemmOp{.lhs = "A", .rhs = "B", .out = "C"}}});
EXPECT_EQ(r.tensor("A").layout, Layout::Row); // default preserved
EXPECT_EQ(r.tensor("B").layout, Layout::Row); // overridden from Col
EXPECT_EQ(r.tensor("C").layout, Layout::Row); // default preserved
}
TEST(Resolve, AllowsMultipleLayoutOverrides)
{
// Override both A and B (C×C layout)
constexpr auto r = resolve( //
Signature{.dtype = DataType::FP16,
.tensors = {Tensor{.name = "A", .layout = Layout::Col},
Tensor{.name = "B", .layout = Layout::Col}},
.ops = {GemmOp{.lhs = "A", .rhs = "B", .out = "C"}}});
EXPECT_EQ(r.tensor("A").layout, Layout::Col);
EXPECT_EQ(r.tensor("B").layout, Layout::Col);
EXPECT_EQ(r.tensor("C").layout, Layout::Row); // default preserved
}
// ============================================================================
// GEMM + Add + Relu chain
// ============================================================================
TEST(Resolve, ResolvesGemmAddReluToSixTensors)
{
constexpr auto r = resolve( //
Signature{.dtype = DataType::FP16,
.ops = {GemmOp{.lhs = "A", .rhs = "B", .out = "C"},
AddOp{.lhs = "C", .rhs = "bias", .out = "D"},
ReluOp{.in = "D", .out = "E"}}});
EXPECT_EQ(r.num_tensors, 6); // A, B, C, bias, D, E
}
TEST(Resolve, PropagatesRankAndLayoutThroughEpilogueChain)
{
constexpr auto r = resolve( //
Signature{.dtype = DataType::FP16,
.ops = {GemmOp{.lhs = "A", .rhs = "B", .out = "C"},
AddOp{.lhs = "C", .rhs = "bias", .out = "D"},
ReluOp{.in = "D", .out = "E"}}});
EXPECT_EQ(r.tensor("C").rank, 2);
EXPECT_EQ(r.tensor("bias").rank, 2);
EXPECT_EQ(r.tensor("bias").layout, Layout::Row);
EXPECT_EQ(r.tensor("D").rank, 2);
EXPECT_EQ(r.tensor("D").layout, Layout::Row);
EXPECT_EQ(r.tensor("E").rank, 2);
EXPECT_EQ(r.tensor("E").layout, Layout::Row);
}
TEST(Resolve, PropagatesRankAndLayoutThroughDiamondDAG)
{
// Diamond: GEMM→C splits into two Add paths, then joins.
// C → Add(C,bias1)→D1 ─→ Add(D1,D2)→E
// C → Add(C,bias2)→D2 ─┘
constexpr auto r = resolve( //
Signature{.dtype = DataType::FP16,
.ops = {GemmOp{.lhs = "A", .rhs = "B", .out = "C"},
AddOp{.lhs = "C", .rhs = "bias1", .out = "D1"},
AddOp{.lhs = "C", .rhs = "bias2", .out = "D2"},
AddOp{.lhs = "D1", .rhs = "D2", .out = "E"}}});
EXPECT_EQ(r.num_tensors, 8); // A, B, C, bias1, D1, bias2, D2, E
EXPECT_EQ(r.tensor("D1").rank, 2);
EXPECT_EQ(r.tensor("D2").rank, 2);
EXPECT_EQ(r.tensor("E").rank, 2);
EXPECT_EQ(r.tensor("bias1").layout, Layout::Row);
EXPECT_EQ(r.tensor("E").layout, Layout::Row);
}
TEST(Resolve, AssignsSequentialIndicesToChainedOps)
{
constexpr auto r = resolve( //
Signature{.dtype = DataType::FP16,
.ops = {GemmOp{.lhs = "A", .rhs = "B", .out = "C"},
AddOp{.lhs = "C", .rhs = "bias", .out = "D"},
ReluOp{.in = "D", .out = "E"}}});
EXPECT_EQ(r.tensorIndex("A"), 0);
EXPECT_EQ(r.tensorIndex("B"), 1);
EXPECT_EQ(r.tensorIndex("C"), 2);
EXPECT_EQ(r.tensorIndex("bias"), 3);
EXPECT_EQ(r.tensorIndex("D"), 4);
EXPECT_EQ(r.tensorIndex("E"), 5);
}
// ============================================================================
// Standalone AddOp
// ============================================================================
TEST(Resolve, ResolvesStandaloneAddWithoutImpliedRank)
{
constexpr auto r = resolve(
Signature{.dtype = DataType::FP32, .ops = {AddOp{.lhs = "A", .rhs = "B", .out = "C"}}});
EXPECT_EQ(r.num_tensors, 3);
EXPECT_EQ(r.tensor("A").rank, 0); // no op implies rank
EXPECT_EQ(r.tensor("A").layout, Layout::Auto); // no op implies layout
}
// ============================================================================
// Conflict detection — redundant identical sets are silent
// ============================================================================
TEST(Resolve, AllowsRedundantIdenticalLayoutFromTwoGemmOps)
{
// GemmOp1 outputs "C" as Row. GemmOp2 uses "C" as lhs (also Row).
// Two ops set the same layout → no conflict.
constexpr auto r = resolve( //
Signature{.dtype = DataType::FP16,
.ops = {GemmOp{.lhs = "A", .rhs = "B", .out = "C"},
GemmOp{.lhs = "C", .rhs = "D", .out = "E"}}});
EXPECT_EQ(r.tensor("C").layout, Layout::Row);
EXPECT_EQ(r.tensor("C").rank, 2);
}
TEST(Resolve, AllowsPropagationThroughAddWithConsistentLayout)
{
// GemmOp sets C=Row. AddOp connects C to bias and D.
// Propagation sets bias and D to Row (matching C) → no conflict.
constexpr auto r = resolve( //
Signature{.dtype = DataType::FP16,
.ops = {GemmOp{.lhs = "A", .rhs = "B", .out = "C"},
AddOp{.lhs = "C", .rhs = "bias", .out = "D"}}});
EXPECT_EQ(r.tensor("C").layout, Layout::Row);
EXPECT_EQ(r.tensor("bias").layout, Layout::Row);
EXPECT_EQ(r.tensor("D").layout, Layout::Row);
}
// ============================================================================
// FMHA pattern: two GemmOps + SoftmaxOp
// ============================================================================
TEST(Resolve, ResolvesFMHATwoGemmSoftmaxPattern)
{
constexpr auto r = resolve( //
Signature{.dtype = DataType::FP16,
.ops = {GemmOp{.lhs = "Q", .rhs = "K", .out = "S"},
SoftmaxOp{.in = "S", .out = "P"},
GemmOp{.lhs = "P", .rhs = "V", .out = "O"}}});
EXPECT_EQ(r.num_tensors, 6); // Q, K, S, P, V, O
EXPECT_EQ(r.tensor("Q").rank, 2);
EXPECT_EQ(r.tensor("S").rank, 2);
EXPECT_EQ(r.tensor("P").rank, 2); // propagated via SoftmaxOp
EXPECT_EQ(r.tensor("O").rank, 2);
}
// ============================================================================
// Scalar tracking
// ============================================================================
TEST(Resolve, PreservesScalarNamesAndDtypes)
{
constexpr auto r = resolve( //
Signature{.dtype = DataType::FP16,
.scalars = {Scalar{.name = "alpha", .dtype = DataType::FP32},
Scalar{.name = "beta", .dtype = DataType::FP32}},
.ops = {GemmOp{.lhs = "A", .rhs = "B", .out = "C"}}});
EXPECT_EQ(r.num_scalars, 2);
EXPECT_EQ(r.scalar("alpha").dtype, DataType::FP32);
EXPECT_EQ(r.scalar("beta").dtype, DataType::FP32);
EXPECT_EQ(r.scalarIndex("alpha"), 0);
EXPECT_EQ(r.scalarIndex("beta"), 1);
}
TEST(Resolve, ReportsZeroScalarsWhenNoneDeclared)
{
constexpr auto r = resolve(
Signature{.dtype = DataType::FP16, .ops = {GemmOp{.lhs = "A", .rhs = "B", .out = "C"}}});
EXPECT_EQ(r.num_scalars, 0);
}
// ============================================================================
// findTensor / findScalar (constexpr, not consteval — returns -1 on miss)
// ============================================================================
TEST(Resolve, FindsTensorByName)
{
constexpr auto r = resolve(
Signature{.dtype = DataType::FP16, .ops = {GemmOp{.lhs = "A", .rhs = "B", .out = "C"}}});
EXPECT_EQ(r.findTensor("A"), 0);
EXPECT_EQ(r.findTensor("C"), 2);
}
TEST(Resolve, ReturnsNegativeOneForUnknownTensor)
{
constexpr auto r = resolve(
Signature{.dtype = DataType::FP16, .ops = {GemmOp{.lhs = "A", .rhs = "B", .out = "C"}}});
EXPECT_EQ(r.findTensor("Z"), -1);
}
TEST(Resolve, ReturnsNegativeOneForUnknownScalar)
{
constexpr auto r = resolve(
Signature{.dtype = DataType::FP16, .ops = {GemmOp{.lhs = "A", .rhs = "B", .out = "C"}}});
EXPECT_EQ(r.findScalar("nonexistent"), -1);
}
// ============================================================================
// Quantized tensors
// ============================================================================
TEST(Resolve, QuantizedBAutoRegistersScaleTensor)
{
constexpr auto r = resolve( //
Signature{.dtype = DataType::FP16,
.tensors = {Tensor{.name = "B",
.dtype = DataType::I4,
.quantize = Quantization{.scale_name = "scale",
.scale_dtype = DataType::FP32,
.group_size = 128}}},
.ops = {GemmOp{.lhs = "A", .rhs = "B", .out = "C"}}});
// A, B, C from GemmOp + scale auto-registered = 4 tensors
EXPECT_EQ(r.num_tensors, 4);
}
TEST(Resolve, ScaleTensorGetsDtypeFromQuantization)
{
constexpr auto r = resolve(
Signature{.dtype = DataType::FP16,
.tensors = {Tensor{.name = "B",
.dtype = DataType::I4,
.quantize = Quantization{.scale_name = "scale",
.scale_dtype = DataType::FP32}}},
.ops = {GemmOp{.lhs = "A", .rhs = "B", .out = "C"}}});
// Scale tensor dtype comes from Quantization, not the signature cascade
EXPECT_EQ(r.tensor("scale").dtype, DataType::FP32);
}
TEST(Resolve, ScaleTensorGetsRank2RowLayout)
{
constexpr auto r = resolve( //
Signature{.dtype = DataType::FP16,
.tensors = {Tensor{.name = "B",
.dtype = DataType::I4,
.quantize = Quantization{.scale_name = "scale"}}},
.ops = {GemmOp{.lhs = "A", .rhs = "B", .out = "C"}}});
EXPECT_EQ(r.tensor("scale").rank, 2);
EXPECT_EQ(r.tensor("scale").layout, Layout::Row);
}
TEST(Resolve, QuantizedTensorKeepsOwnDtype)
{
constexpr auto r = resolve( //
Signature{.dtype = DataType::FP16,
.tensors = {Tensor{.name = "B",
.dtype = DataType::I4,
.quantize = Quantization{.scale_name = "scale"}}},
.ops = {GemmOp{.lhs = "A", .rhs = "B", .out = "C"}}});
EXPECT_EQ(r.tensor("B").dtype, DataType::I4);
EXPECT_EQ(r.tensor("A").dtype, DataType::FP16); // cascade still works
}
TEST(Resolve, QuantizedResolvedTensorCarriesQuantizeInfo)
{
constexpr auto r = resolve( //
Signature{.dtype = DataType::FP16,
.tensors = {Tensor{.name = "B",
.dtype = DataType::I4,
.quantize = Quantization{.scale_name = "scale",
.scale_dtype = DataType::FP32,
.group_size = 64}}},
.ops = {GemmOp{.lhs = "A", .rhs = "B", .out = "C"}}});
EXPECT_TRUE(r.tensor("B").quantize.has_value());
EXPECT_EQ(r.tensor("B").quantize->scale_name, "scale");
EXPECT_EQ(r.tensor("B").quantize->scale_dtype, DataType::FP32);
EXPECT_EQ(r.tensor("B").quantize->group_size, 64);
}
TEST(Resolve, NonQuantizedTensorHasNoQuantizeInfo)
{
constexpr auto r = resolve( //
Signature{.dtype = DataType::FP16,
.tensors = {Tensor{.name = "B",
.dtype = DataType::I4,
.quantize = Quantization{.scale_name = "scale"}}},
.ops = {GemmOp{.lhs = "A", .rhs = "B", .out = "C"}}});
EXPECT_FALSE(r.tensor("A").quantize.has_value());
EXPECT_FALSE(r.tensor("C").quantize.has_value());
EXPECT_FALSE(r.tensor("scale").quantize.has_value());
}
TEST(Resolve, QuantizedGemmWithEpiloguePreservesScaleTensor)
{
constexpr auto r = resolve( //
Signature{.dtype = DataType::FP16,
.tensors = {Tensor{.name = "B",
.dtype = DataType::I4,
.quantize = Quantization{.scale_name = "scale"}}},
.ops = {GemmOp{.lhs = "A", .rhs = "B", .out = "C"},
AddOp{.lhs = "C", .rhs = "bias", .out = "D"},
ReluOp{.in = "D", .out = "E"}}});
// A, B, C, bias, D, E from ops + scale auto-registered = 7
EXPECT_EQ(r.num_tensors, 7);
EXPECT_EQ(r.tensor("scale").dtype, DataType::FP32);
EXPECT_TRUE(r.tensor("B").quantize.has_value());
}
// ============================================================================
// C++20 concepts
// ============================================================================
TEST(Concepts, ClassifiesAddAndMulAsBinaryOpLike)
{
EXPECT_TRUE(BinaryOpLike<AddOp>);
EXPECT_TRUE(BinaryOpLike<MulOp>);
EXPECT_FALSE(BinaryOpLike<ReluOp>);
EXPECT_FALSE(BinaryOpLike<SoftmaxOp>);
}
TEST(Concepts, ClassifiesActivationsAsUnaryOpLike)
{
EXPECT_TRUE(UnaryOpLike<ReluOp>);
EXPECT_TRUE(UnaryOpLike<FastGeluOp>);
EXPECT_TRUE(UnaryOpLike<GeluOp>);
EXPECT_TRUE(UnaryOpLike<SiluOp>);
EXPECT_TRUE(UnaryOpLike<SigmoidOp>);
EXPECT_TRUE(UnaryOpLike<SoftmaxOp>);
EXPECT_FALSE(UnaryOpLike<AddOp>);
EXPECT_FALSE(UnaryOpLike<GemmOp>);
}
TEST(Concepts, ClassifiesGemmOpAsBinaryButNotUnary)
{
// GemmOp has lhs/rhs/out AND is special-cased, not generic BinaryOpLike
// (it has .lhs, .rhs, .out but is handled separately in registerSlots)
EXPECT_TRUE(BinaryOpLike<GemmOp>); // structurally matches, but dispatch special-cases it
EXPECT_FALSE(UnaryOpLike<GemmOp>);
}
TEST(Concepts, ClassifiesScaleOpAsUnaryNotBinary)
{
EXPECT_TRUE(UnaryOpLike<ScaleOp>);
EXPECT_FALSE(BinaryOpLike<ScaleOp>);
}
// ============================================================================
// ScaleOp with explicit Scalar
// ============================================================================
TEST(Resolve, ScaleOpReferencesExplicitScalar)
{
constexpr auto r = resolve( //
Signature{.dtype = DataType::FP16,
.scalars = {Scalar{.name = "alpha", .dtype = DataType::FP32}},
.ops = {GemmOp{.lhs = "A", .rhs = "B", .out = "C"},
ScaleOp{.in = "C", .out = "D", .scale = "alpha"}}});
EXPECT_EQ(r.num_tensors, 4); // A, B, C, D
EXPECT_EQ(r.num_scalars, 1);
EXPECT_EQ(r.scalar("alpha").dtype, DataType::FP32);
EXPECT_EQ(r.scalarIndex("alpha"), 0);
}
TEST(Resolve, ScaleOpPreservesScalarDtype)
{
constexpr auto r = resolve( //
Signature{.dtype = DataType::FP16,
.scalars = {Scalar{.name = "scale_factor", .dtype = DataType::FP16}},
.ops = {GemmOp{.lhs = "A", .rhs = "B", .out = "C"},
ScaleOp{.in = "C", .out = "D", .scale = "scale_factor"}}});
EXPECT_EQ(r.scalar("scale_factor").dtype, DataType::FP16);
}
// ============================================================================
// Boundary: signature at kMaxTensors
// ============================================================================
TEST(Resolve, HandlesSignatureWithManyTensors)
{
// Create a chain of AddOps to generate many tensors (close to kMaxTensors).
// Each AddOp creates 3 tensors (lhs, rhs, out). We'll create a chain that
// approaches the limit.
// kMaxTensors is 16, so a signature with 3 GEMMs (each with 3 tensors = 9)
// plus some adds should get close.
constexpr auto r = resolve( //
Signature{.dtype = DataType::FP16,
.ops = {GemmOp{.lhs = "A1", .rhs = "B1", .out = "C1"},
GemmOp{.lhs = "A2", .rhs = "B2", .out = "C2"},
GemmOp{.lhs = "A3", .rhs = "B3", .out = "C3"},
AddOp{.lhs = "C1", .rhs = "C2", .out = "D1"},
AddOp{.lhs = "D1", .rhs = "C3", .out = "D2"}}});
// A1, B1, C1, A2, B2, C2, A3, B3, C3, D1, D2 = 11 tensors
EXPECT_EQ(r.num_tensors, 11);
EXPECT_EQ(r.tensor("A1").dtype, DataType::FP16);
EXPECT_EQ(r.tensor("D2").dtype, DataType::FP16);
}

176
rocm_ck/tests/unit/unit_signature.cpp Normal file
View File

@@ -0,0 +1,176 @@
// Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
// SPDX-License-Identifier: MIT
#include <rocm_ck/signature.hpp>
#include <gtest/gtest.h>
using ::rocm_ck::AddOp;
using ::rocm_ck::DataType;
using ::rocm_ck::FastGeluOp;
using ::rocm_ck::GemmOp;
using ::rocm_ck::kMaxOps;
using ::rocm_ck::kMaxScalars;
using ::rocm_ck::kMaxTensors;
using ::rocm_ck::Layout;
using ::rocm_ck::MulOp;
using ::rocm_ck::Op;
using ::rocm_ck::Quantization;
using ::rocm_ck::ReluOp;
using ::rocm_ck::Scalar;
using ::rocm_ck::SigmoidOp;
using ::rocm_ck::Signature;
using ::rocm_ck::Tensor;
// ============================================================================
// Signature construction
// ============================================================================
TEST(Signature, DefaultsToNoDtype)
{
constexpr Signature sig{};
EXPECT_FALSE(sig.dtype.has_value());
}
TEST(Signature, StoresExplicitDtype)
{
constexpr Signature sig{.dtype = DataType::FP16};
EXPECT_TRUE(sig.dtype.has_value());
EXPECT_EQ(*sig.dtype, DataType::FP16);
}
// ============================================================================
// Tensor
// ============================================================================
TEST(Tensor, DefaultsToAutoLayoutAndRankZero)
{
constexpr Tensor t{.name = "A"};
EXPECT_EQ(t.name, "A");
EXPECT_FALSE(t.dtype.has_value());
EXPECT_EQ(t.rank, 0);
EXPECT_EQ(t.layout, Layout::Auto);
}
TEST(Tensor, StoresAllExplicitFields)
{
constexpr Tensor t{.name = "Q", .dtype = DataType::FP32, .rank = 3, .layout = Layout::Row};
EXPECT_EQ(t.name, "Q");
EXPECT_EQ(*t.dtype, DataType::FP32);
EXPECT_EQ(t.rank, 3);
EXPECT_EQ(t.layout, Layout::Row);
}
// ============================================================================
// Scalar
// ============================================================================
TEST(Scalar, DefaultsToFP32Dtype)
{
constexpr Scalar s{.name = "alpha"};
EXPECT_EQ(s.name, "alpha");
EXPECT_EQ(s.dtype, DataType::FP32);
}
TEST(Scalar, StoresExplicitDtype)
{
constexpr Scalar s{.name = "scale", .dtype = DataType::FP16};
EXPECT_EQ(s.dtype, DataType::FP16);
}
// ============================================================================
// Op variant
// ============================================================================
TEST(Op, DefaultsToMonostate)
{
constexpr Op op{};
EXPECT_TRUE(std::holds_alternative<std::monostate>(op));
}
TEST(Op, HoldsGemmOp)
{
constexpr Op op = GemmOp{.lhs = "A", .rhs = "B", .out = "C"};
EXPECT_TRUE(std::holds_alternative<GemmOp>(op));
}
TEST(Op, HoldsAllUnaryOpTypes)
{
constexpr Op relu = ReluOp{.in = "X", .out = "Y"};
EXPECT_TRUE(std::holds_alternative<ReluOp>(relu));
constexpr Op gelu = FastGeluOp{.in = "X", .out = "Y"};
EXPECT_TRUE(std::holds_alternative<FastGeluOp>(gelu));
constexpr Op sigmoid = SigmoidOp{.in = "X", .out = "Y"};
EXPECT_TRUE(std::holds_alternative<SigmoidOp>(sigmoid));
}
TEST(Op, HoldsAllBinaryOpTypes)
{
constexpr Op add = AddOp{.lhs = "X", .rhs = "Y", .out = "Z"};
EXPECT_TRUE(std::holds_alternative<AddOp>(add));
constexpr Op mul = MulOp{.lhs = "X", .rhs = "Y", .out = "Z"};
EXPECT_TRUE(std::holds_alternative<MulOp>(mul));
}
// ============================================================================
// GemmOp defaults
// ============================================================================
TEST(GemmOp, DefaultsAccDtypeToFP32)
{
constexpr GemmOp gemm{.lhs = "A", .rhs = "B", .out = "C"};
EXPECT_EQ(gemm.acc_dtype, DataType::FP32);
}
// ============================================================================
// Quantization
// ============================================================================
TEST(Quantization, DefaultsToFP32ScaleAndGroupSize128)
{
constexpr Quantization q{.scale_name = "scale"};
EXPECT_EQ(q.scale_name, "scale");
EXPECT_EQ(q.scale_dtype, DataType::FP32);
EXPECT_EQ(q.group_size, 128);
}
TEST(Quantization, StoresExplicitFields)
{
constexpr Quantization q{.scale_name = "bq", .scale_dtype = DataType::FP16, .group_size = 64};
EXPECT_EQ(q.scale_name, "bq");
EXPECT_EQ(q.scale_dtype, DataType::FP16);
EXPECT_EQ(q.group_size, 64);
}
TEST(Tensor, DefaultsToNoQuantize)
{
constexpr Tensor t{.name = "B"};
EXPECT_FALSE(t.quantize.has_value());
}
TEST(Tensor, StoresQuantizeMetadata)
{
constexpr Tensor t{
.name = "B",
.dtype = DataType::I4,
.quantize =
Quantization{.scale_name = "scale", .scale_dtype = DataType::FP32, .group_size = 128}};
EXPECT_TRUE(t.quantize.has_value());
EXPECT_EQ(t.quantize->scale_name, "scale");
EXPECT_EQ(t.quantize->scale_dtype, DataType::FP32);
EXPECT_EQ(t.quantize->group_size, 128);
}
// ============================================================================
// Capacity constants
// ============================================================================
TEST(Signature, DefinesExpectedCapacityLimits)
{
EXPECT_EQ(kMaxTensors, 16);
EXPECT_EQ(kMaxScalars, 16);
EXPECT_EQ(kMaxOps, 8);
}