IR Refactor Part 2: Reflection, Semantic Categories

Cargo.toml

@@ -29,6 +29,7 @@ type_hash = "0.3"
 dashmap = "5.5.3"
 hashbrown = "0.14.5"
 spin = { version = "0.9.8", features = ["mutex", "spin_mutex"] }
+smallvec = { version = "1", features = ["union", "const_generics"] }
 
 getrandom = { version = "0.2.15", default-features = false }
 rand = { version = "0.8.5", default-features = false, features = [

crates/cubecl-core/Cargo.toml

@@ -10,7 +10,7 @@ keywords = []
 license.workspace = true
 name = "cubecl-core"
 readme.workspace = true
-repository = "https://github.com/tracel-ai/cubecl/tree/main/cubecl-cube"
+repository = "https://github.com/tracel-ai/cubecl/tree/main/crates/cubecl-core"
 version.workspace = true
 
 [features]

crates/cubecl-core/src/frontend/container/array/base.rs

@@ -259,8 +259,10 @@ mod metadata {
 
 /// Module that contains the implementation details of the index functions.
 mod indexation {
+    use cubecl_ir::Operator;
+
     use crate::{
-        ir::{BinaryOperator, Instruction, Operator},
+        ir::{BinaryOperator, Instruction},
         prelude::{CubeIndex, CubeIndexMut},
     };
 

crates/cubecl-core/src/frontend/container/line/base.rs

@@ -1,9 +1,9 @@
 use std::num::NonZero;
 
-use cubecl_ir::ExpandElement;
+use cubecl_ir::{Comparison, ExpandElement};
 
 use crate::{
-    ir::{BinaryOperator, ConstantScalarValue, Elem, Instruction, Item, Operator},
+    ir::{Arithmetic, BinaryOperator, ConstantScalarValue, Elem, Instruction, Item},
     prelude::{binary_expand_fixed_output, CubeContext, Dot, Numeric},
     unexpanded,
 };
@@ -220,7 +220,7 @@ macro_rules! impl_line_comparison {
                         let output = context.create_local_mut(Item::vectorized(bool::as_elem(context), size));
 
                         context.register(Instruction::new(
-                            Operator::$operator(BinaryOperator { lhs, rhs }),
+                            Comparison::$operator(BinaryOperator { lhs, rhs }),
                             output.clone().into(),
                         ));
 
@@ -286,6 +286,6 @@ impl<N: Numeric> Dot for Line<N> {
         let lhs: ExpandElement = lhs.into();
         let mut item = lhs.item;
         item.vectorization = None;
-        binary_expand_fixed_output(context, lhs, rhs.into(), item, Operator::Dot).into()
+        binary_expand_fixed_output(context, lhs, rhs.into(), item, Arithmetic::Dot).into()
     }
 }

crates/cubecl-core/src/frontend/container/shared_memory.rs

@@ -95,8 +95,10 @@ impl<T: CubePrimitive + Clone> SharedMemory<T> {
 
 /// Module that contains the implementation details of the index functions.
 mod indexation {
+    use cubecl_ir::Operator;
+
     use crate::{
-        ir::{BinaryOperator, Instruction, Operator},
+        ir::{BinaryOperator, Instruction},
         prelude::{CubeIndex, CubeIndexMut},
         unexpanded,
     };

crates/cubecl-core/src/frontend/container/slice.rs

@@ -1,10 +1,10 @@
 use std::marker::PhantomData;
 
-use cubecl_ir::ExpandElement;
+use cubecl_ir::{ExpandElement, Operator};
 
 use crate::{
     frontend::{indexation::Index, Tensor},
-    ir::{self, Operator},
+    ir::{self},
     prelude::{CubeContext, IntoRuntime},
     unexpanded,
 };
@@ -123,10 +123,11 @@ mod metadata {
 
 /// Module that contains the implementation details of the index functions.
 mod indexation {
+    use cubecl_ir::Operator;
     use ir::Instruction;
 
     use crate::{
-        ir::{BinaryOperator, Operator},
+        ir::BinaryOperator,
         prelude::{CubeIndex, CubeIndexMut},
     };
 

crates/cubecl-core/src/frontend/container/tensor/base.rs

@@ -23,7 +23,7 @@ mod metadata {
 
     use super::*;
     use crate::{
-        ir::{BinaryOperator, Instruction, Operator},
+        ir::{Arithmetic, BinaryOperator, Instruction},
         prelude::Array,
     };
 
@@ -177,7 +177,7 @@ mod metadata {
             // Compute `num_strides = index / stride`.
             let num_strides = context.create_local(Item::new(u32::as_elem(context)));
             context.register(Instruction::new(
-                Operator::Div(BinaryOperator {
+                Arithmetic::Div(BinaryOperator {
                     lhs: *index,
                     rhs: stride.expand.into(),
                 }),
@@ -187,7 +187,7 @@ mod metadata {
             // Compute `coordinate = num_strides % shape `.
             let coordinate = context.create_local(Item::new(u32::as_elem(context)));
             context.register(Instruction::new(
-                Operator::Modulo(BinaryOperator {
+                Arithmetic::Modulo(BinaryOperator {
                     lhs: *num_strides,
                     rhs: shape.expand.into(),
                 }),
@@ -223,8 +223,10 @@ mod metadata {
 
 /// Module that contains the implementation details of the index functions.
 mod indexation {
+    use cubecl_ir::Operator;
+
     use crate::{
-        ir::{BinaryOperator, Instruction, Operator},
+        ir::{BinaryOperator, Instruction},
         prelude::{CubeIndex, CubeIndexMut},
     };
 

crates/cubecl-core/src/frontend/element/cast.rs

@@ -1,11 +1,8 @@
-use cubecl_ir::ExpandElement;
+use cubecl_ir::{ExpandElement, Operator};
 
+use crate::frontend::{cast, CubeContext, CubePrimitive, CubeType};
 use crate::ir::{Instruction, Item, UnaryOperator, Variable};
 use crate::unexpanded;
-use crate::{
-    frontend::{cast, CubeContext, CubePrimitive, CubeType},
-    ir::Operator,
-};
 
 use super::ExpandElementTyped;