Clean dynamic Triton artifacts and add loader helper

Author: drisspg

examples/export_autograd_triton_basic.py

@@ -1,25 +1,20 @@
 from __future__ import annotations
 
-import importlib.util
 from pathlib import Path
 
 import torch
 
-from transformer_nuggets.export_autograd_triton import Specialization, export_autograd_triton
+from transformer_nuggets.export_autograd_triton import (
+    Specialization,
+    export_autograd_triton,
+    load_exported_module,
+)
 
 
 def affine_relu(x, w):
     return torch.relu(x @ w)
 
 
-def import_generated(path: Path):
-    spec = importlib.util.spec_from_file_location(path.stem, path)
-    module = importlib.util.module_from_spec(spec)
-    assert spec.loader is not None
-    spec.loader.exec_module(module)
-    return module
-
-
 def main():
     if not torch.cuda.is_available():
         raise RuntimeError("CUDA is required for this example")
@@ -33,7 +28,7 @@ def main():
         [Specialization(args=(x, w), name="static_4x8_8x3")],
         output_path,
     )
-    generated = import_generated(output_path)
+    generated = load_exported_module(output_path)
     y = generated.affine_relu_compiled(x, w)
     y.sum().backward()
     print(f"wrote {output_path}")

examples/export_autograd_triton_dynamic.py

@@ -1,25 +1,20 @@
 from __future__ import annotations
 
-import importlib.util
 from pathlib import Path
 
 import torch
 
-from transformer_nuggets.export_autograd_triton import Specialization, export_autograd_triton
+from transformer_nuggets.export_autograd_triton import (
+    Specialization,
+    export_autograd_triton,
+    load_exported_module,
+)
 
 
 def affine_relu(x, w):
     return torch.relu(x @ w)
 
 
-def import_python_file(path: Path):
-    spec = importlib.util.spec_from_file_location(path.stem, path)
-    module = importlib.util.module_from_spec(spec)
-    assert spec.loader is not None
-    spec.loader.exec_module(module)
-    return module
-
-
 def main():
     if not torch.cuda.is_available():
         raise RuntimeError("CUDA is required for this example")
@@ -38,15 +33,24 @@ def main():
             )
         ],
         out=output_path,
-        source_backend="inductor",
+        source_backend="clean_triton",
     )
 
-    generated = import_python_file(output_path)
+    generated = load_exported_module(output_path)
     for batch in (1, 4, 16):
-        dynamic_x = torch.randn(batch, 8, device="cuda", requires_grad=True)
-        eager = affine_relu(dynamic_x, w)
-        compiled = generated.affine_relu_compiled(dynamic_x, w)
+        eager_x = torch.randn(batch, 8, device="cuda", requires_grad=True)
+        compiled_x = eager_x.detach().clone().requires_grad_()
+        eager_w = w.detach().clone().requires_grad_()
+        compiled_w = w.detach().clone().requires_grad_()
+
+        eager = affine_relu(eager_x, eager_w)
+        compiled = generated.affine_relu_compiled(compiled_x, compiled_w)
         torch.testing.assert_close(compiled, eager)
+
+        eager_grads = torch.autograd.grad(eager.sum(), (eager_x, eager_w))
+        compiled_grads = torch.autograd.grad(compiled.sum(), (compiled_x, compiled_w))
+        for compiled_grad, eager_grad in zip(compiled_grads, eager_grads, strict=True):
+            torch.testing.assert_close(compiled_grad, eager_grad)
         print(f"batch={batch}: {compiled.shape}")
 
     print(f"generated file: {output_path}")

examples/export_autograd_triton_minimal.py

@@ -1,11 +1,14 @@
 from __future__ import annotations
 
-import importlib.util
 from pathlib import Path
 
 import torch
 
-from transformer_nuggets.export_autograd_triton import Specialization, export_autograd_triton
+from transformer_nuggets.export_autograd_triton import (
+    Specialization,
+    export_autograd_triton,
+    load_exported_module,
+)
 
 
 def tiny_mlp(x, w1, w2, *, activation="relu"):
@@ -19,14 +22,6 @@ def tiny_mlp(x, w1, w2, *, activation="relu"):
     return hidden @ w2
 
 
-def import_python_file(path: Path):
-    spec = importlib.util.spec_from_file_location(path.stem, path)
-    module = importlib.util.module_from_spec(spec)
-    assert spec.loader is not None
-    spec.loader.exec_module(module)
-    return module
-
-
 def main():
     if not torch.cuda.is_available():
         raise RuntimeError("CUDA is required for this example")
@@ -50,7 +45,7 @@ def main():
         max_autotune=True,
     )
 
-    generated = import_python_file(output_path)
+    generated = load_exported_module(output_path)
     eager_x = x.detach().clone().requires_grad_()
     eager_w1 = w1.detach().clone().requires_grad_()
     eager_w2 = w2.detach().clone().requires_grad_()

test/test_export_autograd_triton.py

@@ -1,14 +1,17 @@
 from __future__ import annotations
 
-import importlib.util
 import inspect
 import os
 from pathlib import Path
 
 import pytest
 import torch
 
-from transformer_nuggets.export_autograd_triton import Specialization, export_autograd_triton
+from transformer_nuggets.export_autograd_triton import (
+    Specialization,
+    export_autograd_triton,
+    load_exported_module,
+)
 
 os.environ.setdefault("TORCHINDUCTOR_COMPILE_THREADS", "1")
 
@@ -98,11 +101,7 @@ def _requires_export_runtime():
 
 
 def _import_generated(path: Path):
