refactor: optimize thread indexing using metal intrinsics and simplify sigmoid activation implementation

Author: apstenku123

cppmega_mlx/nn/_tilelang/mamba3_path_c.py

@@ -981,9 +981,8 @@ def snapshots(
         h0: T.Tensor((BATCH, HEADS, HEADDIM, STATE), h0_dtype),
         h_snap: T.Tensor((BATCH, BLOCKS + 1, HEADS, HEADDIM, STATE), "float32"),
     ):
-        with T.Kernel(T.ceildiv(LANES, THREADS), threads=THREADS) as bx:
-            tid = T.get_thread_binding(0)
-            global_lane = bx * THREADS + tid
+        with T.Kernel(T.ceildiv(LANES, THREADS), threads=THREADS) as _bx:
+            global_lane = T.call_intrin("int32", "tir.metal.thread_position_in_grid_x")
             h_state = T.alloc_local((STATE,), accum_dtype)
             if global_lane < LANES:
                 p = global_lane % HEADDIM
@@ -1120,9 +1119,8 @@ def bwd_simd(
         dD_batch: T.Tensor((BATCH, HEADS), dD_dtype),
         dh0: T.Tensor((BATCH, HEADS, HEADDIM, STATE), dh0_dtype),
     ):
-        with T.Kernel(T.ceildiv(LANES, THREADS), threads=THREADS) as bx:
-            tid = T.get_thread_binding(0)
-            global_lane = bx * THREADS + tid
+        with T.Kernel(T.ceildiv(LANES, THREADS), threads=THREADS) as _bx:
+            global_lane = T.call_intrin("int32", "tir.metal.thread_position_in_grid_x")
             h_state = T.alloc_local((STATE,), accum_dtype)
             dh = T.alloc_local((STATE,), accum_dtype)
             if global_lane < LANES:
@@ -1161,12 +1159,7 @@ def bwd_simd(
                         y_state += h_state[n] * T.cast(C[b, t, h, n], accum_dtype)
                     D_h = T.cast(D[h], accum_dtype)
                     y_skipped = y_state + D_h * x_val
-                    sig_z = T.alloc_var(T.float32, init=0.0)
-                    if z_val >= 0.0:
-                        sig_z = 1.0 / (1.0 + T.exp(-z_val))
-                    else:
-                        sig_z = T.exp(z_val)
-                        sig_z = sig_z / (1.0 + sig_z)
+                    sig_z = 1.0 / (1.0 + T.exp(-z_val))
                     silu_z = z_val * sig_z
                     silu_dz = sig_z * (1.0 + z_val * (1.0 - sig_z))
 
@@ -1329,9 +1322,8 @@ def bwd_snap_simd(
         dD_batch: T.Tensor((BATCH, HEADS), dD_dtype),
         dh0: T.Tensor((BATCH, HEADS, HEADDIM, STATE), dh0_dtype),
     ):
-        with T.Kernel(T.ceildiv(LANES, THREADS), threads=THREADS) as bx:
-            tid = T.get_thread_binding(0)
-            global_lane = bx * THREADS + tid
+        with T.Kernel(T.ceildiv(LANES, THREADS), threads=THREADS) as _bx:
+            global_lane = T.call_intrin("int32", "tir.metal.thread_position_in_grid_x")
             h_state = T.alloc_local((STATE,), accum_dtype)
             dh = T.alloc_local((STATE,), accum_dtype)
             if global_lane < LANES:
@@ -1368,12 +1360,7 @@ def bwd_snap_simd(
                                     accum_dtype,
                                 )
                             y_skipped = y_state + D_h * x_val
-                            sig_z = T.alloc_var(T.float32, init=0.0)
-                            if z_val >= 0.0:
-                                sig_z = 1.0 / (1.0 + T.exp(-z_val))
-                            else:
-                                sig_z = T.exp(z_val)
-                                sig_z = sig_z / (1.0 + sig_z)
+                            sig_z = 1.0 / (1.0 + T.exp(-z_val))
                             silu_z = z_val * sig_z
                             silu_dz = sig_z * (1.0 + z_val * (1.0 - sig_z))
 
@@ -1515,9 +1502,8 @@ def bwd_partial(
         dD_partial: T.Tensor((BATCH, HEADS, HEADDIM), dD_dtype),
         dh0: T.Tensor((BATCH, HEADS, HEADDIM, STATE), dh0_dtype),
     ):
-        with T.Kernel(T.ceildiv(LANES, THREADS), threads=THREADS) as bx:
-            tid = T.get_thread_binding(0)
-            global_lane = bx * THREADS + tid
+        with T.Kernel(T.ceildiv(LANES, THREADS), threads=THREADS) as _bx:
+            global_lane = T.call_intrin("int32", "tir.metal.thread_position_in_grid_x")
             h_state = T.alloc_local((STATE,), accum_dtype)
             dh = T.alloc_local((STATE,), accum_dtype)
             if global_lane < LANES:
@@ -1544,12 +1530,7 @@ def bwd_partial(
                         h_state[n] = h_snap[b, t + 1, h, p, n]
                         y_state += h_state[n] * T.cast(C[b, t, h, n], accum_dtype)
                     y_skipped = y_state + D_h * x_val
-                    sig_z = T.alloc_var(T.float32, init=0.0)
-                    if z_val >= 0.0:
-                        sig_z = 1.0 / (1.0 + T.exp(-z_val))
-                    else:
-                        sig_z = T.exp(z_val)
-                        sig_z = sig_z / (1.0 + sig_z)
+                    sig_z = 1.0 / (1.0 + T.exp(-z_val))
                     silu_z = z_val * sig_z
                     silu_dz = sig_z * (1.0 + z_val * (1.0 - sig_z))
 

tests/test_tilelang_mamba3_path_c.py

@@ -361,13 +361,16 @@ def test_lowered_msl_reuses_hot_scalar_temporaries() -> None:
     assert re.search(r"z_val \* \([^;\n]+exp\(", fwd) is None
 
     bwd = dump_lowered_bwd_msl(batch=1, seq=4, heads=1, headdim=2, state=4)
+    assert "gridThreadIdx [[thread_position_in_grid]]" in bwd
+    assert "blockIdx.x) * 256) + (((int)threadIdx.x)" not in bwd
     assert len(re.findall(r"float decay = exp\(", bwd)) == 1
     assert "h_snap" in bwd
     assert "1.000000e+00 / decay" not in bwd
     assert re.search(r"d_decay\[0\] \* exp\(", bwd) is None
     assert re.search(r"dh\[n_\d+\] = \(dh\[n_\d+\] \* exp\(", bwd) is None
-    assert len(re.findall(r"sig_z = .*exp\(", bwd)) == 2
-    assert "sig_z = exp(z_val)" in bwd
+    assert "if (0.000000e+00f <= z_val)" not in bwd
+    assert "sig_z = exp(z_val)" not in bwd
+    assert "1.000000e+00f + exp(" in bwd
     assert re.search(r"dY \* \(z_val \* \([^;\n]+exp\(", bwd) is None
     assert "y_state = sig_z" not in bwd
     assert "dx_inp = sig_z" not in bwd