It works.

Author: wildarch

compiler/include/garel/GARelAttr.td

@@ -38,6 +38,7 @@ def AggregateFunc : I64EnumAttr<
         I64EnumAttrCase<"MAX", 2>,
         I64EnumAttrCase<"LOR", 3>, /* Logical OR (over i1) */
         I64EnumAttrCase<"ARGMIN", 4>,
+        I64EnumAttrCase<"COUNT", 5>,
     ]
 > {
     let cppNamespace = "::garel";
@@ -49,7 +50,7 @@ def Aggregator : GARel_Attr<"Aggregator", "aggregator"> {
 
     let parameters = (ins
         "AggregateFunc":$func,
-        ArrayRefParameter<"ColumnIdx">:$inputs);
+        OptionalArrayRefParameter<"ColumnIdx">:$inputs);
 
     let assemblyFormat = [{
         `<` $func $inputs `>`

compiler/include/garel/GARelOps.td

@@ -143,7 +143,7 @@ def ForOp : GARel_Op<"for", [InferTypeOpAdaptor]> {
 
     let arguments = (ins
         Variadic<Relation>:$init,
-        I64Attr:$iters,
+        I64Relation:$iters,
         I64Attr:$resultIdx);
 
     let regions = (region

compiler/include/garel/GARelTypes.h

@@ -1,5 +1,6 @@
 #pragma once
 
+#include <mlir/IR/MLIRContext.h>
 #include <mlir/IR/Types.h>
 
 #include "garel/GARelAttr.h"
@@ -11,4 +12,6 @@ namespace garel {
 
 bool isColumnType(mlir::Type t);
 
+RelationType getI64RelationType(mlir::MLIRContext *ctx);
+
 } // namespace garel

compiler/include/garel/GARelTypes.td

@@ -34,4 +34,10 @@ def Tuple : GARel_Type<"Tuple", "tuple"> {
 
 def ColumnType : Type<CPred<"::garel::isColumnType($_self)">, "column type">;
 
+def I64Relation : Type<
+        CPred<"::garel::getI64RelationType($_self.getContext()) == $_self">,
+        "relation with a single i64 column",
+        "RelationType">,
+    BuildableType<"::garel::getI64RelationType($_builder.getContext())">;
+
 #endif // GAREL_TYPES

compiler/src/garel/GARelAttr.cpp

@@ -24,23 +24,32 @@ mlir::Type AggregatorAttr::getResultType(mlir::Type inputRel) {
   case AggregateFunc::ARGMIN:
     // NOTE: argmin(arg, val) also uses first input column as output type.
     return llvm::cast<RelationType>(inputRel).getColumns()[getInputs()[0]];
+  case AggregateFunc::COUNT:
+    return mlir::IntegerType::get(inputRel.getContext(), 64);
+  }
+}
+
+static std::size_t expectedNumInputs(AggregateFunc f) {
+  switch (f) {
+  case AggregateFunc::SUM:
+  case AggregateFunc::MIN:
+  case AggregateFunc::MAX:
+  case AggregateFunc::LOR:
+    return 1;
+  case AggregateFunc::ARGMIN:
+    return 2;
+  case AggregateFunc::COUNT:
+    return 0;
   }
 }
 
 mlir::LogicalResult
 AggregatorAttr::verify(llvm::function_ref<mlir::InFlightDiagnostic()> emitError,
                        AggregateFunc func, llvm::ArrayRef<ColumnIdx> inputs) {
-  if (func == AggregateFunc::ARGMIN) {
-    if (inputs.size() != 2) {
-      return emitError() << stringifyAggregateFunc(func)
-                         << " expects exactly two inputs (arg, val), got "
-                         << inputs.size();
-    }
-  } else {
-    if (inputs.size() != 1) {
-      return emitError() << stringifyAggregateFunc(func)
-                         << " expects exactly one input, got " << inputs.size();
-    }
+  if (inputs.size() != expectedNumInputs(func)) {
+    return emitError() << stringifyAggregateFunc(func) << " expects exactly "
+                       << expectedNumInputs(func) << " inputs, got "
+                       << inputs.size();
   }
 
   return mlir::success();

compiler/src/garel/GARelTypes.cpp

@@ -1,6 +1,7 @@
 #include <llvm/ADT/ArrayRef.h>
 #include <llvm/ADT/TypeSwitch.h>
 #include <mlir/IR/Builders.h>
+#include <mlir/IR/BuiltinTypes.h>
 #include <mlir/IR/DialectImplementation.h>
 #include <mlir/IR/OpImplementation.h>
 
@@ -18,6 +19,11 @@ bool isColumnType(mlir::Type t) {
          t.isIndex();
 }
 
+RelationType getI64RelationType(mlir::MLIRContext *ctx) {
+  return RelationType::get(
+      ctx, mlir::ArrayRef<mlir::Type>{mlir::IntegerType::get(ctx, 64)});
+}
+
 // Need to define this here to avoid depending on IPRTypes in
 // IPRDialect and creating a cycle.
 void GARelDialect::registerTypes() {