feat: Split traits (#868)

Author: wingertge

crates/cubecl-attention/src/components/args.rs

@@ -243,7 +243,10 @@ impl<EI: Numeric, EO: Numeric, MA: AttentionArgs> VirtualTensorOperationsExpand<
     ) -> ExpandElementTyped<TensorMap<EO>> {
         unimplemented!("TensorOutputExpand can't be turned into a tensor map");
     }
+}
 
+impl<EI: Numeric, EO: Numeric, MA: AttentionArgs> Lined for TensorOutput<EI, EO, MA> {}
+impl<EI: Numeric, EO: Numeric, MA: AttentionArgs> LinedExpand for TensorOutputExpand<EI, EO, MA> {
     fn line_size(&self) -> u32 {
         let mut scope = Scope::root(false);
         TensorOutputExpand::__expand_line_size_method(self.clone(), &mut scope)
@@ -312,7 +315,10 @@ impl<EI: Numeric, EO: Numeric, MA: AttentionArgs> VirtualTensorOperationsExpand<
     ) -> ExpandElementTyped<TensorMap<EI>> {
         TensorInputExpand::__expand_as_tensor_map_method(self.clone(), scope)
     }
+}
 
+impl<EI: Numeric, EO: Numeric, MA: AttentionArgs> Lined for TensorInput<EI, EO, MA> {}
+impl<EI: Numeric, EO: Numeric, MA: AttentionArgs> LinedExpand for TensorInputExpand<EI, EO, MA> {
     fn line_size(&self) -> u32 {
         let mut scope = Scope::root(false);
         TensorInputExpand::__expand_line_size_method(self.clone(), &mut scope)

crates/cubecl-attention/src/components/global/dummy/attention.rs

@@ -113,7 +113,7 @@ impl<
         comment!("Global: Init Query Loader");
         let layout =
             SimpleGlobalLayout::new(&query, config.global_memory_config(FlashIdent::Query));
-        DummyQueryLoader::<AP, Self::Config>::new(q_offset, query.view(layout.virt()), config)
+        DummyQueryLoader::<AP, Self::Config>::new(q_offset, query.view(layout), config)
     }
 
     fn init_key_loader(
@@ -122,7 +122,7 @@ impl<
     ) -> Self::KeyLoader {
         comment!("Global: Init Key Loader");
         let layout = SimpleGlobalLayout::new(&key, config.global_memory_config(FlashIdent::Key));
-        DummyKeyLoader::new(key.view(layout.virt()), config)
+        DummyKeyLoader::new(key.view(layout), config)
     }
 
     fn init_value_loader(
@@ -132,7 +132,7 @@ impl<
         comment!("Global: Init Value Loader");
         let layout =
             SimpleGlobalLayout::new(&value, config.global_memory_config(FlashIdent::Value));
-        DummyValueLoader::new(value.view(layout.virt()), config)
+        DummyValueLoader::new(value.view(layout), config)
     }
 
     fn init_writer(
@@ -142,6 +142,6 @@ impl<
     ) -> Self::Writer {
         comment!("Global: Init Writer");
         let layout = SimpleGlobalLayout::new(&out, config.global_memory_config(FlashIdent::Out));
-        SA::init_writer(q_offset, out.view_mut(layout.virt()))
+        SA::init_writer(q_offset, out.view_mut(layout))
     }
 }

crates/cubecl-convolution/src/components/global/layout/base.rs

