Improve bitwidth optimization for FTD circuits

include/dynamatic/Transforms/HandshakeMinimizeCstWidth.h

@@ -0,0 +1,29 @@
+//===- HandshakeMinimizeCstWidth.h - Min. constants bitwidth ----*- C++ -*-===//
+//
+// Dynamatic is 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 declares the --handshake-minimize-cst-width pass.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef DYNAMATIC_TRANSFORMS_HANDSHAKEMINIMIZECSTWIDTH_H
+#define DYNAMATIC_TRANSFORMS_HANDSHAKEMINIMIZECSTWIDTH_H
+
+#include "dynamatic/Dialect/Handshake/HandshakeOps.h"
+#include "dynamatic/Support/DynamaticPass.h"
+#include "dynamatic/Support/LLVM.h"
+
+namespace dynamatic {
+/// Attempts to find an equivalent constant to the one provided in the circuit,
+/// that is a constant with the same control and same value attribute. Returns
+/// such an equivalent constant if it finds one, nullptr otherwise.
+handshake::ConstantOp findEquivalentCst(handshake::ConstantOp cstOp);
+/// Computes the minimum required bitwidth needed to store the provided integer.
+unsigned computeRequiredBitwidth(APInt val);
+} // namespace dynamatic
+
+#endif // DYNAMATIC_TRANSFORMS_HANDSHAKEMINIMIZECSTWIDTH_H

include/dynamatic/Transforms/Passes.td

@@ -197,6 +197,21 @@ def HandshakeOptimizeBitwidths : DynamaticPass< "handshake-optimize-bitwidths",
   ];
 }
 
