[Arc] Add Initial Cost Model (#7360)

* [Arc] Add Initial Cost Model

* [Arc] Add the DummyAnalysisTester pass

* [Arc] Fix clang-format failure

* [Arc] Address @fabianschuiki comments

* [Arc] Remove DummyAnalysisTester.cpp

include/circt/Dialect/Arc/ArcCostModel.h

@@ -0,0 +1,62 @@
+//===- ArcCostModel.h -----------------------------------------------------===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef CIRCT_DIALECT_ARC_ARCCOSTMODEL_H
+#define CIRCT_DIALECT_ARC_ARCCOSTMODEL_H
+
+#include "circt/Dialect/Arc/ArcOps.h"
+#include "mlir/IR/Operation.h"
+#include "mlir/Pass/AnalysisManager.h"
+
+using namespace mlir;
+
+namespace circt {
+namespace arc {
+
+struct OperationCosts {
+  size_t normalCost{0};
+  size_t packingCost{0};
+  size_t shufflingCost{0};
+  size_t vectorizeOpsBodyCost{0};
+  size_t totalCost() const {
+    return normalCost + packingCost + shufflingCost + vectorizeOpsBodyCost;
+  }
+  OperationCosts &operator+=(const OperationCosts &other) {
+    this->normalCost += other.normalCost;
+    this->packingCost += other.packingCost;
+    this->shufflingCost += other.shufflingCost;
+    this->vectorizeOpsBodyCost += other.vectorizeOpsBodyCost;
+    return *this;
+  }
+};
+
+class ArcCostModel {
+public:
+  OperationCosts getCost(Operation *op);
+
+private:
+  OperationCosts computeOperationCost(Operation *op);
+
+  // gets the cost to pack the vectors we have some cases we need to consider:
+  // 1: the input is scalar so we can give it a cost of 1
+  // 2: the input is a result of another vector but with no shuffling so the
+  //    is 0
+  // 3: the input is a result of another vector but with some shuffling so
+  //    the cost is the (number of out of order elements) * 2
+  // 4: the input is a mix of some vectors:
+  //    a) same order we multiply by 2
+  //    b) shuffling we multiply by 3
+  OperationCosts getInputVectorsCost(VectorizeOp vecOp);
+  size_t getShufflingCost(const ValueRange &inputVec, bool isSame = false);
+  DenseMap<Operation *, OperationCosts> opCostCache;
+};
+
+} // namespace arc
+} // namespace circt
+
+#endif // CIRCT_DIALECT_ARC_ARCCOSTMODEL_H

include/circt/Dialect/Arc/ArcPasses.h

@@ -46,6 +46,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> createPrintCostModelPass();
 std::unique_ptr<mlir::Pass> createSimplifyVariadicOpsPass();
 std::unique_ptr<mlir::Pass> createSplitLoopsPass();
 std::unique_ptr<mlir::Pass> createStripSVPass();

include/circt/Dialect/Arc/ArcPasses.td

@@ -58,6 +58,24 @@ def Dedup : Pass<"arc-dedup", "mlir::ModuleOp"> {
   ];
 }
 
