[AIG] Refactor AIG longest-path analysis with lazy mode and max-delay filtering

include/circt-c/Dialect/AIG.h

@@ -48,12 +48,17 @@ DEFINE_C_API_STRUCT(AIGLongestPathCollection, void);
 
 // Create a LongestPathAnalysis for the given module
 MLIR_CAPI_EXPORTED AIGLongestPathAnalysis
-aigLongestPathAnalysisCreate(MlirOperation module, bool traceDebugPoints);
+aigLongestPathAnalysisCreate(MlirOperation module, bool collectDebugInfo,
+                             bool keepOnlyMaxDelayPaths, bool lazyComputation);
 
 // Destroy a LongestPathAnalysis
 MLIR_CAPI_EXPORTED void
 aigLongestPathAnalysisDestroy(AIGLongestPathAnalysis analysis);
 
+MLIR_CAPI_EXPORTED AIGLongestPathCollection
+aigLongestPathAnalysisGetPaths(AIGLongestPathAnalysis analysis, MlirValue value,
+                               int64_t bitPos, bool elaboratePaths);
+
 MLIR_CAPI_EXPORTED AIGLongestPathCollection aigLongestPathAnalysisGetAllPaths(
     AIGLongestPathAnalysis analysis, MlirStringRef moduleName,
     bool elaboratePaths);
@@ -79,6 +84,10 @@ MLIR_CAPI_EXPORTED AIGLongestPathDataflowPath
 aigLongestPathCollectionGetDataflowPath(AIGLongestPathCollection collection,
                                         size_t pathIndex);
 
+MLIR_CAPI_EXPORTED void
+aigLongestPathCollectionMerge(AIGLongestPathCollection dest,
+                              AIGLongestPathCollection src);
+
 //===----------------------------------------------------------------------===//
 // DataflowPath API
 //===----------------------------------------------------------------------===//

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

@@ -29,6 +29,7 @@
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/ImmutableList.h"
 #include "llvm/ADT/SmallVector.h"
+#include "llvm/Support/ErrorHandling.h"
 #include "llvm/Support/JSON.h"
 #include <variant>
 
