Merge pull request #14 from graphcore-research/text-microxscaling

Add microxscaling tests, Add `quantize_block`.

Author: awf

docs/source/api.rst

@@ -5,15 +5,22 @@ API
 
 .. module:: gfloat
 
-Functions
----------
+Scalar Functions
+----------------
 
 .. autofunction:: decode_float
 .. autofunction:: round_float
 .. autofunction:: encode_float
 
+Block format functions
+----------------------
+
 .. autofunction:: decode_block
 .. autofunction:: encode_block
+.. autofunction:: quantize_block
+
+.. autofunction:: compute_scale_amax
+
 
 Classes
 -------

docs/source/conf.py

@@ -24,6 +24,16 @@
     "myst_nb",
 ]
 
+autodoc_typehints = "none"  # We have them in the parameter descriptors
+autodoc_typehints_format = "short"
+python_use_unqualified_type_names = True
+
+autodoc_type_aliases = {
+    "Iterable": "Iterable",
+    "npt.ArrayLike": "ArrayLike",
+    "npt.NDArray": "NDArray",
+}
+
 autodoc_default_options = {
     "member-order": "bysource",
 }

docs/source/formats.rst

@@ -31,3 +31,13 @@ IEEE WG P3109 Formats
 ---------------------
 
 .. autofunction:: format_info_p3109
+
+Block Formats
+---------------------
+
+.. autodata:: format_info_mxfp8_e5m2
+.. autodata:: format_info_mxfp8_e4m3
+.. autodata:: format_info_mxfp6_e3m2
+.. autodata:: format_info_mxfp6_e2m3
+.. autodata:: format_info_mxfp4_e2m1
+.. autodata:: format_info_mxint8

pyproject.toml

@@ -33,3 +33,8 @@ optional-dependencies = {dev = {file = ["requirements-dev.txt"]}}
 [tool.black]
 line-length = 88
 fast = true
+
+[tool.mypy]
+[[tool.mypy.overrides]]
+module = "mx.*"
+ignore_missing_imports = true

requirements-dev.txt

@@ -1,12 +1,14 @@
 # Requirements for tests
 pytest
 ml_dtypes
+mx @ git+https://github.com/microsoft/microxcaling
 
 # Requirements for development
 pre-commit
 black
 mypy
 black[jupyter]
+isort
 
 # Requirements for docs
 sphinx==7.1.2

src/gfloat/__init__.py

@@ -1,6 +1,12 @@
 # Copyright (c) 2024 Graphcore Ltd. All rights reserved.
 
-from .block import BlockFormatInfo, decode_block, encode_block
+from .block import (
+    BlockFormatInfo,
+    compute_scale_amax,
+    decode_block,
+    encode_block,
+    quantize_block,
+)
 from .decode import decode_float
 from .printing import float_pow2str, float_tilde_unless_roundtrip_str
 from .round import encode_float, round_float

src/gfloat/block.py

@@ -4,10 +4,13 @@
 # https://en.wikipedia.org/wiki/Block_floating_point
 
 from dataclasses import dataclass
-from typing import Iterable
+from typing import Callable, Iterable
+
+import numpy as np
+import numpy.typing as npt
 
 from .decode import decode_float
-from .round import encode_float, round_float
+from .round import RoundMode, encode_float, round_float
 from .types import FormatInfo
 
 
@@ -45,8 +48,12 @@ def block_size_bytes(self) -> int:
         assert bits % 8 == 0
         return bits // 8
 
+    @property
+    def __name__(self) -> str:
+        return self.name
+
     def __str__(self) -> str:
-        return f"{self.name}"
+        return f"BlockFormatInfo:{self.name})"
 
 
 def decode_block(fi: BlockFormatInfo, block: Iterable[int]) -> Iterable[float]:
