Verify loops bounds outside of core conversion.

Author: wildarch

compiler/include/graphalg/GraphAlgPasses.td

@@ -33,6 +33,18 @@ def GraphAlgToCore : Pass<"graphalg-to-core", "::mlir::func::FuncOp"> {
     let summary = "Desugar GraphAlg IR into the Core operations";
 }
 
+def GraphAlgVerifyLoopBounds : Pass<"graphalg-verify-loop-bounds", "::mlir::func::FuncOp"> {
+    let summary = "Verifies that loops do not use dynamic ranges";
+
+    let description = [{
+        ForOps can specify dynamic loop bounds that depend on arbitrary values
+        in the program. This is useful during parsing, and it allows using
+        parameters that will later be inlined as constants. GraphAlg Core,
+        however, requires that loop bounds are either a constant integer value
+        or the dimensions of a matrix. This pass enforces that requirement.
+    }];
+}
+
 def GraphAlgExplicateSparsity : Pass<"graphalg-explicate-sparsity", "mlir::func::FuncOp"> {
     let summary = "Makes explicit the sparsity of matrices";
 }

compiler/src/graphalg/CMakeLists.txt

@@ -43,6 +43,7 @@ add_library(GraphAlgPasses
     GraphAlgToCore.cpp
     GraphAlgToCorePipeline.cpp
     GraphAlgVerifyDimensions.cpp
+    GraphAlgVerifyLoopBounds.cpp
 )
 add_dependencies(GraphAlgPasses MLIRGraphAlgPassesIncGen)
 target_link_libraries(

compiler/src/graphalg/GraphAlgToCore.cpp

@@ -254,8 +254,6 @@ 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);
@@ -268,20 +266,6 @@ 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();

compiler/src/graphalg/GraphAlgVerifyLoopBounds.cpp

@@ -0,0 +1,32 @@
+#include <mlir/Dialect/Func/IR/FuncOps.h>
+#include <mlir/Pass/Pass.h>
+
+#include "graphalg/GraphAlgOps.h"
+
+namespace graphalg {
+
+#define GEN_PASS_DEF_GRAPHALGVERIFYLOOPBOUNDS
+#include "graphalg/GraphAlgPasses.h.inc"
+
+namespace {
+
+class GraphAlgVerifyLoopBounds
+    : public impl::GraphAlgVerifyLoopBoundsBase<GraphAlgVerifyLoopBounds> {
+  using impl::GraphAlgVerifyLoopBoundsBase<
+      GraphAlgVerifyLoopBounds>::GraphAlgVerifyLoopBoundsBase;
+
+  void runOnOperation() final;
+};
+
+} // namespace
+
+void GraphAlgVerifyLoopBounds::runOnOperation() {
+  getOperation()->walk([&](ForOp op) {
+    if (op.isDynamicRange()) {
+      op.emitOpError("loop bound is not a constant or matrix dimension");
+      signalPassFailure();
+    }
+  });
+}
+
+} // namespace graphalg

compiler/test/verify-loop-bounds/for.mlir

@@ -0,0 +1,36 @@
+// RUN: graphalg-opt --graphalg-verify-loop-bounds --split-input-file --verify-diagnostics < %s
+#dim = #graphalg.dim<distinct[0]<>>
+
+func.func @OkDim(%arg0: !graphalg.mat<#dim x 1 x i64>) -> !graphalg.mat<1 x 1 x i64> {
+  %0 = graphalg.const_mat 0 : i64 -> <1 x 1 x i64>
+  %1 = graphalg.for begin=0 iters=#dim init(%0) : !graphalg.mat<1 x 1 x i64> -> !graphalg.mat<1 x 1 x i64> body {
+  ^bb0(%arg1: !graphalg.mat<1 x 1 x i64>, %arg2: !graphalg.mat<1 x 1 x i64>):
+    graphalg.yield %arg2 : !graphalg.mat<1 x 1 x i64>
+  } until {
+  }
+  return %1 : !graphalg.mat<1 x 1 x i64>
+}
+
+func.func @OkConst(%arg0: !graphalg.mat<#dim x 1 x i64>) -> !graphalg.mat<1 x 1 x i64> {
+  %0 = graphalg.const_mat 0 : i64 -> <1 x 1 x i64>
+  %1 = graphalg.for begin=0 iters=<42> init(%0) : !graphalg.mat<1 x 1 x i64> -> !graphalg.mat<1 x 1 x i64> body {
+  ^bb0(%arg1: !graphalg.mat<1 x 1 x i64>, %arg2: !graphalg.mat<1 x 1 x i64>):
+    graphalg.yield %arg2 : !graphalg.mat<1 x 1 x i64>
+  } until {
+  }
+  return %1 : !graphalg.mat<1 x 1 x i64>
+}
+
+// -----
+#dim = #graphalg.dim<distinct[0]<>>
+
+func.func @DynDim(%arg0 : !graphalg.mat<1 x 1 x i64>) -> !graphalg.mat<1 x 1 x i64> {
+  %0 = graphalg.const_mat 0 : i64 -> <1 x 1 x i64>
+  // expected-error@below{{'graphalg.for' op loop bound is not a constant or matrix dimension}}
+  %1 = graphalg.for dyn_end=%arg0 begin=0 init(%0) : !graphalg.mat<1 x 1 x i64> -> !graphalg.mat<1 x 1 x i64> body {
+  ^bb0(%arg1: !graphalg.mat<1 x 1 x i64>, %arg2: !graphalg.mat<1 x 1 x i64>):
+    graphalg.yield %arg2 : !graphalg.mat<1 x 1 x i64>
+  } until {
+  }
+  return %1 : !graphalg.mat<1 x 1 x i64>
+}

playground/cpp/graphalg-playground.cpp

@@ -25,6 +25,7 @@
 #include "graphalg/GraphAlgAttr.h"
 #include "graphalg/GraphAlgDialect.h"
 #include "graphalg/GraphAlgPasses.h"
+#include "graphalg/GraphAlgSetConstArg.h"
 #include "graphalg/GraphAlgTypes.h"
 #include "graphalg/SemiringTypes.h"
 #include "graphalg/evaluate/Evaluator.h"
@@ -205,6 +206,31 @@ bool Playground::evaluate() {
   }
 
   _argBuilders.clear();
+
+  // Inline constant arguments.
+  llvm::SmallVector<mlir::TypedAttr> constArgs;
+  for (auto arg : args) {
+    graphalg::MatrixAttrReader reader(arg);
+    if (reader.nRows() == 1 && reader.nCols() == 1) {
+      constArgs.push_back(reader.at(0, 0));
+    } else {
+      // Not constant
+      constArgs.push_back(mlir::TypedAttr());
+    }
+  }
+
+  if (mlir::failed(graphalg::setConstantArguments(_funcOp, constArgs))) {
+    return false;
+  }
+
+  // Verify that loop bounds are OK now.
+  mlir::PassManager pm(&_ctx);
+  pm.addPass(mlir::createCanonicalizerPass()); // To propagate constants
+  pm.addPass(graphalg::createGraphAlgVerifyLoopBounds());
+  if (mlir::failed(pm.run(_funcOp))) {
+    return false;
+  }
+
   _result = graphalg::evaluate(_funcOp, args);
   if (!_result) {
     return false;