Revert "Don't require setting dimensions when lowering to garel (#22)"

This reverts commit 56bf019.

Author: wildarch

compiler/include/garel/GARelAttr.td

@@ -38,7 +38,6 @@ def AggregateFunc : I64EnumAttr<
         I64EnumAttrCase<"MAX", 2>,
         I64EnumAttrCase<"LOR", 3>, /* Logical OR (over i1) */
         I64EnumAttrCase<"ARGMIN", 4>,
-        I64EnumAttrCase<"COUNT", 5>,
     ]
 > {
     let cppNamespace = "::garel";
@@ -50,7 +49,7 @@ def Aggregator : GARel_Attr<"Aggregator", "aggregator"> {
 
     let parameters = (ins
         "AggregateFunc":$func,
-        OptionalArrayRefParameter<"ColumnIdx">:$inputs);
+        ArrayRefParameter<"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,
-        I64Relation:$iters,
+        I64Attr:$iters,
         I64Attr:$resultIdx);
 
     let regions = (region

compiler/include/garel/GARelTypes.h

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

compiler/include/garel/GARelTypes.td

@@ -34,10 +34,4 @@ 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/include/graphalg/GraphAlgOps.td

@@ -403,22 +403,16 @@ def BroadcastOp : Core_Op<"broadcast", [
     let hasVerifier = 1;
 }
 
-def ForOp : Core_Op<"for", [
+// Not core according to spec, but we don't want to unroll in the general case.
+def ForConstOp : Core_Op<"for_const", [
         Pure,
-        AttrSizedOperandSegments,
+        AllTypesMatch<["rangeBegin", "rangeEnd"]>,
         DeclareOpInterfaceMethods<RegionBranchOpInterface, ["getEntrySuccessorOperands"]>]> {
-    let summary = "For loop with dynamic bounds";
+    let summary = "For loop with constant bounds";
 
     let description = [{
-        A loop iterating over one of three ranges:
-        1) `dynBegin` (inclusive) to `dynEnd` (exclusive)
-        2) `begin` to `begin` + `iters`, where `iters` is an integer
-        3) `begin` to `begin` + `iters`, where `iters` is a matrix dimension
-
-        Only instances with range types 2 or 3 are considered part of GraphAlg
-        Core. Constant propagation is expected to transform range type 1 into
-        either 2 or 3.
-
+        A loop iterating over the integer range starting at `rangeBegin`
+        (inclusive) and ending at `rangeEnd` (exclusive).
         The `body` region is executed once for every value in the integer range
         (that value is passed as the first block argument).
         At the first iteration of the loop, the other block arguments take the
@@ -438,39 +432,58 @@ def ForOp : Core_Op<"for", [
 
     let arguments = (ins
         Variadic<Matrix>:$initArgs,
-        Optional<I64Scalar>:$dynBegin,
-        Optional<I64Scalar>:$dynEnd,
-        OptionalAttr<I64Attr>:$begin,
-        OptionalAttr<DimAttr>:$iters);
+        I64Scalar:$rangeBegin,
+        I64Scalar:$rangeEnd);
 
     let results = (outs Variadic<Matrix>:$results);
 
     let regions = (region SizedRegion<1>:$body, MaxSizedRegion<1>:$until);
 
     let assemblyFormat = [{
-        (`dyn_begin` `` `=` `` $dynBegin^)?
-        (`dyn_end` `` `=` `` $dynEnd^)?
-        (`begin` `` `=` `` $begin^)?
-        (`iters` `` `=` `` $iters^)?
+        `range` `(`
+            $rangeBegin `,`
+            $rangeEnd
+        `)` `:` type($rangeEnd)
         `init` `(` $initArgs `)` `:` type($initArgs) `->` type($results) attr-dict
         `body` $body
         `until` $until
     }];
 
-    let hasVerifier = 1;
     let hasRegionVerifier = 1;
-    let hasFolder = 1;
+}
 
-    let extraClassDeclaration = [{
-        /** Whether at least one of `dyn_begin` and `dyn_end` is set. */
-        bool isDynamicRange();
+def ForDimOp : Core_Op<"for_dim", [
+        Pure,
+        DeclareOpInterfaceMethods<RegionBranchOpInterface, ["getEntrySuccessorOperands"]>]> {
+    let summary = "For loop over a matrix dimension";
+
+    let description = [{
+        A loop iterating over the half-open range [0..`dim`).
+
+        This op is otherwise equivalent to `ForConstOp`.
     }];
+
+    let arguments = (ins Variadic<Matrix>:$initArgs, DimAttr:$dim);
+
+    let results = (outs Variadic<Matrix>:$results);
+
+    let regions = (region SizedRegion<1>:$body, MaxSizedRegion<1>:$until);
+
+    let assemblyFormat = [{
+        `range` `(` custom<BareAttr>($dim) `)`
+        `init` `(` $initArgs `)` `:` type($initArgs) `->` type($results) attr-dict
+        `body` $body
+        `until` $until
+    }];
+
+    let hasRegionVerifier = 1;
+    let hasCanonicalizer = 1;
 }
 
 def YieldOp : Core_Op<"yield", [
         Pure,
         Terminator,
-        HasParent<"ForOp">,
+        ParentOneOf<["ForConstOp", "ForDimOp"]>,
         DeclareOpInterfaceMethods<RegionBranchTerminatorOpInterface>]> {
     let summary = "Yield from a loop body";
 

compiler/include/graphalg/GraphAlgPasses.td

@@ -3,16 +3,6 @@
 
 include "mlir/Pass/PassBase.td"
 
-def GraphAlgSetConstArg : Pass<"graphalg-set-const-arg", "::mlir::ModuleOp"> {
-    let summary = "Propagate a constant integer argument into a function";
-
-    let options = [
-        Option<"functionName", "func", "std::string", /*default=*/"\"\"", "Name of the function to call">,
-        Option<"argumentNumber", "argNum", "int", /*default=*/"-1", "The argument number that is constant">,
-        Option<"value", "value", "std::int64_t", /*default=*/"0", "The value to propagate">,
-    ];
-}
-
 def GraphAlgPrepareInline : Pass<"graphalg-prepare-inline", "::mlir::ModuleOp"> {
     let summary = "Prepares the IR for function inlining";
 

compiler/include/graphalg/GraphAlgSetConstArg.h

@@ -24,32 +24,23 @@ 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 (inputs.size() != expectedNumInputs(func)) {
-    return emitError() << stringifyAggregateFunc(func) << " expects exactly "
-                       << expectedNumInputs(func) << " inputs, got "
-                       << inputs.size();
+  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();
+    }
   }
 
   return mlir::success();

compiler/src/garel/GARelAttr.cpp

@@ -1,7 +1,6 @@
 #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>
 
@@ -19,11 +18,6 @@ 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() {