[solvers] Add solver_benchmark tool evaluating Z3 and Bitwuzla

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"],

xls/examples/bitonic_sort.x

@@ -61,6 +61,10 @@ fn bitonic_sort<N: u32, LOG_N: u32 = { std::clog2(N) }>(array: u32[N]) -> u32[N]
     result
 }
 
+pub fn bitonic_sort_u32_8(array: u32[8]) -> u32[8] {
+    bitonic_sort(array)
+}
+
 ////////////////////////////////////////////////////////////////////////////////
 // Tests
 ////////////////////////////////////////////////////////////////////////////////

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<I: u32, P: u32 = {compute_pad_bits(I + u32:65)}, R: u32 = {I + u32:65 + P}>
+    (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
+}

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",
+    ],
+)

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<Term> 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<Term> IrTranslator::TranslateLiteralValue(const Type* type,

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) {