@@ -6,14 +6,14 @@ use cubecl_matmul::components::{
 };
 use cubecl_std::{
     FastDivmod,
-    tensor::layout::{Coords3d, Layout, VirtualLayoutOperations, VirtualLayoutOperationsExpand},
+    tensor::layout::{Coords3d, Layout, LayoutExpand},
 };
 
 use crate::{
     components::{
         ConvolutionConfig,
         global::{
-            layout::{NhwcCoords, unwrap, virtual_layout},
+            layout::{NhwcCoords, unwrap},
             load::im2col_tma::div_mod_seq,
         },
     },
@@ -24,7 +24,7 @@ use crate::{
 /// It first decomposes the `(m, k)` matrix into `((n, out_h, out_w), (k_h, k_w, c))`, then applies
 /// the convolution parameters to calculate the position in the input tensor for that kernel element.
 #[derive(CubeType, Clone)]
-pub struct Im2colGlobalLayout {
+pub struct Im2colLayout {
     /// Shape of output DHW
     pub shape_out: Sequence<FastDivmod>,
     /// Shape of channel, for decomposing k
@@ -54,14 +54,14 @@ pub struct Im2colGlobalLayout {
 }
 
 #[cube]
-impl Im2colGlobalLayout {
+impl Im2colLayout {
     pub fn new<G: GlobalConfig>(
         args: &RuntimeArgs,
         #[comptime] config: ConvolutionConfig<G>,
-    ) -> Im2colGlobalLayout {
+    ) -> Im2colLayout {
         let shape_out = args.shape_out.clone();
 
-        Im2colGlobalLayout {
+        Im2colLayout {
             shape_out,
             shape_channel: args.shape_channel,
             shape_m: args.shape_m,
@@ -76,32 +76,32 @@ impl Im2colGlobalLayout {
 }
 
 #[cube]
-impl Layout for Im2colGlobalLayout {
+impl Layout for Im2colLayout {
     type Coordinates = Coords3d;
     type SourceCoordinates = NhwcCoords;
 
-    fn to_source_pos(this: &Self, pos: Self::Coordinates) -> NhwcCoords {
+    fn to_source_pos(&self, pos: Self::Coordinates) -> NhwcCoords {
         let (_, view_m, view_k) = pos;
 
-        let (batch, out_offs) = div_mod_seq(view_m, &this.shape_out);
+        let (batch, out_offs) = div_mod_seq(view_m, &self.shape_out);
 
-        let (mut rem, channel) = this.shape_channel.div_mod(view_k);
+        let (mut rem, channel) = self.shape_channel.div_mod(view_k);
 
-        let spatial_dims = comptime![this.shape_out.len()];
+        let spatial_dims = comptime![self.shape_out.len()];
         let mut in_pos = Sequence::<i32>::new();
 
         #[unroll]
         for i in 0..spatial_dims {
             let i = unwrap(i);
             let dim = comptime![spatial_dims - i - 1];
-            let ksize = comptime![this.kernel_size[dim as usize]];
+            let ksize = comptime![self.kernel_size[dim as usize]];
             let k_pos = rem % ksize;
             rem /= ksize;
 
             let out_pos = *out_offs.index(dim);
-            let stride = comptime![this.stride[dim as usize]];
-            let dilate = comptime![this.dilation[dim as usize]];
-            let pad = comptime![this.padding[dim as usize]];
+            let stride = comptime![self.stride[dim as usize]];
+            let dilate = comptime![self.dilation[dim as usize]];
+            let pad = comptime![self.padding[dim as usize]];
 
             let pos = (out_pos * stride + k_pos * dilate) as i32 - pad;
             in_pos.push(pos);
@@ -116,21 +116,19 @@ impl Layout for Im2colGlobalLayout {
         }
     }
 
-    fn shape(this: &Self) -> Self::Coordinates {
-        (1, this.shape_m, this.shape_k)
+    fn shape(&self) -> Self::Coordinates {
+        (1, self.shape_m, self.shape_k)
     }
 
-    fn to_source_pos_checked(this: &Self, pos: Self::Coordinates) -> (NhwcCoords, bool) {
-        (this.to_source_pos(pos), this.is_in_bounds(pos))
+    fn to_source_pos_checked(&self, pos: Self::Coordinates) -> (NhwcCoords, bool) {
+        (self.to_source_pos(pos), self.is_in_bounds(pos))
     }
 
-    fn is_in_bounds(this: &Self, pos: Self::Coordinates) -> bool {
+    fn is_in_bounds(&self, pos: Self::Coordinates) -> bool {
         let (_, view_m, view_k) = pos;
         // Shouldn't be relied on because it doesn't check spatial
-        let m_in_bounds = comptime!(!this.config.check_row_bounds) || view_m < this.shape_m;
-        let k_in_bounds = comptime!(!this.config.check_col_bounds) || view_k < this.shape_k;
+        let m_in_bounds = comptime!(!self.config.check_row_bounds) || view_m < self.shape_m;
+        let k_in_bounds = comptime!(!self.config.check_col_bounds) || view_k < self.shape_k;
         m_in_bounds && k_in_bounds
     }
 }
-
-virtual_layout!(Im2colGlobalLayout, Im2colGlobalLayoutExpand);

crates/cubecl-convolution/src/components/global/layout/im2col.rs

@@ -1,10 +1,8 @@
-mod base;
 mod im2col;
 mod spatial;
 mod weight;
 mod write;
 
-pub(crate) use base::*;
 pub use im2col::*;
 pub use spatial::*;
 pub use weight::*;

crates/cubecl-convolution/src/components/global/layout/mod.rs

@@ -1,13 +1,11 @@
 use cubecl::prelude::*;
 use cubecl_core::{self as cubecl, intrinsic};
 use cubecl_std::tensor::{
-    layout::{
-        Coordinates, Coords1d, Layout, VirtualLayoutOperations, VirtualLayoutOperationsExpand,
-    },
+    layout::{Coordinates, Coords1d, Layout, LayoutExpand},
     r#virtual::VirtualTensor,
 };
 
-use crate::components::{Dimensionality, global::layout::virtual_layout};
+use crate::components::Dimensionality;
 
 #[derive(CubeType, Clone)]
 pub struct NhwcCoords {
@@ -93,42 +91,39 @@ impl Layout for NhwcLayout {
     type Coordinates = NhwcCoords;
     type SourceCoordinates = Coords1d;
 
-    fn to_source_pos(this: &Self, pos: Self::Coordinates) -> Self::SourceCoordinates {
+    fn to_source_pos(&self, pos: Self::Coordinates) -> Self::SourceCoordinates {
         let NhwcCoords {
             batch,
             spatial,
             channel,
         } = pos;
 
-        let spatial_dims = this.shapes_spatial.len();
-        let mut read_pos = batch * this.stride_batch + channel * this.stride_channel;
+        let spatial_dims = self.shapes_spatial.len();
+        let mut read_pos = batch * self.stride_batch + channel * self.stride_channel;
 
         #[unroll]
         for i in 0..spatial_dims {
             let i = unwrap(i);
-            read_pos += *spatial.index(i) as u32 * *this.strides_spatial.index(i);
+            read_pos += *spatial.index(i) as u32 * *self.strides_spatial.index(i);
         }
 
-        read_pos / this.line_size
+        read_pos / self.line_size
     }
 
-    fn to_source_pos_checked(
-        this: &Self,
-        pos: Self::Coordinates,
-    ) -> (Self::SourceCoordinates, bool) {
-        (this.to_source_pos(pos.clone()), this.is_in_bounds(pos))
+    fn to_source_pos_checked(&self, pos: Self::Coordinates) -> (Self::SourceCoordinates, bool) {
+        (self.to_source_pos(pos.clone()), self.is_in_bounds(pos))
     }
 
-    fn is_in_bounds(this: &Self, pos: Self::Coordinates) -> bool {
-        if comptime![this.check_spatial] {
-            let spatial_dims = this.shapes_spatial.len();
+    fn is_in_bounds(&self, pos: Self::Coordinates) -> bool {
+        if comptime![self.check_spatial] {
+            let spatial_dims = self.shapes_spatial.len();
             let mut spatial_in_bounds = true;
 
             #[unroll]
             for i in 0..spatial_dims {
                 let i = unwrap(i);
                 let pos = *pos.spatial.index(i);
-                spatial_in_bounds &= pos >= 0 && (pos as u32) < *this.shapes_spatial.index(i);
+                spatial_in_bounds &= pos >= 0 && (pos as u32) < *self.shapes_spatial.index(i);
             }
 
             spatial_in_bounds
@@ -137,17 +132,15 @@ impl Layout for NhwcLayout {
         }
     }
 
-    fn shape(this: &Self) -> Self::Coordinates {
+    fn shape(&self) -> Self::Coordinates {
         NhwcCoords {
-            batch: this.shape_batch,
-            spatial: cast_seq(this.shapes_spatial.clone()),
-            channel: this.shape_channel,
+            batch: self.shape_batch,
+            spatial: cast_seq(self.shapes_spatial.clone()),
+            channel: self.shape_channel,
         }
     }
 }
 
-virtual_layout!(NhwcLayout, NhwcLayoutExpand);
-
 #[allow(unused_variables)]
 #[cube]
 pub(crate) fn unwrap(v: u32) -> comptime_type!(u32) {

crates/cubecl-convolution/src/components/global/layout/spatial.rs

@@ -7,15 +7,15 @@ use cubecl_matmul::components::{
 use cubecl_std::{
     FastDivmod,
     tensor::{
-        layout::{Coords3d, Layout, VirtualLayoutOperations, VirtualLayoutOperationsExpand},
+        layout::{Coords3d, Layout, LayoutExpand},
         r#virtual::VirtualTensor,
     },
 };
 
 use crate::{
     components::{
         ConvGemmConfig, ConvolutionConfig,
-        global::layout::{NhwcCoords, unwrap, virtual_layout},
+        global::layout::{NhwcCoords, unwrap},
     },
     kernels::layered::selector::RuntimeArgs,
 };
@@ -83,19 +83,19 @@ impl Layout for WeightLayout {
     type Coordinates = Coords3d;
     type SourceCoordinates = NhwcCoords;
 
-    fn to_source_pos(this: &Self, coords: Self::Coordinates) -> NhwcCoords {
+    fn to_source_pos(&self, coords: Self::Coordinates) -> NhwcCoords {
         let (_, k, n) = coords;
 
-        let (mut rem, in_c) = this.channels.div_mod(k);
+        let (mut rem, in_c) = self.channels.div_mod(k);
 
-        let spatial_dims = comptime![this.strides_spatial.len()];
+        let spatial_dims = comptime![self.strides_spatial.len()];
         let mut kernel_pos = Sequence::<i32>::new();
 
         #[unroll]
         for i in 0..spatial_dims {
             let i = unwrap(i);
             let dim = comptime![spatial_dims - i - 1];
-            let ksize = comptime![this.kernel_size[dim as usize]];
+            let ksize = comptime![self.kernel_size[dim as usize]];
             let k_pos = rem % ksize;
             rem /= ksize;
 
@@ -111,20 +111,18 @@ impl Layout for WeightLayout {
         }
     }
 
-    fn to_source_pos_checked(this: &Self, coords: Self::Coordinates) -> (NhwcCoords, bool) {
-        (this.to_source_pos(coords), this.is_in_bounds(coords))
+    fn to_source_pos_checked(&self, coords: Self::Coordinates) -> (NhwcCoords, bool) {
+        (self.to_source_pos(coords), self.is_in_bounds(coords))
     }
 
-    fn shape(this: &Self) -> Self::Coordinates {
-        (1, this.shape_k, this.shape_n)
+    fn shape(&self) -> Self::Coordinates {
+        (1, self.shape_k, self.shape_n)
     }
 
-    fn is_in_bounds(this: &Self, pos: Self::Coordinates) -> bool {
+    fn is_in_bounds(&self, pos: Self::Coordinates) -> bool {
         let (_, k, n) = pos;
-        let check_k = comptime![this.config.check_row_bounds];
-        let check_n = comptime![this.config.check_col_bounds];
-        (!check_k || k < this.shape_k) && (!check_n || n < this.shape_n)
+        let check_k = comptime![self.config.check_row_bounds];
+        let check_n = comptime![self.config.check_col_bounds];
+        (!check_k || k < self.shape_k) && (!check_n || n < self.shape_n)
     }
 }
-
-virtual_layout!(WeightLayout, WeightLayoutExpand);