Improve export_autograd_triton generated readability

Author: drisspg

test/test_export_autograd_triton.py

@@ -16,7 +16,15 @@
     CleanTritonUnsupportedError,
     clean_triton_source,
 )
+from transformer_nuggets.export_autograd_triton.codegen import (
+    generate_autograd_source,
+    write_autograd_source,
+)
 from transformer_nuggets.export_autograd_triton.guards import tensor_guard_for
+from transformer_nuggets.export_autograd_triton.specs import (
+    CapturedSpecialization,
+    TensorGuardSpec,
+)
 
 os.environ.setdefault("TORCHINDUCTOR_COMPILE_THREADS", "1")
 
@@ -153,6 +161,37 @@ def _clone_args(args):
     return tuple(_clone_tensor(arg) if isinstance(arg, torch.Tensor) else arg for arg in args)
 
 
+def _fake_captured_specialization(
+    forward_source="def call(args):\n    return ()\n",
+    name="spec_0",
+):
+    return CapturedSpecialization(
+        name=name,
+        runtime_tensor_names=("x",),
+        static_args=(),
+        tensor_guards=(
+            TensorGuardSpec(
+                name="x",
+                rank=2,
+                shape=(4, 8),
+                stride=(8, 1),
+                dtype="torch.float32",
+                device_type="cuda",
+                device_index=0,
+            ),
+        ),
+        forward_source=forward_source,
+        backward_source=None,
+        num_user_outputs=1,
+        output_kind="single",
+        needs_autograd=False,
+        dynamic=False,
+        differentiable_output_mask=(True,),
+        forward_residual_names=(),
+        backward_saved_input_names=(),
+    )
+
+
 def _flatten_tensors(value):
     if isinstance(value, torch.Tensor):
         return [value]
@@ -228,6 +267,121 @@ def test_unbounded_dynamic_dim_guard_uses_none_for_symbolic_max():
     assert guard.shape[0] == {"symbol": "batch", "min": 0, "max": None}
 
 
+def test_generated_wrapper_source_is_labeled_and_splits_specs():
+    source = generate_autograd_source(
+        trig_pointwise,
+        "trig_pointwise_compiled",
+        [_fake_captured_specialization()],
+    )
+
+    assert source.startswith('"""Generated by transformer_nuggets.export_autograd_triton.')
+    assert "# Generated Triton artifact index. Paths are relative to this wrapper file." in source
+    assert "#   spec_0 forward: <this_module>_artifacts/spec_0_forward.py" in source
+    assert "#   spec_0 backward: <this_module>_artifacts/spec_0_backward.py" not in source
+    assert (
+        "# Calls the selected *_forward.py artifact listed in the artifact index above." in source
+    )
+    assert (
+        "# Calls the selected *_backward.py artifact listed in the artifact index above." in source
+    )
+    assert "\n\n\n_SPEC_0" not in source
+    assert "_SPEC_0 = {" in source
+    assert "_SPECS = [\n    _SPEC_0," in source
+    assert "class _TrigPointwiseCompiledAutogradFunction(torch.autograd.Function):" in source
+    assert "return _RUNTIME.autograd_forward(ctx, spec_id, runtime_tensors)" in source
+    assert "return _RUNTIME.autograd_backward(ctx, *grad_outputs)" in source
+    assert (
+        "result = _TrigPointwiseCompiledAutogradFunction.apply(spec_id, *runtime_tensors)"
+        in source
+    )
+    assert "return _RUNTIME.restore_output_container(spec_id, result)" in source
+    assert "Run the exported autograd/Triton specialization" in source
+
+
+def test_write_autograd_source_labels_artifacts_and_removes_stale_files(tmp_path):
+    generated_path = tmp_path / "generated_fake.py"
+    artifact_dir = tmp_path / "generated_fake_artifacts"
+    artifact_dir.mkdir()
+    (artifact_dir / "stale.py").write_text("stale")
+
+    write_autograd_source(
+        generated_path,
+        "wrapper source",
+        [_fake_captured_specialization()],
+        source_backend="inductor",
+    )
+
+    artifact_path = artifact_dir / "spec_0_forward.py"
+    artifact_source = artifact_path.read_text()
+    assert not (artifact_dir / "stale.py").exists()
+    assert artifact_source.startswith('"""Generated Triton artifact.')
+    assert "Runtime tensor order:\n  - x" in artifact_source
+    assert "Tensor guards:\n  - x: shape=(4, 8), stride=(8, 1)" in artifact_source
+    assert "The runtime imports this file and calls call(args)." in artifact_source
+
+
+def test_generated_wrapper_uses_output_path_in_artifact_index(tmp_path):
+    generated_path = tmp_path / "generated_fake.py"
+
+    source = generate_autograd_source(
+        trig_pointwise,
+        "trig_pointwise_compiled",
+        [_fake_captured_specialization()],
+        artifact_dir_name=f"{generated_path.stem}_artifacts",
+    )
+
+    assert "#   spec_0 forward: generated_fake_artifacts/spec_0_forward.py" in source
+
+
+def test_write_clean_triton_artifact_strips_benchmark_tail_without_async_compile(tmp_path):
+    generated_path = tmp_path / "generated_fake.py"
+    write_autograd_source(
+        generated_path,
+        "wrapper source",
+        [
+            _fake_captured_specialization(
+                forward_source="""def call(args):
+    return args
+
+
+def get_args():
+    return []
+
+
+def benchmark_compiled_module(args, times=10, repeat=10):
+    return None
+
+
+if __name__ == "__main__":
+    print(get_args())
+"""
+            )
+        ],
+        source_backend="clean_triton",
+    )
+
+    artifact_source = (tmp_path / "generated_fake_artifacts" / "spec_0_forward.py").read_text()
+    assert "def call(args):" in artifact_source
+    assert "def get_args():" not in artifact_source
+    assert "benchmark_compiled_module" not in artifact_source
+    assert "__main__" not in artifact_source
+
+
+def test_artifact_header_escapes_docstring_metadata(tmp_path):
+    generated_path = tmp_path / "generated_fake.py"
+    write_autograd_source(
+        generated_path,
+        "wrapper source",
+        [_fake_captured_specialization(name='bad"""name\nline')],
+        source_backend="inductor",
+    )
+
+    artifact_path = tmp_path / "generated_fake_artifacts" / "spec_0_bad_name_line_forward.py"
+    artifact_source = artifact_path.read_text()
+    compile(artifact_source, str(artifact_path), "exec")
+    assert 'Specialization: bad\\"\\"\\"name\\nline' in artifact_source
+
+
 def test_export_static_specialization_forward_backward_signature_and_errors(tmp_path):
     _requires_export_runtime()
     x = torch.randn(4, 8, device="cuda", requires_grad=True)
