Add leading zero anticipator to DSLX standard library and use in apfloat

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

xls/dslx/stdlib/apfloat.x

@@ -15,6 +15,7 @@
 // Arbitrary-precision floating point routines.
 import std;
 import abs_diff;
+import lza;
 
 pub struct APFloat<EXP_SZ: u32, FRACTION_SZ: u32> {
     sign: bits[1],  // Sign bit.
@@ -3061,7 +3062,7 @@ fn or_last_bit<WIDTH: u32>(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<EXP_SZ: u32, FRACTION_SZ: u32>
+pub fn add<EXP_SZ: u32, FRACTION_SZ: u32, USE_LZA: bool = {true}>
     (a: APFloat<EXP_SZ, FRACTION_SZ>, b: APFloat<EXP_SZ, FRACTION_SZ>)
     -> APFloat<EXP_SZ, FRACTION_SZ> {
     // WIDE_EXP: Widened exponent to capture a possible carry bit.
@@ -3134,13 +3135,40 @@ pub fn add<EXP_SZ: u32, FRACTION_SZ: u32>
     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_FRACTION>(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<EXP_SZ: u32, FRACTION_SZ: u32>
 //  - 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<EXP_SZ: u32, FRACTION_SZ: u32>
+pub fn sub<EXP_SZ: u32, FRACTION_SZ: u32, USE_LZA: bool = {true}>
     (x: APFloat<EXP_SZ, FRACTION_SZ>, y: APFloat<EXP_SZ, FRACTION_SZ>)
     -> APFloat<EXP_SZ, FRACTION_SZ> {
     let y = APFloat<EXP_SZ, FRACTION_SZ> { sign: !y.sign, bexp: y.bexp, fraction: y.fraction };
-    add(x, y)
+    add<EXP_SZ, FRACTION_SZ, USE_LZA>(x, y)
+}
+
+#[quickcheck(exhaustive)]
+fn add_clz_lza_equivalence(a: APFloat<u32:3, u32:3>, b: APFloat<u32:3, u32:3>) -> bool {
+    add<u32:3, u32:3, true>(a, b) == add<u32:3, u32:3, false>(a, b)
 }
 
 // add is thoroughly tested elsewhere so a few simple tests is sufficient.

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<N: u32, STAGES: u32 = {std::clog2(N)}>
+    (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<N: u32, RESULT_BITS: u32 = {std::clog2(N + u32:1)}>
+    (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
+}