Add linalg.reduce op

mlir/dialects/linalg.py

@@ -5,8 +5,10 @@
 import mlir.astnodes as mast
 from mlir.dialect import Dialect, DialectOp, is_op
 from dataclasses import dataclass
-from typing import Optional, List
+from typing import Optional, List, Tuple, Union
 
+Literal = Union[mast.StringLiteral, float, int, bool]
+SsaUse = Union[mast.SsaId, Literal]
 
 @dataclass
 class LinalgBatchMatmul(DialectOp):
@@ -200,6 +202,22 @@ class LinalgRange(DialectOp):
                  " : {out_type.type}")]
 
 
+@dataclass
+class LinalgReduce(DialectOp):
+    inargs: List[mast.SsaId]
+    in_types: List[mast.Type]
+    outargs: List[mast.SsaId]
+    out_types: List[mast.Type]
+    dimensions: List[SsaUse]
+    region: mast.Region
+    args: List[Tuple[mast.SsaId, mast.Type]]
+
+    _syntax_ = [("linalg.reduce"
+                 " ins( {inargs.ssa_id_list} : {in_types.type_list_no_parens} )"
+                 " outs( {outargs.ssa_id_list} : {out_types.type_list_no_parens} )"
+                 " dimensions = [ {dimensions.ssa_use_list} ]"
+                 " ( {args.argument_list} ) {region.region}")]
+
 @dataclass
 class LinalgReshape(DialectOp):
     src_id: mast.SsaId

tests/test_linalg.py

@@ -106,6 +106,17 @@ def test_indexed_generic():
     return
   }
 }""")
+
+def test_reduce():
+    assert_roundtrip_equivalence("""module {
+  func.func @reduce(%arg0: tensor<16x32x64xf32>, %arg1: tensor<16x64xf32>) {
+    %reduce = linalg.reduce ins( %arg0 : tensor<16x32x64xf32> ) outs( %arg1 : tensor<16x64xf32> ) dimensions = [ 1 ] ( %in: f32, %out: f32 ) {
+      %0 = arith.addf %out, %in : f32
+      linalg.yield %0 : f32
+    }
+    return
+  }
+}""")
 
 
 def test_view():