Fix 32-bit regressions in ndarray encode/decode and stochastic rounding tests (#60)

Author: awf

.github/workflows/ci.yaml

@@ -72,7 +72,7 @@ jobs:
 
       - name: Run full unit tests on linux/386
         run: |
-          bash etc/test-linux-386.sh run
+          bash etc/test-linux-386.sh run -vv test
 
       - name: Save linux/386 docker image to cache path
         if: always()

etc/test-linux-386.sh

@@ -2,6 +2,8 @@
 
 # Copyright (c) 2024 Graphcore Ltd. All rights reserved.
 
+[ -n "${BASH_VERSION:-}" ] || exec bash "$0" "$@"
+
 set -euo pipefail
 
 SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
@@ -10,14 +12,23 @@ REPO_DIR="$(cd "${SCRIPT_DIR}/.." && pwd)"
 IMAGE="gfloat-linux-386:py310-uv"
 PLATFORM="linux/386"
 DOCKERFILE="etc/linux-386.Dockerfile"
-MODE="${1:-all}"
+MODE="all"
+
+if [[ "${1:-}" == "all" || "${1:-}" == "load" || "${1:-}" == "build" || "${1:-}" == "run" ]]; then
+    MODE="$1"
+    shift
+fi
 
 usage() {
-    echo "Usage: bash etc/test-linux-386.sh [load|build|run]"
+    echo "Usage: bash etc/test-linux-386.sh [all|load|build|run] [pytest args...]"
     echo "  load  Build image only if missing"
     echo "  build Force rebuild image"
     echo "  run   Run tests (image must already exist)"
     echo "  (no arg) Build image if missing, then run tests"
+    echo ""
+    echo "Examples:"
+    echo "  bash etc/test-linux-386.sh run test/test_encode.py::test_encode[binary32]"
+    echo "  bash etc/test-linux-386.sh run -k stochastic -q"
 }
 
 build_image() {
@@ -39,16 +50,16 @@ run_tests() {
     bash -lc '
         set -euo pipefail
         export PYTHONPATH="/work/src${PYTHONPATH:+:${PYTHONPATH}}"
-        python -m pytest -vv test
-    '
+        python -m pytest "$@"
+    ' _ "$@"
 }
 
 if [[ "${MODE}" != "all" && "${MODE}" != "load" && "${MODE}" != "build" && "${MODE}" != "run" ]]; then
     usage
     exit 2
 fi
 
-echo "Running full unit tests in Docker (${PLATFORM})"
+echo "Running unit tests in Docker (${PLATFORM})"
 
 if ! docker info >/dev/null 2>&1; then
     echo "Docker daemon is not running; start Docker and rerun this script." >&2
@@ -73,5 +84,5 @@ if [[ "${MODE}" == "run" || "${MODE}" == "all" ]]; then
         exit 1
     fi
 
-    run_tests
+    run_tests "$@"
 fi

src/gfloat/decode_ndarray.py

@@ -78,10 +78,12 @@ def decode_ndarray(
 
     issubnormal = (exp == 0) & (significand != 0) & fi.has_subnormals
     expval = np.where(issubnormal, 1 - bias, exp - bias)
-    fsignificand = np.where(issubnormal, 0.0, 1.0) + np.ldexp(significand, -t)
+    fsignificand = np.where(issubnormal, 0.0, 1.0) + np.ldexp(
+        significand.astype(np.float64), np.int32(-t)
+    )
 
     # Normal/Subnormal/Zero case, other values will be overwritten
-    expval_safe = np.where(isspecial | iszero, 0, expval)
+    expval_safe = np.where(isspecial | iszero, 0, expval).astype(np.int32)
     fval_finite_safe = sign * np.ldexp(fsignificand, expval_safe)
     fval = np.where(~(iszero | isspecial), fval_finite_safe, fval)
 

src/gfloat/encode_ndarray.py

@@ -66,11 +66,11 @@ def encode_ndarray(fi: FormatInfo, v: npt.NDArray) -> npt.NDArray:
         biased_exp = exp.astype(np.int64) + (fi.bias - 1)
         subnormal_mask = (biased_exp < 1) & fi.has_subnormals
 
-        biased_exp_safe = np.where(subnormal_mask, biased_exp, 0)
+        biased_exp_safe = np.where(subnormal_mask, biased_exp, 0).astype(np.int32)
         tsig = np.where(subnormal_mask, np.ldexp(sig, biased_exp_safe), sig * 2 - 1.0)
         biased_exp[subnormal_mask] = 0
 
-        isig = np.floor(np.ldexp(tsig, t)).astype(np.int64)
+        isig = np.floor(np.ldexp(tsig, np.int32(t))).astype(np.int64)
 
         zero_mask = fi.has_zero & (isig == 0) & (biased_exp == 0)
         if not fi.has_nz:
@@ -80,8 +80,9 @@ def encode_ndarray(fi: FormatInfo, v: npt.NDArray) -> npt.NDArray:
         if fi.is_twos_complement:
             isig[finite_sign] = (1 << t) - isig[finite_sign]
 
-        code[finite_mask] = (
-            (finite_sign.astype(int) << (k - 1)) | (biased_exp << t) | (isig << 0)
-        )
+        sign_field = np.left_shift(finite_sign.astype(np.uint64), np.uint64(k - 1))
+        exp_field = np.left_shift(biased_exp.astype(np.uint64), np.uint64(t))
+        sig_field = isig.astype(np.uint64)
+        code[finite_mask] = sign_field | exp_field | sig_field
 
     return code

test/conftest.py

@@ -1,59 +1 @@
 # Copyright (c) 2024 Graphcore Ltd. All rights reserved.
-
-import struct
-
-import pytest
-
-
-def _is_32bit_python() -> bool:
-    return struct.calcsize("P") * 8 == 32
-
-
-def pytest_collection_modifyitems(
-    config: pytest.Config, items: list[pytest.Item]
-) -> None:
-    if not _is_32bit_python():
-        return
-
-    mark = pytest.mark.xfail(
-        reason="Known 32-bit regressions (issue #57)",
-        strict=False,
-    )
-
-    for item in items:
-        nodeid = item.nodeid
-
-        if nodeid.startswith("test/test_array_api.py::"):
-            item.add_marker(mark)
-            continue
-
-        if nodeid.startswith(
-            "test/test_round.py::test_stochastic_rounding_scalar_eq_array"
-        ):
-            item.add_marker(mark)
-            continue
-
-        if nodeid.startswith("test/test_encode.py::test_encode["):
-            item.add_marker(mark)
-            continue
-
-        if nodeid.startswith("test/test_encode.py::test_encode_edges[encode_ndarray-"):
-            item.add_marker(mark)
-            continue
-
-        if (
-            nodeid.startswith("test/test_decode.py::test_spot_check_")
-            and "[array]" in nodeid
-        ):
-            item.add_marker(mark)
-            continue
-
-        if (
-            nodeid.startswith("test/test_decode.py::test_specials_decode[")
-            and "[array-" in nodeid
-        ):
-            item.add_marker(mark)
-            continue
-
-        if nodeid.startswith("test/test_decode.py::test_consistent_decodes_all_values["):
-            item.add_marker(mark)

test/test_array_api.py

@@ -25,7 +25,7 @@ def test_array_api(fi: FormatInfo, rnd: RoundMode, sat: bool) -> None:
     a = xp.asarray(a0)
 
     srnumbits = 32
-    srbits0 = np.random.randint(0, 2**srnumbits, a.shape)
+    srbits0 = np.random.randint(0, 2**srnumbits, a.shape, dtype=np.int64)
     srbits = xp.asarray(srbits0)
 
     round_ndarray(fi, a, rnd, sat, srbits=srbits, srnumbits=srnumbits)  # type: ignore

test/test_round.py

@@ -537,7 +537,7 @@ def test_stochastic_rounding(
     n = 10_000
     expected_up_count = expected_up * n
 
-    srbits = np.random.randint(0, 2**srnumbits, size=(n,))
+    srbits = np.random.randint(0, 2**srnumbits, size=(n,), dtype=np.int64)
     if impl == "scalar":
         count_v1 = 0
         for k in range(n):
@@ -591,7 +591,7 @@ def test_stochastic_rounding_scalar_eq_array(
     for alpha in (0, 0.3, 0.5, 0.6, 0.7, 0.9, 1.25):
         v = _linterp(v0, v1, alpha)
         assert np.isfinite(v).all()
-        srbits = np.random.randint(0, 2**srnumbits, v.shape)
+        srbits = np.random.randint(0, 2**srnumbits, v.shape, dtype=np.int64)
 
         val_array = round_ndarray(
             fi,