Fix (minifloat): restructuring quantizers

Xilinx · Feb 23, 2024 · 8f599cd · 8f599cd
1 parent 254a91f
commit 8f599cd
Show file tree

Hide file tree

Showing 4 changed files with 31 additions and 47 deletions.
diff --git a/src/brevitas/core/quant/float.py b/src/brevitas/core/quant/float.py
@@ -7,13 +7,11 @@
 import torch.nn as nn
 
 import brevitas
-from brevitas.core.function_wrapper import FloatClamp
 from brevitas.core.function_wrapper import RoundSte
 from brevitas.core.scaling import ConstScaling
 from brevitas.core.utils import StatelessBuffer
 from brevitas.function.ops import max_float
 from brevitas.function.ops_ste import floor_ste
-from brevitas.utils.float_quant_utils import get_max_value
 
 
 class FloatQuant(brevitas.jit.ScriptModule):

diff --git a/src/brevitas/quant/experimental/float_base.py b/src/brevitas/quant/experimental/float_base.py
@@ -15,20 +15,37 @@
 from brevitas.utils.float_quant_utils import get_max_value
 
 
-class FloatWeightBase(SolveTensorQuantFloatToIntImplFromEnum):
-    proxy_class = WeightQuantProxyFromInjector
+class FloatBase(SolveTensorQuantFloatToIntImplFromEnum):
     tensor_quant = FloatQuant
     signed = True
     float_to_int_impl_type = 'round'
     scaling_min_val = 1e-10
+    float_clamp_impl = FloatClamp
+    tensor_clamp_impl = TensorClamp
 
+    @value
+    def exponent_bias(exponent_bit_width):
+        return 2 ** (exponent_bit_width - 1) - 1
 
-class FloatActBase(SolveTensorQuantFloatToIntImplFromEnum):
+    @value
+    def max_value(
+            exponent_bit_width, mantissa_bit_width, exponent_bias, nan_values, inf_values,
+            saturating):
+        return get_max_value(
+            exponent_bit_width,
+            mantissa_bit_width,
+            exponent_bias,
+            nan_values,
+            inf_values,
+            saturating)
+
+
+class FloatWeightBase(FloatBase):
+    proxy_class = WeightQuantProxyFromInjector
+
+
+class FloatActBase(FloatBase):
     proxy_class = ActQuantProxyFromInjector
-    tensor_quant = FloatQuant
-    signed = True
-    float_to_int_impl_type = 'round'
-    scaling_min_val = 1e-10
 
 
 class ScaledFloatWeightBase(FloatWeightBase, WeightQuantSolver):
@@ -46,45 +63,19 @@ class ScaledFloatActBase(FloatActBase, ActQuantSolver):
     float_scaling_impl = FloatScaling
 
 
-class ExponentBiasMixin(ExtendedInjector):
-
-    @value
-    def exponent_bias(exponent_bit_width):
-        return 2 ** (exponent_bit_width - 1) - 1
-
-
-class MaxFloatInfNaNMixin(ExtendedInjector):
-
-    @value
-    def max_value(
-            exponent_bit_width, mantissa_bit_width, exponent_bias, nan_values, inf_values,
-            saturating):
-        return get_max_value(
-            exponent_bit_width,
-            mantissa_bit_width,
-            exponent_bias,
-            nan_values,
-            inf_values,
-            saturating)
-
-
-class Fp8e4m3Mixin(ExponentBiasMixin, MaxFloatInfNaNMixin):
+class Fp8e4m3Mixin(ExtendedInjector):
     bit_width = 8
     exponent_bit_width = 4
     mantissa_bit_width = 3
-    float_clamp_impl = FloatClamp
-    tensor_clamp_impl = TensorClamp
     nan_values = (('111',))
     inf_values = None
     saturating = True
 
 
-class Fp8e5m2Mixin(ExponentBiasMixin, MaxFloatInfNaNMixin):
+class Fp8e5m2Mixin(ExtendedInjector):
     bit_width = 8
     exponent_bit_width = 5
     mantissa_bit_width = 2
-    float_clamp_impl = FloatClamp
-    tensor_clamp_impl = TensorClamp
     nan_values = ('01', '11', '10')
     inf_values = (('00',))
     saturating = True
diff --git a/tests/brevitas/core/minifloat_fixtures.py b/tests/brevitas/core/minifloat_fixtures.py
@@ -4,18 +4,14 @@
 import pytest_cases
 from pytest_cases import fixture_union
 
-from brevitas.core.function_wrapper import FloatClamp
 from brevitas.inject.enum import BitWidthImplType
-from brevitas.quant.experimental.float_base import ExponentBiasMixin
-from brevitas.quant.experimental.float_base import MaxFloatInfNaNMixin
 from brevitas.quant.experimental.float_base import ScaledFloatWeightBase
 
 
-class Fp8e4m3Base(ExponentBiasMixin, MaxFloatInfNaNMixin, ScaledFloatWeightBase):
+class Fp8e4m3Base(ScaledFloatWeightBase):
     bit_width = 8
     exponent_bit_width = 4
     mantissa_bit_width = 3
-    float_clamp_impl = FloatClamp
     nan_values = None
     inf_values = None
     saturating = True
@@ -24,11 +20,10 @@ class Fp8e4m3Base(ExponentBiasMixin, MaxFloatInfNaNMixin, ScaledFloatWeightBase)
     hypothesis_internal_is_this_a_mock_check = False
 
 
-class Fp8e5m2Base(ExponentBiasMixin, MaxFloatInfNaNMixin, ScaledFloatWeightBase):
+class Fp8e5m2Base(ScaledFloatWeightBase):
     bit_width = 8
     exponent_bit_width = 5
     mantissa_bit_width = 2
-    float_clamp_impl = FloatClamp
     nan_values = None
     inf_values = None
     saturating = True

diff --git a/tests/brevitas/core/test_minifloat.py b/tests/brevitas/core/test_minifloat.py
@@ -4,13 +4,13 @@
 from hypothesis import given
 import pytest
 
-from brevitas.quant.experimental.float_base import Fp8e4m3Mixin
-from brevitas.quant.experimental.float_base import Fp8e5m2Mixin
+from brevitas.quant.experimental.float import Fp8e4m3Weight
+from brevitas.quant.experimental.float import Fp8e5m2Weight
 from tests.brevitas.hyp_helper import float_tensor_random_shape_st
 
 from .minifloat_fixtures import *
 
-FORMATS = {Fp8e5m2Mixin: 57344., Fp8e4m3Mixin: 448., Fp8e4m3Base: 480., Fp8e5m2Base: 114688.}
+FORMATS = {Fp8e5m2Weight: 57344., Fp8e4m3Weight: 448., Fp8e4m3Base: 480., Fp8e5m2Base: 114688.}
 
 
 @pytest.mark.parametrize(