IR Refactor

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

@@ -4,7 +4,7 @@ use super::Compiler;
 use crate::{
     ir::{
         Binding, CubeDim, Elem, Item, KernelDefinition, Location, ReadingStrategy, Scope, Variable,
-        Vectorization, Visibility,
+        VariableKind, Vectorization, Visibility,
     },
     Runtime,
 };
@@ -399,15 +399,15 @@ impl KernelIntegrator {
                         size: None,
                     });
                     self.expansion.scope.write_global(
-                        Variable::Local {
-                            id: local,
+                        Variable::new(
+                            VariableKind::Local {
+                                id: local,
+
+                                depth: self.expansion.scope.depth,
+                            },
                             item,
-                            depth: self.expansion.scope.depth,
-                        },
-                        Variable::GlobalOutputArray {
-                            id: index,
-                            item: item_adapted,
-                        },
+                        ),
+                        Variable::new(VariableKind::GlobalOutputArray(index), item_adapted),
                         position,
                     );
                     index += 1;
@@ -419,15 +419,14 @@ impl KernelIntegrator {
                     position,
                 } => {
                     self.expansion.scope.write_global(
-                        Variable::Local {
-                            id: local,
+                        Variable::new(
+                            VariableKind::Local {
+                                id: local,
+                                depth: self.expansion.scope.depth,
+                            },
                             item,
-                            depth: self.expansion.scope.depth,
-                        },
-                        Variable::GlobalInputArray {
-                            id: input,
-                            item: bool_item(item),
-                        },
+                        ),
+                        Variable::new(VariableKind::GlobalInputArray(input), bool_item(item)),
                         position,
                     );
                 }

crates/cubecl-core/src/frontend/branch.rs

@@ -396,8 +396,7 @@ pub fn if_else_expr_expand<C: CubePrimitive>(
         None => {
             let mut then_child = context.child();
             let ret = then_block(&mut then_child);
-            let out: ExpandElementTyped<C> =
-                context.create_local_variable(ret.expand.item()).into();
+            let out: ExpandElementTyped<C> = context.create_local_variable(ret.expand.item).into();
             assign::expand(&mut then_child, ret, out.clone());
 
             IfElseExprExpand::Runtime {
@@ -502,7 +501,7 @@ pub fn switch_expand_expr<I: Int, C: CubePrimitive>(
 ) -> SwitchExpandExpr<I, C> {
     let mut default_child = context.child();
     let default = default_block(&mut default_child);
-    let out: ExpandElementTyped<C> = context.create_local_variable(default.expand.item()).into();
+    let out: ExpandElementTyped<C> = context.create_local_variable(default.expand.item).into();
     assign::expand(&mut default_child, default, out.clone());
 
     SwitchExpandExpr {

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

@@ -49,7 +49,7 @@
 use std::marker::PhantomData;
 
 use crate::{
-    ir::{self, Operation},
+    ir::{self, Instruction, Operation},
     unexpanded,
 };
 
@@ -239,10 +239,10 @@ pub mod fill {
         value: ExpandElementTyped<C>,
     ) {
         let value: ExpandElement = value.into();
-        context.register(Operation::CoopMma(ir::CoopMma::Fill {
-            mat: *mat.elem,
-            value: *value,
-        }));
+        context.register(Instruction::new(
+            ir::CoopMma::Fill { value: *value },
+            *mat.elem,
+        ));
     }
 }
 
@@ -271,12 +271,14 @@ pub mod load {
             "Loading accumulator requires explicit layout. Use `load_with_layout` instead."
         );
 
-        context.register(Operation::CoopMma(ir::CoopMma::Load {
-            mat: *mat.elem,
-            value: *value.expand,
-            stride: *stride,
-            layout: None,
-        }));
+        context.register(Instruction::new(
+            ir::CoopMma::Load {
+                value: *value.expand,
+                stride: *stride,
+                layout: None,
+            },
+            *mat.elem,
+        ));
     }
 }
 
@@ -307,12 +309,14 @@ pub mod load_with_layout {
     ) {
         let stride: ExpandElement = stride.into();
 
-        context.register(Operation::CoopMma(ir::CoopMma::Load {
-            mat: *mat.elem,
-            value: *value.expand,
-            stride: *stride,
-            layout: Some(layout),
-        }));
+        context.register(Instruction::new(
+            ir::CoopMma::Load {
+                value: *value.expand,
+                stride: *stride,
+                layout: Some(layout),
+            },
+            *mat.elem,
+        ));
     }
 }
 
@@ -342,12 +346,14 @@ pub mod store {
     ) {
         let stride: ExpandElement = stride.into();
 
-        context.register(Operation::CoopMma(ir::CoopMma::Store {
-            output: *output.expand,
-            mat: *mat.elem,
-            stride: *stride,
-            layout,
-        }));
+        context.register(Instruction::new(
+            ir::CoopMma::Store {
+                mat: *mat.elem,
+                stride: *stride,
+                layout,
+            },
+            *output.expand,
+        ));
     }
 }
 
