Fix/autotuner (#1062)

* Improve error messages when no valid kernel is found

* Fix autotuner

* Clippy

* Cleanup

Author: nathanielsimard

crates/cubecl-common/src/format.rs

@@ -0,0 +1,119 @@
+use alloc::format;
+use alloc::string::String;
+
+/// Format strings for use in identifiers and types.
+pub fn format_str(string: &str, markers: &[(char, char)], include_space: bool) -> String {
+    let mut result = String::new();
+    let mut depth = 0;
+    let indentation = 4;
+
+    let mut prev = ' ';
+    let mut in_string = false;
+
+    for c in string.chars() {
+        if c == ' ' {
+            if in_string {
+                result.push(c);
+            }
+
+            continue;
+        }
+        if c == '"' {
+            in_string = !in_string;
+        }
+
+        let mut found_marker = false;
+
+        for (start, end) in markers {
+            let (start, end) = (*start, *end);
+
+            if c == start {
+                depth += 1;
+                if prev != ' ' && include_space {
+                    result.push(' ');
+                }
+                result.push(start);
+                result.push('\n');
+                result.push_str(&" ".repeat(indentation * depth));
+                found_marker = true;
+            } else if c == end {
+                depth -= 1;
+                if prev != start {
+                    if prev == ' ' {
+                        result.pop();
+                    }
+                    result.push_str(",\n");
+                    result.push_str(&" ".repeat(indentation * depth));
+                    result.push(end);
+                } else {
+                    for _ in 0..(&" ".repeat(indentation * depth).len()) + 1 + indentation {
+                        result.pop();
+                    }
+                    result.push(end);
+                }
+                found_marker = true;
+            }
+        }
+
+        if found_marker {
+            prev = c;
+            continue;
+        }
+
+        if c == ',' && depth > 0 {
+            if prev == ' ' {
+                result.pop();
+            }
+
+            result.push_str(",\n");
+            result.push_str(&" ".repeat(indentation * depth));
+            continue;
+        }
+
+        if c == ':' && include_space {
+            result.push(c);
+            result.push(' ');
+            prev = ' ';
+        } else {
+            result.push(c);
+            prev = c;
+        }
+    }
+
+    result
+}
+
+/// Format a debug type.
+pub fn format_debug<F: core::fmt::Debug>(string: &F) -> String {
+    let string = format!("{string:?}");
+    format_str(&string, &[('(', ')'), ('[', ']'), ('{', '}')], true)
+}
+
+#[cfg(test)]
+mod tests {
+    use hashbrown::HashMap;
+
+    use super::*;
+
+    #[derive(Debug)]
+    #[allow(unused)]
+    struct Test {
+        map: HashMap<String, u32>,
+    }
+
+    #[test]
+    fn test_format_debug() {
+        let test = Test {
+            map: HashMap::from_iter([("Hey with space".to_string(), 8)].into_iter()),
+        };
+
+        let formatted = format_debug(&test);
+        let expected = r#"Test {
+    map: {
+        "Hey with space": 8,
+    },
+}"#;
+
+        assert_eq!(expected, formatted);
+    }
+}

crates/cubecl-common/src/lib.rs

@@ -51,6 +51,9 @@ pub mod future;
 /// Quantization primitives required outside of `cubecl-quant`
 pub mod quant;
 
+/// Format utilities.
+pub mod format;
+
 /// Various utilities to create ID's.
 extern crate alloc;
 

crates/cubecl-core/src/compute/kernel.rs

@@ -5,11 +5,11 @@ use std::{
 };
 
 use crate::{Compiler, KernelOptions};
-use cubecl_common::{CubeDim, ExecutionMode};
+use cubecl_common::{CubeDim, ExecutionMode, format::format_str};
 use cubecl_ir::{Id, Scope, StorageType, Type};
 use cubecl_runtime::{
     config::{GlobalConfig, compilation::CompilationLogLevel},
-    id::{KernelId, format_str},
+    id::KernelId,
     kernel::KernelMetadata,
 };
 use serde::{Deserialize, Serialize};

crates/cubecl-runtime/src/id.rs

@@ -1,13 +1,13 @@
 use alloc::format;
 use alloc::string::String;
-use alloc::string::ToString;
 use alloc::sync::Arc;
 use core::{
     any::{Any, TypeId},
     fmt::Display,
     hash::{BuildHasher, Hash, Hasher},
 };
 use cubecl_common::ExecutionMode;
+use cubecl_common::format::format_str;
 
 #[macro_export(local_inner_macros)]
 /// Create a new storage ID type.
