Add AccumulateOp.

Author: ingomueller-net

experimental/iterators/include/iterators/Dialect/Iterators/IR/IteratorsOps.td

@@ -55,6 +55,121 @@ def Iterators_PrintOp : Iterators_Base_Op<"print"> {
 // High-level iterators
 //===----------------------------------------------------------------------===//
 
+/// Looks up the given symbol, which must refer to a FuncOp, in the scope of the
+/// given op and returns the function type of that symbol.
+class LookupFuncType<string opName, string symbolName>
+  : StrFunc<"::mlir::SymbolTable::lookupNearestSymbolFrom<func::FuncOp>("
+            "    &$" # opName # ","
+            "    " # symbolName # ".dyn_cast<FlatSymbolRefAttr>())"
+            "       .getFunctionType()">;
+
+/// An init function, that is, a function with a signature of the form () -> T.
+class Iterators_InitFunctionType
+  : Type<And<[FunctionType.predicate,
+              CPred<"$_self.dyn_cast<FunctionType>().getInputs().size() == 0">,
+              CPred<"$_self.dyn_cast<FunctionType>().getResults().size() == 1">
+             ]>,
+         "function with signature () -> T",
+         "FunctionType">;
+
+/// A FlatSymbolRef referring to an init function.
+def Iterators_InitFunctionSymbol
+  : Confined<FlatSymbolRefAttr, [
+        ReferToOp<"func::FuncOp">,
+        AttrConstraint<
+          SubstLeaves<"$_self",
+                      LookupFuncType<"_op", "$_self.dyn_cast<FlatSymbolRefAttr>()">.result,
+                      Iterators_InitFunctionType<>.predicate>,
+          "referring to a " # Iterators_InitFunctionType<>.summary>
+      ]>;
+
+/// An accumulate function, that is, a function with a signature of the form
+/// (T1, T2) -> T1.
+class Iterators_AccumulateFunctionType
+  : Type<And<[FunctionType.predicate,
+              CPred<"$_self.dyn_cast<FunctionType>().getInputs().size() == 2">,
+              CPred<"$_self.dyn_cast<FunctionType>().getResults().size() == 1">,
+              AllMatchPred<["$_self.dyn_cast<FunctionType>().getInput(0)",
+                            "$_self.dyn_cast<FunctionType>().getResult(0)"]>
+             ]>,
+         "function with signature (T1, T2) -> T1",
+         "FunctionType">;
+
+/// A FlatSymbolRef referring to an accumulate function.
+def Iterators_AccumulateFunctionSymbol
+  : Confined<FlatSymbolRefAttr, [
+        ReferToOp<"func::FuncOp">,
+        AttrConstraint<
+          SubstLeaves<"$_self",
+                      LookupFuncType<"_op", "$_self.dyn_cast<FlatSymbolRefAttr>()">.result,
+                      Iterators_AccumulateFunctionType<>.predicate>,
+          "referring to a " # Iterators_AccumulateFunctionType<>.summary>
+      ]>;
+
+def Iterators_AccumulateOp
+  : Iterators_Op<"accumulate",
+    [AllMatch<["getInitFunc().getResultTypes().front()",
+               "$result.getType().dyn_cast<StreamType>().getElementType()"],
+              "the return type of the init function must match the result "
+              "element type">,
+     AllMatch<["getAccumulateFunc().getArgumentTypes()[0]",
+               "$result.getType().dyn_cast<StreamType>().getElementType()"],
+              "the type of the first argument of the accumulate function must "
+              "match the result element type">,
+     AllMatch<["getAccumulateFunc().getArgumentTypes()[1]",
+               "$input.getType().dyn_cast<StreamType>().getElementType()"],
+              "the type of the second argument of the accumulate function must "
+              "match the input element type">]> {
+  let summary = "Accumulate the elements of a stream into one element";
+  let description = [{
+    Accumulate the elements of the input stream into a single element, i.e.,
+    compute their generalized sum. This is similar to
+    [`std::accumulate`](https://en.cppreference.com/w/cpp/algorithm/accumulate)
+    in C++ and
+    [`functools.reduce`](https://docs.python.org/3/library/functools.html#functools.reduce)
+    with *initializer* in Python. The accumulator is initialized with the value
+    returned by the provided init function; the logic of the accumulation is
+    given by the provided accumulate function.
+
+    Pseudo-code:
+    ```
+    accumulator = @initFuncRef()
+    while (next = upstream->next()):
+      accumulator = @accumulateFuncRef(accumulator, next->value())
+    return accumulator
+
+    Example:
+    ```mlir
+    %0 = iterators.accumulate(%input, @zero_struct, @sum_struct)
+          : (!iterators.stream<!llvm.struct<(i32)>>)
+            -> !iterators.stream<!llvm.struct<(i32)>>
+    ```
+  }];
+  let arguments = (ins
+      Iterators_Stream:$input,
+      Iterators_InitFunctionSymbol:$initFuncRef,
+      Iterators_AccumulateFunctionSymbol:$accumulateFuncRef
+    );
+  let results = (outs Iterators_Stream:$result);
+  let assemblyFormat = [{
+    `(` $input `,` $initFuncRef `,` $accumulateFuncRef `)`
+      attr-dict `:` `(` type($input) `)` `->` type($result)
+  }];
+  let extraClassDeclaration = [{
+    /// Return the init function op that the initFuncRef refers to.
+    func::FuncOp getInitFunc() {
+      return SymbolTable::lookupNearestSymbolFrom<func::FuncOp>(
+          *this, initFuncRefAttr());
+    }
+
+    /// Return the accumulate function op that the accumulateFuncRef refers to.
+    func::FuncOp getAccumulateFunc() {
+      return SymbolTable::lookupNearestSymbolFrom<func::FuncOp>(
+          *this, accumulateFuncRefAttr());
+    }
+  }];
+}
+
 /// Verifies that the element types of nested arrays in the $value array
 /// correspond to the types of the LLVM-struct element type of the $result
 /// Stream.
@@ -97,14 +212,6 @@ def Iterators_ConstantStreamOp
   let results = (outs Iterators_StreamOfLLVMStructOfNumerics:$result);
 }
 
-/// Looks up the given symbol, which must refer to a FuncOp, in the scope of the
-/// given op and returns the function type of that symbol.
-class LookupFuncType<string opName, string symbolName>
-  : StrFunc<"::mlir::SymbolTable::lookupNearestSymbolFrom<func::FuncOp>("
-            "    &$" # opName # ","
-            "    " # symbolName # ".dyn_cast<FlatSymbolRefAttr>())"
-            "       .getFunctionType()">;
-
 /// A predicate, that is, a function with a signature of the form (T) -> i1.
 class Iterators_PredicateFunctionType
   : Type<And<[FunctionType.predicate,