[Arc] Automatic module partition

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,8 @@ 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

@@ -63,9 +63,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