@@ -78,13 +85,18 @@ def decode_block(fi: BlockFormatInfo, block: Iterable[int]) -> Iterable[float]:
 
 
 def encode_block(
-    fi: BlockFormatInfo, scale: float, vals: Iterable[float]
+    fi: BlockFormatInfo,
+    scale: float,
+    vals: Iterable[float],
+    round: RoundMode = RoundMode.TiesToEven,
 ) -> Iterable[int]:
     """
-    Encode a :paramref:`block` of bytes into block Format descibed by :paramref:`fi`
+    Encode float :paramref:`vals` into block Format described by :paramref:`fi`
 
-    The :paramref:`scale` is explicitly passed, and is converted to `1/(1/scale)`
-    before rounding to the target format.
+    The :paramref:`scale` is explicitly passed, and the :paramref:`vals` are
+    assumed to already be multiplied by `1/scale`.
+    That is, this is pure encoding, scaling is computed and applied elsewhere
+    (see e.g. :func:`quantize_block`).
 
     It is checked for overflow in the target format,
     and will raise an exception if it does.
@@ -93,24 +105,88 @@ def encode_block(
       fi (BlockFormatInfo): Describes the target block format
       scale (float): Scale to be recorded in the block
       vals (Iterable[float]): Input block
+      round (RoundMode): Rounding mode to use, defaults to `TiesToEven`
 
     Returns:
       A sequence of ints representing the encoded values.
 
     Raises:
       ValueError: The scale overflows the target scale encoding format.
     """
-    # TODO: this should not do any multiplication - the scale is to be recorded not applied.
-    recip_scale = 1 / scale
-    scale = 1 / recip_scale
 
-    if scale > fi.stype.max:
-        raise ValueError(f"Scaled {scale} too large for {fi.stype}")
+    if scale > fi.stype.max or scale < fi.stype.min:
+        raise ValueError(f"Scaled {scale} out of range for {fi.stype}")
+
+    sat = True  # Saturate elements if out of range
 
     def enc(ty: FormatInfo, x: float) -> int:
-        return encode_float(ty, round_float(ty, x))
+        return encode_float(ty, round_float(ty, x, round, sat))
 
     yield enc(fi.stype, scale)
 
     for val in vals:
-        yield enc(fi.etype, recip_scale * val)
+        yield enc(fi.etype, val)
+
+
+ComputeScaleCallable = Callable[[float, npt.ArrayLike], float]
+
+
+def compute_scale_amax(emax: float, vals: npt.ArrayLike) -> float:
+    """
+    Compute a scale factor such that :paramref:`vals` can be scaled to the
+    range [0, 2**emax].  That is, `scale` is computed such that the largest
+    exponent in the array `vals * scale` will be `emax`.
+
+    The scale is clipped to the range 2**[-127, 127].
+
+    If all values are zero, any scale value smaller than emax would be accurate,
+    but returning the smallest possible means that quick checks on the magnitude
+    to identify near-zero blocks will also find the all-zero blocks.
+
+    Args:
+      emax (float): Maximum exponent to appear in `vals * scale`
+      vals (ArrayLike): Input block
+
+    Returns:
+      A float such that `vals * scale` has exponents less than or equal to `emax`.
+
+    Note:
+      If all vals are zero, 1.0 is returned.
+    """
+    amax = np.max(np.abs(vals))
+    if amax == 0.0:
+        q_log2scale = -127.0
+    else:
+        q_log2scale = np.floor(np.log2(amax)) - emax
+        q_log2scale = np.clip(q_log2scale, -127.0, 127.0)
+    return 2.0**q_log2scale
+
+
+def quantize_block(
+    fi: BlockFormatInfo,
+    vals: npt.NDArray[np.float64],
+    compute_scale: ComputeScaleCallable,
+    round: RoundMode = RoundMode.TiesToEven,
+) -> npt.NDArray[np.float64]:
+    """
+    Encode and decode a block of :paramref:`vals` of bytes into
+    block format described by :paramref:`fi`
+
+    Args:
+      fi (BlockFormatInfo): Describes the target block format
+      vals (numpy.array): Input block
+      compute_scale ((float, ArrayLike) -> float):
+          Callable to compute the scale, defaults to :func:`compute_scale_amax`
+      round (RoundMode): Rounding mode to use, defaults to `TiesToEven`
+
+    Returns:
+      An array of floats representing the quantized values.
+
+    Raises:
+      ValueError: The scale overflows the target scale encoding format.
+    """
+
+    q_scale = compute_scale(fi.etype.emax, vals)
+    scaled_vals = vals / q_scale
+    enc = encode_block(fi, q_scale, scaled_vals, round)
+    return np.fromiter(decode_block(fi, enc), float)

src/gfloat/formats.py

