[circt-synth] Implemented AIG node balancing algorithm

include/circt/Dialect/AIG/AIGOps.td

@@ -22,7 +22,7 @@ include "mlir/Interfaces/SideEffectInterfaces.td"
 class AIGOp<string mnemonic, list<Trait> traits = []> :
   Op<AIG_Dialect, mnemonic, traits>;
 
-def AndInverterOp : AIGOp<"and_inv", [SameOperandsAndResultType, Pure]> {
+def AndInverterOp : AIGOp<"and_inv", [SameOperandsAndResultType, Pure, Commutative]> {
   let summary = "AIG dialect AND operation";
   let description = [{
     The `aig.and_inv` operation represents an And-Inverter in the AIG dialect.

include/circt/Dialect/AIG/AIGPasses.td

@@ -15,9 +15,17 @@ def LowerVariadic : Pass<"aig-lower-variadic", "hw::HWModuleOp"> {
   let summary = "Lower variadic AndInverter operations to binary AndInverter";
 }
 
+def BalanceVariadic : Pass<"aig-balance-variadic", "hw::HWModuleOp"> {
+  let summary = "Lower variadic AndInverter operations to binary AndInverter";
+}
+
 def LowerWordToBits : Pass<"aig-lower-word-to-bits", "hw::HWModuleOp"> {
   let summary = "Lower multi-bit AndInverter to single-bit ones";
   let dependentDialects = ["comb::CombDialect"];
 }
 
+def MaximumAndCover : Pass<"maximum-and-cover", "hw::HWModuleOp"> {
+  let summary = "Maximum And Cover";
+}
+
 #endif // CIRCT_DIALECT_AIG_AIGPASSES_TD

include/circt/Dialect/AIG/Analysis/OpDepthAnalysis.h

@@ -0,0 +1,50 @@
+//===- OpDepthAnalysis.h - operation depth analyses -----------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// This header file defines AIG operation depth analysis.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef CIRCT_ANALYSIS_OPDEPTH_ANALYSIS_H
+#define CIRCT_ANALYSIS_OPDEPTH_ANALYSIS_H
+
+#include "circt/Dialect/AIG/AIGOps.h"
+#include "circt/Dialect/HW/HWOps.h"
+
+namespace mlir {
+class AnalysisManager;
+} // namespace mlir
+namespace circt {
+namespace aig {
+namespace analysis {
+
+class OpDepthAnalysis {
+public:
+  OpDepthAnalysis(hw::HWModuleOp moduleOp, mlir::AnalysisManager &am);
+
+  /// Get the depth of operations of a specific name
+  size_t getOpDepth(AndInverterOp op) const { return opDepths.at(op); }
+
+  const DenseMap<AndInverterOp, size_t> &getOpDepthMap() const {
+    return opDepths;
+  }
+
+  size_t updateLevel(AndInverterOp op, bool isRoot = false);
+
+  SmallVector<AndInverterOp> getPOs();
+
+private:
+  DenseMap<AndInverterOp, size_t> opDepths;
+  hw::HWModuleOp module;
+};
+
+} // namespace analysis
+} // namespace aig
+} // namespace circt
+
+#endif // CIRCT_ANALYSIS_OPDEPTH_ANALYSIS_H

integration_test/circt-synth/aig-balancing-lec.mlir

@@ -0,0 +1,33 @@
+// REQUIRES: libz3
+// REQUIRES: circt-lec-jit
+
+// RUN: circt-opt %s --cse --convert-aig-to-comb -o %t1.mlir
+// RUN: circt-opt %s --maximum-and-cover --aig-balance-variadic --cse --convert-aig-to-comb -o %t2.mlir
+
+// RUN: circt-lec %t1.mlir %t2.mlir  -c1=aig -c2=aig --shared-libs=%libz3 | FileCheck %s --check-prefix=COMB_AIG
+// COMB_AIG: c1 == c2
+hw.module @aig(in %a: i1, in %b: i1, in %c: i1, in %d: i1, out o1: i1, out o2: i1, out o3: i1) {
+  %0 = aig.and_inv %a, %b : i1
+  %1 = aig.and_inv %0, %c : i1
+  %2 = aig.and_inv %b, %c : i1
+  %3 = aig.and_inv %2, %d : i1
+
+  %4 = aig.and_inv %c, %d : i1
+  %5 = aig.and_inv %b, %4 : i1
+  %6 = aig.and_inv %a, %5 : i1
+
+  hw.output %1, %3, %6 : i1, i1, i1
+}
+
+// RUN: circt-lec %t1.mlir %t2.mlir  -c1=aig2 -c2=aig2 --shared-libs=%libz3 | FileCheck %s --check-prefix=COMB_AIG2
+// COMB_AIG2: c1 == c2
+hw.module @aig2(in %a: i1, in %b: i1, in %c: i1, in %d: i1, in %e: i1, in %f: i1, in %g: i1, out o1: i1) {
+  %1 = aig.and_inv %a, not %b : i1
+  %2 = aig.and_inv %d, not %e : i1
+  %3 = aig.and_inv not %2, %f : i1
+  %4 = aig.and_inv not %c, %3 : i1
+  %5 = aig.and_inv %1, not %4 : i1
+  %6 = aig.and_inv %5, %g : i1
+
+  hw.output %6 : i1
+}

lib/Dialect/AIG/Analysis/CMakeLists.txt

@@ -0,0 +1,6 @@
+add_circt_dialect_library(CIRCTAIGAnalysis
+  OpDepthAnalysis.cpp
+
+  LINK_LIBS PUBLIC
+  CIRCTAIG
+)

lib/Dialect/AIG/Analysis/OpDepthAnalysis.cpp

@@ -0,0 +1,56 @@
+//===----------------------------------------------------------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// This file defines the op depth (level) analysis.
+//
+//===----------------------------------------------------------------------===//
+
+#include "circt/Dialect/AIG/Analysis/OpDepthAnalysis.h"
+
+using namespace circt;
+using namespace aig;
+using namespace analysis;
+
+OpDepthAnalysis::OpDepthAnalysis(hw::HWModuleOp moduleOp,
+                                 mlir::AnalysisManager &am)
+    : module(moduleOp) {}
+
+SmallVector<AndInverterOp> OpDepthAnalysis::getPOs() {
+  SmallVector<AndInverterOp> po;
+  for (auto op : module.getOps<AndInverterOp>()) {
+    bool isPO = false;
+    for (auto *user : op->getUsers()) {
+      if (!isa<AndInverterOp>(user)) {
+        isPO = true;
+        break;
+      }
+    }
+    if (isPO)
+      po.push_back(op);
+  }
+  return po;
+}
+
+size_t OpDepthAnalysis::updateLevel(AndInverterOp op, bool isRoot) {
+  if (!isRoot) {
+    if (auto it = opDepths.find(op); it != opDepths.end()) {
+      return it->second;
+    }
+  }
+
+  /// PI is level 0, so the minimum level of an AndInverterOp is 1
+  size_t maxDepth = 1;
+  for (auto fanin : op.getOperands()) {
+    if (auto faninOp = fanin.getDefiningOp<AndInverterOp>()) {
+      size_t faninDepth = updateLevel(faninOp);
+      maxDepth = std::max(maxDepth, faninDepth + 1);
+    }
+  }
+  opDepths[op] = maxDepth;
+  return maxDepth;
+}

lib/Dialect/AIG/CMakeLists.txt

@@ -14,3 +14,4 @@ add_circt_dialect_library(CIRCTAIG
 )
 
 add_subdirectory(Transforms)
+add_subdirectory(Analysis)

lib/Dialect/AIG/Transforms/BalanceVariadic.cpp

@@ -0,0 +1,177 @@
+//===----------------------------------------------------------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// This pass lowers variadic AndInverter operations to balanced binary
+// AndInverter operations.
+//
+//===----------------------------------------------------------------------===//
+#include "llvm/ADT/PriorityQueue.h"
+
+#include "circt/Dialect/AIG/AIGOps.h"
+#include "circt/Dialect/AIG/AIGPasses.h"
+#include "circt/Dialect/AIG/Analysis/OpDepthAnalysis.h"
+
+#define DEBUG_TYPE "aig-balance-variadic"
+
+namespace circt {
+namespace aig {
+#define GEN_PASS_DEF_BALANCEVARIADIC
+#include "circt/Dialect/AIG/AIGPasses.h.inc"
+} // namespace aig
+} // namespace circt
+
+using namespace circt;
+using namespace aig;
+
+namespace {
+/// For wrapping Value and complement information into one object
+struct Signal {
+  Value value;
+  bool complement;
+
+  Signal() = default;
+  Signal(Value v, bool complement) : value(v), complement(complement) {}
+
+  bool isComplement() const { return complement; }
+  Value getValue() const { return value; }
+};
+
+struct BalanceVariadicDriver {
+  BalanceVariadicDriver(mlir::IRRewriter &rewriter,
+                        aig::analysis::OpDepthAnalysis *opDepthAnalysis)
+      : rewriter(rewriter), opDepthAnalysis(opDepthAnalysis) {}
+
+  struct PairSorter {
+    bool operator()(const std::pair<size_t, Signal> &lhs,
+                    const std::pair<size_t, Signal> &rhs) const {
+      // First compare by level (higher level = lower priority)
+      if (lhs.first != rhs.first)
+        return lhs.first > rhs.first;
+
+      // If levels are equal, compare by argnumber or result number for
+      // deterministic ordering
+      auto *lop = lhs.second.getValue().getDefiningOp();
+      auto *rop = rhs.second.getValue().getDefiningOp();
+      if (lop && rop) {
+        return !lop->isBeforeInBlock(rop);
+      }
+
+      if (!lop && rop)
+        return false;
+
+      if (lop && !rop)
+        return true;
+
+      BlockArgument larg = cast<BlockArgument>(lhs.second.getValue());
+      BlockArgument rarg = cast<BlockArgument>(rhs.second.getValue());
+      return larg.getArgNumber() > rarg.getArgNumber();
+    }
+  };
+
+  using NodeLevelHeap =
+      llvm::PriorityQueue<std::pair<size_t, Signal>,
+                          std::vector<std::pair<size_t, Signal>>, PairSorter>;
+
+  void balanceVariadicAndInverterOp(AndInverterOp op) {
+    rewriter.setInsertionPoint(op);
+
+    NodeLevelHeap sortByLevel;
+    for (auto [fanin, inverted] :
+         llvm::zip(op.getOperands(), op.getInverted())) {
+      auto faninOp = fanin.getDefiningOp<AndInverterOp>();
+      size_t level = faninOp ? opDepthAnalysis->updateLevel(faninOp, true) : 0;
+      sortByLevel.push({level, Signal(fanin, inverted)});
+    }
+
+    // extract the top two elements with minimum level
+    // and replace them with a new AndInverterOp
+    while (sortByLevel.size() > 1) {
+      auto [llv, lhs] = sortByLevel.top();
+      sortByLevel.pop();
+      auto [rlv, rhs] = sortByLevel.top();
+      sortByLevel.pop();
+
+      auto balanced = rewriter.create<AndInverterOp>(
+          op.getLoc(), lhs.getValue(), rhs.getValue(), lhs.isComplement(),
+          rhs.isComplement());
+
+      size_t level = std::max(llv, rlv) + 1;
+      sortByLevel.push({level, Signal(balanced, false)});
+    }
+
+    switch (sortByLevel.size()) {
+    case 0:
+      break;
+    case 1: {
+      auto signal = sortByLevel.top().second;
+      Value val = signal.getValue();
+      if (signal.isComplement()) {
+        // TODO: push inverter to use AndInverterOp (fanout)
+        val = rewriter.create<AndInverterOp>(op.getLoc(), val, true);
+      }
+      rewriter.replaceOp(op, val);
+      break;
+    }
+    default:
+      assert(0); // should not happen
+    }
+  }
+
+  void balanceRecursive(AndInverterOp op) {
+    if (visited.count(op))
+      return;
+
+    visited.insert(op);
+    assert(!op->use_empty());
+
+    for (auto fanin : op.getOperands()) {
+      auto faninOp = fanin.getDefiningOp<AndInverterOp>();
+      if (faninOp) {
+        balanceRecursive(faninOp);
+      }
+    }
+
+    if (op.getOperands().size() <= 2)
+      return;
+
+    balanceVariadicAndInverterOp(op);
+  }
+
+  void balancing() {
+    // Balance each variadic AndInverterOp in reverse topological order
+    // Will ignore dangling internal AIG nodes
+    for (AndInverterOp po : opDepthAnalysis->getPOs()) {
+      balanceRecursive(po);
+    }
+  }
+
+private:
+  DenseSet<Operation *> visited;
+  mlir::IRRewriter &rewriter;
+  aig::analysis::OpDepthAnalysis *opDepthAnalysis;
+};
+
+struct BalanceVariadicPass
+    : public impl::BalanceVariadicBase<BalanceVariadicPass> {
+  void runOnOperation() override;
+};
+} // namespace
+
+//===----------------------------------------------------------------------===//
+// Balance Variadic pass
+//===----------------------------------------------------------------------===//
+void BalanceVariadicPass::runOnOperation() {
+  auto *opDepthAnalysis = &getAnalysis<aig::analysis::OpDepthAnalysis>();
+
+  auto module = getOperation();
+  MLIRContext *ctx = module->getContext();
+  mlir::IRRewriter rewriter(ctx);
+
+  BalanceVariadicDriver driver(rewriter, opDepthAnalysis);
+  driver.balancing();
+}

lib/Dialect/AIG/Transforms/CMakeLists.txt

@@ -1,6 +1,8 @@
 add_circt_dialect_library(CIRCTAIGTransforms
   LowerVariadic.cpp
   LowerWordToBits.cpp
+  BalanceVariadic.cpp
+  MaximumAndCover.cpp
 
   DEPENDS
   CIRCTAIGPassesIncGen
@@ -9,4 +11,5 @@ add_circt_dialect_library(CIRCTAIGTransforms
   CIRCTAIG
   CIRCTComb
   CIRCTHW
+  CIRCTAIGAnalysis
 )