feat(cutile): support scalar specialization overrides

cutile-macro/src/_module.rs

@@ -830,6 +830,10 @@ pub fn kernel_launcher(
             _grid: (u32, u32, u32),
             input: Option<DI>,
             function_generics: Option<Vec<String>>,
+            _scalar_hint_overrides: Vec<(
+                String,
+                cutile_compiler::specialization::DivHint,
+            )>,
             _phantom: std::marker::PhantomData<( #(#ki_phantom_types,)* )>,
             _compile_options: CompileOptions,
             // When true, `execute` skips its launch block (set by `.compile()`).
@@ -843,6 +847,7 @@ pub fn kernel_launcher(
                     _grid: (0, 0, 0),
                     input: Some(input),
                     function_generics: None,
+                    _scalar_hint_overrides: Vec::new(),
                     _phantom: std::marker::PhantomData,
                     _compile_options: CompileOptions::default(),
                     _compile_only: false,
@@ -916,6 +921,14 @@ pub fn kernel_launcher(
                 self.function_generics = Some(generics);
                 self
             }
+            fn scalar_hint(
+                mut self,
+                name: impl Into<String>,
+                hint: cutile_compiler::specialization::DivHint,
+            ) -> Self {
+                self._scalar_hint_overrides.push((name.into(), hint));
+                self
+            }
             fn compile_options(mut self, options: CompileOptions) -> Self {
                 self._compile_options = options;
                 self

cutile-macro/src/kernel_launcher_generator.rs

@@ -877,12 +877,35 @@ pub fn generate_kernel_launcher(
 
     // Emit scalar_hints (populated for integer scalar and raw pointer params).
     let scalar_hints_stmt = parse_stmt(format!(
-        "let scalar_hints: Vec<(String, cutile_compiler::specialization::DivHint)> = vec![{}];",
+        "let mut scalar_hints: Vec<(String, cutile_compiler::specialization::DivHint)> = vec![{}];",
         scalar_hint_exprs.join(",")
     ));
     launcher_method.block.stmts.push(scalar_hints_stmt.clone());
     specialization_method.block.stmts.push(scalar_hints_stmt);
 
+    // Scalar specialization overrides replace or extend inferred hints on both
+    // the launch and the specialization path.
+    let scalar_override_stmt = parse_stmt(
+        r#"
+        for (name, hint) in self._scalar_hint_overrides.drain(..) {
+            if let Some((_, inferred)) = scalar_hints
+                .iter_mut()
+                .find(|entry| entry.0 == name)
+            {
+                *inferred = hint;
+            } else {
+                scalar_hints.push((name, hint));
+            }
+        }
+        "#
+        .to_string(),
+    );
+    launcher_method
+        .block
+        .stmts
+        .push(scalar_override_stmt.clone());
+    specialization_method.block.stmts.push(scalar_override_stmt);
+
     let specialization_stmts = syn::parse2::<ExprBlock>(quote! {{
         let const_grid = if self._const_grid { Some(self._grid) } else { None };
         let compile_options = std::mem::take(&mut self._compile_options);

cutile/src/tile_kernel.rs

@@ -906,6 +906,10 @@ where
     }
     /// Sets the type and const generic arguments for this kernel.
     fn generics(self, generics: Vec<String>) -> Self;
+    /// Overrides the inferred divisibility hint for one scalar or raw-pointer
+    /// parameter. This is useful for truly dynamic scalars whose incidental
+    /// runtime divisibility should not create separate JIT specializations.
+    fn scalar_hint(self, name: impl Into<String>, hint: DivHint) -> Self;
     /// Sets a compile-time constant grid, enabling grid-dependent optimizations.
     fn const_grid(self, grid: (u32, u32, u32)) -> Self;
     /// Sets the runtime launch grid dimensions.

cutile/tests/specialization_bits.rs

@@ -435,7 +435,57 @@ fn raw_pointer_launch_scalar_div_hint_covers_powers_of_two_through_16() {
     });
 }
 
+#[test]
+fn raw_pointer_launch_scalar_hint_override_replaces_inferred_hint() {
+    common::with_test_stack(|| {
+        // n=12 would infer div_by<4>; the override to div_by<8> must win.
+        let mlir = launch_raw_ptr_scalar_kernel_with_overrides(
+            12,
+            &[("_n", dh(8))],
+            CompileOptions::default().occupancy(3),
+        );
+        assert!(
+            mlir.contains("assume div_by<8>"),
+            "Expected scalar hint override to emit div_by<8>.\nMLIR:\n{mlir}"
+        );
+        assert!(
+            !mlir.contains("assume div_by<4>"),
+            "Inferred div_by<4> must be replaced by the override.\nMLIR:\n{mlir}"
+        );
+    });
+}
+
+#[test]
+fn raw_pointer_launch_scalar_hint_default_override_removes_inferred_hint() {
+    common::with_test_stack(|| {
+        // A default DivHint (divisor 1) disables specialization for a truly
+        // dynamic scalar: only the raw pointer assume remains.
+        let mlir = launch_raw_ptr_scalar_kernel_with_overrides(
+            12,
+            &[("_n", DivHint::default())],
+            CompileOptions::default().occupancy(3),
+        );
+        assert!(
+            mlir.contains("assume div_by<16>"),
+            "Expected the raw pointer hint to still emit div_by<16>.\nMLIR:\n{mlir}"
+        );
+        assert_eq!(
+            mlir.matches("assume div_by<").count(),
+            1,
+            "Expected only the raw pointer assume once the scalar hint is overridden to default.\nMLIR:\n{mlir}"
+        );
+    });
+}
+
 fn launch_raw_ptr_scalar_kernel_and_read_mlir(n: i32, options: CompileOptions) -> String {
+    launch_raw_ptr_scalar_kernel_with_overrides(n, &[], options)
+}
+
+fn launch_raw_ptr_scalar_kernel_with_overrides(
+    n: i32,
+    overrides: &[(&str, DivHint)],
+    options: CompileOptions,
+) -> String {
     let _lock = RAW_PTR_DUMP_LOCK.lock().expect("lock raw pointer dump dir");
     let dump_dir = Path::new(RAW_PTR_SCALAR_DUMP_DIR);
     let _ = std::fs::remove_dir_all(dump_dir);
@@ -446,11 +496,13 @@ fn launch_raw_ptr_scalar_kernel_and_read_mlir(n: i32, options: CompileOptions) -
         .expect("alloc backing tensor");
     let ptr = backing.device_pointer();
 
-    unsafe { raw_ptr_scalar_kernel(ptr, n) }
+    let launcher = unsafe { raw_ptr_scalar_kernel(ptr, n) }
         .grid((1, 1, 1))
-        .compile_options(options)
-        .sync()
-        .expect("raw pointer scalar kernel launch");
+        .compile_options(options);
+    let launcher = overrides.iter().fold(launcher, |launcher, (name, hint)| {
+        launcher.scalar_hint(*name, *hint)
+    });
+    launcher.sync().expect("raw pointer scalar kernel launch");
 
     let mut mlir = String::new();
     for entry in std::fs::read_dir(dump_dir).expect("read MLIR dump dir") {