Reimplement macro to make writing kernels more ergonomic

.gitignore

@@ -12,3 +12,5 @@ Cargo.lock
 
 # MSVC Windows builds of rustc generate these, which store debugging information
 *.pdb
+**/out
+.clangd

.vscode/settings.json

@@ -0,0 +1,3 @@
+{
+    "files.eol": "\n"
+}

Cargo.toml

@@ -4,17 +4,13 @@
 # https://doc.rust-lang.org/cargo/reference/resolver.html#feature-resolver-version-2
 resolver = "2"
 
-members = [
-    "crates/*",
-    "examples/*", 
-    "xtask",
-]
+members = ["crates/*", "examples/*", "xtask"]
 
 [workspace.package]
 edition = "2021"
-version = "0.2.0"
 license = "MIT OR Apache-2.0"
 readme = "README.md"
+version = "0.2.0"
 
 
 [workspace.dependencies]
@@ -29,23 +25,24 @@ serde = { version = "1.0.204", default-features = false, features = [
 serde_json = { version = "1.0.119", default-features = false }
 
 dashmap = "5.5.3"
-spin = { version = "0.9.8", features = ["mutex", "spin_mutex"] }
 hashbrown = "0.14.5"
+spin = { version = "0.9.8", features = ["mutex", "spin_mutex"] }
 
+getrandom = { version = "0.2.15", default-features = false }
 rand = { version = "0.8.5", default-features = false, features = [
     "std_rng",
 ] } # std_rng is for no_std
-getrandom = { version = "0.2.15", default-features = false }
 
-pollster = "0.3"
+async-channel = "2.3"
 dirs = "5.0.1"
-web-time = "1.1.0"
 md5 = "0.7.0"
-async-channel = "2.3"
+pollster = "0.3"
+weak-table = "0.3"
+web-time = "1.1.0"
 
 # Testing
-serial_test = "3.1.1"
 rstest = "0.19.0"
+serial_test = "3.1.1"
 
 bytemuck = "1.16.1"
 half = { version = "2.4.1", features = [
@@ -57,15 +54,20 @@ num-traits = { version = "0.2.19", default-features = false, features = [
     "libm",
 ] } # libm is for no_std
 
+darling = "0.20.10"
+ident_case = "1"
 proc-macro2 = "1.0.86"
-syn = { version = "2.0.69", features = ["full", "extra-traits"] }
 quote = "1.0.36"
+syn = { version = "2", features = ["full", "extra-traits", "visit-mut"] }
 
 ### For xtask crate ###
-strum = {version = "0.26.3", features = ["derive"]}
+strum = { version = "0.26.3", features = ["derive"] }
 tracel-xtask = { version = "~1.0" }
 
-portable-atomic-util = { version = "0.2.2", features = ["alloc"] } # alloc is for no_std
+portable-atomic-util = { version = "0.2.2", features = [
+    "alloc",
+] } # alloc is for no_std
+pretty_assertions = "1.4"
 
 [profile.dev]
 opt-level = 2

crates/cubecl-common/Cargo.toml

@@ -1,5 +1,8 @@
 [package]
-authors = ["Dilshod Tadjibaev (@antimora)", "Nathaniel Simard (@nathanielsimard)"]
+authors = [
+    "Dilshod Tadjibaev (@antimora)",
+    "Nathaniel Simard (@nathanielsimard)",
+]
 categories = ["science", "mathematics", "algorithms"]
 description = "Common crate for CubeCL"
 edition.workspace = true
@@ -20,18 +23,22 @@ web-time = { version = "1.1.0" }
 
 [dependencies]
 # ** Please make sure all dependencies support no_std when std is disabled **
-spin = { workspace = true }                      # using in place of use std::sync::Mutex;
 derive-new = { workspace = true }
-serde = { workspace = true }
-rand = { workspace = true }
 pollster = { workspace = true, optional = true }
+rand = { workspace = true }
+serde = { workspace = true }
+spin = { workspace = true }                      # using in place of use std::sync::Mutex;
 
 [target.'cfg(target_has_atomic = "ptr")'.dependencies]
 spin = { workspace = true, features = ["mutex", "spin_mutex"] }
 
 [target.'cfg(not(target_has_atomic = "ptr"))'.dependencies]
 portable-atomic-util = { workspace = true }
-spin = { workspace = true, features = ["mutex", "spin_mutex", "portable_atomic"] }
+spin = { workspace = true, features = [
+    "mutex",
+    "spin_mutex",
+    "portable_atomic",
+] }
 
 [dev-dependencies]
 dashmap = { workspace = true }

crates/cubecl-common/src/lib.rs

@@ -22,6 +22,8 @@ pub mod benchmark;
 /// notation.
 pub mod reader;
 
+/// Operators used by macro and IR
+pub mod operator;
 /// Synchronization type module, used both by ComputeServer and Backends.
 pub mod sync_type;
 

crates/cubecl-common/src/operator.rs

@@ -0,0 +1 @@
+

crates/cubecl-core/Cargo.toml

@@ -15,21 +15,23 @@ version.workspace = true
 
 [features]
 default = ["cubecl-runtime/default"]
+export_tests = []
 std = ["cubecl-runtime/std"]
 template = []
-export_tests = []
 
 [dependencies]
 cubecl-runtime = { path = "../cubecl-runtime", version = "0.2.0", default-features = false }
 
 bytemuck = { workspace = true }
-half = { workspace = true, features = ["bytemuck"] }
-serde = { workspace = true }
+cubecl-common = { path = "../cubecl-common", version = "0.2.0" }
 cubecl-macros = { path = "../cubecl-macros", version = "0.2.0" }
 derive-new = { workspace = true }
+half = { workspace = true, features = ["bytemuck"] }
 num-traits = { workspace = true }
+serde = { workspace = true }
 
 log = { workspace = true }
 
 [dev-dependencies]
+pretty_assertions = { workspace = true }
 trybuild = "1"

crates/cubecl-core/src/codegen/integrator.rs

@@ -1,3 +1,5 @@
+use std::num::NonZero;
+
 use super::Compiler;
 use crate::{
     ir::{
@@ -95,18 +97,22 @@ impl core::fmt::Display for KernelSettings {
         }
 
         match self.vectorization_global {
-            Some(vectorization) => f.write_fmt(format_args!("vg{}", vectorization))?,
+            Some(vectorization) => f.write_fmt(format_args!(
+                "vg{}",
+                vectorization.map(NonZero::get).unwrap_or(1)
+            ))?,
             None => f.write_str("vn")?,
         };
 
         for vectorization in self.vectorization_partial.iter() {
             match vectorization {
-                VectorizationPartial::Input { pos, vectorization } => {
-                    f.write_fmt(format_args!("v{vectorization}i{pos}"))?
-                }
-                VectorizationPartial::Output { pos, vectorization } => {
-                    f.write_fmt(format_args!("v{vectorization}o{pos}"))?
-                }
+                VectorizationPartial::Input { pos, vectorization } => f.write_fmt(format_args!(
+                    "v{}i{pos}",
+                    vectorization.map(NonZero::get).unwrap_or(1)
+                ))?,
+                VectorizationPartial::Output { pos, vectorization } => f.write_fmt(
+                    format_args!("v{}o{pos}", vectorization.map(NonZero::get).unwrap_or(1)),
+                )?,
             };
         }
 
@@ -130,7 +136,7 @@ impl KernelSettings {
     pub fn vectorize_input(mut self, position: usize, vectorization: Vectorization) -> Self {
         // Not setting the vectorization factor when it's the default value reduces the kernel id
         // size.
-        if vectorization == 1 {
+        if vectorization.is_none() {
             return self;
         }
 
@@ -147,7 +153,7 @@ impl KernelSettings {
     pub fn vectorize_output(mut self, position: usize, vectorization: Vectorization) -> Self {
         // Not setting the vectorization factor when it's the default value reduces the kernel id
         // size.
-        if vectorization == 1 {
+        if vectorization.is_none() {
             return self;
         }
 
@@ -173,7 +179,7 @@ impl KernelSettings {
             }
         }
 
-        1
+        None
     }
 
     /// Fetch the vectorization for the provided output position.
@@ -190,7 +196,7 @@ impl KernelSettings {
             }
         }
 
-        1
+        None
     }
 
     /// Compile the shader with inplace enabled by the given [mapping](InplaceMapping).