@@ -5,7 +5,7 @@
 
 #: FormatInfo for IEEE-754 Binary32 format
 format_info_binary32 = FormatInfo(
-    name="binary32",
+    name="format_info_binary32",
     k=32,
     precision=24,
     emax=127,
@@ -19,7 +19,7 @@
 
 #: FormatInfo for IEEE-754 Binary16 format
 format_info_binary16 = FormatInfo(
-    name="binary16",
+    name="format_info_binary16",
     k=16,
     precision=11,
     emax=15,
@@ -33,7 +33,7 @@
 
 #: FormatInfo for Google BFloat16 format
 format_info_bfloat16 = FormatInfo(
-    name="bfloat16",
+    name="format_info_bfloat16",
     k=16,
     precision=8,
     emax=127,
@@ -47,7 +47,7 @@
 
 #: FormatInfo for OCP E5M2 format
 format_info_ocp_e5m2 = FormatInfo(
-    name="ocp_e5m2",
+    name="format_info_ocp_e5m2",
     k=8,
     precision=3,
     emax=15,
@@ -61,7 +61,7 @@
 
 #: FormatInfo for OCP E4M3 format
 format_info_ocp_e4m3 = FormatInfo(
-    name="ocp_e4m3",
+    name="format_info_ocp_e4m3",
     k=8,
     precision=4,
     emax=8,
@@ -75,7 +75,7 @@
 
 #: FormatInfo for OCP MX E2M3 format
 format_info_ocp_e2m3 = FormatInfo(
-    name="ocp_e2m3",
+    name="format_info_ocp_e2m3",
     k=6,
     precision=4,
     emax=2,
@@ -89,7 +89,7 @@
 
 #: FormatInfo for OCP MX E3M2 format
 format_info_ocp_e3m2 = FormatInfo(
-    name="ocp_e3m2",
+    name="format_info_ocp_e3m2",
     k=6,
     precision=3,
     emax=4,
@@ -103,7 +103,7 @@
 
 #: FormatInfo for OCP MX E2M1 format
 format_info_ocp_e2m1 = FormatInfo(
-    name="ocp_e2m1",
+    name="format_info_ocp_e2m1",
     k=4,
     precision=2,
     emax=2,
@@ -117,7 +117,7 @@
 
 #: FormatInfo for OCP MX E8M0 format
 format_info_ocp_e8m0 = FormatInfo(
-    name="ocp_e8m0",
+    name="format_info_ocp_e8m0",
     k=8,
     precision=1,
     emax=127,
@@ -131,7 +131,7 @@
 
 #: FormatInfo for OCP MX INT8 format
 format_info_ocp_int8 = FormatInfo(
-    name="ocp_int8",
+    name="format_info_ocp_int8",
     k=8,
     precision=8,
     emax=0,
@@ -210,11 +210,11 @@ def format_info_p3109(precision: int) -> FormatInfo:
 # Block formats
 
 format_info_mxfp8_e5m2 = BlockFormatInfo(
-    "ocp_mxfp8_e5m2", format_info_ocp_e5m2, 32, format_info_ocp_e8m0
+    "format_info_mxfp8_e5m2", format_info_ocp_e5m2, 32, format_info_ocp_e8m0
 )
 
 format_info_mxfp8_e4m3 = BlockFormatInfo(
-    "ocp_mxfp8_e4m3", format_info_ocp_e4m3, 32, format_info_ocp_e8m0
+    "format_info_mxfp8_e4m3", format_info_ocp_e4m3, 32, format_info_ocp_e8m0
 )
 
 format_info_mxfp6_e3m2 = BlockFormatInfo(
@@ -233,11 +233,15 @@ def format_info_p3109(precision: int) -> FormatInfo:
     "format_info_mxfp4_e2m1", format_info_ocp_e2m1, 32, format_info_ocp_e8m0
 )
 
+format_info_mxint8 = BlockFormatInfo(
+    "format_info_mxint8", format_info_ocp_int8, 32, format_info_ocp_e8m0
+)
+
 all_block_formats = [
     format_info_mxfp8_e5m2,
     format_info_mxfp8_e4m3,
     format_info_mxfp6_e3m2,
     format_info_mxfp6_e2m3,
     format_info_mxfp4_e2m1,
-    format_info_mxfp4_e2m1,
+    format_info_mxint8,
 ]

src/gfloat/types.py

@@ -400,5 +400,9 @@ def is_all_subnormal(self) -> bool:
         """
         return (self.expBits == 0) and self.has_subnormals
 
+    @property
+    def __name__(self) -> str:
+        return self.name
+
     def __str__(self) -> str:
         return f"{self.name}"