60 lines
1.8 KiB
Rust
60 lines
1.8 KiB
Rust
//! Placement types for distributed tensors.
|
|
|
|
use crate::comm::ReduceOp;
|
|
use serde::{Deserialize, Serialize};
|
|
|
|
/// Specifies how a tensor dimension is placed across a mesh dimension.
|
|
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
|
|
pub enum Placement {
|
|
/// The tensor dimension is replicated across this mesh dimension.
|
|
/// All devices in this mesh dimension hold the same data.
|
|
Replicate,
|
|
|
|
/// The tensor dimension is sharded across this mesh dimension.
|
|
/// Each device holds a slice of the data.
|
|
Shard {
|
|
/// Which tensor dimension is sharded (index into tensor shape)
|
|
tensor_dim: usize,
|
|
},
|
|
|
|
/// Partial placement - tensor holds partial results that need reduction.
|
|
/// Used after operations like matmul where results are distributed.
|
|
Partial {
|
|
/// Reduction operation needed to combine partial results
|
|
reduce_op: PartialReduceOp,
|
|
},
|
|
|
|
/// Interleaved sharding for better memory access patterns.
|
|
/// Data is distributed in an interleaved fashion rather than contiguous chunks.
|
|
InterleavedShard {
|
|
/// Which tensor dimension is sharded
|
|
tensor_dim: usize,
|
|
/// Interleave factor (default: 1)
|
|
interleave_factor: usize,
|
|
},
|
|
}
|
|
|
|
/// Reduction operation for partial placements.
|
|
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
|
|
pub enum PartialReduceOp {
|
|
/// Sum reduction
|
|
Sum,
|
|
/// Mean reduction
|
|
Mean,
|
|
/// Max reduction
|
|
Max,
|
|
/// Min reduction
|
|
Min,
|
|
}
|
|
|
|
impl From<PartialReduceOp> for ReduceOp {
|
|
fn from(op: PartialReduceOp) -> Self {
|
|
match op {
|
|
PartialReduceOp::Sum => ReduceOp::Sum,
|
|
PartialReduceOp::Mean => ReduceOp::Sum, // Mean is sum then divide
|
|
PartialReduceOp::Max => ReduceOp::Max,
|
|
PartialReduceOp::Min => ReduceOp::Min,
|
|
}
|
|
}
|
|
}
|