affine transform util: add interoperability with AffineMaps

compiler/ir/autoflow/affine_transform.py

@@ -3,6 +3,7 @@
 import numpy as np
 import numpy.typing as npt
 from typing_extensions import Self
+from xdsl.ir.affine import AffineConstantExpr, AffineDimExpr, AffineExpr, AffineMap
 
 
 @dataclass(frozen=True)
@@ -25,6 +26,38 @@ def __post_init__(self):
         if self.A.shape[0] != self.b.shape[0]:
             raise ValueError("Matrix A and vector b must have compatible dimensions.")
 
+    @classmethod
+    def from_affine_map(cls, map: AffineMap) -> Self:
+        """
+        Return the affine transform representation of the given affine map.
+        """
+
+        def generate_one_list(n: int, d: int):
+            return [1 if x == d else 0 for x in range(n)]
+
+        b = np.array(map.eval(generate_one_list(map.num_dims, -1), []))
+        a = np.zeros((len(map.results), map.num_dims), dtype=np.int_)
+        for dim in range(map.num_dims):
+            temp = np.array(map.eval(generate_one_list(map.num_dims, dim), []))
+            a[:, dim] = temp - b
+
+        return cls(a, b)
+
+    def to_affine_map(self) -> AffineMap:
+        """
+        Return the xDSL AffineMap representation of this AffineTransform
+        """
+        results: list[AffineExpr] = []
+        for result in range(self.num_results):
+            expr = AffineConstantExpr(int(self.b[result]))
+            for dim in range(self.num_dims):
+                if self.A[result, dim] != 0:
+                    expr += AffineConstantExpr(
+                        int(self.A[result, dim])
+                    ) * AffineDimExpr(dim)
+            results.append(expr)
+        return AffineMap(self.num_dims, 0, tuple(results))
+
     @property
     def num_dims(self) -> int:
         return self.A.shape[1]

tests/ir/autoflow/test_affine_transform.py

@@ -1,7 +1,9 @@
 import numpy as np
 import pytest
+from xdsl.ir.affine import AffineMap
 
 from compiler.ir.autoflow import AffineTransform
+from compiler.util.canonicalize_affine import canonicalize_map
 
 
 def test_affine_transform_initialization_valid():
@@ -95,3 +97,21 @@ def test_affine_transform_str():
     transform = AffineTransform(A, b)
     expected = "AffineTransform(A=\n[[1 0]\n [0 1]],\nb=[1 2])"
     assert str(transform) == expected
+
+
+def test_affine_map_interop():
+    map = AffineMap.from_callable(lambda a, b, c: (a + 2 * b, -b + c, 3 * a + c + 4))
+
+    # convert AffineMap to AffineTransform
+    transform = AffineTransform.from_affine_map(map)
+
+    expected_a = np.array([[1, 2, 0], [0, -1, 1], [3, 0, 1]])
+    expected_b = np.array([0, 0, 4])
+
+    assert (transform.A == expected_a).all()
+    assert (transform.b == expected_b).all()
+
+    # convert back to AffineMap
+
+    original_map = transform.to_affine_map()
+    assert canonicalize_map(map) == canonicalize_map(original_map)