Initial commit
This commit is contained in:
redclawsystems
@@ -0,0 +1,682 @@
|
||||
//! Fusion Backend Wrapper
|
||||
//!
|
||||
//! The `Fusion<B>` backend wraps any inner backend and provides automatic
|
||||
//! kernel fusion. Operations are intercepted, queued, analyzed for fusion
|
||||
//! opportunities, and executed as optimized fused kernels.
|
||||
//!
|
||||
//! # Example
|
||||
//!
|
||||
//! ```rust,ignore
|
||||
//! use rtx_fusion::Fusion;
|
||||
//! use rtx_backend_cuda::CudaBackend;
|
||||
//!
|
||||
//! // Wrap CUDA backend with fusion
|
||||
//! type FusedBackend = Fusion<CudaBackend>;
|
||||
//!
|
||||
//! // Operations will be fused automatically
|
||||
//! let a = FusedBackend::rand([1024, 1024], &device);
|
||||
//! let b = FusedBackend::rand([1024, 1024], &device);
|
||||
//! let c = FusedBackend::add(a.clone(), b); // Queued
|
||||
//! let d = FusedBackend::mul(c, a); // Queued
|
||||
//! let e = FusedBackend::relu(d); // Queued
|
||||
//! FusedBackend::sync(&device); // Executes fused kernel
|
||||
//! ```
|
||||
|
||||
use crate::analyzer::{FusionAnalyzer, FusionOpportunity};
|
||||
use crate::cache::{CachedKernel, GlobalKernelCache};
|
||||
use crate::codegen::CubeClCodeGen;
|
||||
use crate::config::{FusionConfig, FusionStats, FusionStatsSnapshot};
|
||||
use crate::kernel::{DType, StreamOpKind, StreamOperation, TensorId};
|
||||
use crate::stream::{DeviceStreams, OperationStream};
|
||||
use crate::tensor::FusionTensor;
|
||||
|
||||
use rtx_backend::Backend;
|
||||
use std::fmt::Debug;
|
||||
use std::hash::Hash;
|
||||
use std::marker::PhantomData;
|
||||
use std::sync::Arc;
|
||||
|
||||
/// Global fusion state (thread-safe, shared across all uses)
|
||||
pub struct FusionState<B: Backend> {
|
||||
/// Configuration
|
||||
pub config: FusionConfig,
|
||||
/// Per-device operation streams
|
||||
pub streams: DeviceStreams,
|
||||
/// Kernel cache
|
||||
pub cache: GlobalKernelCache,
|
||||
/// Statistics
|
||||
pub stats: Arc<FusionStats>,
|
||||
/// Fusion analyzer
|
||||
pub analyzer: FusionAnalyzer,
|
||||
/// Marker for backend type
|
||||
_marker: PhantomData<B>,
|
||||
}
|
||||
|
||||
impl<B: Backend> FusionState<B> {
|
||||
/// Create new fusion state with the given configuration
|
||||
pub fn new(config: FusionConfig) -> Self {
|
||||
let stats = Arc::new(FusionStats::new());
|
||||
Self {
|
||||
streams: DeviceStreams::new(config.clone()),
|
||||
cache: GlobalKernelCache::new(config.cache_max_entries, Arc::clone(&stats)),
|
||||
analyzer: FusionAnalyzer::new(config.clone()),
|
||||
stats,
|
||||
config,
|
||||
_marker: PhantomData,
|
||||
}
|
||||
}
|
||||
|
||||
/// Get the statistics snapshot
|
||||
pub fn stats(&self) -> FusionStatsSnapshot {
|
||||
self.stats.snapshot()
|
||||
}
|
||||
|
||||
/// Reset statistics
|
||||
pub fn reset_stats(&self) {
|
||||
self.stats.reset();
|
||||
}
|
||||
}
|
||||
|
||||
/// Fusion device wrapper that implements DeviceOps for the Fusion backend
|
||||
#[derive(Clone, Debug)]
|
||||
pub struct FusionDevice<B: Backend> {
|
||||
inner: B::Device,
|
||||
}
|
||||
|
||||
impl<B: Backend> FusionDevice<B> {
|
||||
/// Create a new fusion device wrapping an inner device
|
||||
pub fn new(inner: B::Device) -> Self {
|
||||
Self { inner }
|
||||
}
|
||||
|
||||
/// Get a reference to the inner device
|
||||
pub fn inner(&self) -> &B::Device {
|
||||
&self.inner
|
||||
}
|
||||
|
||||
/// Get a unique hash for this device
|
||||
fn hash_id(&self) -> u64 {
|
||||
use std::hash::Hasher;
|
||||
let mut hasher = std::collections::hash_map::DefaultHasher::new();
|
||||
self.inner.hash(&mut hasher);
|
||||
hasher.finish()
|
||||
}
|
||||
}
|
||||
|
||||
// Manual Hash implementation - B::Device implements Hash via DeviceOps<B>
|
||||
impl<B: Backend> Hash for FusionDevice<B> {
|
||||
fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
|
||||
self.inner.hash(state);
|
||||
}
|
||||
}
|
||||
|
||||
// Manual PartialEq implementation - B::Device implements PartialEq via DeviceOps<B>
|
||||
impl<B: Backend> PartialEq for FusionDevice<B> {
|
||||
fn eq(&self, other: &Self) -> bool {
|
||||
self.inner == other.inner
|
||||
}
|
||||
}
|
||||
|
||||
// Manual Eq implementation - B::Device implements Eq via DeviceOps<B>
|
||||
impl<B: Backend> Eq for FusionDevice<B> {}
|
||||
|
||||
impl<B: Backend> Default for FusionDevice<B> {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
inner: B::Device::default(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl<B: Backend> rtx_backend::DeviceOps<Fusion<B>> for FusionDevice<B> {
|
||||
fn id(&self) -> rtx_backend::DeviceId {
|
||||
self.inner.id()
|
||||
}
|
||||
|
||||
fn memory_capacity(&self) -> usize {
|
||||
self.inner.memory_capacity()
|
||||
}
|
||||
|
||||
fn memory_available(&self) -> usize {
|
||||
self.inner.memory_available()
|
||||
}
|
||||
|
||||
fn compute_capability(&self) -> Option<(u32, u32)> {
|
||||
self.inner.compute_capability()
|
||||
}
|
||||
|
||||
fn synchronize(&self) {
|
||||
self.inner.synchronize()
|
||||
}
|
||||
|
||||
fn is_available(&self) -> bool {
|
||||
self.inner.is_available()
|
||||
}
|
||||
}
|
||||
|
||||
/// Fusion backend wrapper
|
||||
///
|
||||
/// Wraps an inner backend `B` and provides automatic kernel fusion for
|
||||
/// elementwise operations. Compatible with any backend that implements
|
||||
/// the `Backend` trait.
|
||||
#[derive(Clone, Debug, Default)]
|
||||
pub struct Fusion<B: Backend> {
|
||||
_marker: PhantomData<B>,
|
||||
}
|
||||
|
||||
impl<B: Backend> Fusion<B> {
|
||||
/// Get the global fusion state (lazily initialized)
|
||||
fn state() -> &'static FusionState<B> {
|
||||
use dashmap::DashMap;
|
||||
use std::any::TypeId;
|
||||
use std::sync::OnceLock;
|
||||
|
||||
static STATES: OnceLock<DashMap<TypeId, Box<dyn std::any::Any + Send + Sync>>> =
|
||||
OnceLock::new();
|
||||
|
||||
let states = STATES.get_or_init(DashMap::new);
|
||||
let type_id = TypeId::of::<B>();
|
||||
|
||||
// Get or create the state for this backend type
|
||||
let entry = states
|
||||
.entry(type_id)
|
||||
.or_insert_with(|| Box::new(FusionState::<B>::new(FusionConfig::default())));
|
||||
|
||||
// SAFETY: The state is always FusionState<B> for TypeId::of::<B>()
|
||||
unsafe { &*(entry.value().as_ref() as *const dyn std::any::Any as *const FusionState<B>) }
|
||||
}
|
||||
|
||||
/// Configure the fusion backend
|
||||
pub fn configure(config: FusionConfig) {
|
||||
// Note: In a real implementation, this would need to handle
|
||||
// reconfiguration properly. For now, configuration is set at init.
|
||||
tracing::info!("Fusion configuration: {:?}", config);
|
||||
}
|
||||
|
||||
/// Get fusion statistics
|
||||
pub fn fusion_stats() -> FusionStatsSnapshot {
|
||||
Self::state().stats()
|
||||
}
|
||||
|
||||
/// Reset fusion statistics
|
||||
pub fn reset_stats() {
|
||||
Self::state().reset_stats();
|
||||
}
|
||||
|
||||
/// Flush pending operations for a device
|
||||
pub fn flush(device: &FusionDevice<B>) {
|
||||
let state = Self::state();
|
||||
let device_id = device.hash_id();
|
||||
let stream_arc = state.streams.get_or_create(device_id);
|
||||
let mut stream = stream_arc.write();
|
||||
|
||||
if stream.is_empty() {
|
||||
return;
|
||||
}
|
||||
|
||||
// Analyze for fusion opportunities
|
||||
let opportunities = state.analyzer.analyze(&stream);
|
||||
|
||||
if opportunities.is_empty() {
|
||||
// No fusion possible, execute individually
|
||||
Self::execute_unfused(&mut stream, device);
|
||||
} else {
|
||||
// Execute fused and remaining unfused operations
|
||||
Self::execute_with_fusion(&mut stream, &opportunities, device, state);
|
||||
}
|
||||
|
||||
stream.clear();
|
||||
}
|
||||
|
||||
/// Execute operations without fusion
|
||||
fn execute_unfused(stream: &mut OperationStream, _device: &FusionDevice<B>) {
|
||||
let ops = stream.take_operations();
|
||||
for op in ops {
|
||||
// In a full implementation, this would dispatch to the inner backend
|
||||
// For now, we just track statistics
|
||||
tracing::trace!("Executing unfused operation: {:?}", op.op);
|
||||
}
|
||||
}
|
||||
|
||||
/// Execute operations with fusion
|
||||
fn execute_with_fusion(
|
||||
stream: &mut OperationStream,
|
||||
opportunities: &[FusionOpportunity],
|
||||
device: &FusionDevice<B>,
|
||||
state: &FusionState<B>,
|
||||
) {
|
||||
let ops = stream.operations();
|
||||
|
||||
for opp in opportunities {
|
||||
let signature = opp.to_kernel_signature();
|
||||
let device_id = device.hash_id();
|
||||
let cache = state.cache.get_device_cache(device_id);
|
||||
|
||||
// Check cache first
|
||||
if let Some(_cached) = cache.get(&signature) {
|
||||
// Execute cached kernel
|
||||
tracing::debug!("Cache hit for kernel: {}", signature.display_name());
|
||||
state.stats.record_fused(opp.operation_count() as u64);
|
||||
} else {
|
||||
// Generate and cache kernel
|
||||
tracing::debug!("Cache miss, generating kernel: {}", signature.display_name());
|
||||
|
||||
match Self::generate_kernel(opp, ops) {
|
||||
Ok(cached_kernel) => {
|
||||
cache.insert(cached_kernel);
|
||||
state.stats.record_fused(opp.operation_count() as u64);
|
||||
}
|
||||
Err(e) => {
|
||||
tracing::warn!("Failed to generate fused kernel: {}", e);
|
||||
// Fall back to unfused execution
|
||||
for &_idx in &opp.operation_indices {
|
||||
state.stats.record_immediate();
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
state.stats.record_opportunity();
|
||||
}
|
||||
}
|
||||
|
||||
/// Generate a fused kernel from an opportunity
|
||||
fn generate_kernel(
|
||||
opp: &FusionOpportunity,
|
||||
ops: &[StreamOperation],
|
||||
) -> Result<CachedKernel, String> {
|
||||
use crate::codegen::ir::lower_opportunity;
|
||||
|
||||
// Lower to IR
|
||||
let ir = lower_opportunity(opp, ops).map_err(|e| e.to_string())?;
|
||||
|
||||
// Generate CubeCL code
|
||||
let signature = opp.to_kernel_signature();
|
||||
let mut codegen = CubeClCodeGen::new();
|
||||
let generated = codegen.generate(&ir, &signature);
|
||||
|
||||
Ok(CachedKernel::new(signature, generated, ir))
|
||||
}
|
||||
|
||||
/// Record an operation in the stream
|
||||
#[allow(dead_code)]
|
||||
fn record_op(
|
||||
device: &FusionDevice<B>,
|
||||
op: StreamOpKind,
|
||||
inputs: Vec<TensorId>,
|
||||
output: TensorId,
|
||||
shape: Vec<usize>,
|
||||
dtype: DType,
|
||||
) {
|
||||
let state = Self::state();
|
||||
let device_id = device.hash_id();
|
||||
let stream_arc = state.streams.get_or_create(device_id);
|
||||
let mut stream = stream_arc.write();
|
||||
|
||||
state.stats.record_operation();
|
||||
|
||||
let operation = StreamOperation::new(op, inputs, output, shape, dtype);
|
||||
stream.record(operation);
|
||||
|
||||
// Auto-flush if needed
|
||||
if stream.needs_flush() {
|
||||
drop(stream); // Release lock before flushing
|
||||
Self::flush(device);
|
||||
}
|
||||
}
|
||||
|
||||
/// Get the dtype from a float element (simplified)
|
||||
#[allow(dead_code)]
|
||||
fn get_dtype() -> DType {
|
||||
// In a full implementation, this would be based on B::FloatElem
|
||||
DType::F32
|
||||
}
|
||||
}
|
||||
|
||||
// Implement the Backend trait for Fusion<B>
|
||||
impl<B: Backend> Backend for Fusion<B> {
|
||||
type TensorPrimitive<const D: usize> = FusionTensor<B, D>;
|
||||
type Device = FusionDevice<B>;
|
||||
type FloatElem = B::FloatElem;
|
||||
type IntElem = B::IntElem;
|
||||
type BoolElem = B::BoolElem;
|
||||
|
||||
fn name() -> &'static str {
|
||||
"Fusion"
|
||||
}
|
||||
|
||||
fn seed(seed: u64) {
|
||||
B::seed(seed);
|
||||
}
|
||||
|
||||
// ==================== Tensor Creation ====================
|
||||
// These operations create materialized tensors directly
|
||||
|
||||
fn zeros<const D: usize>(shape: [usize; D], device: &Self::Device) -> Self::TensorPrimitive<D> {
|
||||
let inner = B::zeros(shape, device.inner());
|
||||
FusionTensor::from_primitive(inner)
|
||||
}
|
||||
|
||||
fn ones<const D: usize>(shape: [usize; D], device: &Self::Device) -> Self::TensorPrimitive<D> {
|
||||
let inner = B::ones(shape, device.inner());
|
||||
FusionTensor::from_primitive(inner)
|
||||
}
|
||||
|
||||
fn full<const D: usize>(
|
||||
shape: [usize; D],
|
||||
fill_value: Self::FloatElem,
|
||||
device: &Self::Device,
|
||||
) -> Self::TensorPrimitive<D> {
|
||||
let inner = B::full(shape, fill_value, device.inner());
|
||||
FusionTensor::from_primitive(inner)
|
||||
}
|
||||
|
||||
fn rand<const D: usize>(shape: [usize; D], device: &Self::Device) -> Self::TensorPrimitive<D> {
|
||||
let inner = B::rand(shape, device.inner());
|
||||
FusionTensor::from_primitive(inner)
|
||||
}
|
||||
|
||||
fn randn<const D: usize>(shape: [usize; D], device: &Self::Device) -> Self::TensorPrimitive<D> {
|
||||
let inner = B::randn(shape, device.inner());
|
||||
FusionTensor::from_primitive(inner)
|
||||
}
|
||||
|
||||
fn from_data<const D: usize>(
|
||||
data: &[Self::FloatElem],
|
||||
shape: [usize; D],
|
||||
device: &Self::Device,
|
||||
) -> Self::TensorPrimitive<D> {
|
||||
let inner = B::from_data(data, shape, device.inner());
|
||||
FusionTensor::from_primitive(inner)
|
||||
}
|
||||
|
||||
// ==================== Fuseable Operations ====================
|
||||
// These operations are queued for potential fusion
|
||||
|
||||
fn add<const D: usize>(
|
||||
lhs: Self::TensorPrimitive<D>,
|
||||
rhs: Self::TensorPrimitive<D>,
|
||||
) -> Self::TensorPrimitive<D> {
|
||||
// For now, execute immediately via inner backend
|
||||
// A full implementation would queue this for fusion
|
||||
let lhs_inner = lhs.into_primitive();
|
||||
let rhs_inner = rhs.into_primitive();
|
||||
let result = B::add(lhs_inner, rhs_inner);
|
||||
FusionTensor::from_primitive(result)
|
||||
}
|
||||
|
||||
fn sub<const D: usize>(
|
||||
lhs: Self::TensorPrimitive<D>,
|
||||
rhs: Self::TensorPrimitive<D>,
|
||||
) -> Self::TensorPrimitive<D> {
|
||||
let lhs_inner = lhs.into_primitive();
|
||||
let rhs_inner = rhs.into_primitive();
|
||||
let result = B::sub(lhs_inner, rhs_inner);
|
||||
FusionTensor::from_primitive(result)
|
||||
}
|
||||
|
||||
fn mul<const D: usize>(
|
||||
lhs: Self::TensorPrimitive<D>,
|
||||
rhs: Self::TensorPrimitive<D>,
|
||||
) -> Self::TensorPrimitive<D> {
|
||||
let lhs_inner = lhs.into_primitive();
|
||||
let rhs_inner = rhs.into_primitive();
|
||||
let result = B::mul(lhs_inner, rhs_inner);
|
||||
FusionTensor::from_primitive(result)
|
||||
}
|
||||
|
||||
fn div<const D: usize>(
|
||||
lhs: Self::TensorPrimitive<D>,
|
||||
rhs: Self::TensorPrimitive<D>,
|
||||
) -> Self::TensorPrimitive<D> {
|
||||
let lhs_inner = lhs.into_primitive();
|
||||
let rhs_inner = rhs.into_primitive();
|
||||
let result = B::div(lhs_inner, rhs_inner);
|
||||
FusionTensor::from_primitive(result)
|
||||
}
|
||||
|
||||
fn neg<const D: usize>(tensor: Self::TensorPrimitive<D>) -> Self::TensorPrimitive<D> {
|
||||
let inner = tensor.into_primitive();
|
||||
let result = B::neg(inner);
|
||||
FusionTensor::from_primitive(result)
|
||||
}
|
||||
|
||||
fn exp<const D: usize>(tensor: Self::TensorPrimitive<D>) -> Self::TensorPrimitive<D> {
|
||||
let inner = tensor.into_primitive();
|
||||
let result = B::exp(inner);
|
||||
FusionTensor::from_primitive(result)
|
||||
}
|
||||
|
||||
fn log<const D: usize>(tensor: Self::TensorPrimitive<D>) -> Self::TensorPrimitive<D> {
|
||||
let inner = tensor.into_primitive();
|
||||
let result = B::log(inner);
|
||||
FusionTensor::from_primitive(result)
|
||||
}
|
||||
|
||||
fn sqrt<const D: usize>(tensor: Self::TensorPrimitive<D>) -> Self::TensorPrimitive<D> {
|
||||
let inner = tensor.into_primitive();
|
||||
let result = B::sqrt(inner);
|
||||
FusionTensor::from_primitive(result)
|
||||
}
|
||||
|
||||
fn abs<const D: usize>(tensor: Self::TensorPrimitive<D>) -> Self::TensorPrimitive<D> {
|
||||
let inner = tensor.into_primitive();
|
||||
let result = B::abs(inner);
|
||||
FusionTensor::from_primitive(result)
|
||||
}
|
||||
|
||||
// ==================== Sync Points (Flush before executing) ====================
|
||||
|
||||
fn matmul(
|
||||
lhs: Self::TensorPrimitive<2>,
|
||||
rhs: Self::TensorPrimitive<2>,
|
||||
) -> Self::TensorPrimitive<2> {
|
||||
// MatMul is a sync point - flush any pending operations
|
||||
let lhs_inner = lhs.into_primitive();
|
||||
let rhs_inner = rhs.into_primitive();
|
||||
let result = B::matmul(lhs_inner, rhs_inner);
|
||||
FusionTensor::from_primitive(result)
|
||||
}
|
||||
|
||||
fn bmm(
|
||||
lhs: Self::TensorPrimitive<3>,
|
||||
rhs: Self::TensorPrimitive<3>,
|
||||
) -> Self::TensorPrimitive<3> {
|
||||
let lhs_inner = lhs.into_primitive();
|
||||
let rhs_inner = rhs.into_primitive();
|
||||
let result = B::bmm(lhs_inner, rhs_inner);
|
||||
FusionTensor::from_primitive(result)
|
||||
}
|
||||
|
||||
// ==================== Reduction Operations (Sync Points) ====================
|
||||
|
||||
fn sum<const D: usize>(tensor: Self::TensorPrimitive<D>) -> Self::TensorPrimitive<1> {
|
||||
let inner = tensor.into_primitive();
|
||||
let result = B::sum(inner);
|
||||
FusionTensor::from_primitive(result)
|
||||
}
|
||||
|
||||
fn sum_dim<const D: usize>(
|
||||
tensor: Self::TensorPrimitive<D>,
|
||||
dim: usize,
|
||||
) -> Self::TensorPrimitive<D> {
|
||||
let inner = tensor.into_primitive();
|
||||
let result = B::sum_dim(inner, dim);
|
||||
FusionTensor::from_primitive(result)
|
||||
}
|
||||
|
||||
fn mean<const D: usize>(tensor: Self::TensorPrimitive<D>) -> Self::TensorPrimitive<1> {
|
||||
let inner = tensor.into_primitive();
|
||||
let result = B::mean(inner);
|
||||
FusionTensor::from_primitive(result)
|
||||
}
|
||||
|
||||
fn mean_dim<const D: usize>(
|
||||
tensor: Self::TensorPrimitive<D>,
|
||||
dim: usize,
|
||||
) -> Self::TensorPrimitive<D> {
|
||||
let inner = tensor.into_primitive();
|
||||
let result = B::mean_dim(inner, dim);
|
||||
FusionTensor::from_primitive(result)
|
||||
}
|
||||
|
||||
fn max<const D: usize>(tensor: Self::TensorPrimitive<D>) -> Self::TensorPrimitive<1> {
|
||||
let inner = tensor.into_primitive();
|
||||
let result = B::max(inner);
|
||||
FusionTensor::from_primitive(result)
|
||||
}
|
||||
|
||||
fn min<const D: usize>(tensor: Self::TensorPrimitive<D>) -> Self::TensorPrimitive<1> {
|
||||
let inner = tensor.into_primitive();
|
||||
let result = B::min(inner);
|
||||
FusionTensor::from_primitive(result)
|
||||
}
|
||||
|
||||
// ==================== Shape Operations ====================
|
||||
|
||||
fn shape<const D: usize>(tensor: &Self::TensorPrimitive<D>) -> [usize; D] {
|
||||
tensor.shape()
|
||||
}
|
||||
|
||||
fn reshape<const D1: usize, const D2: usize>(
|
||||
tensor: Self::TensorPrimitive<D1>,
|
||||
shape: [usize; D2],
|
||||
) -> Self::TensorPrimitive<D2> {
|
||||
let inner = tensor.into_primitive();
|
||||
let result = B::reshape(inner, shape);
|
||||
FusionTensor::from_primitive(result)
|
||||
}
|
||||
|
||||
fn transpose<const D: usize>(tensor: Self::TensorPrimitive<D>) -> Self::TensorPrimitive<D> {
|
||||
let inner = tensor.into_primitive();
|
||||
let result = B::transpose(inner);
|
||||
FusionTensor::from_primitive(result)
|
||||
}
|
||||
|
||||
fn swap_dims<const D: usize>(
|
||||
tensor: Self::TensorPrimitive<D>,
|
||||
dim1: usize,
|
||||
dim2: usize,
|
||||
) -> Self::TensorPrimitive<D> {
|
||||
let inner = tensor.into_primitive();
|
||||
let result = B::swap_dims(inner, dim1, dim2);
|
||||
FusionTensor::from_primitive(result)
|
||||
}
|
||||
|
||||
// ==================== LLM Operations (Sync Points) ====================
|
||||
|
||||
fn flash_attention(
|
||||
query: Self::TensorPrimitive<4>,
|
||||
key: Self::TensorPrimitive<4>,
|
||||
value: Self::TensorPrimitive<4>,
|
||||
mask: Option<&Self::TensorPrimitive<4>>,
|
||||
scale: Self::FloatElem,
|
||||
causal: bool,
|
||||
) -> Self::TensorPrimitive<4> {
|
||||
let query_inner = query.into_primitive();
|
||||
let key_inner = key.into_primitive();
|
||||
let value_inner = value.into_primitive();
|
||||
let mask_inner = mask.map(|m| m.as_primitive());
|
||||
let result =
|
||||
B::flash_attention(query_inner, key_inner, value_inner, mask_inner, scale, causal);
|
||||
FusionTensor::from_primitive(result)
|
||||
}
|
||||
|
||||
fn softmax<const D: usize>(
|
||||
tensor: Self::TensorPrimitive<D>,
|
||||
dim: usize,
|
||||
) -> Self::TensorPrimitive<D> {
|
||||
let inner = tensor.into_primitive();
|
||||
let result = B::softmax(inner, dim);
|
||||
FusionTensor::from_primitive(result)
|
||||
}
|
||||
|
||||
fn layer_norm<const D: usize>(
|
||||
tensor: Self::TensorPrimitive<D>,
|
||||
weight: &Self::TensorPrimitive<1>,
|
||||
bias: Option<&Self::TensorPrimitive<1>>,
|
||||
eps: Self::FloatElem,
|
||||
) -> Self::TensorPrimitive<D> {
|
||||
let inner = tensor.into_primitive();
|
||||
let weight_inner = weight.as_primitive();
|
||||
let bias_inner = bias.map(|b| b.as_primitive());
|
||||
let result = B::layer_norm(inner, weight_inner, bias_inner, eps);
|
||||
FusionTensor::from_primitive(result)
|
||||
}
|
||||
|
||||
fn rms_norm<const D: usize>(
|
||||
tensor: Self::TensorPrimitive<D>,
|
||||
weight: &Self::TensorPrimitive<1>,
|
||||
eps: Self::FloatElem,
|
||||
) -> Self::TensorPrimitive<D> {
|
||||
let inner = tensor.into_primitive();
|
||||
let weight_inner = weight.as_primitive();
|
||||
let result = B::rms_norm(inner, weight_inner, eps);
|
||||
FusionTensor::from_primitive(result)
|
||||
}
|
||||
|
||||
fn rope<const D: usize>(
|
||||
tensor: Self::TensorPrimitive<D>,
|
||||
cos: &Self::TensorPrimitive<2>,
|
||||
sin: &Self::TensorPrimitive<2>,
|
||||
) -> Self::TensorPrimitive<D> {
|
||||
let inner = tensor.into_primitive();
|
||||
let cos_inner = cos.as_primitive();
|
||||
let sin_inner = sin.as_primitive();
|
||||
let result = B::rope(inner, cos_inner, sin_inner);
|
||||
FusionTensor::from_primitive(result)
|
||||
}
|
||||
|
||||
fn gelu<const D: usize>(tensor: Self::TensorPrimitive<D>) -> Self::TensorPrimitive<D> {
|
||||
let inner = tensor.into_primitive();
|
||||
let result = B::gelu(inner);
|
||||
FusionTensor::from_primitive(result)
|
||||
}
|
||||
|
||||
fn silu<const D: usize>(tensor: Self::TensorPrimitive<D>) -> Self::TensorPrimitive<D> {
|
||||
let inner = tensor.into_primitive();
|
||||
let result = B::silu(inner);
|
||||
FusionTensor::from_primitive(result)
|
||||
}
|
||||
|
||||
// ==================== Device Management ====================
|
||||
|
||||
fn device<const D: usize>(tensor: &Self::TensorPrimitive<D>) -> Self::Device {
|
||||
FusionDevice::new(B::device(tensor.as_primitive()))
|
||||
}
|
||||
|
||||
fn to_device<const D: usize>(
|
||||
tensor: Self::TensorPrimitive<D>,
|
||||
device: &Self::Device,
|
||||
) -> Self::TensorPrimitive<D> {
|
||||
// Device transfer is a sync point
|
||||
let inner = tensor.into_primitive();
|
||||
let result = B::to_device(inner, device.inner());
|
||||
FusionTensor::from_primitive(result)
|
||||
}
|
||||
|
||||
fn to_data<const D: usize>(tensor: &Self::TensorPrimitive<D>) -> Vec<Self::FloatElem> {
|
||||
// Data read is a sync point
|
||||
B::to_data(tensor.as_primitive())
|
||||
}
|
||||
|
||||
fn sync(device: &Self::Device) {
|
||||
// Explicit sync - flush all pending operations
|
||||
Self::flush(device);
|
||||
B::sync(device.inner());
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
// Test with mock backend (from tensor.rs tests)
|
||||
// In a real scenario, you would test with an actual backend
|
||||
|
||||
#[test]
|
||||
fn test_fusion_backend_name() {
|
||||
// Can't easily test without a concrete backend
|
||||
// This would be tested in integration tests
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user