[XLS] Apply our area & delay models for our LUT conversion pass

By estimating when it's beneficial to merge logic into an existing select, we can apply our LUT identification & conversion logic more often.

PiperOrigin-RevId: 691566073

Author: ericastor

xls/build_rules/BUILD

@@ -216,10 +216,13 @@ bzl_library(
     visibility = ["//visibility:private"],
     deps = [
         ":xls_codegen_rules_bzl",
+        ":xls_common_rules_bzl",
         ":xls_config_rules_bzl",
+        ":xls_dslx_rules_bzl",
         ":xls_ir_macros_bzl",
         ":xls_ir_rules_bzl",
         ":xls_rules_bzl",
+        ":xls_toolchains_bzl",
         ":xls_type_check_utils_bzl",
         "@bazel_skylib//rules:build_test",
         "@bazel_skylib//rules:diff_test",

xls/build_rules/xls_config_rules.bzl

@@ -21,12 +21,14 @@ This module exposes the configuration parameters for the XLS build rules. It
 load(
     "//xls/build_rules:xls_oss_config_rules.bzl",
     _CONFIG = "CONFIG",
+    _DEFAULT_BENCHMARK_SYNTH_AREA_MODEL = "DEFAULT_BENCHMARK_SYNTH_AREA_MODEL",
     _DEFAULT_BENCHMARK_SYNTH_DELAY_MODEL = "DEFAULT_BENCHMARK_SYNTH_DELAY_MODEL",
     _delay_model_to_standard_cells = "delay_model_to_standard_cells",
     _enable_generated_file_wrapper = "enable_generated_file_wrapper",
 )
 
 CONFIG = _CONFIG
 DEFAULT_BENCHMARK_SYNTH_DELAY_MODEL = _DEFAULT_BENCHMARK_SYNTH_DELAY_MODEL
+DEFAULT_BENCHMARK_SYNTH_AREA_MODEL = _DEFAULT_BENCHMARK_SYNTH_AREA_MODEL
 delay_model_to_standard_cells = _delay_model_to_standard_cells
 enable_generated_file_wrapper = _enable_generated_file_wrapper

xls/build_rules/xls_ir_rules.bzl

@@ -225,6 +225,8 @@ def _optimize_ir(ctx, src, original_input_files):
         "optimize_for_best_case_throughput",
         "enable_resource_sharing",
         "top",
+        "delay_model",
+        "area_model",
     )
 
     is_args_valid(opt_ir_args, IR_OPT_FLAGS)
@@ -462,6 +464,8 @@ def get_benchmark_ir_cmd(ctx, src, append_cmd_line_args = True):
         "use_context_narrowing_analysis",
         "optimize_for_best_case_throughput",
         "enable_resource_sharing",
+        "delay_model",
+        "area_model",
         "run_evaluators",
     ] + _CODEGEN_FLAGS + _SCHEDULING_FLAGS
 

xls/build_rules/xls_macros.bzl

