Skip to content

Commit

Permalink
Control opt_main pipeline with a proto
Browse files Browse the repository at this point in the history
Give the ability to specify a full opt pipeline as a proto and the ability to dump an existing pipeline as a proto. This is useful for allowing one to make tools which can modify a pipeline for experimentation.

PiperOrigin-RevId: 699296142
  • Loading branch information
allight authored and copybara-github committed Nov 22, 2024
1 parent 4718484 commit f47eb95
Show file tree
Hide file tree
Showing 20 changed files with 859 additions and 62 deletions.
67 changes: 67 additions & 0 deletions xls/passes/BUILD
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,9 @@

# Optimization passes, pass managers.

load("@rules_python//python:proto.bzl", "py_proto_library")
# cc_proto_library is used in this file

package(
default_applicable_licenses = ["//:license"],
default_visibility = ["//xls:xls_internal"],
Expand All @@ -24,6 +27,58 @@ package(
licenses = ["notice"], # Apache 2.0
)

proto_library(
name = "pass_pipeline_proto",
srcs = ["pass_pipeline.proto"],
visibility = ["//xls:xls_users"],
)

cc_proto_library(
name = "pass_pipeline_cc_proto",
visibility = ["//xls:xls_users"],
deps = [":pass_pipeline_proto"],
)

py_proto_library(
name = "pass_pipeline_py_pb2",
visibility = ["//xls:xls_users"],
deps = [":pass_pipeline_proto"],
)

cc_binary(
name = "dump_default_optimization_pass_pipeline_main",
srcs = ["dump_default_optimization_pass_pipeline_main.cc"],
visibility = ["//visibility:private"],
deps = [
":optimization_pass",
":optimization_pass_pipeline",
":pass_pipeline_cc_proto",
"//xls/common:exit_status",
"//xls/common:init_xls",
"//xls/common/file:filesystem",
"//xls/common/status:ret_check",
"//xls/common/status:status_macros",
"@com_google_absl//absl/flags:flag",
"@com_google_absl//absl/log",
"@com_google_absl//absl/status",
"@com_google_protobuf//:protobuf",
],
)

genrule(
name = "default_optimization_pipeline_proto",
outs = ["default_opt_pass_pipeline.binarypb"],
cmd = "$(location :dump_default_optimization_pass_pipeline_main) --proto_out $@",
tools = [":dump_default_optimization_pass_pipeline_main"],
)

genrule(
name = "default_optimization_pipeline_textproto",
outs = ["default_opt_pass_pipeline.textpb"],
cmd = "$(location :dump_default_optimization_pass_pipeline_main) --textproto_out $@",
tools = [":dump_default_optimization_pass_pipeline_main"],
)

cc_library(
name = "optimization_pass_pipeline",
srcs = ["optimization_pass_pipeline.cc"],
Expand Down Expand Up @@ -58,6 +113,7 @@ cc_library(
":optimization_pass",
":optimization_pass_registry",
":pass_base",
":pass_pipeline_cc_proto",
":proc_inlining_pass",
":proc_state_array_flattening_pass",
":proc_state_narrowing_pass",
Expand Down Expand Up @@ -573,6 +629,7 @@ cc_library(
hdrs = ["optimization_pass_registry.h"],
deps = [
":optimization_pass",
":pass_pipeline_cc_proto",
"//xls/common:module_initializer",
"@com_google_absl//absl/log:check",
"@com_google_absl//absl/status",
Expand Down Expand Up @@ -1223,12 +1280,14 @@ cc_library(
srcs = ["pass_base.cc"],
hdrs = ["pass_base.h"],
deps = [
":pass_pipeline_cc_proto",
"//xls/common:casts",
"//xls/common/file:filesystem",
"//xls/common/logging:log_lines",
"//xls/common/status:ret_check",
"//xls/common/status:status_macros",
"//xls/ir",
"@com_google_absl//absl/algorithm:container",
"@com_google_absl//absl/container:flat_hash_map",
"@com_google_absl//absl/log",
"@com_google_absl//absl/log:vlog_is_on",
Expand Down Expand Up @@ -1272,6 +1331,7 @@ cc_library(
hdrs = ["pass_registry.h"],
deps = [
":pass_base",
":pass_pipeline_cc_proto",
"@com_google_absl//absl/algorithm:container",
"@com_google_absl//absl/base:core_headers",
"@com_google_absl//absl/container:flat_hash_map",
Expand All @@ -1287,8 +1347,10 @@ cc_library(
hdrs = ["pipeline_generator.h"],
deps = [
":pass_base",
":pass_pipeline_cc_proto",
"//xls/common/status:ret_check",
"//xls/common/status:status_macros",
"@com_google_absl//absl/log",
"@com_google_absl//absl/status",
"@com_google_absl//absl/status:statusor",
"@com_google_absl//absl/strings",
Expand All @@ -1303,6 +1365,7 @@ cc_test(
":dce_pass",
":optimization_pass",
":pass_base",
":pass_pipeline_cc_proto",
"//xls/common:xls_gunit_main",
"//xls/common/status:matchers",
"//xls/common/status:ret_check",
Expand All @@ -1313,12 +1376,14 @@ cc_test(
"//xls/ir:function_builder",
"//xls/ir:ir_matcher",
"//xls/ir:ir_test_base",
"//xls/ir:source_location",
"//xls/ir:value",
"@com_google_absl//absl/status",
"@com_google_absl//absl/status:status_matchers",
"@com_google_absl//absl/status:statusor",
"@com_google_absl//absl/strings",
"@com_google_googletest//:gtest",
"@com_google_protobuf//:protobuf",
"@com_googlesource_code_re2//:re2",
],
)
Expand Down Expand Up @@ -1347,6 +1412,7 @@ cc_library(
":optimization_pass",
":optimization_pass_registry",
":pass_base",
":pass_pipeline_cc_proto",
":predicate_dominator_analysis",
":predicate_state",
":proc_state_range_query_engine",
Expand Down Expand Up @@ -1644,6 +1710,7 @@ cc_library(
":optimization_pass",
":optimization_pass_registry",
":pass_base",
":pass_pipeline_cc_proto",
":query_engine",
":stateless_query_engine",
":union_query_engine",
Expand Down
8 changes: 8 additions & 0 deletions xls/passes/conditional_specialization_pass.cc
Original file line number Diff line number Diff line change
Expand Up @@ -956,6 +956,14 @@ absl::StatusOr<bool> ConditionalSpecializationPass::RunOnFunctionBaseInternal(

return changed;
}
absl::StatusOr<PassPipelineProto::Element>
ConditionalSpecializationPass::ToProto() const {
// TODO(allight): This is not very elegant. Ideally the registry could handle
// this? Doing it there would probably be even more weird though.
PassPipelineProto::Element e;
*e.mutable_pass_name() = use_bdd_ ? "cond_spec(Bdd)" : "cond_spec(noBdd)";
return e;
}

XLS_REGISTER_MODULE_INITIALIZER(cond_spec, {
CHECK_OK(RegisterOptimizationPass<ConditionalSpecializationPass>(
Expand Down
3 changes: 3 additions & 0 deletions xls/passes/conditional_specialization_pass.h
Original file line number Diff line number Diff line change
Expand Up @@ -21,6 +21,7 @@
#include "xls/ir/function_base.h"
#include "xls/passes/optimization_pass.h"
#include "xls/passes/pass_base.h"
#include "xls/passes/pass_pipeline.pb.h"

namespace xls {

Expand All @@ -36,6 +37,8 @@ class ConditionalSpecializationPass : public OptimizationFunctionBasePass {
use_bdd_(use_bdd) {}
~ConditionalSpecializationPass() override = default;

absl::StatusOr<PassPipelineProto::Element> ToProto() const override;

protected:
bool use_bdd_;
absl::StatusOr<bool> RunOnFunctionBaseInternal(
Expand Down
80 changes: 80 additions & 0 deletions xls/passes/dump_default_optimization_pass_pipeline_main.cc
Original file line number Diff line number Diff line change
@@ -0,0 +1,80 @@
// Copyright 2024 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 <memory>
#include <optional>
#include <string>
#include <string_view>
#include <vector>

#include "absl/flags/flag.h"
#include "absl/log/log.h"
#include "absl/status/status.h"
#include "google/protobuf/text_format.h"
#include "xls/common/exit_status.h"
#include "xls/common/file/filesystem.h"
#include "xls/common/init_xls.h"
#include "xls/common/status/ret_check.h"
#include "xls/common/status/status_macros.h"
#include "xls/passes/optimization_pass.h"
#include "xls/passes/optimization_pass_pipeline.h"
#include "xls/passes/pass_pipeline.pb.h"

const char kUsage[] = R"(
Dump the default optimization pipeline as a PassPipelineProto.
)";

ABSL_FLAG(std::optional<std::string>, textproto_out, std::nullopt,
"Location to write the pipeline in textproto format.");
ABSL_FLAG(std::optional<std::string>, proto_out, std::nullopt,
"Location to write the pipeline in binary proto format.");

namespace xls {
namespace {

absl::Status RealMain() {
std::unique_ptr<OptimizationPass> passes = CreateOptimizationPassPipeline();
if (!absl::GetFlag(FLAGS_textproto_out) && !absl::GetFlag(FLAGS_proto_out)) {
return absl::InvalidArgumentError(
"One of --proto_out or --textproto_out is required");
}
PassPipelineProto res;
XLS_ASSIGN_OR_RETURN(*res.mutable_top(), passes->ToProto());
if (absl::GetFlag(FLAGS_textproto_out)) {
std::string out;
XLS_RET_CHECK(google::protobuf::TextFormat::PrintToString(res, &out))
<< "Unable to serialize proto.";
XLS_RETURN_IF_ERROR(
SetFileContents(*absl::GetFlag(FLAGS_textproto_out), out));
}
if (absl::GetFlag(FLAGS_proto_out)) {
XLS_RETURN_IF_ERROR(SetFileContents(*absl::GetFlag(FLAGS_proto_out),
res.SerializeAsString()));
}
return absl::OkStatus();
}

} // namespace
} // namespace xls

int main(int argc, char** argv) {
std::vector<std::string_view> positional_arguments =
xls::InitXls(kUsage, argc, argv);

if (!positional_arguments.empty()) {
LOG(QFATAL) << "Expected invocation: " << argv[0];
}

return xls::ExitStatus(xls::RealMain());
}
21 changes: 21 additions & 0 deletions xls/passes/narrowing_pass.cc
Original file line number Diff line number Diff line change
Expand Up @@ -1880,6 +1880,27 @@ std::ostream& operator<<(std::ostream& os, NarrowingPass::AnalysisType a) {
}
}

absl::StatusOr<PassPipelineProto::Element> NarrowingPass::ToProto() const {
// TODO(allight): This is not very elegant. Ideally the registry could handle
// this? Doing it there would probably be even more weird though.
PassPipelineProto::Element e;
switch (analysis_) {
case AnalysisType::kTernary:
*e.mutable_pass_name() = "narrow(Ternary)";
break;
case AnalysisType::kRange:
*e.mutable_pass_name() = "narrow(Range)";
break;
case AnalysisType::kRangeWithContext:
*e.mutable_pass_name() = "narrow(Context)";
break;
case AnalysisType::kRangeWithOptionalContext:
*e.mutable_pass_name() = "narrow(OptionalContext)";
break;
}
return e;
}

XLS_REGISTER_MODULE_INITIALIZER(narrowing_pass, {
CHECK_OK(RegisterOptimizationPass<NarrowingPass>("narrow"));
CHECK_OK(RegisterOptimizationPass<NarrowingPass>(
Expand Down
3 changes: 3 additions & 0 deletions xls/passes/narrowing_pass.h
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,7 @@
#include "xls/ir/function_base.h"
#include "xls/passes/optimization_pass.h"
#include "xls/passes/pass_base.h"
#include "xls/passes/pass_pipeline.pb.h"

namespace xls {

Expand All @@ -43,6 +44,8 @@ class NarrowingPass : public OptimizationFunctionBasePass {
: OptimizationFunctionBasePass(kName, "Narrowing"), analysis_(analysis) {}
~NarrowingPass() override = default;

absl::StatusOr<PassPipelineProto::Element> ToProto() const override;

protected:
AnalysisType analysis_;

Expand Down
Loading

0 comments on commit f47eb95

Please sign in to comment.