Tests passing again.

Author: wildarch

compiler/include/graphalg/GraphAlgOps.td

@@ -467,87 +467,10 @@ 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,
-        AllTypesMatch<["rangeBegin", "rangeEnd"]>,
-        DeclareOpInterfaceMethods<RegionBranchOpInterface, ["getEntrySuccessorOperands"]>]> {
-    let summary = "For loop with constant bounds";
-
-    let description = [{
-        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
-        values of `initArgs`. For subsequent iterations, results from the
-        previous iteration (produced by `YieldOp`) are taken instead.
-        The `until` region, if present, is executed after `body`, and produces a
-        single boolean scalar indicating whether the loop should terminate
-        early.
-
-        In more imperative terms, `initArgs` can be seen as the set of variables
-        that are updated in the loop body.
-        Within the loop body, those variables can be accessed through the block
-        arguments, and their updated values are set through `YieldOp`.
-        Finally, `results` represents the new state of those variables after the
-        loop terminates.
-    }];
-
-    let arguments = (ins
-        Variadic<Matrix>:$initArgs,
-        I64Scalar:$rangeBegin,
-        I64Scalar:$rangeEnd);
-
-    let results = (outs Variadic<Matrix>:$results);
-
-    let regions = (region SizedRegion<1>:$body, MaxSizedRegion<1>:$until);
-
-    let assemblyFormat = [{
-        `range` `(`
-            $rangeBegin `,`
-            $rangeEnd
-        `)` `:` type($rangeEnd)
-        `init` `(` $initArgs `)` `:` type($initArgs) `->` type($results) attr-dict
-        `body` $body
-        `until` $until
-    }];
-
-    let hasRegionVerifier = 1;
-}
-
-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,
-        ParentOneOf<["ForOp", "ForConstOp", "ForDimOp"]>,
+        HasParent<"ForOp">,
         DeclareOpInterfaceMethods<RegionBranchTerminatorOpInterface>]> {
     let summary = "Yield from a loop body";
 

compiler/src/garel/GraphAlgToRel.cpp

@@ -320,6 +320,11 @@ static mlir::Value createDimRead(mlir::Location loc, graphalg::DimAttr dim,
  */
 static mlir::Value createDimInt(mlir::Location loc, graphalg::DimAttr dim,
                                 mlir::OpBuilder &builder) {
+  if (dim.isConcrete()) {
+    return builder.create<ConstantOp>(
+        loc, builder.getI64IntegerAttr(dim.getConcreteDim()));
+  }
+
   auto input = createDimRead(loc, dim, builder);
   llvm::ArrayRef<ColumnIdx> groupBy;
   std::array<AggregatorAttr, 1> aggregators{
@@ -889,71 +894,48 @@ mlir::LogicalResult OpConversion<graphalg::DiagOp>::matchAndRewrite(
   return mlir::success();
 }
 
-// Sharing logic between ForConstOp and ForDimOp
-static mlir::LogicalResult
-convertFor(mlir::Operation *op, mlir::ValueRange adaptorInitArgs,
-           mlir::Value rangeBegin, mlir::Value rangeEnd, mlir::Region &body,
-           mlir::Region &until, const mlir::TypeConverter *typeConverter,
-           mlir::ConversionPatternRewriter &rewriter) {
-  llvm::SmallVector<mlir::Value> initArgs{rangeBegin};
-  initArgs.append(adaptorInitArgs.begin(), adaptorInitArgs.end());
+template <>
+mlir::LogicalResult OpConversion<graphalg::ForOp>::matchAndRewrite(
+    graphalg::ForOp op, OpAdaptor adaptor,
+    mlir::ConversionPatternRewriter &rewriter) const {
+  auto begin =
+      rewriter.create<ConstantOp>(op.getLoc(), rewriter.getI64IntegerAttr(0));
+  auto iters = createDimInt(op.getLoc(), *op.getIters(), rewriter);
+  llvm::SmallVector<mlir::Value> initArgs{begin};
+  initArgs.append(adaptor.getInitArgs().begin(), adaptor.getInitArgs().end());
 
-  auto blockSignature = typeConverter->convertBlockSignature(&body.front());
+  auto blockSignature =
+      typeConverter->convertBlockSignature(&op.getBody().front());
   if (!blockSignature) {
     return op->emitOpError("Failed to convert iter args");
   }
 
-  mlir::Value iters;
-  if (isConstantZeroI64(rangeBegin)) {
-    iters = rangeEnd;
-  } else {
-    // Subtract rangeBegin from rangeEnd
-    auto joinOp = rewriter.create<JoinOp>(
-        op->getLoc(), mlir::ValueRange{rangeBegin, rangeEnd},
-        rewriter.getAttr<JoinPredicatesAttr>(
-            llvm::ArrayRef<JoinPredicateAttr>{}));
-    auto projOp = rewriter.create<ProjectOp>(
-        op->getLoc(), getI64RelationType(op->getContext()), joinOp);
-
-    auto &block = projOp.createProjectionsBlock();
-    mlir::OpBuilder::InsertionGuard guard{rewriter};
-    rewriter.setInsertionPointToStart(&block);
-
-    auto begin =
-        rewriter.create<ExtractOp>(op->getLoc(), 0, block.getArgument(0));
-    auto end =
-        rewriter.create<ExtractOp>(op->getLoc(), 1, block.getArgument(0));
-    auto res = rewriter.create<mlir::arith::SubIOp>(op->getLoc(), end, begin);
-    rewriter.create<ProjectReturnOp>(op->getLoc(), mlir::ValueRange{res});
-
-    iters = projOp;
-  }
-
   // The relational version of this op can only have a single output value.
   // For loops with multiple results, duplicate.
   llvm::SmallVector<mlir::Value> resultValues;
   for (auto i : llvm::seq(op->getNumResults())) {
     auto result = op->getResult(i);
     if (result.use_empty()) {
       // Not used. Take init arg as a dummy value.
-      resultValues.push_back(adaptorInitArgs[i]);
+      resultValues.push_back(adaptor.getInitArgs()[i]);
       continue;
     }
 
     // We are adding the iteration count variable as a first argument, so offset
     // the result index accordingly.
     std::int64_t resultIdx = i + 1;
-    auto resultType = adaptorInitArgs[i].getType();
+    auto resultType = adaptor.getInitArgs()[i].getType();
     auto forOp = rewriter.create<ForOp>(op->getLoc(), resultType, initArgs,
                                         iters, resultIdx);
     // body block
-    rewriter.cloneRegionBefore(body, forOp.getBody(), forOp.getBody().begin());
+    rewriter.cloneRegionBefore(op.getBody(), forOp.getBody(),
+                               forOp.getBody().begin());
     rewriter.applySignatureConversion(&forOp.getBody().front(),
                                       *blockSignature);
 
     // until block
-    if (!until.empty()) {
-      rewriter.cloneRegionBefore(until, forOp.getUntil(),
+    if (!op.getUntil().empty()) {
+      rewriter.cloneRegionBefore(op.getUntil(), forOp.getUntil(),
                                  forOp.getUntil().begin());
       rewriter.applySignatureConversion(&forOp.getUntil().front(),
                                         *blockSignature);
@@ -966,28 +948,6 @@ convertFor(mlir::Operation *op, mlir::ValueRange adaptorInitArgs,
   return mlir::success();
 }
 
-template <>
-mlir::LogicalResult OpConversion<graphalg::ForConstOp>::matchAndRewrite(
-    graphalg::ForConstOp op, OpAdaptor adaptor,
-    mlir::ConversionPatternRewriter &rewriter) const {
-  return convertFor(op, adaptor.getInitArgs(), adaptor.getRangeBegin(),
-                    adaptor.getRangeEnd(), op.getBody(), op.getUntil(),
-                    typeConverter, rewriter);
-}
-
-template <>
-mlir::LogicalResult OpConversion<graphalg::ForDimOp>::matchAndRewrite(
-    graphalg::ForDimOp op, OpAdaptor adaptor,
-    mlir::ConversionPatternRewriter &rewriter) const {
-  auto ctx = op.getContext();
-  auto rangeBegin =
-      rewriter.create<ConstantOp>(op.getLoc(), rewriter.getI64IntegerAttr(0));
-  auto rangeEnd = createDimInt(op.getLoc(), op.getDim(), rewriter);
-
-  return convertFor(op, adaptor.getInitArgs(), rangeBegin, rangeEnd,
-                    op.getBody(), op.getUntil(), typeConverter, rewriter);
-}
-
 template <>
 mlir::LogicalResult OpConversion<graphalg::YieldOp>::matchAndRewrite(
     graphalg::YieldOp op, OpAdaptor adaptor,
@@ -1481,11 +1441,10 @@ void GraphAlgToRel::runOnOperation() {
       OpConversion<graphalg::TransposeOp>, OpConversion<graphalg::BroadcastOp>,
       OpConversion<graphalg::ConstantMatrixOp>,
       OpConversion<graphalg::DeferredReduceOp>, OpConversion<graphalg::DiagOp>,
-      OpConversion<graphalg::ForConstOp>, OpConversion<graphalg::ForDimOp>,
-      OpConversion<graphalg::YieldOp>, OpConversion<graphalg::MatMulJoinOp>,
-      OpConversion<graphalg::PickAnyOp>, OpConversion<graphalg::TrilOp>,
-      OpConversion<graphalg::UnionOp>, OpConversion<graphalg::CastDimOp>>(
-      matrixTypeConverter, &getContext());
+      OpConversion<graphalg::ForOp>, OpConversion<graphalg::YieldOp>,
+      OpConversion<graphalg::MatMulJoinOp>, OpConversion<graphalg::PickAnyOp>,
+      OpConversion<graphalg::TrilOp>, OpConversion<graphalg::UnionOp>,
+      OpConversion<graphalg::CastDimOp>>(matrixTypeConverter, &getContext());
   patterns.add<ApplyOpConversion>(semiringTypeConverter, matrixTypeConverter,
                                   &getContext());
 

compiler/src/graphalg/GraphAlgCanonicalize.cpp

@@ -272,38 +272,9 @@ ForOp::fold(FoldAdaptor adaptor,
     return mlir::success();
   }
 
-  return mlir::failure();
-}
-
-static mlir::LogicalResult forDimConst(ForDimOp op,
-                                       mlir::PatternRewriter &rewriter) {
-  if (!op.getDim().isConcrete()) {
-    return mlir::failure();
-  }
-
-  // The number of iterations is known, so we can replace with a ForConstOp.
-  auto end = op.getDim().getConcreteDim();
-
-  // Range from 0 to dim.
-  auto intType =
-      MatrixType::scalarOf(SemiringTypes::forInt(rewriter.getContext()));
-  auto beginOp = rewriter.create<ConstantMatrixOp>(
-      op->getLoc(), intType, rewriter.getI64IntegerAttr(0));
-  auto endOp = rewriter.create<ConstantMatrixOp>(
-      op->getLoc(), intType, rewriter.getI64IntegerAttr(end));
+  // TODO: Fold if iters=0
 
-  auto forConstOp = rewriter.create<ForConstOp>(
-      op->getLoc(), op->getResultTypes(), op.getInitArgs(), beginOp, endOp);
-  rewriter.inlineRegionBefore(op.getBody(), forConstOp.getBody(),
-                              forConstOp.getBody().begin());
-  rewriter.replaceOp(op, forConstOp);
-
-  return mlir::success();
-}
-
-void ForDimOp::getCanonicalizationPatterns(mlir::RewritePatternSet &patterns,
-                                           mlir::MLIRContext *context) {
-  patterns.add(forDimConst);
+  return mlir::failure();
 }
 
 mlir::OpFoldResult PickAnyOp::fold(FoldAdaptor adaptor) {

compiler/src/graphalg/GraphAlgLoopAggregate.cpp

@@ -28,19 +28,10 @@ class GraphAlgLoopAggregate
 // If the loop body ends with an aggregation op, ensure the init arg is also an
 // aggregation. Later on in the AvantGraph query pipeline, this will signal to
 // the optimizer that the iteration state can be kept as an aggregate table.
-static void addInitReduce(mlir::Operation *op, mlir::IRRewriter &rewriter) {
-  mlir::Block *body;
-  llvm::SmallVector<mlir::Value> newInitArgs;
-  if (auto constOp = llvm::dyn_cast<ForConstOp>(op)) {
-    body = &constOp.getBody().front();
-    newInitArgs = constOp.getInitArgs();
-  } else {
-    auto dimOp = llvm::cast<ForDimOp>(op);
-    body = &dimOp.getBody().front();
-    newInitArgs = dimOp.getInitArgs();
-  }
+static void addInitReduce(ForOp op, mlir::IRRewriter &rewriter) {
+  llvm::SmallVector<mlir::Value> newInitArgs(op.getInitArgs());
 
-  auto yieldOp = llvm::cast<YieldOp>(body->getTerminator());
+  auto yieldOp = llvm::cast<YieldOp>(op.getBody().front().getTerminator());
   for (auto [i, iterResult] : llvm::enumerate(yieldOp.getInputs())) {
     auto iterLastOp = iterResult.getDefiningOp();
     if (llvm::isa_and_present<PickAnyOp, DeferredReduceOp>(iterLastOp)) {
@@ -50,20 +41,13 @@ static void addInitReduce(mlir::Operation *op, mlir::IRRewriter &rewriter) {
     }
   }
 
-  rewriter.modifyOpInPlace(op, [&]() {
-    if (auto constOp = llvm::dyn_cast<ForConstOp>(op)) {
-      constOp.getInitArgsMutable().assign(newInitArgs);
-    } else {
-      auto dimOp = llvm::cast<ForDimOp>(op);
-      dimOp.getInitArgsMutable().assign(newInitArgs);
-    }
-  });
+  rewriter.modifyOpInPlace(
+      op, [&]() { op.getInitArgsMutable().assign(newInitArgs); });
 }
 
 void GraphAlgLoopAggregate::runOnOperation() {
-  llvm::SmallVector<mlir::Operation *> loopOps;
-  getOperation()->walk([&](ForConstOp op) { loopOps.emplace_back(op); });
-  getOperation()->walk([&](ForDimOp op) { loopOps.emplace_back(op); });
+  llvm::SmallVector<ForOp> loopOps;
+  getOperation()->walk([&](ForOp op) { loopOps.emplace_back(op); });
 
   mlir::IRRewriter rewriter(&getContext());
   for (auto op : loopOps) {