New for loop design.

Author: wildarch

compiler/include/graphalg/GraphAlgOps.td

@@ -403,6 +403,70 @@ def BroadcastOp : Core_Op<"broadcast", [
     let hasVerifier = 1;
 }
 
+def ForOp : Core_Op<"for", [
+        Pure,
+        AttrSizedOperandSegments,
+        DeclareOpInterfaceMethods<RegionBranchOpInterface, ["getEntrySuccessorOperands"]>]> {
+    let summary = "For loop with dynamic 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.
+
+        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,
+        Optional<I64Scalar>:$dynBegin,
+        Optional<I64Scalar>:$dynEnd,
+        OptionalAttr<I64Attr>:$begin,
+        OptionalAttr<DimAttr>:$iters);
+
+    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^)?
+        `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();
+    }];
+}
+
 // Not core according to spec, but we don't want to unroll in the general case.
 def ForConstOp : Core_Op<"for_const", [
         Pure,
@@ -483,7 +547,7 @@ def ForDimOp : Core_Op<"for_dim", [
 def YieldOp : Core_Op<"yield", [
         Pure,
         Terminator,
-        ParentOneOf<["ForConstOp", "ForDimOp"]>,
+        ParentOneOf<["ForOp", "ForConstOp", "ForDimOp"]>,
         DeclareOpInterfaceMethods<RegionBranchTerminatorOpInterface>]> {
     let summary = "Yield from a loop body";
 

compiler/src/graphalg/GraphAlgCanonicalize.cpp

@@ -242,6 +242,39 @@ mlir::OpFoldResult BroadcastOp::fold(FoldAdaptor adaptor) {
   return nullptr;
 }
 
+mlir::LogicalResult
+ForOp::fold(FoldAdaptor adaptor,
+            ::llvm::SmallVectorImpl<::mlir::OpFoldResult> &results) {
+  if (!getBegin() && adaptor.getDynBegin()) {
+    // Can infer a constant begin to the range.
+    auto begin = llvm::cast<mlir::IntegerAttr>(adaptor.getDynBegin());
+    setBeginAttr(begin);
+    getDynBeginMutable().clear();
+    return mlir::success();
+  }
+
+  if (!getIters() && getBegin() && adaptor.getDynEnd()) {
+    // Can infer a constant number of iterations.
+    auto begin = *getBegin();
+    auto end = llvm::cast<mlir::IntegerAttr>(adaptor.getDynEnd())
+                   .getValue()
+                   .getZExtValue();
+    // If end < begin, drop to 0 iterations.
+    std::size_t iters = 0;
+    if (begin < end) {
+      iters = end - begin;
+    }
+
+    // NOTE: number of iterations is encoded as a DimAttr.
+    auto dim = DimAttr::getConcrete(getContext(), iters);
+    setItersAttr(dim);
+    getDynEndMutable().clear();
+    return mlir::success();
+  }
+
+  return mlir::failure();
+}
+
 static mlir::LogicalResult forDimConst(ForDimOp op,
                                        mlir::PatternRewriter &rewriter) {
   if (!op.getDim().isConcrete()) {

compiler/src/graphalg/GraphAlgOps.cpp

@@ -505,6 +505,41 @@ mlir::LogicalResult BroadcastOp::verify() {
   return mlir::success();
 }
 
+// === ForOp ===
+mlir::LogicalResult ForOp::verify() {
+  if (getDynBegin() && getBegin()) {
+    return emitOpError("begin and dyn_begin are mutually exclusive");
+  } else if (!getDynBegin() && !getBegin()) {
+    return emitOpError("no loop start: must have 'begin' or 'dyn_begin'");
+  }
+
+  if (getDynEnd() && getIters()) {
+    return emitOpError("begin and iters are mutually exclusive");
+  } else if (!getDynEnd() && !getIters()) {
+    return emitOpError("no loop end: must have 'iters' or 'dyn_end'");
+  }
+
+  return mlir::success();
+}
+
+mlir::LogicalResult ForOp::verifyRegions() {
+  return verifyLoop(getOperation(), getInitArgs(), getBody(), getUntil());
+}
+
+void ForOp::getSuccessorRegions(
+    mlir::RegionBranchPoint point,
+    llvm::SmallVectorImpl<mlir::RegionSuccessor> &regions) {
+  getLoopSuccessorRegions(getOperation(), getBody(), getUntil(), point,
+                          regions);
+}
+
+mlir::OperandRange
+ForOp::getEntrySuccessorOperands(mlir::RegionBranchPoint point) {
+  return getInitArgs();
+}
+
+bool ForOp::isDynamicRange() { return getDynBegin() || getDynEnd(); }
+
 // === ForConstOp ===
 mlir::LogicalResult ForConstOp::verifyRegions() {
   return verifyLoop(getOperation(), getInitArgs(), getBody(), getUntil());

compiler/src/graphalg/GraphAlgSetDimensions.cpp

@@ -125,20 +125,6 @@ class DimConversionPattern : public mlir::ConversionPattern {
                   mlir::ConversionPatternRewriter &rewriter) const override;
 };
 
-/** Template for rewrites (without type conversion). */
-template <typename T>
-class DimOpRewritePattern : public mlir::OpRewritePattern<T> {
-private:
-  const DimMapper &_mapper;
-
-  mlir::LogicalResult
-  matchAndRewrite(T op, mlir::PatternRewriter &rewriter) const override;
-
-public:
-  DimOpRewritePattern(const DimMapper &mapper, mlir::MLIRContext *ctx)
-      : mlir::OpRewritePattern<T>(ctx), _mapper(mapper) {}
-};
-
 } // namespace
 
 mlir::FailureOr<DimMapper>
@@ -309,31 +295,33 @@ mlir::LogicalResult DimConversionPattern::matchAndRewrite(
   return mlir::success();
 }
 
-template <>
-mlir::LogicalResult DimOpRewritePattern<CastDimOp>::matchAndRewrite(
-    CastDimOp op, mlir::PatternRewriter &rewriter) const {
-  auto dim = _mapper.convertAttr(op.getInput());
-  if (!dim) {
-    return mlir::failure();
+static mlir::LogicalResult updateDim(ForOp op, DimMapper &mapper) {
+  if (!op.getIters() || !op.getIters()->isAbstract()) {
+    // No update needed
+    return mlir::success();
   }
 
-  // The folder on CastDimOp should turn this into a constant.
-  auto newOp = rewriter.createOrFold<CastDimOp>(op->getLoc(), dim);
-  rewriter.replaceOp(op, newOp);
+  auto dim = mapper.convertAttr(*op.getIters());
+  if (!dim) {
+    return op.emitOpError("no mapping for ") << *op.getIters();
+  }
 
+  op.setItersAttr(dim);
   return mlir::success();
 }
 
-template <>
-mlir::LogicalResult DimOpRewritePattern<ForDimOp>::matchAndRewrite(
-    ForDimOp op, mlir::PatternRewriter &rewriter) const {
-  auto dim = _mapper.convertAttr(op.getDim());
-  if (!dim) {
-    return mlir::failure();
+static mlir::LogicalResult updateDim(CastDimOp op, DimMapper &mapper) {
+  if (!op.getInput().isAbstract()) {
+    // No update needed
+    return mlir::success();
   }
 
-  rewriter.modifyOpInPlace(op, [&]() { op.setDimAttr(dim); });
+  auto dim = mapper.convertAttr(op.getInput());
+  if (!dim) {
+    return op.emitOpError("no mapping for ") << op.getInput();
+  }
 
+  op.setInputAttr(dim);
   return mlir::success();
 }
 
@@ -356,6 +344,19 @@ void GraphAlgSetDimensions::runOnOperation() {
     return signalPassFailure();
   }
 
+  // Update direct references to dimensions.
+  bool failedDirectUpdate = false;
+  func->walk([&](ForOp op) {
+    if (mlir::failed(updateDim(op, *dimMapper))) {
+      failedDirectUpdate = true;
+    }
+  });
+  func->walk([&](CastDimOp op) {
+    if (mlir::failed(updateDim(op, *dimMapper))) {
+      failedDirectUpdate = true;
+    }
+  });
+
   mlir::ConversionTarget target(getContext());
   target.addDynamicallyLegalDialect<GraphAlgDialect>(
       doesNotUseAbstractDimensions);
@@ -374,12 +375,6 @@ void GraphAlgSetDimensions::runOnOperation() {
   // Convert all result types and block argument types.
   patterns.add<DimConversionPattern>(typeConverter, &getContext());
 
-  // Convert ops that have a special dependency on DimAttr.
-  patterns.add<DimOpRewritePattern<CastDimOp>, DimOpRewritePattern<ForDimOp>>(
-      *dimMapper, &getContext());
-  // Use the canonicalization pattern to rewrite ForDimOp into ForConstOp.
-  ForDimOp::getCanonicalizationPatterns(patterns, &getContext());
-
   if (mlir::failed(
           mlir::applyPartialConversion(func, target, std::move(patterns)))) {
     return signalPassFailure();

compiler/src/graphalg/GraphAlgToCore.cpp

@@ -254,6 +254,8 @@ void GraphAlgToCore::runOnOperation() {
   target.addIllegalDialect<graphalg::GraphAlgDialect>();
   target.addDynamicallyLegalDialect<graphalg::GraphAlgDialect>(
       [](mlir::Operation *op) { return op->hasTrait<IsCore>(); });
+  target.addDynamicallyLegalOp<ForOp>(
+      [](ForOp op) { return !op.isDynamicRange(); });
 
   mlir::RewritePatternSet patterns(&getContext());
   patterns.add(convertVecMatMul);
@@ -266,6 +268,20 @@ void GraphAlgToCore::runOnOperation() {
   patterns.add(convertTriu);
   patterns.add(convertLiteral);
 
+  // Conversion will give very unclear errors about dynamic range for loops, so
+  // do our own analysis first.
+  bool haveDynamicRangeLoops = false;
+  getOperation()->walk([&](ForOp op) {
+    if (op.isDynamicRange()) {
+      op.emitOpError("loop bound must be a constant in GraphAlg Core");
+      haveDynamicRangeLoops = true;
+    }
+  });
+
+  if (haveDynamicRangeLoops) {
+    return signalPassFailure();
+  }
+
   if (mlir::failed(mlir::applyFullConversion(getOperation(), target,
                                              std::move(patterns)))) {
     signalPassFailure();