fix: Align shift right behaviour to align with most other compilers (#894)

Author: wingertge

crates/cubecl-attention/src/components/tile/dummy/flash_matmul/accelerated/setup.rs

@@ -33,7 +33,7 @@ impl FlashMatmulFamily for AcceleratedFlashMatmul {
             1,
             line_sizes.query as u32,
             line_sizes.key as u32,
-            problem.seq_kv as u32 % selection.attention_tile_size.seq_kv != 0,
+            !(problem.seq_kv as u32).is_multiple_of(selection.attention_tile_size.seq_kv),
         )
     }
 }

crates/cubecl-attention/src/components/tile/dummy/flash_matmul/dummy_register/setup.rs

@@ -35,7 +35,7 @@ impl FlashMatmulFamily for DummyRegisterFlashMatmul {
             1,
             line_sizes.query as u32,
             line_sizes.key as u32,
-            problem.seq_kv as u32 % selection.attention_tile_size.seq_kv != 0,
+            !(problem.seq_kv as u32).is_multiple_of(selection.attention_tile_size.seq_kv),
         )
     }
 }

crates/cubecl-attention/src/components/tile/dummy/writer.rs

@@ -48,7 +48,7 @@ impl<EO: Numeric> DummyWriter<EO> {
 
         let unit_step = config.plane_dim() * output_line_size;
         let num_unit_writes = comptime!(div_ceil(tile_size, unit_step));
-        let balanced_workload = comptime!(tile_size % unit_step == 0);
+        let balanced_workload = comptime!(tile_size.is_multiple_of(unit_step));
 
         #[unroll(num_unit_writes == 1)]
         for i in 0..num_unit_writes {

crates/cubecl-core/src/runtime_tests/cmma.rs

@@ -597,8 +597,8 @@ pub fn test_simple_tf32<R: Runtime>(
         return;
     }
 
-    let lhs: Vec<f32> = (0..128).map(|i| (i as f32)).collect();
-    let rhs: Vec<f32> = (0..128).map(|i| ((i % 8) as f32)).collect();
+    let lhs: Vec<f32> = (0..128).map(|i| i as f32).collect();
+    let rhs: Vec<f32> = (0..128).map(|i| (i % 8) as f32).collect();
 
     let lhs = client.create(f32::as_bytes(&lhs));
     let rhs = client.create(f32::as_bytes(&rhs));

crates/cubecl-cpp/src/shared/item.rs

@@ -51,7 +51,7 @@ impl<D: Dialect> Item<D> {
     }
 
     pub fn optimized(&self) -> Item<D> {
-        if !self.can_be_optimized() || self.vectorization % 2 != 0 {
+        if !self.can_be_optimized() || !self.vectorization.is_multiple_of(2) {
             return *self;
         }
 

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

@@ -44,7 +44,6 @@ use cudarc::driver::sys::{
 use cudarc::driver::sys::{CUfunc_st, CUtensorMapInterleave};
 #[cfg(feature = "cuda-12080")]
 use cudarc::driver::sys::{CUtensorMapIm2ColWideMode, cuTensorMapEncodeIm2colWide};
-use serde::{Deserialize, Serialize};
 use std::collections::HashMap;
 use std::ffi::c_char;
 use std::path::PathBuf;
@@ -78,7 +77,8 @@ pub(crate) struct CudaContext {
     compilation_options: CompilationOptions,
 }
 
-#[derive(Debug, Serialize, Deserialize, PartialEq, Eq, Clone)]
+#[cfg(feature = "compilation-cache")]
+#[derive(Debug, serde::Serialize, serde::Deserialize, PartialEq, Eq, Clone)]
 pub struct PtxCacheEntry {
     entrypoint_name: String,
     cube_dim: (u32, u32, u32),
@@ -336,7 +336,7 @@ impl ComputeServer for CudaServer {
                     .expect("Failed to find resource");
                 let device_ptr = resource.ptr as *mut c_void;
                 debug_assert!(
-                    device_ptr as usize % 16 == 0,
+                    (device_ptr as usize).is_multiple_of(16),
                     "Tensor pointer must be 16 byte aligned"
                 );
                 let mut map_ptr = MaybeUninit::zeroed();

crates/cubecl-matmul/src/components/batch/partitioned_matmul/hypercube/base.rs

@@ -76,13 +76,13 @@ impl HypercubeConfig {
         match self.global_order {
             RowMajor | ColMajor => Ok(()),
 
-            SwizzleRowMajor(w) if m_cubes % w != 0 => {
+            SwizzleRowMajor(w) if !m_cubes.is_multiple_of(w) => {
                 Err(MatmulSetupError::InvalidConfig(Box::new(format!(
                     "In swizzle row major, number of cubes in m {m_cubes:?} must be divisible by swizzle step length {w:?}."
                 ))))
             }
 
-            SwizzleColMajor(w) if n_cubes % w != 0 => {
+            SwizzleColMajor(w) if !n_cubes.is_multiple_of(w) => {
                 Err(MatmulSetupError::InvalidConfig(Box::new(format!(
                     "In swizzle col major, number of cubes in n {n_cubes:?} must be divisible by swizzle step length {w:?}."
                 ))))

crates/cubecl-matmul/src/components/batch/partitioned_matmul/hypercube/sm_allocation.rs

@@ -51,7 +51,7 @@ impl SmAllocation {
                     let mut i = 1;
 
                     while i * i <= n {
-                        if n % i == 0 {
+                        if n.is_multiple_of(i) {
                             divs.push(i);
                             if i != n / i {
                                 divs.push(n / i);

crates/cubecl-matmul/src/components/global/load/strategy/async_full_cyclic.rs

@@ -26,7 +26,7 @@ impl<T: TilingOrder> LoadingValidation for AsyncFullCyclicLoading<T> {
         let num_slices = config.tiling_scheme().elements_in_tile_row(ident)
             * config.tiling_scheme().tiles_in_stage(ident);
 
-        if num_slices >= total_units && num_slices % total_units != 0 {
+        if num_slices >= total_units && !num_slices.is_multiple_of(total_units) {
             return Err(Box::new(format!(
                 "Number of units ({total_units:?}) must divide number of slices ({num_slices:?}). Would require units doing different numbers of slices"
             )));

crates/cubecl-matmul/src/components/global/load/strategy/async_full_maximize_slice_length.rs

@@ -83,7 +83,7 @@ impl<IP: InputPrecision> AsyncLoadingJob<IP, StridedTilingLayout>
         let nth_slice = this.unit_count * task_id + UNIT_POS;
 
         #[allow(clippy::collapsible_else_if)]
-        if comptime!(this.num_slices % this.unit_count == 0) {
+        if comptime!(this.num_slices.is_multiple_of(this.unit_count)) {
             load_nth_slice::<IP::Global, IP::Stage, CM, G>(
                 nth_slice,
                 tensor_reader,