Initial commit

crates/training/rtx-distributed/src/dtensor/placement.rs

@@ -0,0 +1,59 @@
+//! 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,
+        }
+    }
+}