diff --git a/rocm_ck/README.md b/rocm_ck/README.md index 57f7786e44..3a9d95a4f4 100644 --- a/rocm_ck/README.md +++ b/rocm_ck/README.md @@ -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 ``` diff --git a/rocm_ck/include/rocm_ck/arch_properties.hpp b/rocm_ck/include/rocm_ck/arch_properties.hpp new file mode 100644 index 0000000000..6e58ba2900 --- /dev/null +++ b/rocm_ck/include/rocm_ck/arch_properties.hpp @@ -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 +#include + +#include +#include + +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 +/// 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(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(target); + } + + static constexpr GpuTarget targetAt(int index) + { + if(index < 0 || index >= kNumTargets) + throw "TargetSet index out of range [0, kNumTargets)"; + return static_cast(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 + 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 diff --git a/rocm_ck/include/rocm_ck/gpu_target.hpp b/rocm_ck/include/rocm_ck/gpu_target.hpp index 0a0ed83803..9ee6510039 100644 --- a/rocm_ck/include/rocm_ck/gpu_target.hpp +++ b/rocm_ck/include/rocm_ck/gpu_target.hpp @@ -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 diff --git a/rocm_ck/include/rocm_ck/ops.hpp b/rocm_ck/include/rocm_ck/ops.hpp index ca49cc75d9..723683c5c7 100644 --- a/rocm_ck/include/rocm_ck/ops.hpp +++ b/rocm_ck/include/rocm_ck/ops.hpp @@ -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; + ScaleOp>; } // namespace rocm_ck diff --git a/rocm_ck/include/rocm_ck/resolve.hpp b/rocm_ck/include/rocm_ck/resolve.hpp new file mode 100644 index 0000000000..080c9a0565 --- /dev/null +++ b/rocm_ck/include/rocm_ck/resolve.hpp @@ -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 +#include + +#include +#include + +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 +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 +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 +consteval auto visitOp(const Op& op, F&& f) +{ + return std::visit(std::forward(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 tensors = {}; + + int num_scalars = 0; + std::array 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; + + if constexpr(std::is_same_v) + { + return {}; + } + else if constexpr(std::is_same_v) + { + // 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) + { + // 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) + { + 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) + { + 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; + + if constexpr(std::is_same_v || std::is_same_v) + { + // monostate: nothing to do + // GemmOp: rank/layout already set in collectTensorSlotsFromOp + } + else if constexpr(BinaryOpLike) + { + propagate_binary(typed_op.lhs, typed_op.rhs, typed_op.out); + } + else if constexpr(UnaryOpLike) + { + 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 diff --git a/rocm_ck/include/rocm_ck/signature.hpp b/rocm_ck/include/rocm_ck/signature.hpp new file mode 100644 index 0000000000..1e950d3e5b --- /dev/null +++ b/rocm_ck/include/rocm_ck/signature.hpp @@ -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 +#include +#include +#include + +#include +#include +#include + +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 dtype = std::nullopt; + int rank = 0; + Layout layout = Layout::Auto; + std::optional 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 dtype; + std::array tensors = {}; + std::array scalars = {}; + std::array ops = {}; +}; + +} // namespace rocm_ck diff --git a/rocm_ck/tests/CMakeLists.txt b/rocm_ck/tests/CMakeLists.txt index 8c48b6e6ba..64ffc5b355 100644 --- a/rocm_ck/tests/CMakeLists.txt +++ b/rocm_ck/tests/CMakeLists.txt @@ -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) diff --git a/rocm_ck/tests/compile_fail/CMakeLists.txt b/rocm_ck/tests/compile_fail/CMakeLists.txt new file mode 100644 index 0000000000..1ce4c751a3 --- /dev/null +++ b/rocm_ck/tests/compile_fail/CMakeLists.txt @@ -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") diff --git a/rocm_ck/tests/compile_fail/conflicting_layout.cpp b/rocm_ck/tests/compile_fail/conflicting_layout.cpp new file mode 100644 index 0000000000..fbae32317c --- /dev/null +++ b/rocm_ck/tests/compile_fail/conflicting_layout.cpp @@ -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 + +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"}}}); diff --git a/rocm_ck/tests/compile_fail/empty_target_set_wavefront.cpp b/rocm_ck/tests/compile_fail/empty_target_set_wavefront.cpp new file mode 100644 index 0000000000..388338e079 --- /dev/null +++ b/rocm_ck/tests/compile_fail/empty_target_set_wavefront.cpp @@ -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 + +using namespace rocm_ck; + +constexpr int bad = TargetSet{}.wavefront_size(); diff --git a/rocm_ck/tests/compile_fail/empty_tensor_name.cpp b/rocm_ck/tests/compile_fail/empty_tensor_name.cpp new file mode 100644 index 0000000000..ac1aada23f --- /dev/null +++ b/rocm_ck/tests/compile_fail/empty_tensor_name.cpp @@ -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 + +using namespace rocm_ck; + +constexpr auto bad = + resolve(Signature{.dtype = DataType::FP16, .ops = {GemmOp{.lhs = "", .rhs = "B", .out = "C"}}}); diff --git a/rocm_ck/tests/compile_fail/fp8_ocp_unsupported.cpp b/rocm_ck/tests/compile_fail/fp8_ocp_unsupported.cpp new file mode 100644 index 0000000000..1ca5fb4959 --- /dev/null +++ b/rocm_ck/tests/compile_fail/fp8_ocp_unsupported.cpp @@ -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 + +using namespace rocm_ck; + +constexpr bool bad = isValidWaveTile(DataType::FP8_OCP, 16, 16, 16, GpuTarget::gfx942); diff --git a/rocm_ck/tests/compile_fail/layout_auto_stride.cpp b/rocm_ck/tests/compile_fail/layout_auto_stride.cpp new file mode 100644 index 0000000000..947b186ad0 --- /dev/null +++ b/rocm_ck/tests/compile_fail/layout_auto_stride.cpp @@ -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 + +#include + +using namespace rocm_ck; + +constexpr auto bad = leadingDimStride(Layout::Auto, std::array{1, 2}); diff --git a/rocm_ck/tests/compile_fail/no_dtype.cpp b/rocm_ck/tests/compile_fail/no_dtype.cpp new file mode 100644 index 0000000000..43a22ada60 --- /dev/null +++ b/rocm_ck/tests/compile_fail/no_dtype.cpp @@ -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 + +using namespace rocm_ck; + +constexpr auto bad = resolve(Signature{.ops = {GemmOp{.lhs = "A", .rhs = "B", .out = "C"}}}); diff --git a/rocm_ck/tests/compile_fail/quantize_empty_scale_name.cpp b/rocm_ck/tests/compile_fail/quantize_empty_scale_name.cpp new file mode 100644 index 0000000000..b96a6ae15f --- /dev/null +++ b/rocm_ck/tests/compile_fail/quantize_empty_scale_name.cpp @@ -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 + +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"}}}); diff --git a/rocm_ck/tests/compile_fail/scale_scalar_not_found.cpp b/rocm_ck/tests/compile_fail/scale_scalar_not_found.cpp new file mode 100644 index 0000000000..8fbd043cc1 --- /dev/null +++ b/rocm_ck/tests/compile_fail/scale_scalar_not_found.cpp @@ -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 + +using namespace rocm_ck; + +constexpr auto bad = resolve( + Signature{.dtype = DataType::FP16, .ops = {ScaleOp{.in = "A", .out = "B", .scale = "alpha"}}}); diff --git a/rocm_ck/tests/compile_fail/ssa_violation.cpp b/rocm_ck/tests/compile_fail/ssa_violation.cpp new file mode 100644 index 0000000000..7707e805ee --- /dev/null +++ b/rocm_ck/tests/compile_fail/ssa_violation.cpp @@ -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 + +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"}}}); diff --git a/rocm_ck/tests/compile_fail/target_set_mixed_wave.cpp b/rocm_ck/tests/compile_fail/target_set_mixed_wave.cpp new file mode 100644 index 0000000000..a35a23d5d6 --- /dev/null +++ b/rocm_ck/tests/compile_fail/target_set_mixed_wave.cpp @@ -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 + +using namespace rocm_ck; + +constexpr auto mixed = TargetSet::all(); +constexpr int bad_wave = mixed.wavefront_size(); // Must fail: mixed wave64/wave32 diff --git a/rocm_ck/tests/compile_fail/tensor_name_too_long.cpp b/rocm_ck/tests/compile_fail/tensor_name_too_long.cpp new file mode 100644 index 0000000000..6da29d86b2 --- /dev/null +++ b/rocm_ck/tests/compile_fail/tensor_name_too_long.cpp @@ -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 + +using namespace rocm_ck; + +constexpr FixedString<16> bad("ABCDEFGHIJKLMNOP"); // 16 chars, needs 17 with null diff --git a/rocm_ck/tests/unit/unit_arch_properties.cpp b/rocm_ck/tests/unit/unit_arch_properties.cpp new file mode 100644 index 0000000000..7d45be7c27 --- /dev/null +++ b/rocm_ck/tests/unit/unit_arch_properties.cpp @@ -0,0 +1,487 @@ +// Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +// SPDX-License-Identifier: MIT + +#include + +#include + +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())); +} diff --git a/rocm_ck/tests/unit/unit_resolve.cpp b/rocm_ck/tests/unit/unit_resolve.cpp new file mode 100644 index 0000000000..2b90f36596 --- /dev/null +++ b/rocm_ck/tests/unit/unit_resolve.cpp @@ -0,0 +1,535 @@ +// Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +// SPDX-License-Identifier: MIT + +#include + +#include + +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); + EXPECT_TRUE(BinaryOpLike); + EXPECT_FALSE(BinaryOpLike); + EXPECT_FALSE(BinaryOpLike); +} + +TEST(Concepts, ClassifiesActivationsAsUnaryOpLike) +{ + EXPECT_TRUE(UnaryOpLike); + EXPECT_TRUE(UnaryOpLike); + EXPECT_TRUE(UnaryOpLike); + EXPECT_TRUE(UnaryOpLike); + EXPECT_TRUE(UnaryOpLike); + EXPECT_TRUE(UnaryOpLike); + EXPECT_FALSE(UnaryOpLike); + EXPECT_FALSE(UnaryOpLike); +} + +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); // structurally matches, but dispatch special-cases it + EXPECT_FALSE(UnaryOpLike); +} + +TEST(Concepts, ClassifiesScaleOpAsUnaryNotBinary) +{ + EXPECT_TRUE(UnaryOpLike); + EXPECT_FALSE(BinaryOpLike); +} + +// ============================================================================ +// 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); +} diff --git a/rocm_ck/tests/unit/unit_signature.cpp b/rocm_ck/tests/unit/unit_signature.cpp new file mode 100644 index 0000000000..5cb5c2c886 --- /dev/null +++ b/rocm_ck/tests/unit/unit_signature.cpp @@ -0,0 +1,176 @@ +// Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +// SPDX-License-Identifier: MIT + +#include + +#include + +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(op)); +} + +TEST(Op, HoldsGemmOp) +{ + constexpr Op op = GemmOp{.lhs = "A", .rhs = "B", .out = "C"}; + EXPECT_TRUE(std::holds_alternative(op)); +} + +TEST(Op, HoldsAllUnaryOpTypes) +{ + constexpr Op relu = ReluOp{.in = "X", .out = "Y"}; + EXPECT_TRUE(std::holds_alternative(relu)); + + constexpr Op gelu = FastGeluOp{.in = "X", .out = "Y"}; + EXPECT_TRUE(std::holds_alternative(gelu)); + + constexpr Op sigmoid = SigmoidOp{.in = "X", .out = "Y"}; + EXPECT_TRUE(std::holds_alternative(sigmoid)); +} + +TEST(Op, HoldsAllBinaryOpTypes) +{ + constexpr Op add = AddOp{.lhs = "X", .rhs = "Y", .out = "Z"}; + EXPECT_TRUE(std::holds_alternative(add)); + + constexpr Op mul = MulOp{.lhs = "X", .rhs = "Y", .out = "Z"}; + EXPECT_TRUE(std::holds_alternative(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); +}