-    spec = importlib.util.spec_from_file_location(path.stem, path)
-    module = importlib.util.module_from_spec(spec)
-    assert spec.loader is not None
-    spec.loader.exec_module(module)
-    return module
+    return load_exported_module(path)
 
 
 def _clone_tensor(tensor):
@@ -521,13 +520,16 @@ def test_dynamic_batch_specialization_dispatches_across_batch_sizes(tmp_path):
             )
         ],
         generated_path,
-        source_backend="inductor",
+        source_backend="clean_triton",
     )
     module = _import_generated(generated_path)
-    assert "s" in "\n".join(
+    artifact_source = "\n".join(
         path.read_text()
         for path in generated_path.with_name("generated_dynamic_artifacts").glob("*.py")
     )
+    assert "@triton.jit" in artifact_source
+    assert "async_compile.triton" not in artifact_source
+    assert "triton.cdiv" in artifact_source
 
     for batch_size in (1, 7, 16):
         dynamic_x = torch.randn(batch_size, 8, device="cuda", requires_grad=True)
@@ -560,11 +562,12 @@ def test_dynamic_shape_limitations_are_explicitly_guarded(tmp_path):
     x = torch.randn(2, 4, device="cuda", requires_grad=True)
     w = torch.randn(4, 5, device="cuda", requires_grad=True)
 
-    with pytest.raises(ValueError, match="source_backend='inductor'"):
+    with pytest.raises(NotImplementedError, match="forward-only"):
         export_autograd_triton(
-            affine_activation,
-            [Specialization(args=(x, w), dynamic_shapes={"x": {0: "batch"}})],
-            tmp_path / "generated_dynamic_clean.py",
+            integer_tensor_output,
+            [Specialization(args=(x,), dynamic_shapes={"x": {0: "batch"}})],
+            tmp_path / "generated_dynamic_forward_only.py",
+            source_backend="inductor",
         )
 
     with pytest.raises(NotImplementedError, match="dynamic dim 0"):

transformer_nuggets/export_autograd_triton/__init__.py

@@ -1,4 +1,10 @@
 from transformer_nuggets.export_autograd_triton.api import export_autograd_triton
+from transformer_nuggets.export_autograd_triton.loading import load_exported_module
 from transformer_nuggets.export_autograd_triton.specs import ExportedAutogradSource, Specialization
 
-__all__ = ["ExportedAutogradSource", "Specialization", "export_autograd_triton"]
+__all__ = [
+    "ExportedAutogradSource",
+    "Specialization",
+    "export_autograd_triton",
+    "load_exported_module",
+]

transformer_nuggets/export_autograd_triton/api.py

@@ -37,11 +37,6 @@ def export_autograd_triton(
         raise ValueError("source_backend must be 'inductor' or 'clean_triton'")
     if not specializations:
         raise ValueError("At least one specialization is required")
-    if source_backend == "clean_triton" and any(
-        specialization.dynamic_shapes is not None for specialization in specializations
-    ):
-        raise ValueError("dynamic_shapes require source_backend='inductor' for now")
-
     config_patches = dict(inductor_config_patches or {})
     if max_autotune:
         config_patches["max_autotune"] = True

transformer_nuggets/export_autograd_triton/codegen.py

@@ -232,11 +232,40 @@ def _rewrite_module_as_clean_triton(path: Path) -> None:
     from torch.utils._get_clean_triton import get_clean_triton
 
     get_clean_triton(path.resolve(), path.resolve(), auto_generate_params=True)
+    _patch_dynamic_pointwise_grids(path)
     launch_params_path = Path(f"{path}.launch_params")
     if launch_params_path.exists():
         launch_params_path.unlink()
 
 
+def _patch_dynamic_pointwise_grids(path: Path) -> None:
+    source = path.read_text()
+    for xnumel_name in sorted(set(re.findall(r"(\w+_xnumel)\s*=", source))):
+        kernel_name = xnumel_name.removesuffix("_xnumel")
+        source = re.sub(
+            rf"{kernel_name}\[\(\d+, \d+, \d+\)\]\((?P<args>[^\n]*?), \d+, XBLOCK=(?P<xblock>\d+),",
+            lambda match: _dynamic_pointwise_launch_replacement(
+                kernel_name,
+                xnumel_name,
+                match,
+            ),
+            source,
+        )
+    path.write_text(source)
+
+
+def _dynamic_pointwise_launch_replacement(
+    kernel_name: str,
+    xnumel_name: str,
+    match: re.Match[str],
+) -> str:
+    xblock = match.group("xblock")
+    return (
+        f"{kernel_name}[(triton.cdiv({xnumel_name}, {xblock}), 1, 1)]("
+        f"{match.group('args')}, {xnumel_name}, XBLOCK={xblock},"
+    )
+
+
 def _generate_public_wrapper(exported_name: str, signature: inspect.Signature) -> str:
     return f"def {exported_name}{signature}:\n" + indent(
         "return _run_with_bound_args(locals())\n", "    "

transformer_nuggets/export_autograd_triton/loading.py

@@ -0,0 +1,15 @@
+from __future__ import annotations
+
+import importlib.util
+from pathlib import Path
+from types import ModuleType
+
+
+def load_exported_module(path: str | Path) -> ModuleType:
+    path = Path(path)
+    spec = importlib.util.spec_from_file_location(path.stem, path)
+    if spec is None or spec.loader is None:
+        raise ImportError(f"Could not load generated module from {path}")
+    module = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(module)
+    return module