[metal] Add MSL walker, 3D launcher, and copy kernel tests

Add the MSL AST walker that converts Python AST to MSL C++ source,
MetalBackend.post_process_function_def to wire it into codegen,
update the launcher to 3D threadgroup dispatch, and add copy kernel
tests.

stack-info: PR: #1794, branch: aditvenk/stack/13

Author: aditvenk

helion/_compiler/backend.py

@@ -2499,3 +2499,38 @@ def launcher_keyword_args(self, config: Config, *, has_barrier: bool) -> list[st
 
         dims = tuple(DeviceFunction.current().codegen.max_thread_block_dims)
         return [f"_block_dims=({dims[0]}, {dims[1]}, {dims[2]})"]
+
+    # --- MSL code generation ---
+
+    def post_process_function_def(
+        self, stmts: list[ast.stmt], device_fn: DeviceFunction
+    ) -> list[ast.stmt]:
+        """Post-process the standard codegen output into an MSL-returning function.
+
+        Extracts the body from the FunctionDef, creates an ``MslAstWalker``
+        to generate MSL, and replaces the function with a zero-arg function
+        that returns ``(msl_source, kernel_name)``.
+        """
+        import ast as _ast
+
+        from .ast_extension import create
+        from .ast_extension import create_arguments
+        from .ast_extension import statement_from_string
+        from .metal.msl_ast_walker import MslAstWalker
+
+        fn_def = next(s for s in stmts if isinstance(s, _ast.FunctionDef))
+
+        kernel_name = device_fn.name
+        walker = MslAstWalker(device_fn, fn_def.body)
+        msl_source = walker.generate()
+
+        fn_body = statement_from_string(f"return ({msl_source!r}, {kernel_name!r})")
+        msl_fn = create(
+            _ast.FunctionDef,
+            name=kernel_name,
+            args=create_arguments([]),
+            body=[fn_body],
+            decorator_list=[],
+            type_params=[],
+        )
+        return [msl_fn]

helion/_compiler/metal/__init__.py

@@ -0,0 +1 @@
+from __future__ import annotations