@@ -659,11 +813,90 @@ def call(args):
     assert result.patched_launches == ["triton_poi_fused_sin_0", "triton_poi_fused_cos_1"]
     assert "triton.cdiv(triton_poi_fused_sin_0_xnumel, 32)" in result.source
     assert "triton.cdiv(triton_poi_fused_cos_1_xnumel, 64)" in result.source
-    assert "buf0, triton_poi_fused_sin_0_xnumel, XBLOCK=32" in result.source
-    assert "buf1, triton_poi_fused_cos_1_xnumel, XBLOCK=64" in result.source
+    assert "buf0,\n        triton_poi_fused_sin_0_xnumel,\n        XBLOCK=32" in result.source
+    assert "buf1,\n        triton_poi_fused_cos_1_xnumel,\n        XBLOCK=64" in result.source
     assert "[(1, 1, 1)]" not in result.source
 
 
+def test_clean_triton_source_removes_unused_imports_assignments_and_blank_runs():
+    source = """
+import math
+import os
+import torch
+import torch
+from pathlib import Path
+from torch._C import unused_stream, _cuda_getCurrentRawStream as get_raw_stream
+
+unused_alias = torch.ops.aten
+used_alias = torch.ops.aten
+
+
+
+
+def call(args):
+    stream = get_raw_stream(0)
+    return torch.empty(1), used_alias, stream
+"""
+
+    result = clean_triton_source(source, dynamic=False)
+
+    assert "import math" not in result.source
+    assert "import os" not in result.source
+    assert result.source.count("import torch") == 1
+    assert "from pathlib import Path" not in result.source
+    assert "unused_alias" not in result.source
+    assert "used_alias = torch.ops.aten" in result.source
+    assert "unused_stream" not in result.source
+    assert "get_raw_stream" in result.source
+    assert "\n\n\n" not in result.source
+
+
+def test_clean_triton_source_strips_standalone_benchmark_tail():
+    source = """
+def call(args):
+    return args
+
+
+def get_args():
+    return []
+
+
+def benchmark_compiled_module(args, times=10, repeat=10):
+    return None
+
+
+if __name__ == "__main__":
+    print(get_args())
+"""
+
+    result = clean_triton_source(source, dynamic=False)
+
+    assert "def call(args):" in result.source
+    assert "def get_args():" not in result.source
+    assert "benchmark_compiled_module" not in result.source
+    assert "__main__" not in result.source
+
+
+def test_clean_triton_source_does_not_strip_non_benchmark_get_args_helper():
+    source = """
+def call(args):
+    return get_args(args)
+
+
+def get_args(args):
+    return args
+
+
+def still_needed(args):
+    return args
+"""
+
+    result = clean_triton_source(source, dynamic=False)
+
+    assert "def get_args(args):" in result.source
+    assert "def still_needed(args):" in result.source
+
+
 def test_dynamic_clean_triton_source_patches_repeated_kernel_launches_by_occurrence():
     original_source = """
 def call(args):
@@ -684,8 +917,14 @@ def call(args):
 
     result = clean_triton_source(source, dynamic=True, original_source=original_source)
 
-    assert "triton.cdiv(s0, 4), 1, 1)](buf0, s0" in result.source
-    assert "triton.cdiv(s1, 4), 1, 1)](buf1, s1" in result.source
+    assert (
+        "triton.cdiv(s0, 4),\n        1,\n        1,\n    )](\n        buf0,\n        s0,"
+        in result.source
+    )
+    assert (
+        "triton.cdiv(s1, 4),\n        1,\n        1,\n    )](\n        buf1,\n        s1,"
+        in result.source
+    )
 
 
 def test_dynamic_clean_triton_reduction(tmp_path):
@@ -756,7 +995,8 @@ def test_dynamic_clean_triton_matmul_dynamic_m(tmp_path):
     assert "@triton.jit" in artifact_source
     assert "async_compile.triton" not in artifact_source
     assert "triton_tem_" in artifact_source
-    assert "[((" in artifact_source
+    assert "[(\n" in artifact_source
+    assert "((15 + s" in artifact_source
     assert "[(4, 1, 1)]" not in artifact_source
 
     for batch_size in (1, 7, 32):

transformer_nuggets/export_autograd_triton/README.md

@@ -55,18 +55,29 @@ agent_space/generated_rms_norm_artifacts/
   spec_0_backward.py
 ```
 
-The top-level generated wrapper is intentionally small. It contains:
+The top-level generated wrapper is intentionally small and labeled, but keeps the autograd node
+visible. It contains:
 
-- `_SPECS`: specialization metadata and tensor guards
+- a generated artifact index with direct relative paths to each forward/backward artifact
+- one `_SPEC_<n>` block per specialization, with metadata and tensor guards
+- `_SPECS`: the ordered specialization list used for dispatch
 - `_RUNTIME = ExportedAutogradRuntime(...)`
-- one public function such as `rms_norm_compiled(...)`
+- one explicit `torch.autograd.Function` subclass whose `forward` and `backward` call the
+  generated artifact runners through the runtime helper
+- one documented public function such as `rms_norm_compiled(...)` that selects a specialization,
+  gathers runtime tensors, and invokes the generated autograd function
+
+Artifact files start with a short generated-artifact docstring that names the specialization,
+forward/backward direction, runtime tensor order, static arguments, tensor guards, output
+metadata, and saved-residual order. The artifact directory is recreated on each export so stale
+kernels from earlier runs do not hang around.
 
 Shared runtime behavior lives in `runtime.py`, not in every generated file:
 
 - artifact module loading
 - specialization dispatch
 - static/tensor guard checking
-- generated autograd function
+- autograd forward/backward helper plumbing
 - saved residual reordering for backward graphs
 - non-differentiable output marking
 - no-match diagnostics
@@ -89,7 +100,10 @@ Use this when:
 
 Runs `torch.utils._get_clean_triton.get_clean_triton(...)` and then applies the dynamic-clean validation/patching layer in `clean_triton.py`.
 
-The goal is readable inline `@triton.jit` artifacts without `async_compile.triton(...)` while preserving dynamic launch arguments.
+The goal is readable inline `@triton.jit` artifacts without `async_compile.triton(...)` while
+preserving dynamic launch arguments. The cleaner also removes unused imports/top-level aliases,
+standalone benchmark scaffolding (`get_args`, `benchmark_compiled_module`,
+`if __name__ == "__main__"`), and formats patched dynamic launches as multiline Triton calls.
 
 Supported dynamic clean families today:
 
@@ -229,6 +243,5 @@ python examples/export_autograd_triton/dynamic.py
 - Symbolic stride guards for broader dynamic layouts
 - Dynamic CUDA device ordinal support
 - More dynamic clean-Triton families and variants as Inductor output changes
-- Cleaner metadata formatting in generated wrappers if `_SPECS` grows large
+- Semantic names for common generated symbols such as `arg0_1`, `buf*`, and `s*`
 - Optional Nsight Compute-based bandwidth measurement for exact HBM traffic
-- Tests that compare generated wrapper readability/size directly

transformer_nuggets/export_autograd_triton/api.py

@@ -46,7 +46,12 @@ def export_autograd_triton(
     ]
     output_path = Path(out)
     public_name = exported_name or f"{fn.__name__}_compiled"
-    source = generate_autograd_source(fn, public_name, captured_specializations)
+    source = generate_autograd_source(
+        fn,
+        public_name,
+        captured_specializations,
+        artifact_dir_name=f"{output_path.stem}_artifacts",
+    )
     write_autograd_source(output_path, source, captured_specializations, source_backend)
     return ExportedAutogradSource(
         output_path=output_path,