From f1582ed996c5b23ba1e923b982afcee19161318f Mon Sep 17 00:00:00 2001
From: Andrew Young <youngar17@gmail.com>
Date: Thu, 12 Dec 2024 17:10:22 -0800
Subject: [PATCH] [firtool] Add option to disable layer sink

The layer sink pass is showing some performance problems, and we want
the ability to turn it off to work around the issue, until we can
properly fix it.
---
 include/circt-c/Firtool/Firtool.h |  3 +++
 include/circt/Firtool/Firtool.h   |  7 +++++++
 lib/CAPI/Firtool/Firtool.cpp      |  5 +++++
 lib/Firtool/Firtool.cpp           | 19 ++++++++++++++-----
 4 files changed, 29 insertions(+), 5 deletions(-)

diff --git a/include/circt-c/Firtool/Firtool.h b/include/circt-c/Firtool/Firtool.h
index 90d37de97c02..beec1f681d25 100644
--- a/include/circt-c/Firtool/Firtool.h
+++ b/include/circt-c/Firtool/Firtool.h
@@ -103,6 +103,9 @@ circtFirtoolOptionsSetEnableDebugInfo(CirctFirtoolFirtoolOptions options,
 MLIR_CAPI_EXPORTED void
 circtFirtoolOptionsSetBuildMode(CirctFirtoolFirtoolOptions options,
                                 CirctFirtoolBuildMode value);
+MLIR_CAPI_EXPORTED void
+circtFirtoolOptionsSetDisableLayerSink(CirctFirtoolFirtoolOptions options,
+                                       bool value);
 
 MLIR_CAPI_EXPORTED void
 circtFirtoolOptionsSetDisableOptimization(CirctFirtoolFirtoolOptions options,
diff --git a/include/circt/Firtool/Firtool.h b/include/circt/Firtool/Firtool.h
index baed0b88fa67..fb85cbebe4bb 100644
--- a/include/circt/Firtool/Firtool.h
+++ b/include/circt/Firtool/Firtool.h
@@ -96,6 +96,7 @@ class FirtoolOptions {
   }
   bool shouldConvertProbesToSignals() const { return probesToSignals; }
   bool shouldReplaceSequentialMemories() const { return replSeqMem; }
+  bool shouldDisableLayerSink() const { return disableLayerSink; }
   bool shouldDisableOptimization() const { return disableOptimization; }
   bool shouldAdvancedLayerSink() const { return advancedLayerSink; }
   bool shouldLowerMemories() const { return lowerMemories; }
@@ -189,6 +190,11 @@ class FirtoolOptions {
     return *this;
   }
 
+  FirtoolOptions &setDisableLayerSink(bool value) {
+    disableLayerSink = value;
+    return *this;
+  }
+
   FirtoolOptions &setDisableOptimization(bool value) {
     disableOptimization = value;
     return *this;
@@ -381,6 +387,7 @@ class FirtoolOptions {
   firrtl::PreserveValues::PreserveMode preserveMode;
   bool enableDebugInfo;
   BuildMode buildMode;
+  bool disableLayerSink;
   bool disableOptimization;
   bool exportChiselInterface;
   std::string chiselInterfaceOutDirectory;
diff --git a/lib/CAPI/Firtool/Firtool.cpp b/lib/CAPI/Firtool/Firtool.cpp
index 0d96c13bbcde..94808e3c603b 100644
--- a/lib/CAPI/Firtool/Firtool.cpp
+++ b/lib/CAPI/Firtool/Firtool.cpp
@@ -125,6 +125,11 @@ void circtFirtoolOptionsSetBuildMode(CirctFirtoolFirtoolOptions options,
   unwrap(options)->setBuildMode(converted);
 }
 
+void circtFirtoolOptionsSetDisableLayerSink(CirctFirtoolFirtoolOptions options,
+                                            bool value) {
+  unwrap(options)->setDisableLayerSink(value);
+}
+
 void circtFirtoolOptionsSetDisableOptimization(
     CirctFirtoolFirtoolOptions options, bool value) {
   unwrap(options)->setDisableOptimization(value);
diff --git a/lib/Firtool/Firtool.cpp b/lib/Firtool/Firtool.cpp
index 4aa85db1b08d..0517b3b28f85 100644
--- a/lib/Firtool/Firtool.cpp
+++ b/lib/Firtool/Firtool.cpp
@@ -228,11 +228,14 @@ LogicalResult firtool::populateCHIRRTLToLowFIRRTL(mlir::PassManager &pm,
   //
   // TODO: Improve LowerLayers to avoid the need for canonicalization. See:
   //   https://github.com/llvm/circt/issues/7896
-  if (opt.shouldAdvancedLayerSink())
-    pm.nest<firrtl::CircuitOp>().addPass(firrtl::createAdvancedLayerSinkPass());
-  else
-    pm.nest<firrtl::CircuitOp>().nest<firrtl::FModuleOp>().addPass(
-        firrtl::createLayerSinkPass());
+  if (!opt.shouldDisableLayerSink()) {
+    if (opt.shouldAdvancedLayerSink())
+      pm.nest<firrtl::CircuitOp>().addPass(
+          firrtl::createAdvancedLayerSinkPass());
+    else
+      pm.nest<firrtl::CircuitOp>().nest<firrtl::FModuleOp>().addPass(
+          firrtl::createLayerSinkPass());
+  }
   pm.nest<firrtl::CircuitOp>().addPass(firrtl::createLowerLayersPass());
   if (!opt.shouldDisableOptimization())
     pm.nest<firrtl::CircuitOp>().nest<firrtl::FModuleOp>().addPass(
@@ -522,6 +525,11 @@ struct FirtoolCmdOptions {
                                   "release", "Compile with optimizations")),
       llvm::cl::init(firtool::FirtoolOptions::BuildModeDefault)};
 
+  llvm::cl::opt<bool> disableLayerSink{
+      "disable-layer-sink",
+      llvm::cl::desc("Disable layer sink"),
+  };
+
   llvm::cl::opt<bool> disableOptimization{
       "disable-opt",
       llvm::cl::desc("Disable optimizations"),
@@ -789,6 +797,7 @@ circt::firtool::FirtoolOptions::FirtoolOptions()
   preserveMode = clOptions->preserveMode;
   enableDebugInfo = clOptions->enableDebugInfo;
   buildMode = clOptions->buildMode;
+  disableLayerSink = clOptions->disableLayerSink;
   disableOptimization = clOptions->disableOptimization;
   exportChiselInterface = clOptions->exportChiselInterface;
   chiselInterfaceOutDirectory = clOptions->chiselInterfaceOutDirectory;