dialects: implement affine.apply (#1752)

Mirror affine.apply from MLIR --------- Co-authored-by: Chris Vasiladiotis <[email protected]>
xdslproject · Nov 6, 2023 · bdcc260 · bdcc260
1 parent cb75f54
commit bdcc260
Show file tree

Hide file tree

Showing 2 changed files with 24 additions and 0 deletions.
diff --git a/tests/filecheck/dialects/affine/affine_ops.mlir b/tests/filecheck/dialects/affine/affine_ops.mlir
@@ -40,9 +40,11 @@
     // CHECK-NEXT: "affine.store"(%value, %memref) {"map" = affine_map<() -> (0, 0)>} : (f64, memref<2x3xf64>) -> ()
 
     %zero = "test.op"() : () -> index
+    %2 = "affine.apply"(%zero, %zero) <{"map" = affine_map<(d0)[s0] -> (((d0 + (s0 * 42)) + -1))>}> : (index, index) -> index
     %same_value = "affine.load"(%memref, %zero, %zero) {"map" = affine_map<(d0, d1) -> (d0, d1)>} : (memref<2x3xf64>, index, index) -> f64
 
     // CHECK:      %zero = "test.op"() : () -> index
+    // CHECK-NEXT: %{{.*}} = "affine.apply"(%{{.*}}, %{{.*}}) <{"map" = affine_map<(d0)[s0] -> (((d0 + (s0 * 42)) + -1))>}> : (index, index) -> index
     // CHECK-NEXT: %same_value = "affine.load"(%memref, %zero, %zero) {"map" = affine_map<(d0, d1) -> (d0, d1)>} : (memref<2x3xf64>, index, index) -> f64
 
     func.func @empty() {
@@ -59,6 +61,7 @@
     }, {
       "affine.yield"() : () -> ()
     }) {"condition" = affine_set<() : (0 == 0)>, "some_attr" = true} : () -> ()
+
     func.return
   }
 // CHECK:    func.func @empty() {
@@ -75,6 +78,7 @@
 // CHECK-NEXT:      }, {
 // CHECK-NEXT:        "affine.yield"() : () -> ()
 // CHECK-NEXT:      }) {"condition" = affine_set<() : (0 == 0)>, "some_attr" = true} : () -> ()
+
 // CHECK-NEXT:      func.return
 // CHECK-NEXT:    }
   func.func @affine_if() -> f32 {

diff --git a/xdsl/dialects/affine.py b/xdsl/dialects/affine.py
@@ -25,12 +25,31 @@
     irdl_op_definition,
     operand_def,
     opt_attr_def,
+    prop_def,
     region_def,
     result_def,
     var_operand_def,
     var_result_def,
 )
 from xdsl.traits import IsTerminator
+from xdsl.utils.exceptions import VerifyException
+
+
+@irdl_op_definition
+class ApplyOp(IRDLOperation):
+    name = "affine.apply"
+
+    mapOperands = var_operand_def(IndexType)
+    map = prop_def(AffineMapAttr)
+    result = result_def(IndexType)
+
+    def verify_(self) -> None:
+        if len(self.mapOperands) != self.map.data.num_dims + self.map.data.num_symbols:
+            raise VerifyException(
+                f"{self.name} expects {self.map.data.num_dims + self.map.data.num_symbols} operands, but got {len(self.mapOperands)}. The number of map operands must match the sum of the dimensions and symbols of its map."
+            )
+        if len(self.map.data.results) != 1:
+            raise VerifyException("affine.apply expects a unidimensional map.")
 
 
 @irdl_op_definition
@@ -209,6 +228,7 @@ def get(*operands: SSAValue | Operation) -> Yield:
 Affine = Dialect(
     "affine",
     [
+        ApplyOp,
         For,
         If,
         Store,