@@ -172,14 +173,43 @@ class DataflowPath {
 // JSON serialization for DataflowPath
 llvm::json::Value toJSON(const circt::aig::DataflowPath &path);
 
-// Options for the longest path analysis.
-struct LongestPathAnalysisOption {
-  bool traceDebugPoints = false;
-  bool incremental = false;
-
-  LongestPathAnalysisOption(bool traceDebugPoints, bool incremental)
-      : traceDebugPoints(traceDebugPoints), incremental(incremental) {}
-  LongestPathAnalysisOption() = default;
+/// Configuration options for the longest path analysis.
+///
+/// This struct controls various aspects of the analysis behavior, including
+/// debugging features, computation modes, and optimization settings. Different
+/// combinations of options are suitable for different use cases.
+///
+/// Example usage:
+///   // For lazily computing paths with debug info
+///   LongestPathAnalysisOption options(true, true, false);
+///
+///   // For fast critical path identification only
+///   LongestPathAnalysisOption options(false, false, true);
+struct LongestPathAnalysisOptions {
+  /// Enable collection of debug points along timing paths.
+  /// When enabled, records intermediate points with delay values and comments
+  /// for debugging, visualization, and understanding delay contributions.
+  /// Moderate performance impact.
+  bool collectDebugInfo = false;
+
+  /// Enable lazy computation mode for on-demand analysis.
+  /// Performs delay computations lazily and caches results, tracking IR
+  /// changes. Better for iterative workflows where only specific paths
+  /// are queried. Disables parallel processing.
+  bool lazyComputation = false;
+
+  /// Keep only the maximum delay path per fanout point.
+  /// Focuses on finding maximum delays, discarding non-critical paths.
+  /// Significantly faster and uses less memory when only delay bounds
+  /// are needed rather than complete path enumeration.
+  bool keepOnlyMaxDelayPaths = false;
+
+  /// Construct analysis options with the specified settings.
+  LongestPathAnalysisOptions(bool collectDebugInfo = false,
+                             bool lazyComputation = false,
+                             bool keepOnlyMaxDelayPaths = false)
+      : collectDebugInfo(collectDebugInfo), lazyComputation(lazyComputation),
+        keepOnlyMaxDelayPaths(keepOnlyMaxDelayPaths) {}
 };
 
 // This analysis finds the longest paths in the dataflow graph across modules.
@@ -192,22 +222,27 @@ class LongestPathAnalysis {
 public:
   // Entry points for analysis.
   LongestPathAnalysis(Operation *moduleOp, mlir::AnalysisManager &am,
-                      const LongestPathAnalysisOption &option = {});
+                      const LongestPathAnalysisOptions &option = {});
   ~LongestPathAnalysis();
 
   // Return all longest paths to each Fanin for the given value and bit
   // position.
-  LogicalResult getResults(Value value, size_t bitPos,
-                           SmallVectorImpl<DataflowPath> &results) const;
+  LogicalResult computeGlobalPaths(Value value, size_t bitPos,
+                                   SmallVectorImpl<DataflowPath> &results);
+
+  // Compute local paths for specified value and bit. Local paths are paths
+  // that are fully contained within a module.
+  FailureOr<ArrayRef<OpenPath>> computeLocalPaths(Value value, size_t bitPos);
 
-  // Return the maximum delay to the given value for all bit positions.
-  int64_t getMaxDelay(Value value) const;
+  // Return the maximum delay to the given value and bit position. If bitPos is
+  // negative, then return the maximum delay across all bits.
+  FailureOr<int64_t> getMaxDelay(Value value, int64_t bitPos = -1);
 
   // Return the average of the maximum delays across all bits of the given
   // value, which is useful approximation for the delay of the value. For each
   // bit position, finds all paths and takes the maximum delay. Then averages
   // these maximum delays across all bits of the value.
-  int64_t getAverageMaxDelay(Value value) const;
+  FailureOr<int64_t> getAverageMaxDelay(Value value);
 
   // Return paths that are closed under the given module. Closed paths are
   // typically register-to-register paths. A closed path is a path that starts
@@ -264,6 +299,7 @@ class LongestPathAnalysis {
 
 private:
   mlir::MLIRContext *ctx;
+  bool isAnalysisValid = true;
 };
 
 // Incremental version of longest path analysis that supports on-demand
@@ -272,14 +308,20 @@ class IncrementalLongestPathAnalysis : private LongestPathAnalysis,
                                        public mlir::PatternRewriter::Listener {
 public:
   IncrementalLongestPathAnalysis(Operation *moduleOp, mlir::AnalysisManager &am)
-      : LongestPathAnalysis(moduleOp, am,
-                            LongestPathAnalysisOption(false, true)) {}
-
-  // Compute maximum delay for specified value and bit.
-  FailureOr<int64_t> getOrComputeMaxDelay(Value value, size_t bitPos);
+      : LongestPathAnalysis(
+            moduleOp, am,
+            LongestPathAnalysisOptions(/*collectDebugInfo=*/false,
+                                       /*lazyComputation=*/true,
+                                       /*keepOnlyMaxDelayPaths=*/true)) {}
+
+  IncrementalLongestPathAnalysis(Operation *moduleOp, mlir::AnalysisManager &am,
+                                 const LongestPathAnalysisOptions &option)
+      : LongestPathAnalysis(moduleOp, am, option) {
+    assert(option.lazyComputation && "Lazy computation must be enabled");
+  }
 
-  // Compute all paths for specified value and bit.
-  FailureOr<ArrayRef<OpenPath>> getOrComputePaths(Value value, size_t bitPos);
+  using LongestPathAnalysis::getAverageMaxDelay;
+  using LongestPathAnalysis::getMaxDelay;
 
   // Check if operation can be safely modified without invalidating analysis.
   bool isOperationValidToMutate(Operation *op) const;
@@ -288,17 +330,6 @@ class IncrementalLongestPathAnalysis : private LongestPathAnalysis,
   void notifyOperationModified(Operation *op) override;
   void notifyOperationReplaced(Operation *op, ValueRange replacement) override;
   void notifyOperationErased(Operation *op) override;
-
-private:
-  bool isAnalysisValid = true;
-};
-
-// A wrapper class for the longest path analysis that also traces debug points.
-// This is necessary for analysis manager to cache the analysis results.
-class LongestPathAnalysisWithTrace : public LongestPathAnalysis {
-public:
-  LongestPathAnalysisWithTrace(Operation *moduleOp, mlir::AnalysisManager &am)
-      : LongestPathAnalysis(moduleOp, am, {true, false}) {}
 };
 
 // A collection of longest paths. The data structure owns the paths, and used
@@ -316,6 +347,9 @@ class LongestPathCollection {
   // Sort and drop all paths except the longest path per fanout point.
   void sortAndDropNonCriticalPathsPerFanOut();
 
+  // Merge another collection into this one.
+  void merge(const LongestPathCollection &other);
+
 private:
   MLIRContext *ctx;
 };

integration_test/Bindings/Python/dialects/aig.py

@@ -102,3 +102,24 @@ def build_top(module):
     # CHECK: top:test_aig;c[0] 0
     # CHECK: top:test_aig;child:test_child;a[0] 2
     print(collection.longest_path.to_flamegraph())
+
+    test_child = m.body.operations[0]
+    body_block = test_child.regions[0].blocks[0]
+    result0 = body_block.operations[0].results[0]
+    result1 = body_block.operations[1].results[0]
+
+    analysis = LongestPathAnalysis(test_child,
+                                   collect_debug_info=True,
+                                   keep_only_max_delay_paths=True,
+                                   lazy_computation=True)
+    c1 = analysis.get_paths(result0, 0)
+    c2 = analysis.get_paths(result1, 0)
+    # CHECK-NEXT: len(c1) = 1
+    # CHECK-NEXT: len(c2) = 1
+    print("len(c1) =", len(c1))
+    print("len(c2) =", len(c2))
+
+    s1, s2 = len(c1), len(c2)
+    c1.merge(c2)
+    # CHECK-NEXT: merge: True
+    print("merge:", len(c1) == s1 + s2)

lib/Bindings/Python/AIGModule.cpp

@@ -32,15 +32,32 @@ void circt::python::populateDialectAIGSubmodule(nb::module_ &m) {
       .def(
           "__init__",
           [](AIGLongestPathAnalysis *self, MlirOperation module,
-             bool traceDebugPoints) {
-            new (self) AIGLongestPathAnalysis(
-                aigLongestPathAnalysisCreate(module, traceDebugPoints));
+             bool collectDebugInfo, bool keepOnlyMaxDelayPaths,
+             bool lazyComputation) {
+            new (self) AIGLongestPathAnalysis(aigLongestPathAnalysisCreate(
+                module, collectDebugInfo, keepOnlyMaxDelayPaths,
+                lazyComputation));
           },
-          nb::arg("module"), nb::arg("trace_debug_points") = true)
+          nb::arg("module"), nb::arg("collect_debug_info") = false,
+          nb::arg("keep_only_max_delay_paths") = false,
+          nb::arg("lazy_computation") = false)
       .def("__del__",
            [](AIGLongestPathAnalysis &self) {
              aigLongestPathAnalysisDestroy(self);
            })
+      .def("get_paths",
+           [](AIGLongestPathAnalysis *self, MlirValue value, int64_t bitPos,
+              bool elaboratePaths) -> AIGLongestPathCollection {
+             auto collection =
+                 AIGLongestPathCollection(aigLongestPathAnalysisGetPaths(
+                     *self, value, bitPos, elaboratePaths));
+
+             if (aigLongestPathCollectionIsNull(collection))
+               throw nb::value_error(
+                   "Failed to get all paths, see previous error(s).");
+
+             return collection;
+           })
       .def("get_all_paths",
            [](AIGLongestPathAnalysis *self, const std::string &moduleName,
               bool elaboratePaths) -> AIGLongestPathCollection {
@@ -71,6 +88,10 @@ void circt::python::populateDialectAIGSubmodule(nb::module_ &m) {
            [](AIGLongestPathCollection &self,
               int pathIndex) -> AIGLongestPathDataflowPath {
              return aigLongestPathCollectionGetDataflowPath(self, pathIndex);
+           })
+      .def("merge",
+           [](AIGLongestPathCollection &self, AIGLongestPathCollection &src) {
+             aigLongestPathCollectionMerge(self, src);
            });
 
   nb::class_<AIGLongestPathDataflowPath>(m, "_LongestPathDataflowPath")

lib/Bindings/Python/dialects/aig.py

@@ -321,6 +321,16 @@ def print_summary(self) -> None:
     print(f"99th percentile delay: {self.get_by_delay_ratio(0.99).delay}")
     print(f"99.9th percentile delay: {self.get_by_delay_ratio(0.999).delay}")
 
+  def merge(self, src: "LongestPathCollection"):
+    """
+        Merge another collection into this one.
+        Args:
+            src: The collection to merge into this one
+    """
+    self.collection.merge(src.collection)
+    # Re-initialize to reflect the merged collection
+    self.__init__(self.collection)
+
 
 # ============================================================================
 # Main Analysis Interface
@@ -337,22 +347,48 @@ class LongestPathAnalysis:
         analysis: The underlying C++ analysis object
     """
 
-  def __init__(self, module, trace_debug_points: bool = True):
+  def __init__(self,
+               module,
+               collect_debug_info: bool = True,
+               keep_only_max_delay_paths: bool = False,
+               lazy_computation: bool = False):
     """
         Initialize the longest path analysis for a given module.
         Args:
             module: The MLIR module to analyze
-            trace_debug_points: Whether to include debug points in the analysis.
-                              The debug points provide additional information about the path,
-                              but increase the analysis time and memory usage.
+            collect_debug_info: Whether to include debug points in the analysis.
+                                Debug points provide additional information about the path,
+                                but increase analysis time and memory usage.
+            keep_only_max_delay_paths: Keep only maximum-delay paths in collections.
+            lazy_computation: Enable lazy (on-demand) computation.
         """
-    self.analysis = aig._LongestPathAnalysis(module, trace_debug_points)
+    self.analysis = aig._LongestPathAnalysis(module, collect_debug_info,
+                                             keep_only_max_delay_paths,
+                                             lazy_computation)
+
+  def get_paths(self,
+                value,
+                bit_pos: int,
+                elaborate_paths: bool = True) -> LongestPathCollection:
+    """
+        Perform longest path analysis and return all timing paths to the
+        specified value and bit position.
+        Args:
+            value: The value to analyze
+            bit_pos: The bit position to analyze
+            elaborate_paths: Whether to elaborate the paths with detailed information
+        Returns:
+            LongestPathCollection containing all paths sorted by delay
+      """
+    return LongestPathCollection(
+        self.analysis.get_paths(value, bit_pos, elaborate_paths))
 
   def get_all_paths(self,
                     module_name: str,
                     elaborate_paths: bool = True) -> LongestPathCollection:
     """
-        Perform longest path analysis and return all timing paths.
+        Perform longest path analysis and return all timing paths inside
+        the module hierarchy.
         This method analyzes the specified module and returns a collection
         of all timing paths, sorted by delay in descending order.
         Args:

lib/CAPI/Dialect/AIG.cpp

@@ -75,23 +75,39 @@ AIGLongestPathObject wrap(const DataflowPath::OutputPort *object) {
 //===----------------------------------------------------------------------===//
 
 AIGLongestPathAnalysis aigLongestPathAnalysisCreate(MlirOperation module,
-                                                    bool traceDebugPoints) {
+                                                    bool collectDebugInfo,
+                                                    bool keepOnlyMaxDelayPaths,
+                                                    bool lazyComputation) {
   auto *op = unwrap(module);
   auto *wrapper = new LongestPathAnalysisWrapper();
   wrapper->analysisManager =
       std::make_unique<mlir::ModuleAnalysisManager>(op, nullptr);
   mlir::AnalysisManager am = *wrapper->analysisManager;
-  if (traceDebugPoints)
-    wrapper->analysis = std::make_unique<LongestPathAnalysisWithTrace>(op, am);
-  else
-    wrapper->analysis = std::make_unique<LongestPathAnalysis>(op, am);
+  wrapper->analysis = std::make_unique<LongestPathAnalysis>(
+      op, am,
+      LongestPathAnalysisOptions(collectDebugInfo, lazyComputation,
+                                 keepOnlyMaxDelayPaths));
   return wrap(wrapper);
 }
 
 void aigLongestPathAnalysisDestroy(AIGLongestPathAnalysis analysis) {
   delete unwrap(analysis);
 }
 
+AIGLongestPathCollection
+aigLongestPathAnalysisGetPaths(AIGLongestPathAnalysis analysis, MlirValue value,
+                               int64_t bitPos, bool elaboratePaths) {
+  auto *wrapper = unwrap(analysis);
+  auto *lpa = wrapper->analysis.get();
+  auto *collection = new LongestPathCollection(lpa->getContext());
+  auto result =
+      lpa->computeGlobalPaths(unwrap(value), bitPos, collection->paths);
+  if (failed(result))
+    return {nullptr};
+  collection->sortInDescendingOrder();
+  return wrap(collection);
+}
+
 AIGLongestPathCollection
 aigLongestPathAnalysisGetAllPaths(AIGLongestPathAnalysis analysis,
                                   MlirStringRef moduleName,
@@ -136,6 +152,13 @@ aigLongestPathCollectionGetDataflowPath(AIGLongestPathCollection collection,
   return wrap(&path);
 }
 
+void aigLongestPathCollectionMerge(AIGLongestPathCollection dest,
+                                   AIGLongestPathCollection src) {
+  auto *destWrapper = unwrap(dest);
+  auto *srcWrapper = unwrap(src);
+  destWrapper->merge(*srcWrapper);
+}
+
 //===----------------------------------------------------------------------===//
 // DataflowPath
 //===----------------------------------------------------------------------===//