465 lines
13 KiB
Rust
465 lines
13 KiB
Rust
//! 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());
|
|
}
|
|
}
|