Merge pull request #34 from graphcore-research/stochastic-rounding

Add stochastic rounding

Author: awf

docs/source/05-stochastic-rounding.ipynb

@@ -12,7 +12,12 @@ def _isodd(v: int) -> bool:
 
 
 def round_float(
-    fi: FormatInfo, v: float, rnd: RoundMode = RoundMode.TiesToEven, sat: bool = False
+    fi: FormatInfo,
+    v: float,
+    rnd: RoundMode = RoundMode.TiesToEven,
+    sat: bool = False,
+    srbits: int = -1,
+    srnumbits: int = 0,
 ) -> float:
     """
     Round input to the given :py:class:`FormatInfo`, given rounding mode and saturation flag
@@ -27,6 +32,8 @@ def round_float(
       v (float): Input value to be rounded
       rnd (RoundMode): Rounding mode to use
       sat (bool): Saturation flag: if True, round overflowed values to `fi.max`
+      srbits (int): Bits to use for stochastic rounding if rnd == Stochastic.
+      srnumbits (int): How many bits are in srbits.  Implies srbits < 2**srnumbits.
 
     Returns:
       A float which is one of the values in the format.
@@ -35,12 +42,17 @@ def round_float(
        ValueError: The target format cannot represent the input
              (e.g. converting a `NaN`, or an `Inf` when the target has no
              `NaN` or `Inf`, and :paramref:`sat` is false)
+       ValueError: Inconsistent arguments, e.g. srnumbits >= 2**srnumbits
     """
 
     # Constants
     p = fi.precision
     bias = fi.expBias
 
+    if rnd == RoundMode.Stochastic:
+        if srbits >= 2**srnumbits:
+            raise ValueError(f"srnumbits={srnumbits} >= 2**srnumbits={2**srnumbits}")
+
     if np.isnan(v):
         if fi.num_nans == 0:
             raise ValueError(f"No NaN in format {fi}")
@@ -75,12 +87,15 @@ def round_float(
         # Round
         isignificand = math.floor(fsignificand)
         delta = fsignificand - isignificand
+
+        # fmt: off
         if (
             (rnd == RoundMode.TowardPositive and not sign and delta > 0)
             or (rnd == RoundMode.TowardNegative and sign and delta > 0)
             or (rnd == RoundMode.TiesToAway and delta >= 0.5)
             or (rnd == RoundMode.TiesToEven and delta > 0.5)
             or (rnd == RoundMode.TiesToEven and delta == 0.5 and _isodd(isignificand))
+            or (rnd == RoundMode.Stochastic and delta > (0.5 + srbits) * 2.0**-srnumbits)
         ):
             isignificand += 1
 
@@ -95,6 +110,7 @@ def round_float(
                 or (rnd == RoundMode.TiesToAway and delta >= 0.5)
                 or (rnd == RoundMode.TiesToEven and delta > 0.5)
                 or (rnd == RoundMode.TiesToEven and delta == 0.5 and code_is_odd)
+                or (rnd == RoundMode.Stochastic and delta > (0.5 + srbits) * 2.0**-srnumbits)
             ):
                 # Go to nextUp.
                 # Increment isignificand if zero,
@@ -105,6 +121,7 @@ def round_float(
                     assert isignificand == 1
                     expval += 1
             ## End special case for Precision=1.
+        # fmt: on
 
         # Reconstruct rounded result to float
         result = isignificand * (2.0**expval)

src/gfloat/round.py

@@ -16,6 +16,7 @@ class RoundMode(Enum):
     TowardPositive = 3  #: :math:`\min \{ r ~ s.t. ~ r \ge v \}`
     TiesToEven = 4  #: Round to nearest, ties to even
     TiesToAway = 5  #: Round to nearest, ties away from zero
+    Stochastic = 6  #: Stochastic rounding
 
 
 class FloatClass(Enum):

src/gfloat/types.py

@@ -488,3 +488,45 @@ def test_round_roundtrip(round_float: Callable, fi: FormatInfo) -> None:
         fv = decode_float(fi, i)
         fval2 = round_float(fi, fv.fval)
         np.testing.assert_equal(fval2, fv.fval)
+
+
+@pytest.mark.parametrize(
+    "v, srnumbits, expected_up",
+    (
+        (259, 3, 0.0 / 8),
+        (259, 5, 1.0 / 32),
+        (277, 3, 3.0 / 8),
+        (288, 3, 0.5),
+        (311, 3, 7.0 / 8),
+    ),
+)
+def test_stochastic_rounding(v: float, srnumbits: int, expected_up: float) -> None:
+    fi = format_info_ocp_e5m2
+
+    v0 = round_float(fi, v, RoundMode.TowardNegative)
+    v1 = round_float(fi, v, RoundMode.TowardPositive)
+
+    n = 10_000
+    expected_up_count = expected_up * n
+
+    srbits = np.random.randint(0, 2**srnumbits, size=(n,))
+    count_v1 = 0
+    for k in range(n):
+        r = round_float(
+            fi,
+            v,
+            RoundMode.Stochastic,
+            sat=False,
+            srbits=srbits[k],
+            srnumbits=srnumbits,
+        )
+        if r == v1:
+            count_v1 += 1
+        else:
+            assert r == v0
+
+    print(f"SRBits={srnumbits}, observed = {count_v1}, expected = {expected_up_count} ")
+    # e.g. if expected is 1250/10000, want to be within 0.5,1.5
+    # this is loose, but should still catch logic errors
+    atol = n * 2.0 ** (-1 - srnumbits)
+    np.testing.assert_allclose(count_v1, expected_up_count, atol=atol)