dialects: (memref_stream) Add some more verification for generic op

tests/filecheck/dialects/memref_stream/verify.mlir

@@ -19,3 +19,87 @@ memref_stream.generic {
 
 // CHECK: Operation does not verify: Unexpected order of iterator types: ['parallel', 'reduction', 'parallel']
 
+// -----
+
+%A, %B, %C = "test.op"() : () -> (memref<4x2xf64>, memref<2x3xf64>, memref<4x3xf64>)
+
+memref_stream.generic {
+    bounds = [#builtin.int<4>, #builtin.int<2>, #builtin.int<3>],
+    indexing_maps = [
+        affine_map<(d0, d1, d2) -> (d0, d1)>,
+        affine_map<(d0, d1, d2) -> (d1, d2)>
+    ],
+    iterator_types = ["parallel", "parallel", "reduction"]
+} ins(%A, %B : memref<4x2xf64>, memref<2x3xf64>) outs(%C : memref<4x3xf64>) {
+^0(%a : f64, %b : f64, %acc_old : f64):
+    %prod = arith.mulf %a, %b : f64
+    %acc_new = arith.addf %acc_old, %prod : f64
+    memref_stream.yield %acc_new : f64
+}
+
+// CHECK: Operation does not verify: The number of affine maps must match the number of operands
+
+// -----
+
+%A, %B, %C = "test.op"() : () -> (memref<4x2xf64>, memref<2x3xf64>, memref<4x3xf64>)
+
+memref_stream.generic {
+    bounds = [#builtin.int<4>, #builtin.int<2>, #builtin.int<3>],
+    indexing_maps = [
+        affine_map<(d0, d1, d2, d3) -> (d0, d1)>,
+        affine_map<(d0, d1, d2) -> (d1, d2)>,
+        affine_map<(d0, d1, d2, d3) -> (d0, d2)>
+    ],
+    iterator_types = ["parallel", "parallel", "reduction"]
+} ins(%A, %B : memref<4x2xf64>, memref<2x3xf64>) outs(%C : memref<4x3xf64>) {
+^0(%a : f64, %b : f64, %acc_old : f64):
+    %prod = arith.mulf %a, %b : f64
+    %acc_new = arith.addf %acc_old, %prod : f64
+    memref_stream.yield %acc_new : f64
+}
+
+// CHECK: Operation does not verify: Invalid number of dims in indexing map 0
+
+// -----
+
+%A, %B, %C = "test.op"() : () -> (memref<4x2xf64>, memref<2x3xf64>, memref<4x3xf64>)
+
+memref_stream.generic {
+    bounds = [#builtin.int<4>, #builtin.int<2>, #builtin.int<3>],
+    indexing_maps = [
+        affine_map<(d0, d1, d2) -> (d0, d1)>,
+        affine_map<(d0, d1, d2) -> (d1, d2)>,
+        affine_map<(d0, d1, d2, d3) -> (d0, d2)>
+    ],
+    iterator_types = ["parallel", "parallel", "reduction"]
+} ins(%A, %B : memref<4x2xf64>, memref<2x3xf64>) outs(%C : memref<4x3xf64>) {
+^0(%a : f64, %b : f64, %acc_old : f64):
+    %prod = arith.mulf %a, %b : f64
+    %acc_new = arith.addf %acc_old, %prod : f64
+    memref_stream.yield %acc_new : f64
+}
+
+// CHECK: Operation does not verify: The number of dims in output indexing maps must be 3 or 2
+
+// -----
+
+%A, %B, %C, %D = "test.op"() : () -> (memref<4x2xf64>, memref<2x3xf64>, memref<4x3xf64>, memref<4x3xf64>)
+
+memref_stream.generic {
+    bounds = [#builtin.int<4>, #builtin.int<2>, #builtin.int<3>],
+    indexing_maps = [
+        affine_map<(d0, d1, d2) -> (d0, d1)>,
+        affine_map<(d0, d1, d2) -> (d1, d2)>,
+        affine_map<(d0, d1, d2) -> (d0, d2)>,
+        affine_map<(d0, d1) -> (d0, d1)>
+    ],
+    iterator_types = ["parallel", "parallel", "reduction"]
+} ins(%A, %B : memref<4x2xf64>, memref<2x3xf64>) outs(%C, %D : memref<4x3xf64>, memref<4x3xf64>) {
+^0(%a : f64, %b : f64, %acc_old0 : f64, %acc_old1 : f64):
+    %prod = arith.mulf %a, %b : f64
+    %acc_new0 = arith.addf %acc_old0, %prod : f64
+    %acc_new1 = arith.addf %acc_old1, %prod : f64
+    memref_stream.yield %acc_new0, %acc_new1 : f64, f64
+}
+
+// CHECK: Operation does not verify: The number of dims in output indexing maps must all be the same

xdsl/dialects/memref_stream.py

@@ -336,6 +336,11 @@ def __init__(
         iterator_types: ArrayAttr[Attribute],
         bounds: ArrayAttr[IntAttr],
     ) -> None:
+        for m in indexing_maps:
+            if m.data.num_symbols:
+                raise NotImplementedError(
+                    f"Symbols currently not implemented in {self.name} indexing maps"
+                )
         super().__init__(
             operands=[inputs, outputs],
             properties={
@@ -530,6 +535,42 @@ def verify_(self) -> None:
                 f"Unexpected order of iterator types: {[it.data.value for it in iterator_types]}"
             )
 
+        if len(self.operands) != len(self.indexing_maps):
+            raise VerifyException(
+                "The number of affine maps must match the number of operands"
+            )
+
+        # Whether or not the operation represents an imperfect loop nest, verify that the
+        # bounds of the outer + inner nests match the domain of the input affine maps
+        input_count = len(self.inputs)
+        input_maps = self.indexing_maps.data[:input_count]
+
+        for i, m in enumerate(input_maps):
+            if len(iterator_types) != m.data.num_dims:
+                raise VerifyException(f"Invalid number of dims in indexing map {i}")
+
+        # If the operation represents an imperfect loop nest, the bounds must match the
+        # number of parallel iterators; otherwise they must match the total number of
+        # iterators. In either case, they must all be the same.
+        output_maps = self.indexing_maps.data[input_count:]
+
+        min_dims = min(m.data.num_dims for m in output_maps)
+        max_dims = max(m.data.num_dims for m in output_maps)
+
+        if min_dims != max_dims:
+            raise VerifyException(
+                "The number of dims in output indexing maps must all be the same"
+            )
+
+        if min_dims not in (len(iterator_types), num_parallel):
+            # To signify that the output is imperfectly nested, the output affine map has
+            # as many dims as parallel iterators. Otherwise, it has as many dims as
+            # the total number of iterators.
+            raise VerifyException(
+                "The number of dims in output indexing maps must be "
+                f"{len(iterator_types)} or {num_parallel}"
+            )
+
 
 @irdl_op_definition
 class YieldOp(AbstractYieldOperation[Attribute]):