Initial commit
This commit is contained in:
redclawsystems
@@ -0,0 +1,464 @@
|
||||
//! TMA-Accelerated Operations
|
||||
//!
|
||||
//! Provides Tensor Memory Accelerator (TMA) accelerated operations.
|
||||
|
||||
use std::fmt;
|
||||
use tracing::{debug, instrument};
|
||||
|
||||
use super::tile_kernel::TileConfig;
|
||||
use crate::error::{IntegrationError, Result};
|
||||
|
||||
#[cfg(all(target_os = "macos", feature = "tma-metal"))]
|
||||
use tensor_accelerator::metal::{MetalTensor, MetalTensorDevice};
|
||||
|
||||
/// Configuration for TMA operations
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct TmaConfig {
|
||||
/// Enable async memory copies
|
||||
pub async_copy: bool,
|
||||
/// Tile size for M dimension
|
||||
pub tile_m: usize,
|
||||
/// Tile size for N dimension
|
||||
pub tile_n: usize,
|
||||
/// Tile size for K dimension
|
||||
pub tile_k: usize,
|
||||
/// Use FP16 accumulation
|
||||
pub use_fp16_accumulator: bool,
|
||||
/// Number of stages for software pipelining
|
||||
pub pipeline_stages: usize,
|
||||
}
|
||||
|
||||
impl Default for TmaConfig {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
async_copy: true,
|
||||
tile_m: 128,
|
||||
tile_n: 128,
|
||||
tile_k: 64,
|
||||
use_fp16_accumulator: false,
|
||||
pipeline_stages: 3,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl TmaConfig {
|
||||
/// Configuration optimized for NVIDIA Blackwell (sm_110)
|
||||
pub fn blackwell() -> Self {
|
||||
Self {
|
||||
async_copy: true,
|
||||
tile_m: 128,
|
||||
tile_n: 128,
|
||||
tile_k: 64,
|
||||
use_fp16_accumulator: true,
|
||||
pipeline_stages: 4,
|
||||
}
|
||||
}
|
||||
|
||||
/// Configuration optimized for Apple M3/M4 (Apple9)
|
||||
pub fn apple_silicon() -> Self {
|
||||
Self {
|
||||
async_copy: true,
|
||||
tile_m: 64,
|
||||
tile_n: 64,
|
||||
tile_k: 32,
|
||||
use_fp16_accumulator: true,
|
||||
pipeline_stages: 2,
|
||||
}
|
||||
}
|
||||
|
||||
/// Convert to TileConfig
|
||||
pub fn to_tile_config(&self) -> TileConfig {
|
||||
TileConfig {
|
||||
tile_m: self.tile_m,
|
||||
tile_n: self.tile_n,
|
||||
tile_k: self.tile_k,
|
||||
warp_tile_m: self.tile_m / 4,
|
||||
warp_tile_n: self.tile_n / 4,
|
||||
pipeline_stages: self.pipeline_stages,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Trait for TMA-accelerated tensor operations
|
||||
///
|
||||
/// Provides optimized implementations of common tensor operations
|
||||
/// using Tensor Memory Accelerator on supported hardware.
|
||||
pub trait TmaOps: Sized {
|
||||
/// Matrix multiply with TMA optimization
|
||||
///
|
||||
/// Uses tile-based memory access patterns for better cache utilization.
|
||||
///
|
||||
/// # Example
|
||||
/// ```ignore
|
||||
/// let c = a.matmul_tma(&b)?;
|
||||
/// ```
|
||||
fn matmul_tma(&self, other: &Self) -> Result<Self>;
|
||||
|
||||
/// Matrix multiply with custom configuration
|
||||
fn matmul_tma_with_config(&self, other: &Self, config: &TmaConfig) -> Result<Self>;
|
||||
|
||||
/// Batched matrix multiply with TMA
|
||||
fn batched_matmul_tma(&self, other: &Self) -> Result<Self>;
|
||||
|
||||
/// Attention operation with TMA optimization
|
||||
///
|
||||
/// Computes: softmax(Q @ K.T / sqrt(d_k)) @ V
|
||||
fn attention_tma(&self, key: &Self, value: &Self, scale: f32) -> Result<Self>;
|
||||
|
||||
/// Flash attention variant with memory-efficient implementation
|
||||
fn flash_attention_tma(
|
||||
&self,
|
||||
key: &Self,
|
||||
value: &Self,
|
||||
scale: f32,
|
||||
causal: bool,
|
||||
) -> Result<Self>;
|
||||
|
||||
/// Element-wise addition with tiled memory access
|
||||
fn add_tma(&self, other: &Self) -> Result<Self>;
|
||||
|
||||
/// Get the shape of this tensor
|
||||
fn shape(&self) -> &[usize];
|
||||
|
||||
/// Check if TMA is available for this tensor
|
||||
fn tma_available(&self) -> bool;
|
||||
}
|
||||
|
||||
/// Wrapper for TMA-capable tensors
|
||||
pub struct TmaTensor {
|
||||
/// Shape of the tensor
|
||||
shape: Vec<usize>,
|
||||
/// Data storage
|
||||
data: Vec<f32>,
|
||||
/// TMA configuration
|
||||
config: TmaConfig,
|
||||
/// Whether tensor is on GPU
|
||||
on_device: bool,
|
||||
}
|
||||
|
||||
impl TmaTensor {
|
||||
/// Create a new TMA tensor
|
||||
pub fn new(shape: Vec<usize>, config: TmaConfig) -> Self {
|
||||
let size: usize = shape.iter().product();
|
||||
Self {
|
||||
shape,
|
||||
data: vec![0.0; size],
|
||||
config,
|
||||
on_device: false,
|
||||
}
|
||||
}
|
||||
|
||||
/// Create from existing data
|
||||
pub fn from_vec(shape: Vec<usize>, data: Vec<f32>) -> Result<Self> {
|
||||
let expected_size: usize = shape.iter().product();
|
||||
if data.len() != expected_size {
|
||||
return Err(IntegrationError::ShapeMismatch {
|
||||
expected: format!("{} elements", expected_size),
|
||||
actual: format!("{} elements", data.len()),
|
||||
});
|
||||
}
|
||||
|
||||
Ok(Self {
|
||||
shape,
|
||||
data,
|
||||
config: TmaConfig::default(),
|
||||
on_device: false,
|
||||
})
|
||||
}
|
||||
|
||||
/// Get tensor data
|
||||
pub fn data(&self) -> &[f32] {
|
||||
&self.data
|
||||
}
|
||||
|
||||
/// Get mutable tensor data
|
||||
pub fn data_mut(&mut self) -> &mut [f32] {
|
||||
&mut self.data
|
||||
}
|
||||
|
||||
/// Set TMA configuration
|
||||
pub fn with_config(mut self, config: TmaConfig) -> Self {
|
||||
self.config = config;
|
||||
self
|
||||
}
|
||||
|
||||
/// Move tensor to device
|
||||
#[instrument(skip(self))]
|
||||
pub fn to_device(&mut self) -> Result<()> {
|
||||
if self.on_device {
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
debug!("Moving tensor {:?} to device", self.shape);
|
||||
|
||||
#[cfg(all(target_os = "macos", feature = "tma-metal"))]
|
||||
{
|
||||
// Actual device transfer would happen here
|
||||
}
|
||||
|
||||
self.on_device = true;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Move tensor to host
|
||||
pub fn to_host(&mut self) -> Result<()> {
|
||||
if !self.on_device {
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
debug!("Moving tensor {:?} to host", self.shape);
|
||||
|
||||
#[cfg(all(target_os = "macos", feature = "tma-metal"))]
|
||||
{
|
||||
// Actual host transfer would happen here
|
||||
}
|
||||
|
||||
self.on_device = false;
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
impl TmaOps for TmaTensor {
|
||||
#[instrument(skip(self, other))]
|
||||
fn matmul_tma(&self, other: &Self) -> Result<Self> {
|
||||
self.matmul_tma_with_config(other, &self.config)
|
||||
}
|
||||
|
||||
fn matmul_tma_with_config(&self, other: &Self, config: &TmaConfig) -> Result<Self> {
|
||||
// Validate shapes for matmul
|
||||
if self.shape.len() < 2 || other.shape.len() < 2 {
|
||||
return Err(IntegrationError::ShapeMismatch {
|
||||
expected: "2D or higher tensors".to_string(),
|
||||
actual: format!("{}D and {}D", self.shape.len(), other.shape.len()),
|
||||
});
|
||||
}
|
||||
|
||||
let m = self.shape[self.shape.len() - 2];
|
||||
let k1 = self.shape[self.shape.len() - 1];
|
||||
let k2 = other.shape[other.shape.len() - 2];
|
||||
let n = other.shape[other.shape.len() - 1];
|
||||
|
||||
if k1 != k2 {
|
||||
return Err(IntegrationError::ShapeMismatch {
|
||||
expected: format!("K dimensions to match: {}", k1),
|
||||
actual: format!("got {}", k2),
|
||||
});
|
||||
}
|
||||
|
||||
debug!(
|
||||
"TMA matmul: ({}, {}) x ({}, {}) with tiles {}x{}x{}",
|
||||
m, k1, k2, n, config.tile_m, config.tile_n, config.tile_k
|
||||
);
|
||||
|
||||
// Output shape
|
||||
let mut output_shape = self.shape.clone();
|
||||
*output_shape.last_mut().unwrap() = n;
|
||||
|
||||
let output_size: usize = output_shape.iter().product();
|
||||
let mut output_data = vec![0.0f32; output_size];
|
||||
|
||||
// Naive implementation (actual TMA would use hardware acceleration)
|
||||
// In production, this would dispatch to GPU kernels
|
||||
for i in 0..m {
|
||||
for j in 0..n {
|
||||
let mut sum = 0.0f32;
|
||||
for k in 0..k1 {
|
||||
sum += self.data[i * k1 + k] * other.data[k * n + j];
|
||||
}
|
||||
output_data[i * n + j] = sum;
|
||||
}
|
||||
}
|
||||
|
||||
Ok(TmaTensor {
|
||||
shape: output_shape,
|
||||
data: output_data,
|
||||
config: config.clone(),
|
||||
on_device: self.on_device,
|
||||
})
|
||||
}
|
||||
|
||||
fn batched_matmul_tma(&self, other: &Self) -> Result<Self> {
|
||||
// For batched matmul, iterate over batch dimensions
|
||||
self.matmul_tma(other)
|
||||
}
|
||||
|
||||
fn attention_tma(&self, key: &Self, value: &Self, scale: f32) -> Result<Self> {
|
||||
self.flash_attention_tma(key, value, scale, false)
|
||||
}
|
||||
|
||||
fn flash_attention_tma(
|
||||
&self,
|
||||
key: &Self,
|
||||
value: &Self,
|
||||
scale: f32,
|
||||
_causal: bool,
|
||||
) -> Result<Self> {
|
||||
debug!(
|
||||
"Flash attention TMA: Q{:?} K{:?} V{:?} scale={}",
|
||||
self.shape, key.shape, value.shape, scale
|
||||
);
|
||||
|
||||
// Compute Q @ K.T
|
||||
// For now, transpose K and multiply
|
||||
let k_t = Self::transpose_2d(key)?;
|
||||
let mut scores = self.matmul_tma(&k_t)?;
|
||||
|
||||
// Scale
|
||||
for v in scores.data.iter_mut() {
|
||||
*v *= scale;
|
||||
}
|
||||
|
||||
// Softmax (simplified - actual would be numerically stable)
|
||||
let seq_len = scores.shape[scores.shape.len() - 1];
|
||||
for i in 0..(scores.data.len() / seq_len) {
|
||||
let start = i * seq_len;
|
||||
let end = start + seq_len;
|
||||
|
||||
let max = scores.data[start..end]
|
||||
.iter()
|
||||
.cloned()
|
||||
.fold(f32::NEG_INFINITY, f32::max);
|
||||
let mut sum = 0.0f32;
|
||||
|
||||
for v in &mut scores.data[start..end] {
|
||||
*v = (*v - max).exp();
|
||||
sum += *v;
|
||||
}
|
||||
|
||||
for v in &mut scores.data[start..end] {
|
||||
*v /= sum;
|
||||
}
|
||||
}
|
||||
|
||||
// Multiply by V
|
||||
scores.matmul_tma(value)
|
||||
}
|
||||
|
||||
fn add_tma(&self, other: &Self) -> Result<Self> {
|
||||
if self.shape != other.shape {
|
||||
return Err(IntegrationError::ShapeMismatch {
|
||||
expected: format!("{:?}", self.shape),
|
||||
actual: format!("{:?}", other.shape),
|
||||
});
|
||||
}
|
||||
|
||||
let data: Vec<f32> = self
|
||||
.data
|
||||
.iter()
|
||||
.zip(other.data.iter())
|
||||
.map(|(a, b)| a + b)
|
||||
.collect();
|
||||
|
||||
Ok(TmaTensor {
|
||||
shape: self.shape.clone(),
|
||||
data,
|
||||
config: self.config.clone(),
|
||||
on_device: self.on_device,
|
||||
})
|
||||
}
|
||||
|
||||
fn shape(&self) -> &[usize] {
|
||||
&self.shape
|
||||
}
|
||||
|
||||
fn tma_available(&self) -> bool {
|
||||
#[cfg(all(target_os = "macos", feature = "tma-metal"))]
|
||||
{
|
||||
true
|
||||
}
|
||||
#[cfg(not(all(target_os = "macos", feature = "tma-metal")))]
|
||||
{
|
||||
false
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl TmaTensor {
|
||||
/// Transpose a 2D tensor
|
||||
fn transpose_2d(tensor: &Self) -> Result<Self> {
|
||||
if tensor.shape.len() != 2 {
|
||||
return Err(IntegrationError::ShapeMismatch {
|
||||
expected: "2D tensor".to_string(),
|
||||
actual: format!("{}D tensor", tensor.shape.len()),
|
||||
});
|
||||
}
|
||||
|
||||
let rows = tensor.shape[0];
|
||||
let cols = tensor.shape[1];
|
||||
let mut data = vec![0.0f32; rows * cols];
|
||||
|
||||
for i in 0..rows {
|
||||
for j in 0..cols {
|
||||
data[j * rows + i] = tensor.data[i * cols + j];
|
||||
}
|
||||
}
|
||||
|
||||
Ok(TmaTensor {
|
||||
shape: vec![cols, rows],
|
||||
data,
|
||||
config: tensor.config.clone(),
|
||||
on_device: tensor.on_device,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
impl fmt::Debug for TmaTensor {
|
||||
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||||
f.debug_struct("TmaTensor")
|
||||
.field("shape", &self.shape)
|
||||
.field("on_device", &self.on_device)
|
||||
.field("tma_available", &self.tma_available())
|
||||
.finish()
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_tma_config_default() {
|
||||
let config = TmaConfig::default();
|
||||
assert!(config.async_copy);
|
||||
assert_eq!(config.tile_m, 128);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_tma_config_presets() {
|
||||
let blackwell = TmaConfig::blackwell();
|
||||
assert_eq!(blackwell.pipeline_stages, 4);
|
||||
|
||||
let apple = TmaConfig::apple_silicon();
|
||||
assert_eq!(apple.tile_m, 64);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_tma_tensor_creation() {
|
||||
let tensor = TmaTensor::new(vec![32, 64], TmaConfig::default());
|
||||
assert_eq!(tensor.shape(), &[32, 64]);
|
||||
assert_eq!(tensor.data.len(), 32 * 64);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_tma_matmul() {
|
||||
let a = TmaTensor::from_vec(vec![2, 3], vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0]).unwrap();
|
||||
|
||||
let b = TmaTensor::from_vec(vec![3, 2], vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0]).unwrap();
|
||||
|
||||
let c = a.matmul_tma(&b).unwrap();
|
||||
assert_eq!(c.shape(), &[2, 2]);
|
||||
|
||||
// Check result: [1,2,3] @ [[1,2],[3,4],[5,6]] = [22, 28]
|
||||
assert!((c.data[0] - 22.0).abs() < 1e-5);
|
||||
assert!((c.data[1] - 28.0).abs() < 1e-5);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_shape_mismatch() {
|
||||
let a = TmaTensor::new(vec![2, 3], TmaConfig::default());
|
||||
let b = TmaTensor::new(vec![4, 2], TmaConfig::default());
|
||||
|
||||
let result = a.matmul_tma(&b);
|
||||
assert!(result.is_err());
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user