@@ -208,81 +208,6 @@ impl PartialEq for KernelId {
 
 impl Eq for KernelId {}
 
-/// Format strings for use in identifiers and types.
-pub fn format_str(kernel_id: &str, markers: &[(char, char)], include_space: bool) -> String {
-    let kernel_id = kernel_id.to_string();
-    let mut result = String::new();
-    let mut depth = 0;
-    let indentation = 4;
-
-    let mut prev = ' ';
-
-    for c in kernel_id.chars() {
-        if c == ' ' {
-            continue;
-        }
-
-        let mut found_marker = false;
-
-        for (start, end) in markers {
-            let (start, end) = (*start, *end);
-
-            if c == start {
-                depth += 1;
-                if prev != ' ' && include_space {
-                    result.push(' ');
-                }
-                result.push(start);
-                result.push('\n');
-                result.push_str(&" ".repeat(indentation * depth));
-                found_marker = true;
-            } else if c == end {
-                depth -= 1;
-                if prev != start {
-                    if prev == ' ' {
-                        result.pop();
-                    }
-                    result.push_str(",\n");
-                    result.push_str(&" ".repeat(indentation * depth));
-                    result.push(end);
-                } else {
-                    for _ in 0..(&" ".repeat(indentation * depth).len()) + 1 + indentation {
-                        result.pop();
-                    }
-                    result.push(end);
-                }
-                found_marker = true;
-            }
-        }
-
-        if found_marker {
-            prev = c;
-            continue;
-        }
-
-        if c == ',' && depth > 0 {
-            if prev == ' ' {
-                result.pop();
-            }
-
-            result.push_str(",\n");
-            result.push_str(&" ".repeat(indentation * depth));
-            continue;
-        }
-
-        if c == ':' && include_space {
-            result.push(c);
-            result.push(' ');
-            prev = ' ';
-        } else {
-            result.push(c);
-            prev = c;
-        }
-    }
-
-    result
-}
-
 impl Display for KernelId {
     fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
         match &self.info {

crates/cubecl-runtime/src/tune/base.rs

@@ -48,6 +48,12 @@ pub struct TuneGroup<K> {
     pub(crate) priority: PriorityFunc<K>,
 }
 
+impl<K> core::fmt::Debug for TuneGroup<K> {
+    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
+        f.debug_struct("TuneGroup").field("id", &self.id).finish()
+    }
+}
+
 impl<K> Clone for TuneGroup<K> {
     fn clone(&self) -> Self {
         Self {
@@ -83,9 +89,12 @@ struct GroupPlan {
     indices: HashMap<i8, Vec<usize>>,
 }
 
+#[derive(Debug)]
 struct Cleanup {
     groups: Vec<i8>,
     tunables: Vec<(i8, i8)>,
+    /// Within group priority is too low to even try.
+    skipped: bool,
 }
 
 impl TunePlan {
@@ -154,13 +163,22 @@ impl TunePlan {
         };
 
         let (mut group_indices, cleanup) = self.group_plan_next(priority);
+        // Some entries are skipped for this round of prioritizing.
+        let skipped = cleanup.skipped || priority < 0;
+
         self.cleanup(cleanup);
 
         if priority >= 0 {
             indices.append(&mut group_indices);
         }
 
-        indices
+        // The indices list is empty, but it doesn't mean we should stop
+        // autotuning, since some entries were skipped.
+        if indices.is_empty() && skipped {
+            self.next()
+        } else {
+            indices
+        }
     }
 
     fn cleanup(&mut self, cleanup: Cleanup) {
@@ -231,6 +249,7 @@ impl TunePlan {
             Cleanup {
                 groups: cleanup_groups,
                 tunables: cleanup_tunables,
+                skipped: within_group_prio < 0,
             },
         )
     }
@@ -320,6 +339,24 @@ mod tests {
         assert!(plan.next().is_empty());
     }
 
+    #[test]
+    fn test_plan_negative_priority() {
+        let group0 = TuneGroup::<FakeAutotuneKey>::new(|_| 2);
+        let group1 = TuneGroup::<FakeAutotuneKey>::new(|_| 1);
+
+        let tunable0 = Tunable::<FakeAutotuneKey, (), ()>::new(fake_kernel);
+        let tunable1 = Tunable::<FakeAutotuneKey, (), ()>::new(fake_kernel).group(&group0, |_| -1);
+        let tunable2 = Tunable::<FakeAutotuneKey, (), ()>::new(fake_kernel).group(&group0, |_| 2);
+        let tunable3 = Tunable::<FakeAutotuneKey, (), ()>::new(fake_kernel).group(&group1, |_| 2);
+
+        let key = FakeAutotuneKey;
+        let mut plan = TunePlan::new(&key, &[tunable0, tunable1, tunable2, tunable3]);
+
+        assert_eq!(plan.next(), vec![0, 2]);
+        assert_eq!(plan.next(), vec![3]);
+        assert!(plan.next().is_empty());
+    }
+
     #[test]
     fn test_plan_no_group() {
         let tunable0 = Tunable::<FakeAutotuneKey, (), ()>::new(fake_kernel);