From ccb0f875097beb434cdd469b9ecd20b5fe4dd588 Mon Sep 17 00:00:00 2001 From: Eric Astor Date: Tue, 16 Jun 2026 16:27:39 -0700 Subject: [PATCH] [solvers] Add solver_benchmark tool evaluating Z3 and Bitwuzla Adds a benchmarking binary to measure and compare translation and solve times across XLS SMT solver backends (Z3 and Bitwuzla) on standard XLS benchmark designs (fp32_fmac, bitonic_sort, riscv_simple, sha256). Also introduces scaled/weakened example benchmarks suitable for solver analysis. PiperOrigin-RevId: 933361496 --- xls/examples/BUILD | 16 ++ xls/examples/bitonic_sort.x | 4 + xls/examples/sha256_scaled.x | 138 ++++++++++++ xls/solvers/BUILD | 30 +++ xls/solvers/bitwuzla_ir_translator.cc | 26 ++- xls/solvers/bitwuzla_ir_translator_test.cc | 114 ++++++++++ xls/solvers/solver_benchmark.cc | 237 +++++++++++++++++++++ 7 files changed, 559 insertions(+), 6 deletions(-) create mode 100644 xls/examples/sha256_scaled.x create mode 100644 xls/solvers/solver_benchmark.cc diff --git a/xls/examples/BUILD b/xls/examples/BUILD index bd8dacdf1b..8cce57b968 100644 --- a/xls/examples/BUILD +++ b/xls/examples/BUILD @@ -427,6 +427,16 @@ xls_dslx_opt_ir( tags = ["optonly"], ) +xls_dslx_opt_ir( + name = "sha256_scaled", + srcs = ["sha256_scaled.x"], + dslx_top = "main", + ir_conv_args = {"lower_to_proc_scoped_channels": "true"}, + ir_file = "sha256_scaled.ir", + opt_ir_file = "sha256_scaled.opt.ir", + tags = ["optonly"], +) + xls_dslx_test( name = "sha256_dslx_test", srcs = ["sha256.x"], @@ -1206,6 +1216,12 @@ xls_dslx_test( dslx_test_args = {"compare": "jit"}, ) +xls_dslx_opt_ir( + name = "bitonic_sort", + dslx_top = "bitonic_sort_u32_8", + library = ":bitonic_sort_dslx", +) + xls_dslx_library( name = "hack_cpu_dslx", srcs = ["hack_cpu.x"], diff --git a/xls/examples/bitonic_sort.x b/xls/examples/bitonic_sort.x index e84041a110..c7c3b185c6 100644 --- a/xls/examples/bitonic_sort.x +++ b/xls/examples/bitonic_sort.x @@ -61,6 +61,10 @@ fn bitonic_sort(array: u32[N]) -> u32[N] result } +pub fn bitonic_sort_u32_8(array: u32[8]) -> u32[8] { + bitonic_sort(array) +} + //////////////////////////////////////////////////////////////////////////////// // Tests //////////////////////////////////////////////////////////////////////////////// diff --git a/xls/examples/sha256_scaled.x b/xls/examples/sha256_scaled.x new file mode 100644 index 0000000000..44d2f34d03 --- /dev/null +++ b/xls/examples/sha256_scaled.x @@ -0,0 +1,138 @@ +#![feature(type_inference_v2)] + +// Copyright 2020 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. + +import std; + +// --- +// +// SHA algorithm based on the description in: +// +// https://en.wikipedia.org/wiki/SHA-2#Pseudocode +// +// We attempt to mirror the pseudocode presented there fairly directly for ease +// of reproducing correct results. +// +// This version is deliberately weakened (just 8 rounds) for easier experimentation with LEC and +// property checking. + +pub type Digest = (u32, u32, u32, u32, u32, u32, u32, u32); + +fn sha256_chunk_w_table(chunk: bits[512]) -> u32[64] { + // Seed the "w" table with the message chunk. + let w_init: u32[64] = (chunk ++ bits[1536]:0) as u32[64]; + + // Build up the remaining values of the "w" table. + // TODO(b/149962183): Make range go from 16 - 64 once counted for + // ranges can start at values other than 0. + let w: u32[64] = for (i, w): (u32, u32[64]) in u32:0..u32:48 { + let w_im15: u32 = w[i + u32:16 - u32:15]; + let s_0: u32 = std::rotr(w_im15, u32:7) ^ std::rotr(w_im15, u32:18) ^ (w_im15 >> u32:3); + let w_im2: u32 = w[i + u32:16 - u32:2]; + let s_1: u32 = std::rotr(w_im2, u32:17) ^ std::rotr(w_im2, u32:19) ^ (w_im2 >> u32:10); + let value: u32 = w[i + u32:16 - u32:16] + s_0 + w[i + u32:16 - u32:7] + s_1; + update(w, i + u32:16, value) + }(w_init); + w +} + +// Evolves the digest for a single chunk in the overall message being +// SHA256-hashed. +fn sha256_chunk_weak(chunk: bits[512], digest_init: Digest) -> Digest { + let w: u32[64] = sha256_chunk_w_table(chunk); + + // The constant "K" table of addends. + const K = u32[64]:[ + u32:0x428a2f98, u32:0x71374491, u32:0xb5c0fbcf, u32:0xe9b5dba5, u32:0x3956c25b, + u32:0x59f111f1, u32:0x923f82a4, u32:0xab1c5ed5, u32:0xd807aa98, u32:0x12835b01, + u32:0x243185be, u32:0x550c7dc3, u32:0x72be5d74, u32:0x80deb1fe, u32:0x9bdc06a7, + u32:0xc19bf174, u32:0xe49b69c1, u32:0xefbe4786, u32:0x0fc19dc6, u32:0x240ca1cc, + u32:0x2de92c6f, u32:0x4a7484aa, u32:0x5cb0a9dc, u32:0x76f988da, u32:0x983e5152, + u32:0xa831c66d, u32:0xb00327c8, u32:0xbf597fc7, u32:0xc6e00bf3, u32:0xd5a79147, + u32:0x06ca6351, u32:0x14292967, u32:0x27b70a85, u32:0x2e1b2138, u32:0x4d2c6dfc, + u32:0x53380d13, u32:0x650a7354, u32:0x766a0abb, u32:0x81c2c92e, u32:0x92722c85, + u32:0xa2bfe8a1, u32:0xa81a664b, u32:0xc24b8b70, u32:0xc76c51a3, u32:0xd192e819, + u32:0xd6990624, u32:0xf40e3585, u32:0x106aa070, u32:0x19a4c116, u32:0x1e376c08, + u32:0x2748774c, u32:0x34b0bcb5, u32:0x391c0cb3, u32:0x4ed8aa4a, u32:0x5b9cca4f, + u32:0x682e6ff3, u32:0x748f82ee, u32:0x78a5636f, u32:0x84c87814, u32:0x8cc70208, + u32:0x90befffa, u32:0xa4506ceb, u32:0xbef9a3f7, u32:0xc67178f2, + ]; + + // Compute the digest using the "w" table over 8 "rounds". + let (a, b, c, d, e, f, g, h): Digest = + for (i, (a, b, c, d, e, f, g, h)): (u32, Digest) in u32:0..u32:8 { + let S1 = std::rotr(e, u32:6) ^ std::rotr(e, u32:11) ^ std::rotr(e, u32:25); + let ch = (e & f) ^ ((!e) & g); + let temp1 = h + S1 + ch + K[i] + w[i]; + let S0 = std::rotr(a, u32:2) ^ std::rotr(a, u32:13) ^ std::rotr(a, u32:22); + let maj = (a & b) ^ (a & c) ^ (b & c); + let temp2 = S0 + maj; + let (h, g, f) = (g, f, e); + let e = d + temp1; + let (d, c, b) = (c, b, a); + let a = temp1 + temp2; + (a, b, c, d, e, f, g, h) + }(digest_init); + + // The new digest mixes together values from the original digest with the + // derived values (a, b, c, ...) we've computed. + let (h0, h1, h2, h3, h4, h5, h6, h7): Digest = digest_init; + (h0 + a, h1 + b, h2 + c, h3 + d, h4 + e, h5 + f, h6 + g, h7 + h) +} + +// Returns the number of bits required to add on to turn bit_count into a +// multiple of 512. +fn compute_pad_bits(bit_count: u32) -> u32 { + std::round_up_to_nearest(bit_count, u32:512) - bit_count +} + +#[test] +fn compute_pad_bits_test() { + assert_eq(u32:511, compute_pad_bits(u32:1)); + assert_eq(u32:1, compute_pad_bits(u32:511)); + assert_eq(u32:0, compute_pad_bits(u32:512)); + assert_eq(u32:511, compute_pad_bits(u32:513)); + assert_eq(u32:0, compute_pad_bits(u32:1024)); +} + +// The SHA algorithm tells us to precondition our input by tacking on a +// trailing "stop" bit, padding out with zeros, and appending the length as a +// 64-bit quantity such that the resulting number of bits is a multiple of 512. +fn pad_to_512b_chunk + (x: bits[I]) -> bits[R] { + let stop_bit: bits[1] = bits[1]:1; + x ++ stop_bit ++ bits[P]:0 ++ I as bits[64] +} + +pub fn sha256_weak(message: bits[512]) -> Digest { + let digest_init: Digest = ( + u32:0x6a09e667, u32:0xbb67ae85, u32:0x3c6ef372, u32:0xa54ff53a, u32:0x510e527f, + u32:0x9b05688c, u32:0x1f83d9ab, u32:0x5be0cd19, + ); + + // TODO(leary): 2019-03-19 Commenting this out for now to avoid needing a + // 'structural' for loop in IR conversion. + // + //for (chunk, digest): (bits[512], Digest) in message { + // let new_digest: Digest = sha256_chunk_weak(chunk, digest); + // new_digest + //}(digest_init) + sha256_chunk_weak(message, digest_init) +} + +pub fn main(message: bits[512]) -> bits[256] { + let digest = sha256_weak(message); + digest.0 ++ digest.1 ++ digest.2 ++ digest.3 ++ digest.4 ++ digest.5 ++ digest.6 ++ digest.7 +} diff --git a/xls/solvers/BUILD b/xls/solvers/BUILD index 38a470181b..51a38eca02 100644 --- a/xls/solvers/BUILD +++ b/xls/solvers/BUILD @@ -354,6 +354,7 @@ cc_test( "//xls/common/status:matchers", "//xls/ir", "//xls/ir:bits", + "//xls/ir:bits_ops", "//xls/ir:function_builder", "//xls/ir:ir_test_base", "//xls/ir:value", @@ -605,3 +606,32 @@ cc_test( "@googletest//:gtest", ], ) + +cc_binary( + name = "solver_benchmark", + testonly = True, + srcs = ["solver_benchmark.cc"], + data = [ + "//xls/examples:bitonic_sort.opt.ir", + "//xls/examples:fp32_fmac.opt.ir", + "//xls/examples:riscv_simple.opt.ir", + "//xls/examples:sha256.opt.ir", + "//xls/examples:sha256_scaled.opt.ir", + ], + deps = [ + ":bitwuzla_ir_translator", + ":solver", + ":z3_ir_translator", + "//xls/common:init_xls", + "//xls/common/file:filesystem", + "//xls/common/file:get_runfile_path", + "//xls/common/status:status_macros", + "//xls/ir", + "//xls/ir:ir_parser", + "@abseil-cpp//absl/log:check", + "@abseil-cpp//absl/status", + "@abseil-cpp//absl/status:statusor", + "@abseil-cpp//absl/strings", + "@googletest//:gtest", + ], +) diff --git a/xls/solvers/bitwuzla_ir_translator.cc b/xls/solvers/bitwuzla_ir_translator.cc index 3285ae32fc..4e15f9f406 100644 --- a/xls/solvers/bitwuzla_ir_translator.cc +++ b/xls/solvers/bitwuzla_ir_translator.cc @@ -256,12 +256,26 @@ Term IrTranslator::TranslateLiteralBits(const Bits& bits) { if (bc == 0) { return tm_.mk_bv_value_uint64(tm_.mk_bv_sort(1), 0); } - std::string str; - str.reserve(bc); - for (int64_t i = bc - 1; i >= 0; --i) { - str.push_back(bits.Get(i) ? '1' : '0'); - } - return tm_.mk_bv_value(tm_.mk_bv_sort(bc), str, 2); + if (bits.IsZero()) { + return tm_.mk_bv_zero(tm_.mk_bv_sort(bc)); + } + if (bits.IsAllOnes()) { + return tm_.mk_bv_ones(tm_.mk_bv_sort(bc)); + } + if (bits.FitsInUint64()) { + return tm_.mk_bv_value_uint64(tm_.mk_bv_sort(bc), bits.ToUint64().value()); + } + int64_t num_words = bits.bitmap().word_count(); + std::vector words; + words.reserve(num_words); + int64_t top_width = bc - (num_words - 1) * 64; + words.push_back(tm_.mk_bv_value_uint64(tm_.mk_bv_sort(top_width), + bits.bitmap().GetWord(num_words - 1))); + for (int64_t i = num_words - 2; i >= 0; --i) { + words.push_back( + tm_.mk_bv_value_uint64(tm_.mk_bv_sort(64), bits.bitmap().GetWord(i))); + } + return ConcatN(words); } absl::StatusOr IrTranslator::TranslateLiteralValue(const Type* type, diff --git a/xls/solvers/bitwuzla_ir_translator_test.cc b/xls/solvers/bitwuzla_ir_translator_test.cc index 0be4e83477..d118a2eff0 100644 --- a/xls/solvers/bitwuzla_ir_translator_test.cc +++ b/xls/solvers/bitwuzla_ir_translator_test.cc @@ -29,6 +29,7 @@ #include "absl/strings/str_format.h" #include "xls/common/status/matchers.h" #include "xls/ir/bits.h" +#include "xls/ir/bits_ops.h" #include "xls/ir/function.h" #include "xls/ir/function_builder.h" #include "xls/ir/ir_test_base.h" @@ -575,6 +576,119 @@ TEST_F(BitwuzlaIrTranslatorTest, TupleWithArrayLiteral) { IsOkAndHolds(IsProvenTrue())); } +TEST_F(BitwuzlaIrTranslatorTest, WideLiteralTranslation) { + // Test that literals wider than 64 bits with multiple non-zero words + // and non-power-of-two widths preserve exact bit order and endianness. + + // 1. 128-bit literal: two 64-bit words + { + auto package = CreatePackage(); + FunctionBuilder fb("wide_128", package.get()); + Bits hi_64 = UBits(0x1122334455667788ULL, 64); + Bits lo_64 = UBits(0x99aabbccddeeff01ULL, 64); + Bits val_128 = bits_ops::Concat({hi_64, lo_64}); + BValue lit_128 = fb.Literal(val_128); + + BValue slice_lo = fb.BitSlice(lit_128, /*start=*/0, /*width=*/64); + BValue slice_hi = fb.BitSlice(lit_128, /*start=*/64, /*width=*/64); + + BValue eq_lo = fb.Eq(slice_lo, fb.Literal(lo_64)); + BValue eq_hi = fb.Eq(slice_hi, fb.Literal(hi_64)); + BValue swapped = fb.Eq(slice_lo, fb.Literal(hi_64)); + BValue eq_concat = + fb.Eq(lit_128, fb.Concat({fb.Literal(hi_64), fb.Literal(lo_64)})); + + XLS_ASSERT_OK_AND_ASSIGN(Function * f, fb.BuildWithReturnValue(eq_lo)); + EXPECT_THAT(solver_->TryProve(f, eq_lo.node(), Predicate::NotEqualToZero(), + SolverLimit()), + IsOkAndHolds(IsProvenTrue())); + EXPECT_THAT(solver_->TryProve(f, eq_hi.node(), Predicate::NotEqualToZero(), + SolverLimit()), + IsOkAndHolds(IsProvenTrue())); + EXPECT_THAT(solver_->TryProve(f, swapped.node(), + Predicate::NotEqualToZero(), SolverLimit()), + IsOkAndHolds(IsProvenFalse())); + EXPECT_THAT(solver_->TryProve(f, eq_concat.node(), + Predicate::NotEqualToZero(), SolverLimit()), + IsOkAndHolds(IsProvenTrue())); + } + + // 2. 100-bit literal: 36-bit top word (MSB) + 64-bit bottom word (LSB) + { + auto package = CreatePackage(); + FunctionBuilder fb("wide_100", package.get()); + Bits hi_36 = UBits(0x123456789ULL, 36); + Bits lo_64 = UBits(0xfedcba9876543210ULL, 64); + Bits val_100 = bits_ops::Concat({hi_36, lo_64}); + BValue lit_100 = fb.Literal(val_100); + + BValue slice_lo = fb.BitSlice(lit_100, /*start=*/0, /*width=*/64); + BValue slice_hi = fb.BitSlice(lit_100, /*start=*/64, /*width=*/36); + // Cross-boundary slice: bits 60..67 (8 bits spanning the 64-bit word + // boundary) Low 4 bits come from lo_64 bits 60..63 (top nibble of + // 0xfedcba9876543210: 0xf) High 4 bits come from hi_36 bits 0..3 (bottom + // nibble of 0x123456789: 0x9) Combined byte: (0x9 << 4) | 0xf = 0x9f. + BValue slice_cross = fb.BitSlice(lit_100, /*start=*/60, /*width=*/8); + BValue expected_cross = fb.Literal(UBits(0x9f, 8)); + + BValue eq_lo = fb.Eq(slice_lo, fb.Literal(lo_64)); + BValue eq_hi = fb.Eq(slice_hi, fb.Literal(hi_36)); + BValue eq_cross = fb.Eq(slice_cross, expected_cross); + BValue eq_concat = + fb.Eq(lit_100, fb.Concat({fb.Literal(hi_36), fb.Literal(lo_64)})); + + XLS_ASSERT_OK_AND_ASSIGN(Function * f, fb.BuildWithReturnValue(eq_lo)); + EXPECT_THAT(solver_->TryProve(f, eq_lo.node(), Predicate::NotEqualToZero(), + SolverLimit()), + IsOkAndHolds(IsProvenTrue())); + EXPECT_THAT(solver_->TryProve(f, eq_hi.node(), Predicate::NotEqualToZero(), + SolverLimit()), + IsOkAndHolds(IsProvenTrue())); + EXPECT_THAT(solver_->TryProve(f, eq_cross.node(), + Predicate::NotEqualToZero(), SolverLimit()), + IsOkAndHolds(IsProvenTrue())); + EXPECT_THAT(solver_->TryProve(f, eq_concat.node(), + Predicate::NotEqualToZero(), SolverLimit()), + IsOkAndHolds(IsProvenTrue())); + } + + // 3. 256-bit literal: 4 distinct 64-bit words + { + auto package = CreatePackage(); + FunctionBuilder fb("wide_256", package.get()); + Bits w3 = UBits(0x0123456789abcdefULL, 64); + Bits w2 = UBits(0xfedcba9876543210ULL, 64); + Bits w1 = UBits(0xa5a5a5a55a5a5a5aULL, 64); + Bits w0 = UBits(0x5a5a5a5aa5a5a5a5ULL, 64); + Bits val_256 = bits_ops::Concat({w3, w2, w1, w0}); + BValue lit_256 = fb.Literal(val_256); + + BValue s0 = fb.BitSlice(lit_256, /*start=*/0, /*width=*/64); + BValue s1 = fb.BitSlice(lit_256, /*start=*/64, /*width=*/64); + BValue s2 = fb.BitSlice(lit_256, /*start=*/128, /*width=*/64); + BValue s3 = fb.BitSlice(lit_256, /*start=*/192, /*width=*/64); + + BValue eq_s0 = fb.Eq(s0, fb.Literal(w0)); + BValue eq_s1 = fb.Eq(s1, fb.Literal(w1)); + BValue eq_s2 = fb.Eq(s2, fb.Literal(w2)); + BValue eq_s3 = fb.Eq(s3, fb.Literal(w3)); + + XLS_ASSERT_OK_AND_ASSIGN(Function * f, fb.BuildWithReturnValue(eq_s0)); + EXPECT_THAT(solver_->TryProve(f, eq_s0.node(), Predicate::NotEqualToZero(), + SolverLimit()), + IsOkAndHolds(IsProvenTrue())); + EXPECT_THAT(solver_->TryProve(f, eq_s1.node(), Predicate::NotEqualToZero(), + SolverLimit()), + IsOkAndHolds(IsProvenTrue())); + EXPECT_THAT(solver_->TryProve(f, eq_s2.node(), Predicate::NotEqualToZero(), + SolverLimit()), + IsOkAndHolds(IsProvenTrue())); + EXPECT_THAT(solver_->TryProve(f, eq_s3.node(), Predicate::NotEqualToZero(), + SolverLimit()), + IsOkAndHolds(IsProvenTrue())); + } +} + // Microbenchmarks comparing Bitwuzla against Z3 on representative IR operations static void BM_OneHotLsb(benchmark::State& state, SolverKind kind) { diff --git a/xls/solvers/solver_benchmark.cc b/xls/solvers/solver_benchmark.cc new file mode 100644 index 0000000000..adc6d44dcf --- /dev/null +++ b/xls/solvers/solver_benchmark.cc @@ -0,0 +1,237 @@ +// 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 +#include +#include +#include +#include +#include +#include +#include + +#include "benchmark/benchmark.h" +#include "absl/log/check.h" +#include "absl/status/status.h" +#include "absl/status/statusor.h" +#include "absl/strings/str_cat.h" +#include "xls/common/file/filesystem.h" +#include "xls/common/file/get_runfile_path.h" +#include "xls/common/init_xls.h" +#include "xls/common/status/status_macros.h" +#include "xls/ir/function.h" +#include "xls/ir/ir_parser.h" +#include "xls/ir/node.h" +#include "xls/ir/topo_sort.h" +#include "xls/solvers/solver.h" + +namespace xls::solvers { +namespace { + +struct BenchmarkInput { + std::string name; + std::string path; + bool translate_only = false; + bool skip_z3_solve = false; + bool skip_z3_ne_zero = false; +}; + +const std::vector& GetInputs() { + static const std::vector* inputs = + new std::vector{ + {.name = "fp32_fmac", .path = "xls/examples/fp32_fmac.opt.ir"}, + {.name = "bitonic_sort", .path = "xls/examples/bitonic_sort.opt.ir"}, + {.name = "riscv_simple", .path = "xls/examples/riscv_simple.opt.ir"}, + {.name = "sha256_scaled", + .path = "xls/examples/sha256_scaled.opt.ir", + .skip_z3_ne_zero = true}, + {.name = "sha256_full", + .path = "xls/examples/sha256.opt.ir", + .translate_only = true}}; + return *inputs; +} + +FunctionBase* GetBenchmarkSubject(Package* package) { + std::optional top = package->GetTop(); + if (top.has_value() && (*top)->node_count() > 10) { + return *top; + } + FunctionBase* best = nullptr; + int64_t max_nodes = -1; + for (const auto& f : package->functions()) { + if (f->node_count() > max_nodes) { + max_nodes = f->node_count(); + best = f.get(); + } + } + for (const auto& p : package->procs()) { + if (p->node_count() > max_nodes) { + max_nodes = p->node_count(); + best = p.get(); + } + } + return best; +} + +absl::StatusOr FindSubjectNode(FunctionBase* f, SolverKind solver_kind) { + if (f->IsFunction()) { + Node* rv = f->AsFunctionOrDie()->return_value(); + if (rv->GetType()->IsBits()) { + return rv; + } + } + XLS_ASSIGN_OR_RETURN(std::vector nodes, TopoSort(f)); + for (auto it = nodes.rbegin(); it != nodes.rend(); ++it) { + if ((*it)->GetType()->IsBits()) { + return *it; + } + } + return absl::NotFoundError("No suitable subject node found"); +} + +void RunTranslateBenchmark(::testing::benchmark::State& state, + std::string_view ir_relative_path, + SolverKind solver_kind) { + absl::StatusOr runfile_path = + GetXlsRunfilePath(ir_relative_path); + CHECK(runfile_path.ok()) << runfile_path.status().message(); + + absl::StatusOr contents = GetFileContents(*runfile_path); + CHECK(contents.ok()) << contents.status().message(); + + absl::StatusOr> package = + Parser::ParsePackage(*contents, ir_relative_path); + CHECK(package.ok()) << package.status().message(); + + FunctionBase* subject_entity = GetBenchmarkSubject(package->get()); + CHECK(subject_entity != nullptr) << "No benchmark subject found in package"; + + absl::StatusOr> solver = CreateSolver(solver_kind); + CHECK(solver.ok()) << solver.status().message(); + + for (auto _ : state) { + auto instance = (*solver)->CreateSolverInstance(subject_entity, + /*allow_unsupported=*/true); + CHECK(instance.ok()) << instance.status().message(); + ::testing::DoNotOptimize(instance); + } +} + +enum class SolverTestKind { + kEqZero, + kNeZero, +}; + +void RunSolveBenchmark(::testing::benchmark::State& state, + std::string_view ir_relative_path, + SolverKind solver_kind, SolverTestKind test_kind) { + absl::StatusOr runfile_path = + GetXlsRunfilePath(ir_relative_path); + CHECK(runfile_path.ok()) << runfile_path.status().message(); + + absl::StatusOr contents = GetFileContents(*runfile_path); + CHECK(contents.ok()) << contents.status().message(); + + absl::StatusOr> package = + Parser::ParsePackage(*contents, ir_relative_path); + CHECK(package.ok()) << package.status().message(); + + FunctionBase* subject_entity = GetBenchmarkSubject(package->get()); + CHECK(subject_entity != nullptr) << "No benchmark subject found in package"; + + absl::StatusOr subject = FindSubjectNode(subject_entity, solver_kind); + if (!subject.ok()) { + std::cerr << "Skipping solve benchmark for " << ir_relative_path + << " because no suitable subject node was found: " + << subject.status().message() << "\n"; + return; + } + + absl::StatusOr> solver = CreateSolver(solver_kind); + CHECK(solver.ok()) << solver.status().message(); + + absl::StatusOr> instance = + (*solver)->CreateSolverInstance(subject_entity, + /*allow_unsupported=*/true); + CHECK(instance.ok()) << instance.status().message(); + + Predicate predicate = (test_kind == SolverTestKind::kEqZero) + ? Predicate::EqualToZero() + : Predicate::NotEqualToZero(); + + for (auto _ : state) { + auto result_or = (*instance)->TryProve(*subject, predicate); + CHECK(result_or.ok()) << result_or.status().message(); + ::testing::DoNotOptimize(result_or); + } +} + +void BM_Translate(::testing::benchmark::State& state) { + SolverKind kind = static_cast(state.range(0)); + const BenchmarkInput& input = GetInputs()[state.range(1)]; + RunTranslateBenchmark(state, input.path, kind); +} + +void BM_Solve(::testing::benchmark::State& state) { + SolverKind kind = static_cast(state.range(0)); + const BenchmarkInput& input = GetInputs()[state.range(1)]; + SolverTestKind test_kind = static_cast(state.range(2)); + RunSolveBenchmark(state, input.path, kind, test_kind); +} + +void RegisterBenchmarks(::testing::Benchmark* b, bool translate) { + const auto& inputs = GetInputs(); + for (SolverKind kind : {SolverKind::kZ3, SolverKind::kBitwuzla}) { + std::string kind_str = (kind == SolverKind::kZ3) ? "Z3" : "Bitwuzla"; + for (int i = 0; i < inputs.size(); ++i) { + if ((!translate && inputs[i].translate_only) || + (kind == SolverKind::kZ3 && inputs[i].skip_z3_solve)) { + continue; + } + if (translate) { + b->Args({static_cast(kind), i}, + absl::StrCat("solver=", kind_str, "/input=", inputs[i].name)); + } else { + for (SolverTestKind test_kind : + {SolverTestKind::kEqZero, SolverTestKind::kNeZero}) { + if (kind == SolverKind::kZ3 && test_kind == SolverTestKind::kNeZero && + inputs[i].skip_z3_ne_zero) { + continue; + } + std::string test_kind_str = + (test_kind == SolverTestKind::kEqZero) ? "eq_zero" : "ne_zero"; + b->Args({static_cast(kind), i, static_cast(test_kind)}, + absl::StrCat("solver=", kind_str, "/input=", inputs[i].name, + "/test=", test_kind_str)); + } + } + } + } +} + +BENCHMARK(BM_Translate)->Apply([](::testing::Benchmark* b) { + RegisterBenchmarks(b, /*translate=*/true); +}); +BENCHMARK(BM_Solve)->Apply([](::testing::Benchmark* b) { + RegisterBenchmarks(b, /*translate=*/false); +}); + +} // namespace +} // namespace xls::solvers + +int main(int argc, char** argv) { + xls::InitXls(argv[0], argc, argv); + RunSpecifiedBenchmarks(); + return 0; +}