+def PrintCostModel : Pass<"arc-print-cost-model", "mlir::ModuleOp"> {
+  let summary = "A dymmy pass to test analysis passes";
+  let constructor = "circt::arc::createPrintCostModelPass()";
+  let dependentDialects = ["arc::ArcDialect"];
+  let statistics = [
+    Statistic<"moduleCost", "Operation(s)",
+      "Number of operations in the module">,
+    Statistic<"packingCost", "Pack operations(s)",
+      "Number of scalar to vector packking in the module">,
+    Statistic<"shufflingCost", "Shuffle operation(s)",
+      "Number of shuffles done to set up the VectorizeOps">,
+    Statistic<"vectoroizeOpsBodyCost", "VectorizeOps Body Cost",
+      "Number of operations inside the body of the VectorizeOps">,
+    Statistic<"allVectorizeOpsCost", "All VectorizeOps Cost",
+      "Total Cost of all VectorizeOps in the module">
+  ];
+}
+
 def FindInitialVectors : Pass<"arc-find-initial-vectors", "mlir::ModuleOp"> {
   let summary = "Find initial groups of vectorizable ops";
   let constructor = "circt::arc::createFindInitialVectorsPass()";

lib/Dialect/Arc/ArcCostModel.cpp

@@ -0,0 +1,132 @@
+//===- ArcCostModel.cpp ---------------------------------------------------===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+
+#include "circt/Dialect/Arc/ArcCostModel.h"
+#include "circt/Dialect/Comb/CombOps.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"
+#include <algorithm>
+
+using namespace llvm;
+using namespace circt;
+using namespace arc;
+using namespace std;
+
+// FIXME: May be refined and we have more accurate operation costs
+enum class OperationCost : size_t {
+  NOCOST,
+  NORMALCOST,
+  PACKCOST = 2,
+  EXTRACTCOST = 3,
+  CONCATCOST = 3,
+  SAMEVECTORNOSHUFFLE = 0,
+  SAMEVECTORSHUFFLECOST = 2,
+  DIFFERENTVECTORNOSHUFFLE = 2,
+  DIFFERENTVECTORSHUFFLECOST = 3
+};
+
+OperationCosts ArcCostModel::getCost(Operation *op) {
+  return computeOperationCost(op);
+}
+
+OperationCosts ArcCostModel::computeOperationCost(Operation *op) {
+  if (auto it = opCostCache.find(op); it != opCostCache.end())
+    return it->second;
+
+  OperationCosts costs;
+
+  if (isa<circt::comb::ConcatOp>(op))
+    costs.normalCost = size_t(OperationCost::CONCATCOST);
+  else if (isa<circt::comb::ExtractOp>(op))
+    costs.normalCost = size_t(OperationCost::EXTRACTCOST);
+  else if (auto vecOp = dyn_cast<arc::VectorizeOp>(op)) {
+    // VectorizeOpCost = packingCost + shufflingCost + bodyCost
+    OperationCosts inputVecCosts = getInputVectorsCost(vecOp);
+    costs.packingCost += inputVecCosts.packingCost;
+    costs.shufflingCost += inputVecCosts.shufflingCost;
+
+    for (auto &region : op->getRegions()) {
+      for (auto &block : region) {
+        for (auto &innerOp : block) {
+          OperationCosts innerCosts = computeOperationCost(&innerOp);
+          costs.vectorizeOpsBodyCost += innerCosts.totalCost();
+        }
+      }
+    }
+  } else if (auto callableOp = dyn_cast<CallOpInterface>(op)) {
+    // Callable Op? then resolve!
+    if (auto *calledOp = callableOp.resolveCallable())
+      return opCostCache[callableOp] = computeOperationCost(calledOp);
+  } else if (isa<func::FuncOp, arc::DefineOp, mlir::ModuleOp>(op)) {
+    // Get the body cost
+    for (auto &region : op->getRegions())
+      for (auto &block : region)
+        for (auto &innerOp : block)
+          costs += computeOperationCost(&innerOp);
+  } else
+    costs.normalCost = size_t(OperationCost::NORMALCOST);
+
+  return opCostCache[op] = costs;
+}
+
+OperationCosts ArcCostModel::getInputVectorsCost(VectorizeOp vecOp) {
+  OperationCosts costs;
+  for (auto inputVec : vecOp.getInputs()) {
+    if (auto otherVecOp = inputVec[0].getDefiningOp<VectorizeOp>();
+        all_of(inputVec.begin(), inputVec.end(), [&](auto element) {
+          return element.template getDefiningOp<VectorizeOp>() == otherVecOp;
+        })) {
+      // This means that they came from the same vector or
+      // VectorizeOp == null so they are all scalars
+
+      // Check if they all scalars we multiply by the PACKCOST (SHL/R + OR)
+      if (!otherVecOp)
+        costs.packingCost += inputVec.size() * size_t(OperationCost::PACKCOST);
+      else
+        costs.shufflingCost += inputVec == otherVecOp.getResults()
+                                   ? size_t(OperationCost::SAMEVECTORNOSHUFFLE)
+                                   : getShufflingCost(inputVec, true);
+    } else
+      // inputVector consists of elements from different vectotrize ops and
+      // may have scalars as well.
+      costs.shufflingCost += getShufflingCost(inputVec);
+  }
+
+  return costs;
+}
+
+size_t ArcCostModel::getShufflingCost(const ValueRange &inputVec, bool isSame) {
+  size_t totalCost = 0;
+  if (isSame) {
+    auto vecOp = inputVec[0].getDefiningOp<VectorizeOp>();
+    for (auto [elem, orig] : llvm::zip(inputVec, vecOp.getResults()))
+      if (elem != orig)
+        ++totalCost;
+
+    return totalCost * size_t(OperationCost::SAMEVECTORSHUFFLECOST);
+  }
+
+  for (size_t i = 0; i < inputVec.size(); ++i) {
+    auto otherVecOp = inputVec[i].getDefiningOp<VectorizeOp>();
+    // If the element is not a result of a vector operation then it's a result
+    // of a scalar operation, then it just needs to be packed into the vector.
+    if (!otherVecOp)
+      totalCost += size_t(OperationCost::PACKCOST);
+    else {
+      // If it's a result of a vector operation, then we have two cases:
+      // (1) Its order in `inputVec` is the same as its order in the result of
+      //     the defining op.
+      // (2) the order is different.
+      size_t idx = find(otherVecOp.getResults().begin(),
+                        otherVecOp.getResults().end(), inputVec[i]) -
+                   otherVecOp.getResults().begin();
+      totalCost += i == idx ? size_t(OperationCost::DIFFERENTVECTORNOSHUFFLE)
+                            : size_t(OperationCost::DIFFERENTVECTORSHUFFLECOST);
+    }
+  }
+  return totalCost;
+}

lib/Dialect/Arc/CMakeLists.txt

@@ -3,6 +3,7 @@ set(CIRCT_Arc_Sources
   ArcFolds.cpp
   ArcOps.cpp
   ArcTypes.cpp
+  ArcCostModel.cpp
   ModelInfo.cpp
 )
 
@@ -26,11 +27,13 @@ add_circt_dialect_library(CIRCTArc
   Support
 
   LINK_LIBS PUBLIC
+  CIRCTComb
   CIRCTHW
   CIRCTSeq
   MLIRIR
   MLIRInferTypeOpInterface
   MLIRSideEffectInterfaces
+  MLIRFuncDialect
 )
 
 add_circt_library(CIRCTArcReductions

lib/Dialect/Arc/Transforms/CMakeLists.txt

@@ -18,6 +18,7 @@ add_circt_dialect_library(CIRCTArcTransforms
   LowerVectorizations.cpp
   MakeTables.cpp
   MuxToControlFlow.cpp
+  PrintCostModel.cpp
   SimplifyVariadicOps.cpp
   SplitFuncs.cpp
   SplitLoops.cpp

lib/Dialect/Arc/Transforms/PrintCostModel.cpp

@@ -0,0 +1,56 @@
+//===- DummyAnalysisTester.cpp --------------------------------------------===//
+//
+// 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 is a dummy pass to test the cost model results it doesn't do any thing.
+// just walks over the ops to compute some statistics.
+//
+//===----------------------------------------------------------------------===//
+
+#include "circt/Dialect/Arc/ArcCostModel.h"
+#include "circt/Dialect/Arc/ArcPasses.h"
+#include "circt/Dialect/HW/HWOps.h"
+#include "mlir/IR/MLIRContext.h"
+#include "mlir/Pass/Pass.h"
+
+#define DEBUG_TYPE "arc-print-cost-model"
+
+namespace circt {
+namespace arc {
+#define GEN_PASS_DEF_PRINTCOSTMODEL
+#include "circt/Dialect/Arc/ArcPasses.h.inc"
+} // namespace arc
+} // namespace circt
+
+using namespace circt;
+using namespace arc;
+
+namespace {
+struct PrintCostModelPass
+    : public arc::impl::PrintCostModelBase<PrintCostModelPass> {
+  void runOnOperation() override;
+};
+} // namespace
+
+void PrintCostModelPass::runOnOperation() {
+  OperationCosts statVars;
+  ArcCostModel arcCostModel;
+  for (auto moduleOp : getOperation().getOps<hw::HWModuleOp>()) {
+    moduleOp.walk([&](Operation *op) { statVars += arcCostModel.getCost(op); });
+  }
+
+  moduleCost = statVars.totalCost();
+  packingCost = statVars.packingCost;
+  shufflingCost = statVars.shufflingCost;
+  vectoroizeOpsBodyCost = statVars.vectorizeOpsBodyCost;
+  allVectorizeOpsCost = statVars.packingCost + statVars.shufflingCost +
+                        statVars.vectorizeOpsBodyCost;
+}
+
+std::unique_ptr<Pass> arc::createPrintCostModelPass() {
+  return std::make_unique<PrintCostModelPass>();
+}