From edf1c02bdaf1e87981da800f5bd1175665476220 Mon Sep 17 00:00:00 2001 From: Angelo Matni Date: Thu, 30 Jul 2026 21:50:41 -0700 Subject: [PATCH] Add leading zero anticipator to DSLX standard library and use in apfloat PiperOrigin-RevId: 956910380 --- xls/dslx/stdlib/BUILD | 22 +++++ xls/dslx/stdlib/apfloat.x | 53 ++++++++--- xls/dslx/stdlib/lza.x | 179 ++++++++++++++++++++++++++++++++++++++ 3 files changed, 244 insertions(+), 10 deletions(-) create mode 100644 xls/dslx/stdlib/lza.x diff --git a/xls/dslx/stdlib/BUILD b/xls/dslx/stdlib/BUILD index 9e0c2865aa..c325bd478d 100644 --- a/xls/dslx/stdlib/BUILD +++ b/xls/dslx/stdlib/BUILD @@ -104,6 +104,28 @@ xls_dslx_library( deps = [":abs_diff_dslx"], ) +xls_dslx_library( + name = "lza_dslx", + srcs = ["lza.x"], +) + +xls_dslx_fmt_test( + name = "lza_dslx_fmt_test", + src = "lza.x", + opportunistic_postcondition = True, +) + +xls_dslx_test( + name = "lza_dslx_test", + srcs = ["lza.x"], + dslx_test_args = {"compare": "jit"}, +) + +xls_dslx_prove_quickcheck_test( + name = "lza_dslx_prove_quickcheck_test", + srcs = ["lza.x"], +) + xls_dslx_test( name = "apfloat_dslx_test", dslx_test_args = {"compare": "jit"}, diff --git a/xls/dslx/stdlib/apfloat.x b/xls/dslx/stdlib/apfloat.x index 0e8181fe7d..ecec8d6aec 100644 --- a/xls/dslx/stdlib/apfloat.x +++ b/xls/dslx/stdlib/apfloat.x @@ -15,6 +15,7 @@ // Arbitrary-precision floating point routines. import std; import abs_diff; +import lza; pub struct APFloat { sign: bits[1], // Sign bit. @@ -3061,7 +3062,7 @@ fn or_last_bit(value: bits[WIDTH], lsb: u1) -> bits[WIDTH] { // The bit widths of different float components are given // in comments throughout this implementation, listed // relative to the widths of a standard float32. -pub fn add +pub fn add (a: APFloat, b: APFloat) -> APFloat { // WIDE_EXP: Widened exponent to capture a possible carry bit. @@ -3134,13 +3135,40 @@ pub fn add let carry_fraction = or_last_bit(carry_fraction, abs_fraction[0:1]); // If we cancelled higher bits, then we'll need to shift left. - // Leading zeroes will be 1 if there's no carry or cancellation. - let leading_zeroes = std::clzt(abs_fraction); - - // Manually apply https://github.com/google/xls/issues/1274 - let cancel_fraction = abs_fraction as uN[WIDE_FRACTION + u32:1] << leading_zeroes; - let cancel_fraction = (cancel_fraction >> u32:1) as uN[NORMALIZED_FRACTION]; - let shifted_fraction = if carry_bit { carry_fraction } else { cancel_fraction }; + const CLZ_WIDTH: u32 = std::clog2(WIDE_FRACTION + u32:1); + let (shifted_fraction, leading_zeroes) = if USE_LZA { + let (approx_lz, _) = + lza::lza(wide_x as uN[WIDE_FRACTION], addend_y as uN[WIDE_FRACTION]); + + // For subtraction, lza is off-by-one at most and will be corrected later by examining the + // shifted_fraction. For addition, when there is carry, this is discarded, and when there + // is no carry, clz is exactly 1, because fraction sum is in range [1.0, 2.0) + let cancel_shift = if x.sign != y.sign { approx_lz } else { uN[CLZ_WIDTH]:1 }; + + // Manually apply https://github.com/google/xls/issues/1274 + let cancel_fraction = abs_fraction as uN[WIDE_FRACTION + u32:1] << cancel_shift; + let cancel_fraction = (cancel_fraction >> u32:1) as uN[NORMALIZED_FRACTION]; + let shifted_fraction = if carry_bit { carry_fraction } else { cancel_fraction }; + + // Leading zeroes will be off-by-one at most if there is cancellation. As an optimization, + // we correct this in a separate shift-left instead of using LZA detection. + let lz_off_by_one = shifted_fraction[-1:] == u1:0; + let shifted_fraction = + if lz_off_by_one { shifted_fraction << u32:1 } else { shifted_fraction }; + let corrected_lz = + if lz_off_by_one { cancel_shift + uN[CLZ_WIDTH]:1 } else { cancel_shift }; + let leading_zeroes = if carry_bit { uN[CLZ_WIDTH]:0 } else { corrected_lz }; + (shifted_fraction, leading_zeroes) + } else { + // Leading zeroes will be 1 if there's no carry or cancellation. + let leading_zeroes = std::clzt(abs_fraction); + + // Manually apply https://github.com/google/xls/issues/1274 + let cancel_fraction = abs_fraction as uN[WIDE_FRACTION + u32:1] << leading_zeroes; + let cancel_fraction = (cancel_fraction >> u32:1) as uN[NORMALIZED_FRACTION]; + let shifted_fraction = if carry_bit { carry_fraction } else { cancel_fraction }; + (shifted_fraction, leading_zeroes) + }; // Step 4: Rounding. // Rounding down is a no-op, since we eventually have to shift off @@ -3233,11 +3261,16 @@ pub fn add // - No exception flags are raised/reported. // In all other cases, results should be identical to other // conforming implementations (modulo exact fraction values in the NaN case). -pub fn sub +pub fn sub (x: APFloat, y: APFloat) -> APFloat { let y = APFloat { sign: !y.sign, bexp: y.bexp, fraction: y.fraction }; - add(x, y) + add(x, y) +} + +#[quickcheck(exhaustive)] +fn add_clz_lza_equivalence(a: APFloat, b: APFloat) -> bool { + add(a, b) == add(a, b) } // add is thoroughly tested elsewhere so a few simple tests is sufficient. diff --git a/xls/dslx/stdlib/lza.x b/xls/dslx/stdlib/lza.x new file mode 100644 index 0000000000..129565b379 --- /dev/null +++ b/xls/dslx/stdlib/lza.x @@ -0,0 +1,179 @@ +// 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. + +// Leading Zero Anticipator (LZA) +// Based on Bruguera and Lang +// Leading-One Prediction Scheme for Latency Improvement in Single Datapath Floating-Point Adders +// https://ieeexplore.ieee.org/document/727065 + +import abs_diff; +import std; + +// Merges (B, M, E, Y) labels as defined in the paper. +// The Y, or yes, indicates a correction pattern has been seen, which amounts to the following: +// Case 1: 0s ++ 1 ++ 0s ++ -1 ++ ... +// Case 2: 0s ++ 1 ++ -1s ++ 0s ++ -1 ++ ... +// +// Case 1 requires correction because 1 ++ 0s ++ -1 annihilates to 0 ++ 1s which shifts the +// leading one to the right by one index. +// Case 2 requires correction because 1 ++ -1s ++ 0s ++ -1 repeatedly annihilates until you're left +// with case 1, i.e. 1 ++ -1 becomes 0 ++ 1 over and over until you reduce to case 1. +fn merge_branch_labels(left: (u1, u1, u1, u1), right: (u1, u1, u1, u1)) -> (u1, u1, u1, u1) { + let (left_begin, left_mid, left_end, left_yes) = left; + let (right_begin, right_mid, right_end, right_yes) = right; + + // begin is unimpacted by mid (zeroes around it) + let begin = (left_begin & right_mid) | (left_mid & right_begin); + // mid is a sequence of 0s + let mid = left_mid & right_mid; + // end is the sequence of 0s ++ -1 + let end = left_end | (left_mid & right_end); + + // begin ++ end -> yes, the full correction pattern has been seen. + // yes ++ any -> yes , mid ++ yes -> yes + let yes = left_yes | (left_mid & right_yes) | (left_begin & right_end); + + (begin, mid, end, yes) +} + +fn tree_reduce + (labels: (u1, u1, u1, u1)[N]) -> (u1, u1, u1, u1) { + const MID_LABEL = (u1:0, u1:1, u1:0, u1:0); + + let final_labels = + for (stage, current_labels): (uN[N], (u1, u1, u1, u1)[N]) in uN[N]:0..STAGES as uN[N] { + // Used to merge a left and right node in the tree; the stride is the offset we add to + // the left node's index in the array to get the right node's index in the array. + // + // For the first stage, stride == 1 means we merge labels[0] with labels[1], labels[2] + // with labels[3], etc. For the second stage, stride == 2 means we merge + // labels[0] with labels[2], labels[1] with labels[3], etc... + let stride = uN[N]:1 << stage; + for (i, next_labels): (uN[N], (u1, u1, u1, u1)[N]) in uN[N]:0..N as uN[N] { + let right_idx = i + stride; + let right_label = + if right_idx < N as uN[N] { current_labels[right_idx] } else { MID_LABEL }; + + // Double the stride is the index periodicity with which we merge labels. So when + // stride == 1, we perform the merge operation every 2nd element; when stride == 2, + // we merge every 4th element, etc... + let merge_period = stride << uN[N]:1; + // When unrolled, computing `is_active` should be optimized to a constant. + let is_active = (i % merge_period) == uN[N]:0; + let merged = merge_branch_labels(current_labels[i], right_label); + let updated_label = if is_active { merged } else { current_labels[i] }; + update(next_labels, i, updated_label) + }(current_labels) + }(labels); + + if N as uN[N] > uN[N]:0 { final_labels[0] } else { MID_LABEL } +} + +fn make_label(begin: u1, end: u1) -> (u1, u1, u1, u1) { (begin, !(begin | end), end, u1:0) } + +// Predicts leading zero count for |a - b| and whether shift needs correction. +pub fn lza + (a: uN[N], b: uN[N]) -> (uN[RESULT_BITS], u1) { + let MID_LABEL = (u1:0, u1:1, u1:0, u1:0); + let init_pos_labels = (u1, u1, u1, u1)[N]:[MID_LABEL, ...]; + let init_neg_labels = (u1, u1, u1, u1)[N]:[MID_LABEL, ...]; + + // Shifted 'a' and 'b' are used to examine bits at index bit_idx-1 and bit_idx+1 below + let a_early = a >> uN[N]:1; + let b_early = b >> uN[N]:1; + let a_late = a << uN[N]:1; + let b_late = b << uN[N]:1; + + let (indicator_vector, pos_labels, neg_labels) = + for (k, (vec, p_labels, n_labels)): ( + uN[N], (uN[N], (u1, u1, u1, u1)[N], (u1, u1, u1, u1)[N]) + ) in uN[N]:0..N as uN[N] { + let bit_idx = N as uN[N] - uN[N]:1 - k; + let a_left = (a_early >> bit_idx) as u1; + let b_left = (b_early >> bit_idx) as u1; + let a_center = (a >> bit_idx) as u1; + let b_center = (b >> bit_idx) as u1; + let a_right = (a_late >> bit_idx) as u1; + let b_right = (b_late >> bit_idx) as u1; + + // a_i-1 == b_i-1 becomes label 0 + let left_e = !(a_left ^ b_left); + // a_i > b_i becomes label 1 + let center_g = a_center & !b_center; + // a_i < b_i becomes label -1 + let center_s = !a_center & b_center; + // a_i+1 > b_i+1 + let right_g = a_right & !b_right; + // a_i+1 < b_i+1 + let right_s = !a_right & b_right; + + // Computing F from the paper: + // For example, ne_s_ns in terms of labels is [1|-1] ++ -1 ++ [0|1] + let ne_s_ns = !left_e & center_s & !right_s; + let e_g_ns = left_e & center_g & !right_s; + let e_s_ng = left_e & center_s & !right_g; + let ne_g_ng = !left_e & center_g & !right_g; + let indicator_bit = (e_g_ns | ne_s_ns) | (e_s_ng | ne_g_ng); + let next_vec = vec | ((indicator_bit as uN[N]) << bit_idx); + + // Computing detection trees in W > 0: + let pos_begin = (center_g & !right_s) | ne_s_ns; + let pos_end = left_e & center_s; + let pos_label = make_label(pos_begin, pos_end); + + // Computing detection trees in W < 0: + let neg_begin = (center_s & !right_g) | ne_g_ng; + let neg_end = left_e & center_g; + let neg_label = make_label(neg_begin, neg_end); + + (next_vec, update(p_labels, k, pos_label), update(n_labels, k, neg_label)) + }((uN[N]:0, init_pos_labels, init_neg_labels)); + + let pred_clz = std::clzt(indicator_vector); + let (_, _, _, pos_yes) = tree_reduce(pos_labels); + let (_, _, _, neg_yes) = tree_reduce(neg_labels); + let has_error = pos_yes | neg_yes; + (pred_clz, has_error) +} + +#[test] +fn test_lza() { + let (shift, err) = lza(u8:1, u8:0); + assert_eq(shift, u4:7); + assert_eq(err, u1:0); + + let (shift, err) = lza(u8:4, u8:1); + assert_eq(shift, u4:5); + assert_eq(err, u1:1); +} + +#[quickcheck(exhaustive)] +fn lza_u8_nonnegative(a: u8, b: u8) -> bool { + let (pred_clz, has_error) = lza(a, b); + let ab_absdiff = abs_diff::to_corrected(abs_diff::abs_diff(a, b)); + let actual_clz = std::clzt(ab_absdiff); + + let eq_or_one_less = (pred_clz == actual_clz) || (pred_clz + u4:1 == actual_clz); + (has_error || (pred_clz == actual_clz)) && eq_or_one_less +} + +#[quickcheck] +fn lza_u16_nonnegative(a: u16, b: u16) -> bool { + let (pred_clz, has_error) = lza(a, b); + let ab_absdiff = abs_diff::to_corrected(abs_diff::abs_diff(a, b)); + let actual_clz = std::clzt(ab_absdiff); + + let eq_or_one_less = (pred_clz == actual_clz) || (pred_clz + u5:1 == actual_clz); + (has_error || (pred_clz == actual_clz)) && eq_or_one_less +}