+def HandshakeMinimizeCstWidth : DynamaticPass<"handshake-minimize-cst-width",
+                                              ["mlir::arith::ArithDialect"]> {
+  let summary = "Minimizes the bitwidth of all Handshake constants.";
+  let description = [{
+    Rewrites constant operations with the minimum required bitwidth to support
+    the constants' values. The pass inserts extension operations as needed to
+    ensure consistency with users of constant operations. The pass also pays
+    attention to not create duplicated constants indirectly due to the
+    minimization process.
+  }];
+  let options =
+    [Option<"optNegatives", "opt-negatives", "bool", "false",
+            "If true, allows bitwidth optimization of negative values.">];
+}
+
 def HandshakeReplaceMemoryInterfaces : DynamaticPass<
   "handshake-replace-memory-interfaces"
 > {

lib/Transforms/CMakeLists.txt

@@ -10,6 +10,7 @@ add_dynamatic_library(DynamaticTransforms
   HandshakeDeactivateMemDependencies.cpp
   HandshakeHoistExtInstances.cpp
   HandshakeMaterialize.cpp
+  HandshakeMinimizeCstWidth.cpp
   HandshakeOptimizeBitwidths.cpp
   HandshakeInferBasicBlocks.cpp
   HandshakeReplaceMemoryInterfaces.cpp

lib/Transforms/HandshakeMinimizeCstWidth.cpp

@@ -0,0 +1,208 @@
+//===- HandshakeMinimizeCstWidth.cpp - Min. constants bitwidth --*- C++ -*-===//
+//
+// Dynamatic is 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
+//
+//===----------------------------------------------------------------------===//
+//
+// Implements the handshake-minimize-cst-width pass, which minimizes the
+// bitwidth of all constants. The pass matches on all Handshake constants in the
+// IR, determines the minimum bitwidth necessary to hold their value, and
+// updates their result/attribute type match to this minimal value.
+//
+//===----------------------------------------------------------------------===//
+
+#include "dynamatic/Transforms/HandshakeMinimizeCstWidth.h"
+#include "dynamatic/Dialect/Handshake/HandshakeOps.h"
+#include "dynamatic/Dialect/Handshake/HandshakeTypes.h"
+#include "dynamatic/Support/CFG.h"
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
+#include "llvm/Support/Debug.h"
+
+#define DEBUG_TYPE "BITWIDTH"
+
+// [START Boilerplate code for the MLIR pass]
+#include "dynamatic/Transforms/Passes.h" // IWYU pragma: keep
+namespace dynamatic {
+#define GEN_PASS_DEF_HANDSHAKEMINIMIZECSTWIDTH
+#include "dynamatic/Transforms/Passes.h.inc"
+} // namespace dynamatic
+// [END Boilerplate code for the MLIR pass]
+
+STATISTIC(savedBits, "Number of saved bits");
+
+using namespace mlir;
+using namespace dynamatic;
+
+/// Determines whether the control value of two constants can be considered
+/// equivalent.
+static bool areCstCtrlEquivalent(Value ctrl, Value otherCtrl) {
+  if (ctrl == otherCtrl)
+    return true;
+
+  // Both controls are equivalent if they originate from sources in the same
+  // block
+  Operation *defOp = ctrl.getDefiningOp();
+  if (!defOp || !isa<handshake::SourceOp>(defOp))
+    return false;
+  Operation *otherDefOp = otherCtrl.getDefiningOp();
+  if (!otherDefOp || !isa<handshake::SourceOp>(otherDefOp))
+    return false;
+  std::optional<unsigned> block = getLogicBB(defOp);
+  std::optional<unsigned> otherBlock = getLogicBB(otherDefOp);
+  return block.has_value() && otherBlock.has_value() &&
+         block.value() == otherBlock.value();
+}
+
+handshake::ConstantOp static findEquivalentCst(TypedAttr valueAttr,
+                                               Value ctrl) {
+  auto funcOp = cast<handshake::FuncOp>(ctrl.getParentBlock()->getParentOp());
+  for (auto cstOp : funcOp.getOps<handshake::ConstantOp>()) {
+    // The constant operation needs to have the same value attribute and the
+    // same control
+    auto cstAttr = cstOp.getValue();
+    if (cstAttr == valueAttr && areCstCtrlEquivalent(ctrl, cstOp.getCtrl()))
+      return cstOp;
+  }
+  return nullptr;
+}
+
+handshake::ConstantOp
+dynamatic::findEquivalentCst(handshake::ConstantOp cstOp) {
+  auto cstAttr = cstOp.getValue();
+  auto funcOp = cstOp->getParentOfType<handshake::FuncOp>();
+  assert(funcOp && "constant should have parent function");
+
+  for (auto otherCstOp : funcOp.getOps<handshake::ConstantOp>()) {
+    // Don't match ourself
+    if (cstOp == otherCstOp)
+      continue;
+
+    // The constant operation needs to have the same value attribute and the
+    // same control
+    auto otherCstAttr = otherCstOp.getValue();
+    if (otherCstAttr == cstAttr &&
+        areCstCtrlEquivalent(cstOp.getCtrl(), otherCstOp.getCtrl()))
+      return otherCstOp;
+  }
+
+  return nullptr;
+}
+
+/// Inserts an extension op after the constant op that extends the constant's
+/// integer result to a provided destination type. The function assumes that it
+/// makes sense to extend the former type into the latter type.
+static handshake::ExtSIOp insertExtOp(handshake::ConstantOp toExtend,
+                                      handshake::ConstantOp toReplace,
+                                      PatternRewriter &rewriter) {
+  rewriter.setInsertionPointAfter(toExtend);
+  auto extOp = rewriter.create<handshake::ExtSIOp>(
+      toExtend.getLoc(), toReplace.getResult().getType(), toExtend.getResult());
+  inheritBB(toExtend, extOp);
+  return extOp;
+}
+
+unsigned dynamatic::computeRequiredBitwidth(APInt val) {
+  bool isNegative = false;
+  if (val.isNegative()) {
+    isNegative = true;
+    int64_t negVal = val.getSExtValue();
+    if (negVal - 1 == 0)
+      // The value is the minimum number representable on 64 bits
+      return APInt::APINT_BITS_PER_WORD;
+
+    // Flip the sign to make it positive
+    val = APInt(APInt::APINT_BITS_PER_WORD, -negVal);
+  }
+
+  unsigned log = val.logBase2();
+  return val.isPowerOf2() && isNegative ? log + 1 : log + 2;
+}
+
+namespace {
+
+/// Minimizes the bitwidth used by Handshake constants as much as possible. If
+/// the bitwidth is reduced, inserts an extension operation after the constant
+/// so that users of the constant result can keep using the same value type.
+struct MinimizeConstantBitwidth
+    : public OpRewritePattern<handshake::ConstantOp> {
+  using OpRewritePattern<handshake::ConstantOp>::OpRewritePattern;
+
+  MinimizeConstantBitwidth(bool optNegatives, MLIRContext *ctx)
+      : OpRewritePattern<handshake::ConstantOp>(ctx),
+        optNegatives(optNegatives) {};
+
+  LogicalResult matchAndRewrite(handshake::ConstantOp cstOp,
+                                PatternRewriter &rewriter) const override {
+    // Only consider integer attributes
+    auto intAttr = dyn_cast<mlir::IntegerAttr>(cstOp.getValue());
+    if (!intAttr)
+      return failure();
+    handshake::ChannelType channelType = cstOp.getResult().getType();
+    IntegerType dataType = cast<IntegerType>(channelType.getDataType());
+
+    // We only support reducing signless values that fit on 64 bits or less
+    APInt val = intAttr.getValue();
+    if (dataType.getSignedness() !=
+            IntegerType::SignednessSemantics::Signless ||
+        !val.isSingleWord())
+      return failure();
+
+    // Do not optimize negative values
+    if (val.isNegative() && !optNegatives)
+      return failure();
+
+    // Check if we can reduce the bitwidth
+    unsigned newWidth = computeRequiredBitwidth(val);
+    if (newWidth >= dataType.getWidth())
+      return failure();
+
+    // Create the new constant value
+    IntegerAttr newAttr = IntegerAttr::get(
+        IntegerType::get(getContext(), newWidth, dataType.getSignedness()),
+        val.trunc(newWidth));
+
+    if (auto otherCstOp = findEquivalentCst(newAttr, cstOp.getCtrl())) {
+      // Use the other constant's result and simply erase the matched constant
+      rewriter.replaceOp(cstOp, insertExtOp(otherCstOp, cstOp, rewriter));
+      return success();
+    }
+
+    // Create a new constant to replace the matched one with
+    auto newCstOp = rewriter.create<handshake::ConstantOp>(
+        cstOp->getLoc(), newAttr, cstOp.getCtrl());
+    rewriter.replaceOp(cstOp, insertExtOp(newCstOp, cstOp, rewriter));
+    return success();
+  }
+
+private:
+  /// Whether to allow optimization of negative values.
+  bool optNegatives;
+};
+
+/// Driver for the constant bitwidth reduction pass. A greedy pattern rewriter
+/// matches on all Handshake constants and minimizes their bitwidth.
+struct HandshakeMinimizeCstWidthPass
+    : public dynamatic::impl::HandshakeMinimizeCstWidthBase<
+          HandshakeMinimizeCstWidthPass> {
+
+  using HandshakeMinimizeCstWidthBase::HandshakeMinimizeCstWidthBase;
+
+  void runDynamaticPass() override {
+    auto *ctx = &getContext();
+    mlir::ModuleOp mod = getOperation();
+
+    mlir::GreedyRewriteConfig config;
+    config.useTopDownTraversal = true;
+    config.enableRegionSimplification = false;
+    RewritePatternSet patterns{ctx};
+    patterns.add<MinimizeConstantBitwidth>(optNegatives, ctx);
+    if (failed(applyPatternsAndFoldGreedily(mod, std::move(patterns), config)))
+      return signalPassFailure();
+
+    LLVM_DEBUG(llvm::dbgs() << "Number of saved bits is " << savedBits << "\n");
+  };
+};
+
+} // namespace