@@ -33,6 +33,7 @@ load(
 )
 load(
     "//xls/build_rules:xls_config_rules.bzl",
+    "DEFAULT_BENCHMARK_SYNTH_AREA_MODEL",
     "DEFAULT_BENCHMARK_SYNTH_DELAY_MODEL",
     "delay_model_to_standard_cells",
     "enable_generated_file_wrapper",
@@ -557,12 +558,14 @@ Examples:
     # Setup shared arguments
     SHARED_FLAGS = (
         "top",
+        "delay_model",
     )
     IR_OPT_FLAGS = (
         "ir_dump_path",
         "passes",
         "skip_passes",
         "opt_level",
+        "area_model",
         "convert_array_index_to_select",
         "use_context_narrowing_analysis",
         "optimize_for_best_case_throughput",
@@ -579,8 +582,14 @@ Examples:
         if k not in IR_OPT_FLAGS or k in SHARED_FLAGS
     }
 
-    # Add default opt args (currently empty)
-    full_opt_args = dict()
+    # Add default opt args
+    full_opt_args = {
+        "delay_model": DEFAULT_BENCHMARK_SYNTH_DELAY_MODEL,
+        "area_model": DEFAULT_BENCHMARK_SYNTH_AREA_MODEL,
+    }
+    if "delay_model" in opt_ir_args and "area_model" not in opt_ir_args:
+        # Default to the area model matching the delay model.
+        opt_ir_args["area_model"] = opt_ir_args["delay_model"]
     full_opt_args.update(opt_ir_args)
 
     # Add default codegen args

xls/build_rules/xls_oss_config_rules.bzl

@@ -23,6 +23,7 @@ CONFIG = {
 }
 
 DEFAULT_BENCHMARK_SYNTH_DELAY_MODEL = "asap7"
+DEFAULT_BENCHMARK_SYNTH_AREA_MODEL = "asap7"
 
 def enable_generated_file_wrapper(**kwargs):  # @unused
     """The function is a placeholder for enable_generated_file_wrapper.

xls/dev_tools/BUILD

@@ -621,6 +621,8 @@ cc_binary(
         "//xls/common/status:ret_check",
         "//xls/common/status:status_macros",
         "//xls/data_structures:binary_decision_diagram",
+        "//xls/estimators/area_model:area_estimator",
+        "//xls/estimators/area_model:area_estimators",
         "//xls/estimators/delay_model:analyze_critical_path",
         "//xls/estimators/delay_model:delay_estimator",
         "//xls/estimators/delay_model:delay_estimators",

xls/dev_tools/benchmark_main.cc

@@ -50,6 +50,8 @@
 #include "xls/common/status/status_macros.h"
 #include "xls/data_structures/binary_decision_diagram.h"
 #include "xls/dev_tools/pipeline_metrics.h"
+#include "xls/estimators/area_model/area_estimator.h"
+#include "xls/estimators/area_model/area_estimators.h"
 #include "xls/estimators/delay_model/analyze_critical_path.h"
 #include "xls/estimators/delay_model/delay_estimator.h"
 #include "xls/estimators/delay_model/delay_estimators.h"
@@ -129,6 +131,7 @@ ABSL_FLAG(
     "but at the cost of constraining the schedule and thus increasing area.");
 ABSL_FLAG(bool, enable_resource_sharing, false,
           "Enable the resource sharing optimization to save area.");
+ABSL_FLAG(std::string, area_model, "", "Area model name to use from registry.");
 ABSL_FLAG(bool, run_evaluators, true,
           "Whether to run the JIT and interpreter.");
 ABSL_FLAG(bool, compare_delay_to_synthesis, false,
@@ -176,7 +179,9 @@ int64_t DurationToMs(absl::Duration duration) {
 
 // Run the standard pipeline on the given package and prints stats about the
 // passes and execution time.
-absl::Status RunOptimizationAndPrintStats(Package* package) {
+absl::Status RunOptimizationAndPrintStats(Package* package,
+                                          DelayEstimator* delay_estimator,
+                                          AreaEstimator* area_estimator) {
   std::unique_ptr<OptimizationCompoundPass> pipeline =
       CreateOptimizationPassPipeline();
 
@@ -197,6 +202,8 @@ absl::Status RunOptimizationAndPrintStats(Package* package) {
           : std::make_optional(split_next_value_selects);
   pass_options.use_context_narrowing_analysis =
       absl::GetFlag(FLAGS_use_context_narrowing_analysis);
+  pass_options.delay_estimator = delay_estimator;
+  pass_options.area_estimator = area_estimator;
   PassResults pass_results;
   OptimizationContext context;
   XLS_RETURN_IF_ERROR(
@@ -285,6 +292,18 @@ absl::Status PrintTotalDelay(FunctionBase* f,
   return absl::OkStatus();
 }
 
+absl::Status PrintTotalArea(FunctionBase* f,
+                            const AreaEstimator& area_estimator) {
+  int64_t total_area = 0;
+  for (Node* node : f->nodes()) {
+    XLS_ASSIGN_OR_RETURN(double op_area,
+                         area_estimator.GetOperationAreaInSquareMicrons(node));
+    total_area += op_area;
+  }
+  std::cout << absl::StrFormat("Total area: %.4f um2\n", total_area);
+  return absl::OkStatus();
+}
+
 // Returns the critical-path delay through each pipeline stage.
 absl::StatusOr<std::vector<int64_t>> GetDelayPerStageInPs(
     FunctionBase* f, const PipelineSchedule& schedule,
@@ -695,8 +714,9 @@ absl::Status RunInterpreterAndJit(FunctionBase* function_base,
 
 absl::Status AnalyzeAndPrintCriticalPath(
     FunctionBase* f, std::optional<int64_t> effective_clock_period_ps,
-    const DelayEstimator& delay_estimator, const QueryEngine& query_engine,
-    PipelineScheduleOrGroup* schedules, synthesis::Synthesizer* synthesizer) {
+    const DelayEstimator& delay_estimator, const AreaEstimator* area_estimator,
+    const QueryEngine& query_engine, PipelineScheduleOrGroup* schedules,
+    synthesis::Synthesizer* synthesizer) {
   XLS_ASSIGN_OR_RETURN(
       std::vector<CriticalPathEntry> critical_path,
       AnalyzeCriticalPath(f, effective_clock_period_ps, delay_estimator));
@@ -718,6 +738,9 @@ absl::Status AnalyzeAndPrintCriticalPath(
   }
   XLS_RETURN_IF_ERROR(PrintCriticalPath(f, query_engine, delay_diff));
   XLS_RETURN_IF_ERROR(PrintTotalDelay(f, delay_estimator));
+  if (area_estimator != nullptr) {
+    XLS_RETURN_IF_ERROR(PrintTotalArea(f, *area_estimator));
+  }
   return absl::OkStatus();
 }
 
@@ -747,7 +770,19 @@ absl::Status RealMain(std::string_view path) {
     XLS_RETURN_IF_ERROR(
         RunInterpreterAndJit(package->GetTop().value(), "unoptimized"));
   }
-  XLS_RETURN_IF_ERROR(RunOptimizationAndPrintStats(package.get()));
+  DelayEstimator* delay_estimator = nullptr;
+  if (delay_model_flag_passed) {
+    XLS_ASSIGN_OR_RETURN(
+        delay_estimator,
+        GetDelayEstimator(scheduling_options_flags_proto.delay_model()));
+  }
+  AreaEstimator* area_estimator = nullptr;
+  if (std::string area_model = absl::GetFlag(FLAGS_area_model);
+      !area_model.empty()) {
+    XLS_ASSIGN_OR_RETURN(area_estimator, GetAreaEstimator(area_model));
+  }
+  XLS_RETURN_IF_ERROR(RunOptimizationAndPrintStats(
+      package.get(), delay_estimator, area_estimator));
 
   FunctionBase* f = package->GetTop().value();
   BddQueryEngine query_engine(BddQueryEngine::kDefaultPathLimit);
@@ -768,15 +803,8 @@ absl::Status RealMain(std::string_view path) {
               100;
     }
   }
-  const DelayEstimator* pdelay_estimator;
-  if (!delay_model_flag_passed) {
-    pdelay_estimator = &GetStandardDelayEstimator();
-  } else {
-    XLS_ASSIGN_OR_RETURN(
-        pdelay_estimator,
-        GetDelayEstimator(scheduling_options_flags_proto.delay_model()));
-  }
-  const auto& delay_estimator = *pdelay_estimator;
+  const DelayEstimator& defaulted_delay_estimator =
+      delay_estimator ? *delay_estimator : GetStandardDelayEstimator();
   std::unique_ptr<synthesis::Synthesizer> synthesizer;
   if (absl::GetFlag(FLAGS_compare_delay_to_synthesis)) {
     synthesis::GrpcSynthesizerParameters parameters(
@@ -794,7 +822,8 @@ absl::Status RealMain(std::string_view path) {
       scheduling_options_flags_proto.pipeline_stages() > 0;
   if (!f->IsProc() && !benchmark_codegen) {
     XLS_RETURN_IF_ERROR(AnalyzeAndPrintCriticalPath(
-        f, effective_clock_period_ps, delay_estimator, query_engine,
+        f, effective_clock_period_ps, defaulted_delay_estimator, area_estimator,
+        query_engine,
         /*schedules=*/nullptr, synthesizer.get()));
   } else if (benchmark_codegen) {
     PipelineScheduleOrGroup schedules = PackagePipelineSchedules();
@@ -803,17 +832,17 @@ absl::Status RealMain(std::string_view path) {
                            SetUpSchedulingOptions(
                                scheduling_options_flags_proto, package.get()));
       absl::Duration scheduling_time;
-      XLS_ASSIGN_OR_RETURN(schedules,
-                           Schedule(package.get(), scheduling_options,
-                                    &delay_estimator, &scheduling_time));
+      XLS_ASSIGN_OR_RETURN(
+          schedules, Schedule(package.get(), scheduling_options,
+                              &defaulted_delay_estimator, &scheduling_time));
       std::cout << absl::StreamFormat("Scheduling time: %dms\n",
                                       scheduling_time / absl::Milliseconds(1));
       XLS_RETURN_IF_ERROR(AnalyzeAndPrintCriticalPath(
-          f, effective_clock_period_ps, delay_estimator, query_engine,
-          &schedules, synthesizer.get()));
+          f, effective_clock_period_ps, defaulted_delay_estimator,
+          area_estimator, query_engine, &schedules, synthesizer.get()));
 
       XLS_RETURN_IF_ERROR(PrintScheduleInfo(
-          f, schedules, query_engine, delay_estimator,
+          f, schedules, query_engine, defaulted_delay_estimator,
           scheduling_options_flags_proto.has_clock_period_ps()
               ? std::make_optional(
                     scheduling_options_flags_proto.clock_period_ps())

xls/ir/BUILD

@@ -1008,6 +1008,7 @@ cc_library(
         ":verifier",
         "//xls/common/status:ret_check",
         "//xls/common/status:status_macros",
+        "//xls/estimators/area_model:area_estimator",
         "//xls/estimators/delay_model:delay_estimator",
         "@com_google_absl//absl/container:flat_hash_map",
         "@com_google_absl//absl/log",

xls/ir/function_base.cc

@@ -135,7 +135,8 @@ absl::StatusOr<Node*> FunctionBase::GetNode(
 }
 
 absl::Status FunctionBase::RemoveNode(Node* node) {
-  XLS_RET_CHECK(node->users().empty()) << node->GetName();
+  XLS_RET_CHECK(node->users().empty())
+      << node->GetName() << ", users " << node->GetUsersString();
   XLS_RET_CHECK(!HasImplicitUse(node)) << node->GetName();
   VLOG(4) << absl::StrFormat("Removing node from FunctionBase %s: %s", name(),
                              node->ToString());

xls/ir/ir_test_base.h

@@ -27,6 +27,7 @@
 #include "absl/status/statusor.h"
 #include "absl/strings/str_format.h"
 #include "absl/strings/str_replace.h"
+#include "xls/estimators/area_model/area_estimator.h"
 #include "xls/estimators/delay_model/delay_estimator.h"
 #include "xls/ir/function.h"
 #include "xls/ir/function_base.h"
@@ -148,6 +149,41 @@ class TestDelayEstimator : public DelayEstimator {
   int64_t base_delay_;
 };
 
+class TestAreaEstimator : public AreaEstimator {
+ public:
+  explicit TestAreaEstimator(double base_area = 1)
+      : AreaEstimator("test"), base_area_(base_area) {}
+
+  absl::StatusOr<double> GetOperationAreaInSquareMicrons(
+      Node* node) const override {
+    switch (node->op()) {
+      case Op::kAfterAll:
+      case Op::kMinDelay:
+      case Op::kBitSlice:
+      case Op::kConcat:
+      case Op::kLiteral:
+      case Op::kParam:
+      case Op::kNext:
+      case Op::kReceive:
+      case Op::kSend:
+      case Op::kTupleIndex:
+        return 0.0;
+      case Op::kUDiv:
+      case Op::kSDiv:
+        return 2 * base_area_;
+      default:
+        return base_area_;
+    }
+  }
+
+ private:
+  absl::StatusOr<double> GetOneBitRegisterAreaInSquareMicrons() const override {
+    return base_area_;
+  }
+
+  double base_area_;
+};
+
 // Helper to record IR before and after some test event which changes it.
 struct ScopedRecordIr {
  public: