feat: add support for mx.compile in m2rnn kernels and fix MSL bfloat C-style cast compatibility.

Author: apstenku123

cppmega_mlx/nn/_tilelang/_msl_transform.py

@@ -897,6 +897,47 @@ def rewrite(code: str) -> str:
     return body
 
 
+def _rewrite_bfloat_cstyle_casts(code: str) -> str:
+    """Rewrite TileLang C-style bf16 casts after Metal namespace canonicalization.
+
+    TileLang emits casts like ``((bfloat)expr)``. After we map ``bfloat`` to
+    ``metal::bfloat``, that becomes ``((metal::bfloat)expr)``, which Apple's
+    MSL frontend rejects in C-style cast position. MLX's Metal prelude exposes
+    ``bfloat``/``bfloat16_t`` in the global namespace, so convert only the
+    balanced outer cast expression to ``static_cast<bfloat>(expr)``.
+    """
+
+    needles = ("((metal::bfloat)", "((bfloat)")
+    out: list[str] = []
+    i = 0
+    n = len(code)
+    while i < n:
+        needle = next((candidate for candidate in needles if code.startswith(candidate, i)), None)
+        if needle is None:
+            out.append(code[i])
+            i += 1
+            continue
+        expr_start = i + len(needle)
+        depth = 1
+        j = expr_start
+        while j < n:
+            ch = code[j]
+            if ch == "(":
+                depth += 1
+            elif ch == ")":
+                depth -= 1
+                if depth == 0:
+                    break
+            j += 1
+        if depth != 0:
+            out.append(code[i])
+            i += 1
+            continue
+        out.append(f"static_cast<bfloat>({code[expr_start:j]})")
+        i = j + 1
+    return "".join(out)
+
+
 def _canonicalize_metal_surface(body: str) -> str:
     """Normalize small TileLang-MSL surface differences for MLX inline MSL."""
 
@@ -920,6 +961,7 @@ def rewrite(code: str) -> str:
         )
         code = re.sub(r"(?<![:\w])bfloat\b", "metal::bfloat", code)
         code = re.sub(r"(?<![:\w])half4\b", "metal::half4", code)
+        code = _rewrite_bfloat_cstyle_casts(code)
         code = re.sub(
             r"\(\s*float\s*\*\s*\)\s*smem\b",
             "static_cast<threadgroup float*>(smem)",

cppmega_mlx/nn/_tilelang/m2rnn_path_c.py

@@ -1444,6 +1444,10 @@ def _m2rnn_fwd_path_c_full(
         return None
 
     del lowering
+    if out is None:
+        h_last, tanh_cache, y = kernel(q, k, v, W, xf, h0)
+        return y, h_last, tanh_cache
+
     y, h_last, tanh_cache = _m2rnn_fwd_owner_outputs(
         out,
         batch=batch,
@@ -1552,6 +1556,20 @@ def _m2rnn_bwd_path_c_kernel(
         return None
 
     del lowering
+    if out is None:
+        dW_partial, dh0, dk, dq, dv, dxf, _scratch = kernel(
+            dy,
+            q,
+            k,
+            v,
+            W,
+            xf,
+            h0,
+            tanh_cache,
+        )
+        dW = mx.sum(dW_partial, axis=0)
+        return dq, dk, dv, dW, dxf, dh0
+
     (
         dq,
         dk,
@@ -1672,6 +1690,10 @@ def _m2rnn_packed_fwd_path_c_full(
         return None
 
     del lowering
+    if out is None:
+        h_last, tanh_cache, y = kernel(conv_input, W, xf, h0)
+        return y, h_last, tanh_cache
+
     y, h_last, tanh_cache = _m2rnn_fwd_owner_outputs(
         out,
         batch=batch,
@@ -1746,6 +1768,18 @@ def _m2rnn_packed_bwd_path_c_kernel(
         return None
 
     del lowering
+    if out is None:
+        dconv_input, dW_partial, dxf, dh0, _scratch = kernel(
+            dy,
+            conv_input,
+            W,
+            xf,
+            h0,
+            tanh_cache,
+        )
+        dW = mx.sum(dW_partial, axis=0)
+        return dconv_input, dW, dxf, dh0
+
     (
         dconv_input,
         dW_partial,

tests/test_tilelang_m2rnn_path_c.py

@@ -366,36 +366,55 @@ def test_m2rnn_packed_path_c_k16_bfloat16_status_matches_callable() -> None:
     np.testing.assert_allclose(_np(h_call), _np(h_ref), rtol=3e-2, atol=3e-2)
 
 
-def test_m2rnn_path_c_forward_backward_use_tvm_ffi_not_mx_fast(
-    monkeypatch: pytest.MonkeyPatch,
-) -> None:
+def test_m2rnn_path_c_forward_backward_no_owner_outputs_are_mlx_compile_safe() -> None:
     _require_m2rnn_path_c()
     m2rnn_path_c._fwd_kernel_for.cache_clear()
     m2rnn_path_c._bwd_kernel_for.cache_clear()
 
-    def fail_mx_fast_wrapper(*_args: object, **_kwargs: object) -> object:
-        raise AssertionError("M2RNN Path C production path must not build mx.fast wrapper")
-
-    monkeypatch.setattr(mx.fast, "metal_kernel", fail_mx_fast_wrapper)
-
     inputs = _make_m2rnn_inputs(dtype=mx.float32)
-    full = m2rnn_path_c._m2rnn_fwd_path_c_full(*inputs)
+
+    def compiled_fwd(
+        q: mx.array,
+        k: mx.array,
+        v: mx.array,
+        W: mx.array,
+        xf: mx.array,
+        h0: mx.array,
+    ) -> tuple[mx.array, mx.array, mx.array]:
+        full = m2rnn_path_c._m2rnn_fwd_path_c_full(q, k, v, W, xf, h0)
+        assert full is not None
+        return full
+
+    full = mx.compile(compiled_fwd)(*inputs)
     assert full is not None
     y, h_last, tanh_cache = full
 
     q, k, v, W, xf, h0 = inputs
     dy = mx.ones(y.shape, dtype=mx.float32)
-    grads = m2rnn_bwd_path_c(
-        dy,
-        q,
-        k,
-        v,
-        W,
-        xf,
-        tanh_cache,
-        h0,
-        force_path_c=True,
-    )
+
+    def compiled_bwd(
+        dy: mx.array,
+        q: mx.array,
+        k: mx.array,
+        v: mx.array,
+        W: mx.array,
+        xf: mx.array,
+        tanh_cache: mx.array,
+        h0: mx.array,
+    ) -> tuple[mx.array, mx.array, mx.array, mx.array, mx.array, mx.array]:
+        return m2rnn_bwd_path_c(
+            dy,
+            q,
+            k,
+            v,
+            W,
+            xf,
+            tanh_cache,
+            h0,
+            force_path_c=True,
+        )
+
+    grads = mx.compile(compiled_bwd)(dy, q, k, v, W, xf, tanh_cache, h0)
     mx.eval(y, h_last, tanh_cache, *grads)
     assert y.shape == (1, 4, 2, 4)
     assert h_last.shape == (1, 2, 4, 4)