/home/noah/src/trueno/src/tiling/geometry.rs

Source
//! TCB geometry types and level definitions.

use serde::{Deserialize, Serialize};
use std::fmt;

// ============================================================================
// TILE-001: TcbGeometry Struct
// ============================================================================

/// Dimensions for a Tiling Compute Block
///
/// Represents the (M, N, K) dimensions of a tile in matrix operations:
/// - M: Output rows
/// - N: Output columns
/// - K: Reduction dimension (inner product)
///
/// # Alignment Constraints
///
/// Per the TCB-03 pattern (Tile Quantization Alignment), K must align with
/// the quantization superblock size:
/// - Q4_0: K % 32 == 0
/// - Q4_K: K % 256 == 0
/// - Q8_0: K % 32 == 0
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub struct TcbGeometry {
    /// Items processed in M dimension (rows)
    pub m: u32,
    /// Items processed in N dimension (columns)
    pub n: u32,
    /// Reduction dimension (inner product)
    pub k: u32,
    /// Alignment requirement in bytes (typically 16 for SIMD, 32 for AVX2, 64 for AVX-512)
    pub alignment: u32,
}

impl TcbGeometry {
    /// Create a new TCB geometry
    ///
    /// # Panics
    /// Panics if any dimension is zero.
    #[must_use]
    pub fn new(m: u32, n: u32, k: u32) -> Self {
        assert!(m > 0 && n > 0 && k > 0, "TCB dimensions must be non-zero");
        Self {
            m,
            n,
            k,
            alignment: 16, // Default to SSE/NEON alignment
        }
    }

    /// Create geometry with explicit alignment
    #[must_use]
    pub fn with_alignment(m: u32, n: u32, k: u32, alignment: u32) -> Self {
        assert!(m > 0 && n > 0 && k > 0, "TCB dimensions must be non-zero");
        assert!(
            alignment.is_power_of_two(),
            "Alignment must be power of 2"
        );
        Self { m, n, k, alignment }
    }

    /// Calculate arithmetic intensity (FLOPS per byte loaded)
    ///
    /// For GEMM: AI = (2 * M * N * K) / (M*K + K*N) * sizeof(f32)
    ///
    /// Higher AI means compute-bound; lower means memory-bound.
    #[must_use]
    pub fn arithmetic_intensity(&self) -> f32 {
        let flops = 2.0 * self.m as f64 * self.n as f64 * self.k as f64;
        let bytes = (self.m as f64 * self.k as f64 + self.k as f64 * self.n as f64) * 4.0;
        (flops / bytes) as f32
    }

    /// Calculate total elements in the tile
    #[must_use]
    pub fn total_elements(&self) -> u64 {
        self.m as u64 * self.n as u64
    }

    /// Calculate total FLOPs for this tile
    #[must_use]
    pub fn total_flops(&self) -> u64 {
        2 * self.m as u64 * self.n as u64 * self.k as u64
    }

    /// Check if K dimension aligns with Q4_K superblock (256)
    #[must_use]
    pub fn is_q4k_aligned(&self) -> bool {
        self.k % 256 == 0
    }

    /// Check if K dimension aligns with Q4_0/Q8_0 block (32)
    #[must_use]
    pub fn is_q4_0_aligned(&self) -> bool {
        self.k % 32 == 0
    }

    /// Calculate bytes needed for A tile (M × K × sizeof(f32))
    #[must_use]
    pub fn a_tile_bytes(&self) -> usize {
        self.m as usize * self.k as usize * 4
    }

    /// Calculate bytes needed for B tile (K × N × sizeof(f32))
    #[must_use]
    pub fn b_tile_bytes(&self) -> usize {
        self.k as usize * self.n as usize * 4
    }

    /// Calculate bytes needed for C tile (M × N × sizeof(f32))
    #[must_use]
    pub fn c_tile_bytes(&self) -> usize {
        self.m as usize * self.n as usize * 4
    }

