Simplify allocator #388

maxtremblay · 2024-12-20T16:42:56Z

Highlights

The allocator implementation is simpler.
The benchmark results are similar.
I discovered and fixed a few bugs.

Allocator refactor

Remove the allocator trait and its two implementations in favor of a single allocator struct.
Rename the local variable creations to new named closer to rust naming scheme (example, use mut for mutable variables).
Rename the variant in VariableKind to match the changes.
Propagate the changes. This explains the large amount of modified files. A lot of it is renaming and changing to the new allocator.
Document to allocator struct.

Changing the backends

Rename the local variable variant of the Variable enum to match cubecl-core.
Add the depth field to LocalVariable.
All local immutable variables are now named l_{depth}_{id}.
All local mutable variables (including restricted) are now named l_mut_{depth}_{id}.

Allocation strategy

In a few expand functions, I replaced the use of a mutable variable with an immutable one.

Discovered bugs

With all the above changes, I discovered a few bugs in cubecl that I fixed.

Because the reusing allocator (used by cuda) was not doing any distinction between mutable and immutable variables, we often used immutable variables when we should have used mutable one. Mostly for accumulators in plane operations.
Fix an issue with line indexing where elements of the line where tagged with a vectorization > 1 while there are scalar. Before, we were lucky enough for this bug to have no effect (wgpu: inference let us avoid the type during declaration. cpp: we were using mutables everywhere before which were not causing an error, but were wastefully declared as vectorized).
There were some missing fmt_left in the cpp backend.
I added some tests to reproduce the bugs more easily.

Benchmarks

There are no statistically meaningful changes.

bench_cuda_main.txt
bench_cuda_new_alloc.txt
bench_wgpu_main.txt
bench_wgpu_new_alloc.txt

nathanielsimard

I'm not sure if the frontend tests are necessary anymore. They are impacted by the new allocator and I don't think it's worth updating them.

nathanielsimard · 2024-12-20T17:53:23Z

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

@@ -206,7 +206,7 @@ pub mod plane_any {
        elem: ExpandElementTyped<bool>,
    ) -> ExpandElementTyped<bool> {
        let elem: ExpandElement = elem.into();
-        let output = context.create_local_binding(elem.item);
+        let output = context.create_local_mut(elem.item);


Is it only for cuda? Does mutability work differently in wgsl?

In fact, it is an issue with the planeAny implementation of cuda. I will fix it.

crates/cubecl-core/src/ir/allocator.rs

nathanielsimard · 2024-12-20T17:56:03Z

crates/cubecl-core/src/ir/operation.rs

+    let array_len = scope.create_local_mut(Item::new(Elem::UInt(UIntKind::U32)));
+    let inside_bound = scope.create_local_mut(Item::new(Elem::Bool));
+    let item = scope.create_local_mut(out.item);


Is mut necessary?

nathanielsimard · 2024-12-20T17:56:10Z

crates/cubecl-core/src/ir/operation.rs

+    let array_len = scope.create_local_mut(Item::new(Elem::UInt(UIntKind::U32)));
+    let inside_bound = scope.create_local_mut(Item::new(Elem::Bool));


Is mut necessary?

nathanielsimard · 2024-12-20T17:56:32Z

crates/cubecl-core/src/ir/scope.rs

@@ -67,7 +67,7 @@ impl Scope {

    /// Create a variable initialized at zero.
    pub fn zero<I: Into<Item>>(&mut self, item: I) -> Variable {
-        let local = self.create_local(item);
+        let local = self.create_local_mut(item);


Is mut necessary?

nathanielsimard · 2024-12-20T17:56:48Z

crates/cubecl-core/src/ir/scope.rs

@@ -88,7 +88,7 @@ impl Scope {
            Elem::AtomicUInt(kind) => ConstantScalarValue::UInt(value.to_u64().unwrap(), kind),
            Elem::Bool => ConstantScalarValue::Bool(value.to_u32().unwrap() == 1),
        };
-        let local = self.create_local(item);
+        let local = self.create_local_mut(item);


Is mut necessary?

nathanielsimard · 2024-12-21T21:00:15Z

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

@@ -208,7 +208,7 @@ macro_rules! impl_line_comparison {
                        let lhs = self.expand.into();
                        let rhs = rhs.expand.into();

-                        let output = context.create_local_binding(Item::vectorized(bool::as_elem(), size));
+                        let output = context.create_local_mut(Item::vectorized(bool::as_elem(), size));


I don't think this needs to be mut?

nathanielsimard · 2024-12-21T21:00:55Z

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

@@ -79,7 +79,7 @@ mod fill {
            value: ExpandElementTyped<P>,
        ) -> Self {
            let length = self.expand.item.vectorization;
-            let output = context.create_local_binding(Item::vectorized(P::as_elem(), length));
+            let output = context.create_local(Item::vectorized(P::as_elem(), length));


Is mut required?

nathanielsimard · 2024-12-21T21:01:30Z

crates/cubecl-core/src/frontend/operation/base.rs

@@ -170,7 +170,7 @@ where
    F: Fn(UnaryOperator) -> Operator,
 {
    let input = input.consume();
-    let output = context.create_local_binding(out_item);
+    let output = context.create_local_mut(out_item);


Is mut required?

nathanielsimard · 2024-12-21T21:03:12Z

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

@@ -694,11 +694,12 @@ impl<D: Dialect> Remainder<D> {

            write_op(&lhs, &rhs, &out_tmp, item_out_optimized)?;

+            let maybe_const = if out.is_const() { " cosnt" } else { "" };


typo: cosnt

nathanielsimard · 2024-12-21T21:04:36Z

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

+
+fn cast<D: Dialect>(input: &Variable<D>, target: Item<D>) -> String {
+    if target != input.item() {
+        let maybe_const = if input.is_const() { " const" } else { "" };


We could have a method for that, it's repeated at multiple places.

nathanielsimard · 2024-12-21T21:05:45Z