First draft to automate the rewriting of loops (#101)

* first draft

* scan

* fix rewriting

* fix issues

* extend test

* fix

* fix

Author: xadupre

_unittests/ut_torch_export_patches/test_patch_module.py

@@ -2,12 +2,16 @@
 import inspect
 import textwrap
 import unittest
+import numpy as np
+from scipy.spatial.distance import cdist
 import torch
-from onnx_diagnostic.ext_test_case import ExtTestCase, hide_stdout
+from onnx_diagnostic.ext_test_case import ExtTestCase, hide_stdout, has_torch, requires_torch
 from onnx_diagnostic.torch_export_patches import torch_export_patches, torch_export_rewrite
 from onnx_diagnostic.torch_export_patches.patch_module import (
     transform_method,
     inplace_add_parent,
+    ShapeFinder,
+    RewriteControlFlow,
 )
 
 
@@ -42,7 +46,7 @@ def forward(self, x, y):
             hasattr(node, "parent") for node in ast.walk(tree)
         ), f"Missing parent in {ast.dump(tree, indent=2)}"
 
-    def test_rewrite_forward_return1(self):
+    def test_rewrite_test_in_forward_return1(self):
 
         class Model(torch.nn.Module):
             def forward(self, x, y):
@@ -69,7 +73,7 @@ def forward(self, x, y):
         self.assertEqualAny(expected_, ep.module()(-x, y))
 
     @hide_stdout()
-    def test_rewrite_forward_return2(self):
+    def test_rewrite_test_in_forward_return2(self):
 
         class Model(torch.nn.Module):
             def forward(self, x, y):
@@ -95,7 +99,7 @@ def forward(self, x, y):
         self.assertEqualAny(expected, ep.module()(x, y))
         self.assertEqualAny(expected_, ep.module()(-x, y))
 
-    def test_rewrite_forward_assign1(self):
+    def test_rewrite_test_in_forward_assign1(self):
 
         class Model(torch.nn.Module):
             def forward(self, x, y):
@@ -122,7 +126,7 @@ def forward(self, x, y):
         self.assertEqualAny(expected, ep.module()(x, y))
         self.assertEqualArray(expected_, ep.module()(-x, y))
 
-    def test_rewrite_forward_assign2(self):
+    def test_rewrite_test_in_forward_assign2(self):
 
         class Model(torch.nn.Module):
             def forward(self, x, y):
@@ -149,7 +153,7 @@ def forward(self, x, y):
         self.assertEqualAny(expected, ep.module()(x, y))
         self.assertEqualAny(expected_, ep.module()(-x, y))
 
-    def test_rewrite_forward_assign_noelse(self):
+    def test_rewrite_test_in_forward_assign_noelse(self):
 
         class Model(torch.nn.Module):
             def forward(self, x, y):
@@ -174,7 +178,7 @@ def forward(self, x, y):
         self.assertEqualAny(expected, ep.module()(x, y))
         self.assertEqualAny(expected_, ep.module()(-x, y))
 
-    def test_rewrite_forward_return_noelse(self):
+    def test_rewrite_test_in_forward_return_noelse(self):
 
         class Model(torch.nn.Module):
             def forward(self, x, y):
@@ -186,7 +190,7 @@ def forward(self, x, y):
             lambda: transform_method(Model.forward, verbose=self.verbose), NotImplementedError
         )
 
-    def test_rewrite_forward_assign2_in_2(self):
+    def test_rewrite_test_in_forward_assign2_in_2(self):
 
         class Model(torch.nn.Module):
             def forward(self, x, y):
@@ -215,7 +219,7 @@ def forward(self, x, y):
         self.assertEqualAny(expected, ep.module()(x, y))
         self.assertEqualAny(expected_, ep.module()(-x, y))
 
-    def test_rewrite_forward_assign2_in_3(self):
+    def test_rewrite_test_in_forward_assign2_in_3(self):
 
         class Model(torch.nn.Module):
             def forward(self, x, y):
@@ -292,7 +296,7 @@ def torch_cond_else_2(y):
         x, y = torch.rand((3, 4)), torch.rand((3, 4))
         Model()(x, y)
 
-    def test_rewrite_forward_assign_nested(self):
+    def test_rewrite_test_in_forward_assign_nested(self):
 
         class Model(torch.nn.Module):
             def forward(self, x, y):
@@ -341,7 +345,7 @@ def forward(self, x, y):
         self.assertEqualAny(expected_0, ep.module()(x, -y))
         self.assertEqualAny(expected_1, ep.module()(-x, -y))
 
-    def test_rewrite_forward_none(self):
+    def test_rewrite_test_in_forward_none(self):
 
         class Model(torch.nn.Module):
             def forward(self, x, y):
@@ -365,7 +369,7 @@ def forward(self, x, y):
         self.assertEqualAny(expected, ep.module()(x, y))
         self.assertEqualAny(expected_, ep.module()(-x, y))
 
-    def test_rewrite_PLBartEncoderLayer(self):
+    def test_rewrite_test_in_PLBartEncoderLayer(self):
         from transformers.models.plbart.modeling_plbart import PLBartEncoderLayer
 
         rewritten = transform_method(PLBartEncoderLayer.forward, verbose=self.verbose)
