[metal] Add MSL AST walker for Python-to-C++ translation

[aditvenk]

Stacked PRs:

• [metal] Auto-cap threadgroup size via CuteND/CuteFlattenedTileStrategy #1855
• Add backend-agnostic lane loop APIs to tile strategies #1798
• Add Metal test job to CI test matrix #1862
• [metal] Wire @metal_jit into MetalBackend and simplify launcher #1992
• [metal] Add @metal_jit decorator for AST-to-MSL compilation #1991
• ->[metal] Add MSL AST walker for Python-to-C++ translation #1794

---

[metal] Add MSL AST walker for Python-to-C++ translation

Add msl_ast_walker.py which translates Python AST to MSL C++ source.
Handles statement-level translation (assignments, if/for, etc.),
tl.load/tl.store → pointer dereferences, and C++ namespace restoration
(metal.precise.sin → metal::precise::sin).

This is a standalone library module — not yet wired into the backend.

[aditvenk]

Example:

@helion.kernel(backend="metal", configs=[helion.Config(block_sizes=[256], num_warps=4)])
  def masked_copy(x: torch.Tensor, mask: torch.Tensor) -> torch.Tensor:
      out = torch.zeros_like(x)
      for tile in hl.tile(x.size(0)):
          out[tile] = torch.where(mask[tile], x[tile], 0.0)
      return out

translated to

 kernel void _helion_masked_copy(device bool* mask [[buffer(0)]], device float* x [[buffer(1)]], device float* out [[buffer(2)]], uint3 tgid [[threadgroup_position_in_grid]], uint3 tid [[thread_position_in_threadgroup]]) {
      constexpr int _BLOCK_SIZE_0 = 256;
      auto pid_0 = tgid[0];
      auto offset_0 = (pid_0 * _BLOCK_SIZE_0);
      auto indices_0 = (offset_0 + tid[0]);
      auto mask_0 = (indices_0 < 1000);
      auto load = (mask_0 ? *((mask + (indices_0 * 1))) : (0));
      auto load_1 = (mask_0 ? *((x + (indices_0 * 1))) : (0));
      auto v_0 = ((float)(0.0));
      auto v_1 = select(v_0, load_1, load);
      if (mask_0) {
          *((out + (indices_0 * 1))) = v_1;
      }
  }

[jansel]

Hrm, so this design still doesn't seem ideal to me. Some alternate ideas:

1. We could extend Helion to be able to codegen C++ directly and not force things to go via a fake-python-AST. This seems the cleanest, though I would require more work.

2. Alternately, we could split the Python -> C++ translation into its own separate layer -- so you could do something like:

@metal_python_kernel
def my_metal_kernel(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
   ...

The @metal_python_kernel decorator would handle:

1. Translating the fake AST into C++
2. Compiling that C++
3. Generating the launcher wrapper calls C++ from Python

This might be useful as its own standalone thing. I'm sure there are people out there who would rather write metal using a Python interface.

It will also be easier to debug, since less is happening in a single step.

[aditvenk]

Hrm, so this design still doesn't seem ideal to me. Some alternate ideas:

1. We could extend Helion to be able to codegen C++ directly and not force things to go via a fake-python-AST. This seems the cleanest, though I would require more work.
2. Alternately, we could split the Python -> C++ translation into its own separate layer -- so you could do something like:

@metal_python_kernel
def my_metal_kernel(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
   ...

The @metal_python_kernel decorator would handle:

1. Translating the fake AST into C++
2. Compiling that C++
3. Generating the launcher wrapper calls C++ from Python

This might be useful as its own standalone thing. I'm sure there are people out there who would rather write metal using a Python interface.

It will also be easier to debug, since less is happening in a single step.

Thanks, this is an interesting idea. Let me play around a bit

[aditvenk]

@jansel -- I prototyped Option 2. Helion AST is lowered (via MetalOverrides and MetalBackend) into a fake Metal Python AST + a decorator. The decorator does the last mile AST walking to convert to C++, compile MSL, and generate the launcher call.

Here's an example of the fake Metal Python

  @metal_python_kernel
  def _helion_silu(x, out):
      pid_flat = tgid[0]
      offsets_0 = pid_flat * _BLOCK_SIZE_0 + tid[0]
      indices_0 = offsets_0
      mask_0 = tid[0] < _BLOCK_SIZE_0 and offsets_0 < 1024
      load = tl.load(x + indices_0 * 1, mask_0, other=0)
      v_0 = static_cast < decltype(load) > -load
      v_1 = metal.precise.exp(v_0)
      v_2 = 1
      v_3 = v_1 + v_2
      v_4 = load / v_3
      tl.store(out + indices_0 * 1, v_4, mask_0)

This is closer to "pseudo-code" and can help with understanding/debugging, but I will not consider this as a useful way to write Metal kernel from Python. (PS: The tl.load/tl.store comes from PointerIndexingStrategy which currently hardcodes tl.load/store)

The decorator internally uses the AST walker in this PR to lower to pure C++.

Wdyt about this approach?

[jansel]

@aditvenk let's go with that option, that seems cleaner.

[jansel]

Re-request review once you want me to look at this again.

[aditvenk]

Re-request review once you want me to look at this again.

Absolutely. I am still refining the metal_jit decorator (follow-up PR) -- will re-request review once I am happy with it.

[aditvenk]

@jansel -- ready for re-review.

• This PR introduces the AST walker to C++, it is stripped down as most of the heavy-lifting is done by Metal inductor overrides.
• The next PR in stack ([metal] Add @metal_jit decorator for AST-to-MSL compilation #1991) introduces the metal_jit decorator which invokes the AST walking to generate C++, compile MSL shader
• The PR on top of that ([metal] Wire @metal_jit into MetalBackend and simplify launcher #1992) wires metal_jit into the MetalBackend and adds e2e tests for elementwise, copy kernels.