[StaticLogic] Add optional trip count attribute. (#2657)

This is a useful piece of high-level information that is intrinsic to
the pipeline. This could be used by the Calyx lowering to easily
generate a @bound annotation when possible. It could also be useful to
generalize the Affine lowering to nested loops, by iteratively
applying the current analysis and transformation (#2659).

include/circt/Dialect/StaticLogic/StaticLogic.td

@@ -117,6 +117,7 @@ def PipelineWhileOp : Op<StaticLogic_Dialect, "pipeline.while", []> {
 
   let arguments = (ins
     I64Attr:$II,
+    OptionalAttr<I64Attr>:$tripCount,
     Variadic<AnyType>:$iterArgs
   );
 
@@ -138,7 +139,8 @@ def PipelineWhileOp : Op<StaticLogic_Dialect, "pipeline.while", []> {
   let skipDefaultBuilders = 1;
   let builders = [
     OpBuilder<(ins "mlir::TypeRange":$resultTypes, "mlir::IntegerAttr":$II,
-                   "mlir::ValueRange":$iterArgs)>
+                   "llvm::Optional<IntegerAttr>": $tripCount,
+		   "mlir::ValueRange":$iterArgs)>
   ];
 
   let extraClassDeclaration = [{

lib/Conversion/AffineToStaticLogic/AffineToStaticLogic.cpp

@@ -15,6 +15,7 @@
 #include "circt/Scheduling/Problems.h"
 #include "mlir/Conversion/AffineToStandard/AffineToStandard.h"
 #include "mlir/Dialect/Affine/Analysis/AffineAnalysis.h"
+#include "mlir/Dialect/Affine/Analysis/LoopAnalysis.h"
 #include "mlir/Dialect/Affine/IR/AffineMemoryOpInterfaces.h"
 #include "mlir/Dialect/Affine/IR/AffineOps.h"
 #include "mlir/Dialect/Affine/LoopUtils.h"
@@ -24,13 +25,13 @@
 #include "mlir/Dialect/SCF/SCF.h"
 #include "mlir/Dialect/StandardOps/IR/Ops.h"
 #include "mlir/IR/BlockAndValueMapping.h"
+#include "mlir/IR/BuiltinDialect.h"
 #include "mlir/IR/Dominance.h"
 #include "mlir/IR/ImplicitLocOpBuilder.h"
 #include "mlir/Transforms/DialectConversion.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/TypeSwitch.h"
 #include "llvm/Support/Debug.h"
-#include <mlir/IR/BuiltinDialect.h>
 
 #define DEBUG_TYPE "affine-to-staticlogic"
 
@@ -372,7 +373,14 @@ LogicalResult AffineToStaticLogic::createStaticLogicPipeline(
   iterArgs.append(innerLoop.getIterOperands().begin(),
                   innerLoop.getIterOperands().end());
 
-  auto pipeline = builder.create<PipelineWhileOp>(resultTypes, ii, iterArgs);
+  // If possible, attach a constant trip count attribute. This could be
+  // generalized to support non-constant trip counts by supporting an AffineMap.
+  Optional<IntegerAttr> tripCountAttr;
+  if (auto tripCount = getConstantTripCount(forOp))
+    tripCountAttr = builder.getI64IntegerAttr(*tripCount);
+
+  auto pipeline =
+      builder.create<PipelineWhileOp>(resultTypes, ii, tripCountAttr, iterArgs);
 
   // Create the condition, which currently just compares the induction variable
   // to the upper bound.

lib/Dialect/StaticLogic/StaticLogicOps.cpp

@@ -34,6 +34,14 @@ ParseResult PipelineWhileOp::parse(OpAsmParser &parser,
     return failure();
   result.addAttribute("II", ii);
 
+  // Parse optional trip count.
+  if (succeeded(parser.parseOptionalKeyword("trip_count"))) {
+    IntegerAttr tripCount;
+    if (parser.parseEqual() || parser.parseAttribute(tripCount))
+      return failure();
+    result.addAttribute("tripCount", tripCount);
+  }
+
   // Parse iter_args assignment list.
   SmallVector<OpAsmParser::OperandType> regionArgs, operands;
   if (succeeded(parser.parseOptionalKeyword("iter_args"))) {
@@ -71,6 +79,10 @@ void PipelineWhileOp::print(OpAsmPrinter &p) {
   // Print the initiation interval.
   p << " II = " << ' ' << II();
 
+  // Print the optional tripCount.
+  if (tripCount())
+    p << " trip_count = " << ' ' << *tripCount();
+
   // Print iter_args assignment list.
   p << " iter_args(";
   llvm::interleaveComma(
@@ -94,6 +106,11 @@ void PipelineWhileOp::print(OpAsmPrinter &p) {
 }
 
 static LogicalResult verifyPipelineWhileOp(PipelineWhileOp op) {
+  // Verify trip count is not negative.
+  if (op.tripCount() && *op.tripCount() < 0)
+    return op.emitOpError("trip count must not be negative, found ")
+           << *op.tripCount();
+
   // Verify the condition block is "combinational" based on an allowlist of
   // Arithmetic ops.
   Block &conditionBlock = op.condition().front();
@@ -167,11 +184,14 @@ static LogicalResult verifyPipelineWhileOp(PipelineWhileOp op) {
 
 void PipelineWhileOp::build(OpBuilder &builder, OperationState &state,
                             TypeRange resultTypes, IntegerAttr ii,
+                            Optional<IntegerAttr> tripCount,
                             ValueRange iterArgs) {
   OpBuilder::InsertionGuard g(builder);
 
   state.addTypes(resultTypes);
   state.addAttribute("II", ii);
+  if (tripCount)
+    state.addAttribute("tripCount", *tripCount);
   state.addOperands(iterArgs);
 
   Region *condRegion = state.addRegion();

test/Conversion/AffineToStaticLogic/loops.mlir

@@ -4,11 +4,11 @@
 func @minimal(%arg0 : memref<10xindex>) {
   // Setup constants.
   // CHECK: %[[LB:.+]] = arith.constant 0 : [[ITER_TYPE:.+]]
-  // CHECK: %[[UB:.+]] = arith.constant 10 : [[ITER_TYPE]]
+  // CHECK: %[[UB:.+]] = arith.constant [[TRIP_COUNT:.+]] : [[ITER_TYPE]]
   // CHECK: %[[STEP:.+]] = arith.constant 1 : [[ITER_TYPE]]
 
   // Pipeline header.
-  // CHECK: staticlogic.pipeline.while II = 1 iter_args(%[[ITER_ARG:.+]] = %[[LB]]) : ([[ITER_TYPE]]) -> ()
+  // CHECK: staticlogic.pipeline.while II = 1 trip_count = [[TRIP_COUNT]] iter_args(%[[ITER_ARG:.+]] = %[[LB]]) : ([[ITER_TYPE]]) -> ()
 
   // Condition block.
   // CHECK: %[[COND_RESULT:.+]] = arith.cmpi ult, %[[ITER_ARG]]

test/Dialect/StaticLogic/round-trip.mlir

@@ -180,3 +180,17 @@ func @test5(%arg0: memref<?xi32>) {
   }
   return
 }
+
+func @trip_count_attr() {
+  %false = arith.constant 0 : i1
+  // CHECK: staticlogic.pipeline.while II = 1 trip_count = 3
+  staticlogic.pipeline.while II = 1 trip_count = 3 iter_args(%arg0 = %false) : (i1) -> () {
+    staticlogic.pipeline.register %arg0 : i1
+  } do {
+    %0 = staticlogic.pipeline.stage start = 0 {
+      staticlogic.pipeline.register %arg0 : i1
+    } : i1
+    staticlogic.pipeline.terminator iter_args(%0), results() : (i1) -> ()
+  }
+  return
+}