[xls][mlir] Enable populating codegen metrics.

Plumb through callback to enable populating using the returned package and BlockMetricsProto.

PiperOrigin-RevId: 714680081

Author: jpienaar

xls/contrib/mlir/BUILD

@@ -541,6 +541,8 @@ cc_library(
     ],
     deps = [
         ":mlir_xls",
+        "//xls/codegen:module_signature",
+        "//xls/codegen:xls_metrics_cc_proto",
         "//xls/codegen/vast",
         "//xls/common/file:filesystem",
         "//xls/common/file:get_runfile_path",

xls/contrib/mlir/tools/xls_translate/xls_translate.cc

@@ -58,6 +58,7 @@
 #include "mlir/include/mlir/Support/LLVM.h"
 #include "mlir/include/mlir/Support/LogicalResult.h"
 #include "mlir/include/mlir/Transforms/Passes.h"
+#include "xls/codegen/module_signature.h"
 #include "xls/codegen/vast/vast.h"
 #include "xls/common/file/filesystem.h"
 #include "xls/common/file/get_runfile_path.h"
@@ -1363,7 +1364,8 @@ LogicalResult setTop(Operation* op, std::string_view name, Package* package) {
 }  // namespace
 
 LogicalResult MlirXlsToXlsTranslate(Operation* op, llvm::raw_ostream& output,
-                                    MlirXlsToXlsTranslateOptions options) {
+                                    MlirXlsToXlsTranslateOptions options,
+                                    MetricsReporter metrics_reporter) {
   DslxPackageCache maybe_cache;
   if (options.dslx_cache == nullptr) {
     options.dslx_cache = &maybe_cache;
@@ -1416,6 +1418,20 @@ LogicalResult MlirXlsToXlsTranslate(Operation* op, llvm::raw_ostream& output,
     return failure();
   }
 
+  if (metrics_reporter) {
+    const ::xls::verilog::ModuleSignature signature =
+        xls_codegen_results.value().module_generator_result.signature;
+    if (signature.proto().has_metrics()) {
+      const ::xls::verilog::XlsMetricsProto& metrics =
+          signature.proto().metrics();
+      if (metrics.has_block_metrics()) {
+        const ::xls::verilog::BlockMetricsProto& block_metrics =
+            metrics.block_metrics();
+        metrics_reporter(**package, block_metrics);
+      }
+    }
+  }
+
   output << xls_codegen_results->module_generator_result.verilog_text;
   return success();
 }

xls/contrib/mlir/tools/xls_translate/xls_translate.h

@@ -1,4 +1,3 @@
-#include "xls/tools/opt.h"
 // Copyright 2024 The XLS Authors
 //
 // Licensed under the Apache License, Version 2.0 (the "License");
@@ -22,10 +21,13 @@
 #include "absl/container/flat_hash_map.h"
 #include "absl/log/check.h"
 #include "absl/status/statusor.h"
+#include "llvm/include/llvm/ADT/STLFunctionalExtras.h"
 #include "llvm/include/llvm/ADT/StringRef.h"
 #include "mlir/include/mlir/Support/LLVM.h"
+#include "xls/codegen/xls_metrics.pb.h"
 #include "xls/tools/codegen_flags.h"
 #include "xls/tools/codegen_flags.pb.h"
+#include "xls/tools/opt.h"
 #include "xls/tools/scheduling_options_flags.h"
 #include "xls/tools/scheduling_options_flags.pb.h"
 
@@ -89,9 +91,14 @@ struct MlirXlsToXlsTranslateOptions {
       DieUnlessOk(::xls::GetSchedulingOptionsFlagsProto());
 };
 
+// Callback for reporting codegen metrics.
+using MetricsReporter = llvm::function_ref<void(
+    const ::xls::Package&, const ::xls::verilog::BlockMetricsProto&)>;
+
 // Translates an operation with XLS dialect to DSLX.
 LogicalResult MlirXlsToXlsTranslate(Operation* op, llvm::raw_ostream& output,
-                                    MlirXlsToXlsTranslateOptions options = {});
+                                    MlirXlsToXlsTranslateOptions options = {},
+                                    MetricsReporter metrics_reporter = nullptr);
 
 }  // namespace mlir::xls
 

xls/contrib/mlir/tools/xls_translate/xls_translate_main.cc

@@ -64,6 +64,12 @@ llvm::cl::opt<bool> privatizeAndDceFunctions(
                    "SymbolDCE first"),
     llvm::cl::init(false));
 
+// NOLINTNEXTLINE
+llvm::cl::opt<bool> dumpCodegenMetrics(
+    "dump-codegen-metrics",
+    llvm::cl::desc("Whether to dump XLS codegen metric"),
+    llvm::cl::init(false));
+
 void registerInputDialects(DialectRegistry& registry) {
   // TODO(jpienaar): Registering all as start/prototyping.
   mlir::registerAllDialects(registry);
@@ -73,6 +79,23 @@ void registerInputDialects(DialectRegistry& registry) {
                   mlir::tensor::TensorDialect>();
 }
 
+static void printCodegenMetrics(
+    const ::xls::Package& package,
+    const ::xls::verilog::BlockMetricsProto& metrics) {
+  llvm::errs() << "Generated XLS metrics:\n";
+  llvm::errs() << "  xls_node count: " << package.GetFunctionNodeCount()
+               << "\n";
+  llvm::errs() << "  xls_flop count: " << metrics.flop_count() << "\n";
+  llvm::errs() << "  xls_max_input_to_reg_delay_ps: "
+               << metrics.max_input_to_reg_delay_ps() << "\n";
+  llvm::errs() << "  xls_max_reg_to_reg_delay_ps: "
+               << metrics.max_reg_to_reg_delay_ps() << "\n";
+  llvm::errs() << "  xls_max_reg_to_output_delay_ps: "
+               << metrics.max_reg_to_output_delay_ps() << "\n";
+  llvm::errs() << "  xls_max_feedthrough_path_delay_ps: "
+               << metrics.max_feedthrough_path_delay_ps() << "\n";
+}
+
 LogicalResult mlirXlsToXlsTranslate(Operation* op, llvm::raw_ostream& output) {
   MlirXlsToXlsTranslateOptions options;
   options.main_function = mainFunction;
@@ -90,6 +113,9 @@ LogicalResult mlirXlsToVerilogTranslate(Operation* op,
   options.optimize_ir = optimizeIr;
   options.dslx_search_path = dslxSearchPath;
   options.generate_verilog = true;
+  if (dumpCodegenMetrics) {
+    return MlirXlsToXlsTranslate(op, output, options, printCodegenMetrics);
+  }
   return MlirXlsToXlsTranslate(op, output, options);
 }