[Arc] Initial implementation for Arc-based state-directed partitioning

include/circt/Dialect/Arc/ArcOps.td

@@ -464,6 +464,21 @@ def FinalOp : ClockTreeLikeOp<"final"> {
   }];
 }
 
+//===----------------------------------------------------------------------===//
+// Partition
+//===----------------------------------------------------------------------===//
+
+def TaskOp : ArcOp<"task",  [
+    RecursiveMemoryEffects, NoTerminator, NoRegionArguments, SingleBlock,
+    // ParentOneOf<["ClockTreeOp", "ModelOp"]>
+]> {
+  let summary = "A region reserved after the computation of different partitions";
+  let arguments = (ins
+    OptionalAttr<SymbolNameAttr>:$task_name
+  );
+  let regions = (region SizedRegion<1>:$body);
+}
+
 //===----------------------------------------------------------------------===//
 // Storage Allocation
 //===----------------------------------------------------------------------===//

include/circt/Dialect/Arc/ArcPasses.h

@@ -43,6 +43,7 @@ std::unique_ptr<mlir::Pass> createLowerVectorizationsPass(
     LowerVectorizationsModeEnum mode = LowerVectorizationsModeEnum::Full);
 std::unique_ptr<mlir::Pass> createMakeTablesPass();
 std::unique_ptr<mlir::Pass> createMuxToControlFlowPass();
+std::unique_ptr<mlir::Pass> createPartitionPass(const PartitionOptions &opts = {});
 std::unique_ptr<mlir::Pass> createPrintCostModelPass();
 std::unique_ptr<mlir::Pass> createSimplifyVariadicOpsPass();
 std::unique_ptr<mlir::Pass> createSplitLoopsPass();

include/circt/Dialect/Arc/ArcPasses.td

@@ -318,6 +318,18 @@ def SimplifyVariadicOps : Pass<"arc-simplify-variadic-ops", "mlir::ModuleOp"> {
   ];
 }
 
+def Partition: Pass<"arc-partition", "arc::ModelOp"> {
+  let summary = "Partition the entire circuit into multiple chunks";
+  let options = [
+    Option<"chunks", "chunks", "unsigned", "1",
+      "Number of resulting chunks">
+  ];
+}
+
+def PartitionClone: Pass<"arc-partition-clone", "arc::ModelOp"> {
+  let summary = "Commit the planned partition by cloning models";
+}
+
 def SplitFuncs : Pass<"arc-split-funcs", "mlir::ModuleOp"> {
   let summary = "Split large funcs into multiple smaller funcs";
   let dependentDialects = ["mlir::func::FuncDialect"];

lib/Dialect/Arc/Transforms/AllocateState.cpp

@@ -62,10 +62,27 @@ void AllocateStatePass::allocateBlock(Block *block) {
   }
   LLVM_DEBUG(llvm::dbgs() << "- Visiting block in "
                           << block->getParentOp()->getName() << "\n");
+
+  // The allocation should run in (reversed?) topo-sorted order for preexisting
+  // substorages to work correctly.
+
+  DenseSet<Value> allocated;
+  std::function<void(Value)> visit;
+  visit = [&](Value parent) {
+    if(allocated.contains(parent)) return;
+    for(auto childAlloc : opsByStorage[parent])
+      if(auto childStorageAlloc = dyn_cast<AllocStorageOp>(childAlloc)) {
+        auto child = childStorageAlloc.getResult();
+        if(opsByStorage.contains(child)) // Used by allocations in this block
+          visit(child);
+      }
+    allocateOps(parent, block, opsByStorage[parent]);
+    allocated.insert(parent);
+  };
 
   // Actually allocate each operation.
-  for (auto &[storage, ops] : opsByStorage)
-    allocateOps(storage, block, ops);
+  for (auto &[storage, _] : opsByStorage)
+    visit(storage);
 }
 
 void AllocateStatePass::allocateOps(Value storage, Block *block,

lib/Dialect/Arc/Transforms/CMakeLists.txt

@@ -17,6 +17,7 @@ add_circt_dialect_library(CIRCTArcTransforms
   MakeTables.cpp
   MergeIfs.cpp
   MuxToControlFlow.cpp
+  Partition.cpp
   PrintCostModel.cpp
   SimplifyVariadicOps.cpp
   SplitFuncs.cpp

lib/Dialect/Arc/Transforms/LowerClocksToFuncs.cpp

@@ -42,11 +42,9 @@ struct LowerClocksToFuncsPass
 
   void runOnOperation() override;
   LogicalResult lowerModel(ModelOp modelOp);
-  LogicalResult lowerClock(Operation *clockOp, Value modelStorageArg,
+  LogicalResult lowerClock(Operation *clockOp, SmallVector<Value> modelStorageArgs,
                            OpBuilder &funcBuilder);
-  LogicalResult isolateClock(Operation *clockOp, Value modelStorageArg,
-                             Value clockStorageArg);
-
+  LogicalResult isolateClock(Operation *clockOp, IRMapping storageArgMapping);
   SymbolTable *symbolTable;
 
   Statistic numOpsCopied{this, "ops-copied", "Ops copied into clock trees"};
@@ -100,14 +98,18 @@ LogicalResult LowerClocksToFuncsPass::lowerModel(ModelOp modelOp) {
   // Perform the actual extraction.
   OpBuilder funcBuilder(modelOp);
   for (auto *op : clocks)
-    if (failed(lowerClock(op, modelOp.getBody().getArgument(0), funcBuilder)))
+    if (failed(
+            lowerClock(op,
+                       llvm::map_to_vector(modelOp.getBody().getArguments(),
+                                           [](auto v) -> Value { return v; }),
+                       funcBuilder)))
       return failure();
 
   return success();
 }
 
 LogicalResult LowerClocksToFuncsPass::lowerClock(Operation *clockOp,
-                                                 Value modelStorageArg,
+                                                 SmallVector<Value> modelStorageArgs,
                                                  OpBuilder &funcBuilder) {
   LLVM_DEBUG(llvm::dbgs() << "- Lowering clock " << clockOp->getName() << "\n");
   assert((isa<InitialOp, FinalOp>(clockOp)));
@@ -116,11 +118,17 @@ LogicalResult LowerClocksToFuncsPass::lowerClock(Operation *clockOp,
   // going to use to pass the storage pointer to the clock once it has been
   // pulled out into a separate function.
   Region &clockRegion = clockOp->getRegion(0);
-  Value clockStorageArg = clockRegion.addArgument(modelStorageArg.getType(),
-                                                  modelStorageArg.getLoc());
+
+  IRMapping mapping;
+
+  for(const auto &storage : modelStorageArgs) {
+    Value mapped = clockRegion.addArgument(storage.getType(),
+                                           storage.getLoc());
+    mapping.map(storage, mapped);
+  }
 
   // Ensure the clock tree does not use any values defined outside of it.
-  if (failed(isolateClock(clockOp, modelStorageArg, clockStorageArg)))
+  if (failed(isolateClock(clockOp, mapping)))
     return failure();
 
   // Add a return op to the end of the body.
@@ -137,9 +145,12 @@ LogicalResult LowerClocksToFuncsPass::lowerClock(Operation *clockOp,
   else if (isa<FinalOp>(clockOp))
     funcName.append("_final");
 
+  SmallVector<Type> fnArgTypes = llvm::map_to_vector(
+      modelStorageArgs, [](Value v) { return v.getType(); });
+
   auto funcOp = funcBuilder.create<func::FuncOp>(
       clockOp->getLoc(), funcName,
-      builder.getFunctionType({modelStorageArg.getType()}, {}));
+      builder.getFunctionType({fnArgTypes}, {}));
   symbolTable->insert(funcOp); // uniquifies the name
   LLVM_DEBUG(llvm::dbgs() << "  - Created function `" << funcOp.getSymName()
                           << "`\n");
@@ -174,16 +185,15 @@ LogicalResult LowerClocksToFuncsPass::lowerClock(Operation *clockOp,
 /// body. Anything besides constants should no longer exist after a proper run
 /// of the pipeline.
 LogicalResult LowerClocksToFuncsPass::isolateClock(Operation *clockOp,
-                                                   Value modelStorageArg,
-                                                   Value clockStorageArg) {
+                                                   IRMapping storageMapping) {
   auto *clockRegion = &clockOp->getRegion(0);
   auto builder = OpBuilder::atBlockBegin(&clockRegion->front());
   DenseMap<Value, Value> copiedValues;
   auto result = clockRegion->walk([&](Operation *op) {
     for (auto &operand : op->getOpOperands()) {
       // Block arguments are okay, since there's nothing we can move.
-      if (operand.get() == modelStorageArg) {
-        operand.set(clockStorageArg);
+      if (storageMapping.contains(operand.get())) {
+        operand.set(storageMapping.lookup(operand.get()));
         continue;
       }
       if (isa<BlockArgument>(operand.get())) {