wgpu: resolve merge; import stride::contiguous_strides; format

Author: ariawisp

crates/cubecl-wgpu/src/compute/server.rs

@@ -16,6 +16,7 @@ use cubecl_core::{
 };
 use cubecl_runtime::logging::ServerLogger;
 use cubecl_runtime::memory_management::offset_handles;
+use cubecl_runtime::stride::contiguous_strides;
 use cubecl_runtime::{
     memory_management::MemoryDeviceProperties, server::ComputeServer, storage::BindingResource,
 };
@@ -166,15 +167,14 @@ impl ComputeServer for WgpuServer {
                 continue;
             }
 
-            // 2D pitched rows: support rank==2 with inner-most contiguous dimension.
-            // Note: contiguous path unchanged; pitched path uses per-row queue.write_buffer with small overhead.
-            let shape = desc.shape;
-            if shape.len() == 2 && desc.strides[1] == 1 {
-                let rows = shape[0] as u64;
-                let cols = shape[1] as u64;
+            // Inner-contiguous pitched rows: rank>=2, inner-most contiguous
+            if desc.shape.len() >= 2 && desc.strides[desc.shape.len() - 1] == 1 {
+                let last = desc.shape.len() - 1;
+                let rows = desc.shape[..last].iter().product::<usize>() as u64;
+                let cols = desc.shape[last] as u64;
                 let elem = desc.elem_size as u64;
                 let row_bytes = cols * elem;
-                let row_pitch = desc.strides[0] as u64 * elem;
+                let row_pitch = desc.strides[last - 1] as u64 * elem;
 
                 let resource = self.stream.mem_manage.get_resource(desc.binding);
                 self.stream
@@ -245,11 +245,4 @@ fn compiler(backend: wgpu::Backend) -> AutoCompiler {
     }
 }
 
-pub(crate) fn contiguous_strides(shape: &[usize]) -> Vec<usize> {
-    let rank = shape.len();
-    let mut strides = vec![1; rank];
-    for i in (0..rank - 1).rev() {
-        strides[i] = strides[i + 1] * shape[i + 1];
-    }
-    strides
-}
+// Note: use cubecl_runtime::stride::contiguous_strides for canonical row-major strides.

crates/cubecl-wgpu/src/compute/stream.rs

@@ -189,8 +189,8 @@ impl WgpuStream {
             let resource = self.mem_manage.get_resource(binding);
 
             // Contiguous path: copy entire resource range
-            let is_contiguous = super::super::compute::server::contiguous_strides(descriptor.shape)
-                == descriptor.strides;
+            let is_contiguous =
+                cubecl_runtime::stride::contiguous_strides(descriptor.shape) == descriptor.strides;
 
             if is_contiguous {
                 let size = descriptor.shape.iter().product::<usize>() * descriptor.elem_size;
@@ -214,12 +214,13 @@ impl WgpuStream {
                 continue;
             }
 
-            // 2D pitched rows: rank==2, innermost contiguous
-            if descriptor.shape.len() == 2 && descriptor.strides[1] == 1 {
-                let rows = descriptor.shape[0] as u64;
-                let cols = descriptor.shape[1] as u64;
+            // Inner-contiguous pitched rows: rank>=2, innermost contiguous
+            if descriptor.shape.len() >= 2 && descriptor.strides[descriptor.shape.len() - 1] == 1 {
+                let last = descriptor.shape.len() - 1;
+                let rows = descriptor.shape[..last].iter().product::<usize>() as u64;
+                let cols = descriptor.shape[last] as u64;
                 let row_bytes = cols * elem;
-                let row_pitch = descriptor.strides[0] as u64 * elem;
+                let row_pitch = descriptor.strides[last - 1] as u64 * elem;
                 let total = rows * row_pitch;
                 let align = wgpu::COPY_BUFFER_ALIGNMENT;
                 let aligned_total = total.div_ceil(align) * align;