    /// Check if tile fits in given cache size (bytes)
    #[must_use]
    pub fn fits_in_cache(&self, cache_bytes: usize) -> bool {
        self.a_tile_bytes() + self.b_tile_bytes() <= cache_bytes
    }
}

impl Default for TcbGeometry {
    fn default() -> Self {
        // Sensible default: 4×4 micro-tile for SIMD
        Self {
            m: 4,
            n: 4,
            k: 4,
            alignment: 16,
        }
    }
}

impl fmt::Display for TcbGeometry {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(
            f,
            "TCB({}×{}×{}, align={}, AI={:.2})",
            self.m,
            self.n,
            self.k,
            self.alignment,
            self.arithmetic_intensity()
        )
    }
}
// ============================================================================
// TILE-001: Tiling Levels
// ============================================================================

/// Tiling hierarchy level
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum TcbLevel {
    /// Macro-tile: L3 cache / GPU global memory partitioning
    Macro,
    /// Midi-tile: L2 cache / GPU shared memory
    Midi,
    /// Micro-tile: Registers / SIMD lanes
    Micro,
}

impl TcbLevel {
    /// Get typical cache size for this level (x86_64)
    #[must_use]
    pub fn typical_cache_bytes(&self) -> usize {
        match self {
            TcbLevel::Macro => 32 * 1024 * 1024, // 32 MB L3
            TcbLevel::Midi => 256 * 1024,         // 256 KB L2
            TcbLevel::Micro => 32 * 1024,         // 32 KB L1
        }
    }
}