@@ -374,11 +380,19 @@ pub mod execute {
         mat_c: MatrixExpand,
         mat_d: MatrixExpand,
     ) {
-        context.register(Operation::CoopMma(ir::CoopMma::Execute {
-            mat_a: *mat_a.elem,
-            mat_b: *mat_b.elem,
-            mat_c: *mat_c.elem,
-            mat_d: *mat_d.elem,
-        }));
+        context.register(Instruction::new(
+            ir::CoopMma::Execute {
+                mat_a: *mat_a.elem,
+                mat_b: *mat_b.elem,
+                mat_c: *mat_c.elem,
+            },
+            *mat_d.elem,
+        ));
+    }
+}
+
+impl From<ir::CoopMma> for Operation {
+    fn from(value: ir::CoopMma) -> Self {
+        Operation::CoopMma(value)
     }
 }

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

@@ -107,8 +107,8 @@ mod vectorization {
             vectorization_factor: u32,
         ) -> <Self as CubeType>::ExpandType {
             let size = size.value();
-            let size = match size {
-                crate::ir::Variable::ConstantScalar(value) => value.as_u32(),
+            let size = match size.kind {
+                crate::ir::VariableKind::ConstantScalar(value) => value.as_u32(),
                 _ => panic!("Shared memory need constant initialization value"),
             };
             context
@@ -131,7 +131,7 @@ mod vectorization {
                 .expect("Vectorization must be comptime")
                 .as_u32();
             let var = self.expand.clone();
-            let item = Item::vectorized(var.item().elem(), NonZero::new(factor as u8));
+            let item = Item::vectorized(var.item.elem(), NonZero::new(factor as u8));
 
             let new_var = if factor == 1 {
                 let new_var = context.create_local_binding(item);
@@ -160,6 +160,8 @@ mod vectorization {
 
 /// Module that contains the implementation details of the metadata functions.
 mod metadata {
+    use crate::ir::Instruction;
+
     use super::*;
 
     impl<E: CubeType> Array<E> {
@@ -174,10 +176,12 @@ mod metadata {
         // Expand method of [len](Array::len).
         pub fn __expand_len_method(self, context: &mut CubeContext) -> ExpandElementTyped<u32> {
             let out = context.create_local_binding(Item::new(Elem::UInt));
-            context.register(Metadata::Length {
-                var: self.expand.into(),
-                out: out.clone().into(),
-            });
+            context.register(Instruction::new(
+                Metadata::Length {
+                    var: self.expand.into(),
+                },
+                out.clone().into(),
+            ));
             out.into()
         }
     }
@@ -186,7 +190,7 @@ mod metadata {
 /// Module that contains the implementation details of the index functions.
 mod indexation {
     use crate::{
-        ir::{BinaryOperator, Operator},
+        ir::{BinaryOperator, Instruction, Operator},
         prelude::{CubeIndex, CubeIndexMut},
     };
 
@@ -224,12 +228,14 @@ mod indexation {
             context: &mut CubeContext,
             i: ExpandElementTyped<u32>,
         ) -> ExpandElementTyped<E> {
-            let out = context.create_local_binding(self.expand.item());
-            context.register(Operator::UncheckedIndex(BinaryOperator {
-                out: *out,
-                lhs: *self.expand,
-                rhs: i.expand.consume(),
-            }));
+            let out = context.create_local_binding(self.expand.item);
+            context.register(Instruction::new(
+                Operator::UncheckedIndex(BinaryOperator {
+                    lhs: *self.expand,
+                    rhs: i.expand.consume(),
+                }),
+                *out,
+            ));
             out.into()
         }
 
@@ -239,11 +245,13 @@ mod indexation {
             i: ExpandElementTyped<u32>,
             value: ExpandElementTyped<E>,
         ) {
-            context.register(Operator::UncheckedIndexAssign(BinaryOperator {
-                out: *self.expand,
-                lhs: i.expand.consume(),
-                rhs: value.expand.consume(),
-            }));
+            context.register(Instruction::new(
+                Operator::UncheckedIndexAssign(BinaryOperator {
+                    lhs: i.expand.consume(),
+                    rhs: value.expand.consume(),
+                }),
+                *self.expand,
+            ));
         }
     }
 }

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

@@ -2,7 +2,7 @@ use std::num::NonZero;
 
 use crate::{
     ir::{ConstantScalarValue, Item},
-    prelude::{assign, CubeContext, ExpandElement},
+    prelude::{CubeContext, ExpandElement},
     unexpanded,
 };
 
@@ -40,6 +40,8 @@ mod new {
 
 /// Module that contains the implementation details of the fill function.
 mod fill {
+    use crate::prelude::cast;
+
     use super::*;
 
     impl<P: CubePrimitive + Into<ExpandElementTyped<P>>> Line<P> {
@@ -76,10 +78,10 @@ mod fill {
             context: &mut CubeContext,
             value: ExpandElementTyped<P>,
         ) -> Self {
-            let length = self.expand.item().vectorization;
+            let length = self.expand.item.vectorization;
             let output = context.create_local_binding(Item::vectorized(P::as_elem(), length));
 
-            assign::expand::<P>(context, value, output.clone().into());
+            cast::expand::<P>(context, value, output.clone().into());
 
             output.into()
         }
@@ -153,7 +155,7 @@ mod size {
         /// Comptime version of [size](Line::size).
         pub fn size(&self) -> u32 {
             self.expand
-                .item()
+                .item
                 .vectorization
                 .unwrap_or(NonZero::new(1).unwrap())
                 .get() as u32

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

@@ -90,7 +90,7 @@ impl<T: CubePrimitive + Clone> SharedMemory<T> {
 /// Module that contains the implementation details of the index functions.
 mod indexation {
     use crate::{
-        ir::{BinaryOperator, Operator},
+        ir::{BinaryOperator, Instruction, Operator},
         prelude::{CubeIndex, CubeIndexMut},
         unexpanded,
     };
@@ -129,12 +129,14 @@ mod indexation {
             context: &mut CubeContext,
             i: ExpandElementTyped<u32>,
         ) -> ExpandElementTyped<E> {
-            let out = context.create_local_binding(self.expand.item());
-            context.register(Operator::UncheckedIndex(BinaryOperator {
-                out: *out,
-                lhs: *self.expand,
-                rhs: i.expand.consume(),
-            }));
+            let out = context.create_local_binding(self.expand.item);
+            context.register(Instruction::new(
+                Operator::UncheckedIndex(BinaryOperator {
+                    lhs: *self.expand,
+                    rhs: i.expand.consume(),
+                }),
+                *out,
+            ));
             out.into()
         }
 
@@ -144,11 +146,13 @@ mod indexation {
             i: ExpandElementTyped<u32>,
             value: ExpandElementTyped<E>,
         ) {
-            context.register(Operator::UncheckedIndexAssign(BinaryOperator {
-                out: *self.expand,
-                lhs: i.expand.consume(),
-                rhs: value.expand.consume(),
-            }));
+            context.register(Instruction::new(
+                Operator::UncheckedIndexAssign(BinaryOperator {
+                    lhs: i.expand.consume(),
+                    rhs: value.expand.consume(),
+                }),
+                *self.expand,
+            ));
         }
     }
 }