@@ -443,6 +447,135 @@ def forward(self, x, y):
             got = ep.module()(x, y)
             self.assertEqualArray(expected, got)
 
+    def test_shape_finder(self):
+        expr = "range(x.shape[0])"
+        node = ast.parse(expr)
+        sh = ShapeFinder()
+        sh.visit(node)
+        self.assertEqual({"x"}, sh.found_shape)
+
+    def test__find_loop_vars(self):
+        code = textwrap.dedent(
+            """
+            for i in range(x.shape[0]):
+                z[i, :] = ((x[i : i + 1, :] - y) ** 2).sum(dim=-1)
+            """
+        )
+        node = ast.parse(code)
+        tr = RewriteControlFlow()
+        vars = tr._find_loop_vars(node.body[0])
+        self.assertEqual(
+            {
+                "init": ["z"],
+                "input": ["y"],
+                "loop": ["i"],
+                "output": [],
+                "scan": [],
+                "scan_shape": ["x"],
+            },
+            vars,
+        )
+
+    @requires_torch("2.8")
+    def test_rewrite_loop(self):
+
+        class Model(torch.nn.Module):
+            def forward(self, x, y):
+                z = torch.empty((x.shape[0], y.shape[0]))
+                for i in range(x.shape[0]):
+                    z[i, :] = ((x[i, :] - y) ** 2).sum(dim=-1)
+                return z
+
+        x, y = torch.rand((3, 4)), torch.rand((5, 4))
+        expected = Model()(x, y)
+        self.assertEqualArray(
+            expected.numpy(),
+            cdist(x.numpy(), y.numpy(), metric="sqeuclidean").astype(np.float32),
+            atol=1e-5,
+        )
+
+        class RewrittenModel(torch.nn.Module):
+            def forward(self, x, y):
+                def loop_body_0(x, y):
+                    x = x.reshape((-1, *x.shape))
+                    z = ((x - y) ** 2).sum(dim=-1)
+                    return (z,)
+
+                z = torch.ops.higher_order.scan(loop_body_0, [], [x], [y])
+                return z[0]
+
+        rewritten_expected = RewrittenModel()(x, y)
+        self.assertEqualArray(expected, rewritten_expected)
+
+        DYN = torch.export.Dim.DYNAMIC
+        ds = ({0: DYN, 1: DYN}, {0: DYN, 1: DYN})
+        torch.export.export(RewrittenModel(), (x, y), dynamic_shapes=ds)
+
+        class RewrittenModelLoop(torch.nn.Module):
+            def forward(self, z, iv, x, y):
+                z = z.clone()
+                i = iv.item()
+                z[i, :] = ((x[i, :] - y) ** 2).sum(dim=-1)
+                return (z, iv)
+
+        inputs = (
+            torch.empty((x.shape[0], y.shape[0])),
+            torch.tensor([2], dtype=torch.int64),
+            x,
+            y,
+        )
+        RewrittenModelLoop()(*inputs)
+        try:
+            from experimental_experiment.torch_interpreter.tracing import CustomTracer
+        except ImportError:
+            CustomTracer = None
+        if CustomTracer:
+            graph = CustomTracer().trace(RewrittenModelLoop())
+            self.assertNotEmpty(graph)
+
+        # does not wiork
+        # dsl = ({0: DYN, 1: DYN}, {}, {0: DYN, 1: DYN}, {0: DYN, 1: DYN})
+        # torch.export.export(RewrittenModelLoop(), inputs, dynamic_shapes=dsl)
+
+        class RewrittenModel2(torch.nn.Module):
+            def forward(self, x, y):
+                def loop_body_1(z, iv, x, y):
+                    z = z.clone()
+                    i = iv.item()
+                    z[i, :] = ((x[i, :] - y) ** 2).sum(dim=-1)
+                    return (z, iv)
+
+                z = torch.empty((x.shape[0], y.shape[0]))
+                r = torch.ops.higher_order.scan(
+                    loop_body_1,
+                    [z],
+                    [torch.arange(x.shape[0], dtype=torch.int64).reshape((-1, 1))],
+                    [x, y],
+                )
+                return r[0]
+
+        rewritten = transform_method(Model.forward, verbose=self.verbose)
+        self.assertIn("torch.ops.higher_order.scan(", rewritten.code)
+        Model.forward = rewritten.func
+        self.assertEqualAny(expected, Model()(x, y))
+
+        rewritten_expected2 = RewrittenModel2()(x, y)
+        self.assertEqualArray(expected, rewritten_expected2)
+
+        if not has_torch("2.9"):
+            raise unittest.SkipTest("skipped export, torch must be >= 2.9")
+
+        torch.export.export(RewrittenModel2(), (x, y), dynamic_shapes=ds, strict=False)
+        ep = torch.export.export(Model(), (x, y), dynamic_shapes=ds, strict=False)
+        self.assertEqualAny(expected, ep.module()(x, y))
+
+        """
+                    position_encodings = torch.cat(
+                [weight[:, :required_pos_encodings_columns]
+                 for weight in broadcasted_weights], dim=-1
+            )
+        """
+
 
 if __name__ == "__main__":
     unittest.main(verbosity=2)