Line	Count	Source
1		//! TCB geometry types and level definitions.
2
3		use serde::{Deserialize, Serialize};
4		use std::fmt;
5
6		// ============================================================================
7		// TILE-001: TcbGeometry Struct
8		// ============================================================================
9
10		/// Dimensions for a Tiling Compute Block
11		///
12		/// Represents the (M, N, K) dimensions of a tile in matrix operations:
13		/// - M: Output rows
14		/// - N: Output columns
15		/// - K: Reduction dimension (inner product)
16		///
17		/// # Alignment Constraints
18		///
19		/// Per the TCB-03 pattern (Tile Quantization Alignment), K must align with
20		/// the quantization superblock size:
21		/// - Q4_0: K % 32 == 0
22		/// - Q4_K: K % 256 == 0
23		/// - Q8_0: K % 32 == 0
24		#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
25		pub struct TcbGeometry {
26		/// Items processed in M dimension (rows)
27		pub m: u32,
28		/// Items processed in N dimension (columns)
29		pub n: u32,
30		/// Reduction dimension (inner product)
31		pub k: u32,
32		/// Alignment requirement in bytes (typically 16 for SIMD, 32 for AVX2, 64 for AVX-512)
33		pub alignment: u32,
34		}
35
36		impl TcbGeometry {
37		/// Create a new TCB geometry
38		///
39		/// # Panics
40		/// Panics if any dimension is zero.
41		#[must_use]
42	0	pub fn new(m: u32, n: u32, k: u32) -> Self {
43	0	assert!(m > 0 && n > 0 && k > 0, "TCB dimensions must be non-zero");
44	0	Self {
45	0	m,
46	0	n,
47	0	k,
48	0	alignment: 16, // Default to SSE/NEON alignment
49	0	}
50	0	}
51
52		/// Create geometry with explicit alignment
53		#[must_use]
54	0	pub fn with_alignment(m: u32, n: u32, k: u32, alignment: u32) -> Self {
55	0	assert!(m > 0 && n > 0 && k > 0, "TCB dimensions must be non-zero");
56	0	assert!(
57	0	alignment.is_power_of_two(),
58	0	"Alignment must be power of 2"
59		);
60	0	Self { m, n, k, alignment }
61	0	}
62
63		/// Calculate arithmetic intensity (FLOPS per byte loaded)
64		///
65		/// For GEMM: AI = (2 * M * N * K) / (MK + KN) * sizeof(f32)
66		///
67		/// Higher AI means compute-bound; lower means memory-bound.
68		#[must_use]
69	0	pub fn arithmetic_intensity(&self) -> f32 {
70	0	let flops = 2.0 * self.m as f64 * self.n as f64 * self.k as f64;
71	0	let bytes = (self.m as f64 * self.k as f64 + self.k as f64 * self.n as f64) * 4.0;
72	0	(flops / bytes) as f32
73	0	}
74
75		/// Calculate total elements in the tile
76		#[must_use]
77	0	pub fn total_elements(&self) -> u64 {
78	0	self.m as u64 * self.n as u64
79	0	}
80
81		/// Calculate total FLOPs for this tile
82		#[must_use]
83	0	pub fn total_flops(&self) -> u64 {
84	0	2 * self.m as u64 * self.n as u64 * self.k as u64
85	0	}
86
87		/// Check if K dimension aligns with Q4_K superblock (256)
88		#[must_use]
89	0	pub fn is_q4k_aligned(&self) -> bool {
90	0	self.k % 256 == 0
91	0	}
92
93		/// Check if K dimension aligns with Q4_0/Q8_0 block (32)
94		#[must_use]
95	0	pub fn is_q4_0_aligned(&self) -> bool {
96	0	self.k % 32 == 0
97	0	}
98
99		/// Calculate bytes needed for A tile (M × K × sizeof(f32))
100		#[must_use]
101	0	pub fn a_tile_bytes(&self) -> usize {
102	0	self.m as usize * self.k as usize * 4
103	0	}
104
105		/// Calculate bytes needed for B tile (K × N × sizeof(f32))
106		#[must_use]
107	0	pub fn b_tile_bytes(&self) -> usize {
108	0	self.k as usize * self.n as usize * 4
109	0	}
110
111		/// Calculate bytes needed for C tile (M × N × sizeof(f32))
112		#[must_use]
113	0	pub fn c_tile_bytes(&self) -> usize {
114	0	self.m as usize * self.n as usize * 4
115	0	}
116
117		/// Check if tile fits in given cache size (bytes)
118		#[must_use]
119	0	pub fn fits_in_cache(&self, cache_bytes: usize) -> bool {
120	0	self.a_tile_bytes() + self.b_tile_bytes() <= cache_bytes
121	0	}
122		}
123
124		impl Default for TcbGeometry {
125	0	fn default() -> Self {
126		// Sensible default: 4×4 micro-tile for SIMD
127	0	Self {
128	0	m: 4,
129	0	n: 4,
130	0	k: 4,
131	0	alignment: 16,
132	0	}
133	0	}
134		}
135
136		impl fmt::Display for TcbGeometry {
137	0	fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
138	0	write!(
139	0	f,
140	0	"TCB({}×{}×{}, align={}, AI={:.2})",
141		self.m,
142		self.n,
143		self.k,
144		self.alignment,
145	0	self.arithmetic_intensity()
146		)
147	0	}
148		}
149		// ============================================================================
150		// TILE-001: Tiling Levels
151		// ============================================================================
152
153		/// Tiling hierarchy level
154		#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
155		pub enum TcbLevel {
156		/// Macro-tile: L3 cache / GPU global memory partitioning
157		Macro,
158		/// Midi-tile: L2 cache / GPU shared memory
159		Midi,
160		/// Micro-tile: Registers / SIMD lanes
161		Micro,
162		}
163
164		impl TcbLevel {
165		/// Get typical cache size for this level (x86_64)
166		#[must_use]
167	0	pub fn typical_cache_bytes(&self) -> usize {
168	0	match self {
169	0	TcbLevel::Macro => 32 * 1024 * 1024, // 32 MB L3
170	0	TcbLevel::Midi => 256 * 1024, // 256 KB L2
171	0	TcbLevel::Micro => 32 * 1024, // 32 KB L1
172		}
173	0	}
174		}

Coverage Report

Created: 2026-01-25 15:05