[codegen] Add extension mechanism to CodegenOptions and CodegenFlagsProto

xls/codegen/codegen_options.h

@@ -16,6 +16,7 @@
 #define XLS_CODEGEN_CODEGEN_OPTIONS_H_
 
 #include <cstdint>
+#include <memory>
 #include <optional>
 #include <string>
 #include <string_view>
@@ -37,11 +38,32 @@
 
 namespace xls::verilog {
 
+class CodegenOptionExtension {
+ public:
+  virtual ~CodegenOptionExtension() = default;
+  virtual std::string_view extension_name() const = 0;
+};
+
 // Options describing how codegen should be performed.
 class CodegenOptions {
  public:
   explicit CodegenOptions() = default;
 
+  template <typename T>
+  const T* get_extension() const {
+    for (const auto& ext : extensions_) {
+      if (ext->extension_name() == T::kExtensionName) {
+        return static_cast<const T*>(ext.get());
+      }
+    }
+    return nullptr;
+  }
+
+  CodegenOptions& add_extension(std::unique_ptr<CodegenOptionExtension> ext) {
+    extensions_.push_back(std::move(ext));
+    return *this;
+  }
+
   // Enum to describe which codegen version to use.
   enum class Version : uint8_t {
     kDefault = 0,
@@ -461,6 +483,7 @@ class CodegenOptions {
   int64_t max_trace_verbosity_ = 0;
   RegisterMergeStrategy register_merge_strategy_ =
       RegisterMergeStrategy::kDefault;
+
   SourceAnnotationStrategy source_annotation_strategy_ =
       SourceAnnotationStrategy::kNone;
   std::optional<PackageInterfaceProto> package_interface_;
@@ -474,6 +497,8 @@ class CodegenOptions {
   std::vector<int32_t> randomize_order_seed_;
   std::optional<CodegenResidualData> residual_data_;
   std::optional<std::string> ir_dump_path_;
+
+  std::vector<std::shared_ptr<CodegenOptionExtension>> extensions_;
 };
 
 template <typename Sink>

xls/tools/BUILD

@@ -725,6 +725,7 @@ cc_library(
     hdrs = ["codegen_flags.h"],
     deps = [
         ":codegen_flags_cc_proto",
+        ":codegen_flags_handler_registry",
         "//xls/codegen:codegen_residual_data_cc_proto",
         "//xls/common/file:filesystem",
         "//xls/common/status:ret_check",
@@ -765,12 +766,25 @@ cc_library(
     ],
 )
 
+cc_library(
+    name = "codegen_flags_handler_registry",
+    srcs = ["codegen_flags_handler_registry.cc"],
+    hdrs = ["codegen_flags_handler_registry.h"],
+    deps = [
+        ":codegen_flags_cc_proto",
+        "//xls/codegen:codegen_options",
+        "//xls/common/status:status_macros",
+        "@com_google_absl//absl/status",
+    ],
+)
+
 cc_library(
     name = "codegen",
     srcs = ["codegen.cc"],
     hdrs = ["codegen.h"],
     deps = [
         ":codegen_flags_cc_proto",
+        ":codegen_flags_handler_registry",
         ":schedule",
         ":scheduling_options_flags_cc_proto",
         "//xls/codegen:block_metrics",

xls/tools/codegen.cc

@@ -53,6 +53,7 @@
 #include "xls/scheduling/scheduling_options.h"
 #include "xls/scheduling/scheduling_result.h"
 #include "xls/tools/codegen_flags.pb.h"
+#include "xls/tools/codegen_flags_handler_registry.h"
 #include "xls/tools/schedule.h"
 #include "xls/tools/scheduling_options_flags.pb.h"
 
@@ -455,6 +456,7 @@ absl::StatusOr<verilog::CodegenOptions> CodegenOptionsFromProto(
     options.set_ir_dump_path(p.ir_dump_path());
   }
 
+  XLS_RETURN_IF_ERROR(CodegenFlagsHandlerRegistry::Process(p, options));
   return options;
 }
 

xls/tools/codegen_flags.cc

@@ -37,6 +37,7 @@
 #include "xls/common/status/status_macros.h"
 #include "xls/ir/xls_ir_interface.pb.h"
 #include "xls/tools/codegen_flags.pb.h"
+#include "xls/tools/codegen_flags_handler_registry.h"
 
 // LINT.IfChange
 ABSL_FLAG(
@@ -438,6 +439,7 @@ static absl::StatusOr<bool> SetOptionsFromFlags(CodegenFlagsProto& proto) {
   POPULATE_FLAG(array_index_bounds_checking);
   POPULATE_FLAG(fifo_module);
   POPULATE_FLAG(nodata_fifo_module);
+
   if (absl::GetFlag(FLAGS_materialize_internal_fifos)) {
     any_flags_set = true;
     if (!FLAGS_fifo_module.IsSpecifiedOnCommandLine()) {
@@ -451,6 +453,7 @@ static absl::StatusOr<bool> SetOptionsFromFlags(CodegenFlagsProto& proto) {
              "--materialize_internal_fifos.";
     }
   }
+
   XLS_ASSIGN_OR_RETURN(
       RegisterMergeStrategyProto merge_strategy,
       MergeStrategyFromString(absl::GetFlag(FLAGS_register_merge_strategy)));
@@ -514,6 +517,7 @@ absl::StatusOr<CodegenFlagsProto> GetCodegenFlags() {
     XLS_RETURN_IF_ERROR(xls::ParseTextProtoFile(
         absl::GetFlag(FLAGS_codegen_options_proto), &proto));
   }
+  XLS_RETURN_IF_ERROR(CodegenFlagsHandlerRegistry::ParseFlags(proto));
   if (absl::GetFlag(FLAGS_codegen_options_used_textproto_file)) {
     XLS_RETURN_IF_ERROR(SetTextProtoFile(
         *absl::GetFlag(FLAGS_codegen_options_used_textproto_file), proto));

xls/tools/codegen_flags.proto

@@ -12,7 +12,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-syntax = "proto3";
+edition = "2023";
 
 package xls;
 
@@ -56,55 +56,55 @@ enum CodegenVersionProto {
 //
 // See codegen_flags.cc ABSL_FLAG() definitions for the meaning of these fields.
 message CodegenFlagsProto {
-  optional string top = 1;
-  optional GeneratorKind generator = 2;
-  optional string input_valid_signal = 3;
-  optional string output_valid_signal = 4;
-  optional string manual_load_enable_signal = 5;
-  optional bool flop_inputs = 6;
-  optional bool flop_outputs = 7;
-  optional IOKindProto flop_inputs_kind = 8;
-  optional IOKindProto flop_outputs_kind = 9;
-  optional bool flop_single_value_channels = 10;
-  optional bool add_idle_output = 11;
-  optional string module_name = 12;
-  optional string output_port_name = 13;
-  optional string reset = 14;
-  optional bool reset_active_low = 15;
-  optional bool reset_asynchronous = 16;
-  optional bool reset_data_path = 17;
-  optional bool use_system_verilog = 18;
-  optional bool separate_lines = 19;
-  optional int64 max_inline_depth = 35;
-  optional string gate_format = 20;
-  optional string assert_format = 21;
-  optional string smulp_format = 22;
-  optional string umulp_format = 23;
-  optional string streaming_channel_data_suffix = 24;
-  optional string streaming_channel_valid_suffix = 25;
-  optional string streaming_channel_ready_suffix = 26;
+  string top = 1;
+  GeneratorKind generator = 2;
+  string input_valid_signal = 3;
+  string output_valid_signal = 4;
+  string manual_load_enable_signal = 5;
+  bool flop_inputs = 6;
+  bool flop_outputs = 7;
+  IOKindProto flop_inputs_kind = 8;
+  IOKindProto flop_outputs_kind = 9;
+  bool flop_single_value_channels = 10;
+  bool add_idle_output = 11;
+  string module_name = 12;
+  string output_port_name = 13;
+  string reset = 14;
+  bool reset_active_low = 15;
+  bool reset_asynchronous = 16;
+  bool reset_data_path = 17;
+  bool use_system_verilog = 18;
+  bool separate_lines = 19;
+  int64 max_inline_depth = 35;
+  string gate_format = 20;
+  string assert_format = 21;
+  string smulp_format = 22;
+  string umulp_format = 23;
+  string streaming_channel_data_suffix = 24;
+  string streaming_channel_valid_suffix = 25;
+  string streaming_channel_ready_suffix = 26;
   repeated string ram_configurations = 27;
-  optional bool gate_recvs = 28;
-  optional bool array_index_bounds_checking = 29;
-  optional RegisterMergeStrategyProto register_merge_strategy = 30;
-  optional int64 max_trace_verbosity = 31;
+  bool gate_recvs = 28;
+  bool array_index_bounds_checking = 29;
+  RegisterMergeStrategyProto register_merge_strategy = 30;
+  int64 max_trace_verbosity = 31;
   // If present details about the interface requested. Eg specific sv types to
   // use for arguments etc. Only the 'top' and channels are interpreted. Unknown
   // interface elements are ignored.
-  optional PackageInterfaceProto package_interface = 32;
+  PackageInterfaceProto package_interface = 32;
   // Should annotated arguments be emitted with the sv_types they are annotated
   // with.
-  optional bool emit_sv_types = 33;
+  bool emit_sv_types = 33;
 
-  optional string simulation_macro_name = 34;
+  string simulation_macro_name = 34;
   repeated string assertion_macro_names = 39;
 
-  optional CodegenVersionProto codegen_version = 36;
+  CodegenVersionProto codegen_version = 36;
 
   // Which module to use for FIFOs. If empty, will materialize an internal
   // implementation.
-  optional string fifo_module = 40;
-  optional string nodata_fifo_module = 41;
+  string fifo_module = 40;
+  string nodata_fifo_module = 41;
 
   // If present, the seed used to randomize the order of lines in the output. If
   // empty, will use a default order. This can be useful for creating multiple
@@ -113,14 +113,16 @@ message CodegenFlagsProto {
 
   // If false, runtime invariant assertions (e.g. one-hot selector checks)
   // are omitted from generated RTL.  Default is true.
-  optional bool add_invariant_assertions = 42;
+  bool add_invariant_assertions = 42;
 
   // Parsed reference residual data. When present, codegen uses this to
   // influence emission order; preferred over reading from a path at runtime.
-  optional verilog.CodegenResidualData reference_residual_data = 43;
+  verilog.CodegenResidualData reference_residual_data = 43;
 
-  optional SourceAnnotationStrategyProto source_annotation_strategy = 44;
+  SourceAnnotationStrategyProto source_annotation_strategy = 44;
 
   // Debugging flag to cause the block conversion to dump its pass results.
-  optional string ir_dump_path = 45;
+  string ir_dump_path = 45;
+
+  extensions 20000 to max;
 }

xls/tools/codegen_flags_handler_registry.cc

@@ -0,0 +1,65 @@
+// Copyright 2026 The XLS Authors
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "xls/tools/codegen_flags_handler_registry.h"
+
+#include <string_view>
+#include <utility>
+#include <vector>
+
+#include "absl/status/status.h"
+#include "xls/codegen/codegen_options.h"
+#include "xls/common/status/status_macros.h"
+#include "xls/tools/codegen_flags.pb.h"
+
+namespace xls {
+namespace {
+
+struct RegisteredHandler {
+  std::string_view name;
+  CodegenFlagsHandler handler;
+  CodegenFlagsParser parser;
+};
+
+std::vector<RegisteredHandler>& GetRegistry() {
+  static auto* registry = new std::vector<RegisteredHandler>();
+  return *registry;
+}
+
+}  // namespace
+
+void CodegenFlagsHandlerRegistry::Register(std::string_view name,
+                                           CodegenFlagsHandler handler,
+                                           CodegenFlagsParser parser) {
+  GetRegistry().push_back({name, std::move(handler), std::move(parser)});
+}
+
+absl::Status CodegenFlagsHandlerRegistry::ParseFlags(CodegenFlagsProto& proto) {
+  for (const auto& registered : GetRegistry()) {
+    if (registered.parser != nullptr) {
+      XLS_RETURN_IF_ERROR(registered.parser(proto));
+    }
+  }
+  return absl::OkStatus();
+}
+
+absl::Status CodegenFlagsHandlerRegistry::Process(
+    const CodegenFlagsProto& proto, verilog::CodegenOptions& options) {
+  for (const auto& registered : GetRegistry()) {
+    XLS_RETURN_IF_ERROR(registered.handler(proto, options));
+  }
+  return absl::OkStatus();
+}
+
+}  // namespace xls