Better handle types and overflows in softmax (especially lagacy and stable)

hls4ml/backends/fpga/fpga_backend.py

@@ -44,7 +44,6 @@
     IntegerPrecisionType,
     PrecisionType,
     RoundingMode,
-    SaturationMode,
     StandardFloatPrecisionType,
     UnspecifiedPrecisionType,
     XnorPrecisionType,
@@ -129,22 +128,25 @@ def __init__(self, name):
             ConfigurableAttribute('skip', value_type=bool, default=False, description=descriptions.softmax_skip),
             TypeAttribute(
                 'exp_table',
-                default=FixedPrecisionType(18, 8, rounding_mode=RoundingMode.RND, saturation_mode=SaturationMode.SAT),
+                default=UnspecifiedPrecisionType(),
                 description=descriptions.table_type,
             ),
             TypeAttribute(
                 'inv_table',
-                default=FixedPrecisionType(18, 8, rounding_mode=RoundingMode.RND, saturation_mode=SaturationMode.SAT),
+                default=UnspecifiedPrecisionType(),
                 description=descriptions.table_type,
             ),
             TypeAttribute(
                 'inv_inp',
-                default=FixedPrecisionType(18, 8, rounding_mode=RoundingMode.RND, saturation_mode=SaturationMode.SAT),
+                default=UnspecifiedPrecisionType(),
+                description='What the accumulated value is cast to before accessing the inversion table (only in stable)',
             ),
             TypeAttribute(
-                'accum',
-                default=FixedPrecisionType(18, 8, rounding_mode=RoundingMode.RND, saturation_mode=SaturationMode.SAT),
+                'inp_norm',
+                default=UnspecifiedPrecisionType(),
+                description='The internal width used for the exp table lookup (only in stable)',
             ),
+            TypeAttribute('accum', description=descriptions.accum_type),
         ]
         self.attribute_map[Softmax] = softmax_attrs
 

hls4ml/backends/fpga/passes/fix_softmax_table_size.py

@@ -1,5 +1,3 @@
-import warnings
-
 from hls4ml.model.layers import Layer, Softmax
 from hls4ml.model.optimizer import OptimizerPass
 
@@ -8,61 +6,41 @@ class FixSoftmaxTableSize(OptimizerPass):
     def match(self, node):
         if not isinstance(node, Softmax):
             return False
-        if 'inv_table_size' in node.attributes:
-            return False  # handler generating inv_table_size sets it properly
+        if node.get_attr('table_sizes_checked'):
+            return False  # This optimizer has already run
         return True
 
     def transform(self, model, node: Layer):
         inp_layer = node.get_input_node()  # type: ignore
         if not isinstance(inp_layer, Layer):
             raise RuntimeError(f'Softmax layer {node.name} does not have an input layer')
 
-        input_bw: int = inp_layer.get_attr('result_t').precision.width  # type: ignore
-        table_bw: int = node.get_attr('inv_table_t').precision.width  # type: ignore
         table_size = int(node.get_attr('table_size'))  # type: ignore
+        exp_table_size = node.get_attr('exp_table_size', table_size)
+        inv_table_size = node.get_attr('inv_table_size', table_size)
 
-        backend = model.config.config['Backend']
+        implemenation = node.get_attr('implementation')
 
-        # Somehow, Intel want one extra bits for the table.
-        # I don't know why but if not simulation will crash with segmentation fault.
-        backend_limitation = -1 if backend == 'Quartus' else 0
+        if implemenation == 'stable':
+            inp_norm_bw = node.get_attr('inp_norm_t').precision.width
+            node.set_attr('exp_table_size', min(2**inp_norm_bw, exp_table_size))
 
-        if 2 ** (min(input_bw, table_bw) + backend_limitation) < table_size:
-            # If table size is too large w.r.t. input bitwidth and table bitwidth,
-            # reduce table size to avoid undefined behavior when cutting indices from,
-            # fixed point number.
-            node.set_attr('table_size', str(2 ** (min(input_bw, table_bw) + backend_limitation)))
-            if 2**input_bw < table_size:
-                # The warning message does not have to be looking like this, but you are asking
-                # 125 characters long line.
-                warnings.warn(
-                    (
-                        f'Softmax layer {node.name} table size is too large for input'
-                        f'bitwidth {input_bw}. Setting table size to {2**input_bw}.'
-                        'To avoid this warning, please increase input bitwidth or'
-                        'decrease table size.'
-                    ),
-                    stacklevel=1,
-                )
-            if 2**table_bw < table_size:
-                warnings.warn(
-                    (
-                        f'Softmax layer {node.name} table size is too large for input'
-                        f'bitwidth {input_bw}. Setting table size to {2**input_bw}.'
-                        'To avoid this warning, please increase input bitwidth or'
-                        'decrease table size.'
-                    ),
-                    stacklevel=1,
-                )
-            if backend == 'Quartus':
-                warnings.warn(
-                    (
-                        "Quartus backend's table size is half of 2^min(input_bw-1,table_bw-1)"
-                        ' instead of 2^min(input_bw,table_bw).'
-                    ),
-                    stacklevel=1,
-                )
-            return False
+            inv_inp_bw = node.get_attr('inv_inp_t').precision.width
+            node.set_attr('inv_table_size', min(2**inv_inp_bw, inv_table_size))
+
+        elif implemenation == 'latency':
+            input_bw = inp_layer.get_attr('result_t').precision.width
+            node.set_attr('exp_table_size', min(2**input_bw, exp_table_size))
+
+            accum_bw = node.get_attr('accum_t').precision.width
+            node.set_attr('inv_table_size', min(2**accum_bw, inv_table_size))
+
+        # One could try shrinking the size of legacy, but ignore it for now.
+        # Argmax doesn't use tables so the size is irrelevant in that case.
+
+        node.set_attr('table_sizes_checked', True)
+
+        return False
 
 
 def register_softmax__table_size_fix(backend):

hls4ml/converters/keras/core.py

@@ -82,8 +82,13 @@ def parse_activation_layer(keras_layer, input_names, input_shapes, data_reader):
             raise Exception('PReLU with shared_axes other than None is not supported in hsl4ml')
         layer['param_data'] = get_weights_data(data_reader, layer['name'], 'alpha')
 
-    if layer['class_name'] == 'Activation' and layer['activation'] == 'softmax':
+    if (layer['class_name'] == 'Activation' and layer['activation'] == 'softmax') or layer['class_name'] == 'Softmax':
         layer['class_name'] = 'Softmax'
+        ax = len(input_shapes[0]) - 1
+        n_outer: int = math.prod(input_shapes[0][1:ax])  # type: ignore
+        n_inner: int = math.prod(input_shapes[0][ax + 1 :])  # type: ignore
+        layer['n_outer'] = n_outer
+        layer['n_inner'] = n_inner
     if layer['class_name'] == 'Activation' and layer['activation'] == 'hard_sigmoid':
         layer['class_name'] = 'HardActivation'
     if layer['class_name'] == 'Softmax':

hls4ml/converters/keras_v3/core.py

@@ -66,7 +66,13 @@ def handle(
         match activation:
             case keras.activations.softmax:
                 class_name = 'Softmax'
+                ax = len(in_tensors[0].shape) - 1
+                n_outer: int = prod(in_tensors[0].shape[1:ax])  # type: ignore
+                n_inner: int = prod(in_tensors[0].shape[ax + 1 :])  # type: ignore
                 config['axis'] = -1
+                config['activation'] = 'softmax'
+                config['n_outer'] = n_outer
+                config['n_inner'] = n_inner
             case keras.activations.hard_sigmoid:
                 class_name = 'HardActivation'
             case keras.activations.leaky_relu:

hls4ml/model/layers.py

@@ -965,6 +965,9 @@ def initialize(self):
         if 'n_in' not in self.attributes:
             self.set_attr('n_in', self.get_input_variable().size())
 
+        # set the needed types if needed
+        self._set_type_t('table')
+
 
 class ParametrizedActivation(Activation):
     _expected_attributes = [
@@ -1031,6 +1034,12 @@ class Softmax(Activation):
     def initialize(self):
         super().initialize()
 
+        # set the needed types if needed
+        self._set_type_t('exp_table')
+        self._set_type_t('inv_table')
+        self._set_type_t('inv_inp')
+        self._set_type_t('inp_norm')
+
 
 class TernaryTanh(Activation):
     def initialize(self):

hls4ml/model/optimizer/passes/infer_precision.py

@@ -90,6 +90,9 @@ def _infer_precision(self, node, types_to_infer):
 
         if node_class in ['PReLU']:
             return self._infer_prelu_act_precision(node, types_to_infer)
+
+        if node_class in ['Softmax']:
+            return self._infer_softmax_precision(node, types_to_infer)
         # What about quantized activation layer? Setting it to 'auto' manually will break it here. We should prevent
         # this in config_from_* functions
 
@@ -605,6 +608,58 @@ def _infer_prelu_act_precision(self, node, types_to_infer):
 
         return inferred_types
 
+    def _infer_softmax_precision(self, node, types_to_infer):
+        inferred_types = []
+
+        if 'exp_table_t' in types_to_infer:
+            if node.get_attr('implementation') == 'stable':
+                # this is <= 1
+                prec = FixedPrecisionType(
+                    16, 1, signed=False, rounding_mode=RoundingMode.RND_CONV, saturation_mode=SaturationMode.SAT
+                )
+            else:
+                prec = FixedPrecisionType(
+                    18, 8, signed=False, rounding_mode=RoundingMode.RND_CONV, saturation_mode=SaturationMode.SAT
+                )
+            node.types['exp_table_t'].precision = prec
+            inferred_types.append('exp_table_t')
+
+        if 'inv_table_t' in types_to_infer:
+            # this is <= 1
+            prec = FixedPrecisionType(
+                16, 1, signed=False, rounding_mode=RoundingMode.RND_CONV, saturation_mode=SaturationMode.SAT
+            )
+            node.types['inv_table_t'].precision = prec
+            inferred_types.append('inv_table_t')
+
+        if 'accum_t' in types_to_infer:
+            exp_w = node.types['exp_table_t'].precision.width
+            exp_i = node.types['exp_table_t'].precision.integer
+            exp_s = node.types['exp_table_t'].precision.signed
+            n_slice = node.get_attr('n_in') // node.get_attr('n_inner') // node.get_attr('n_outer')
+            ceillog = math.ceil(np.log2(n_slice))
+            node.types['accum_t'].precision = FixedPrecisionType(exp_w + ceillog, exp_i + ceillog, signed=exp_s)
+            inferred_types.append('accum_t')
+
+        if 'inv_inp_t' in types_to_infer:
+            # if not set, just choose the accumulator type
+            node.types['inv_inp_t'].precision = node.types['accum_t'].precision
+            inferred_types.append('inv_inp_t')
+
+        if 'inp_norm_t' in types_to_infer:
+            in_type = node.get_input_variable().type.precision
+            inp_norm_width = in_type.width - in_type.signed
+            inp_norm_int = in_type.integer - in_type.signed
+            node.types['inp_norm_t'].precision = FixedPrecisionType(inp_norm_width, inp_norm_int, signed=False)
+            inferred_types.append('inp_norm_t')
+
+        if 'result_t' in types_to_infer:
+            # if not set, just choose the inv_table_t type
+            node.types['result_t'].precision = node.types['inv_table_t'].precision
+            inferred_types.append('result_t')
+
+        return inferred_types
+
 
 def _get_precision_from_constant(value: int | float, max_width=8):
     """A utility function to find a fixed type to store the constant

hls4ml/model/types.py

@@ -434,6 +434,9 @@ class UnspecifiedPrecisionType(PrecisionType):
     def __init__(self):
         super().__init__(width=0, signed=False)
 
+    def __str__(self):
+        return 'auto'
+
 
 def find_minimum_width(data, signed=True):
     """

hls4ml/templates/vivado/nnet_utils/nnet_activation.h

@@ -189,14 +189,13 @@ void softmax_latency(data_T data[CONFIG_T::n_slice], res_T res[CONFIG_T::n_slice
         // Note we are exponentiating the inputs, which have type data_T
         init_exp_table<data_T, CONFIG_T>(exp_table);
         // Note we are inverting the exponentials, which have type exp_table_t
-        init_invert_table<typename CONFIG_T::inv_inp_t, CONFIG_T>(invert_table);
+        init_invert_table<typename CONFIG_T::accum_t, CONFIG_T>(invert_table);
         initialized = true;
     }
 
     // Calculate all the e^x's
     typename CONFIG_T::accum_t exp_res[CONFIG_T::n_slice];
     #pragma HLS array_partition variable=exp_res complete
-    typename CONFIG_T::inv_inp_t exp_sum(0);
     for (unsigned i = 0; i < CONFIG_T::n_slice; i++) {
         #pragma HLS unroll
         unsigned x = softmax_idx_from_real_val<data_T, CONFIG_T::exp_table_size>(data[i]);
@@ -206,10 +205,11 @@ void softmax_latency(data_T data[CONFIG_T::n_slice], res_T res[CONFIG_T::n_slice
     // Explicitly sum the results with an adder tree.
     // Rounding & Saturation mode, which improve accuracy, prevent Vivado from expression balancing
     Op_add<typename CONFIG_T::accum_t> op_add;
-    exp_sum = reduce<typename CONFIG_T::accum_t, CONFIG_T::n_slice, Op_add<typename CONFIG_T::accum_t>>(exp_res, op_add);
+    typename CONFIG_T::accum_t exp_sum =
+        reduce<typename CONFIG_T::accum_t, CONFIG_T::n_slice, Op_add<typename CONFIG_T::accum_t>>(exp_res, op_add);
 
     typename CONFIG_T::inv_table_t inv_exp_sum =
-        invert_table[softmax_idx_from_real_val<typename CONFIG_T::inv_inp_t, CONFIG_T::inv_table_size>(exp_sum)];
+        invert_table[softmax_idx_from_real_val<typename CONFIG_T::accum_t, CONFIG_T::inv_table_size>(exp_sum)];
     for (unsigned i = 0; i < CONFIG_T::n_slice; i++) {
         #pragma HLS unroll
         res[i] = exp_res[i] * inv_exp_sum;
@@ -251,7 +251,6 @@ void softmax_stable(data_T data[CONFIG_T::n_slice], res_T res[CONFIG_T::n_slice]
     // Calculate all the e^x's
     typename CONFIG_T::accum_t exp_res[CONFIG_T::n_slice];
     #pragma HLS array_partition variable=exp_res complete
-    typename CONFIG_T::inv_inp_t exp_sum(0);
     for (unsigned i = 0; i < CONFIG_T::n_slice; i++) {
         #pragma HLS unroll
         unsigned x = softmax_idx_from_real_val<typename CONFIG_T::inp_norm_t, CONFIG_T::exp_table_size>(d_xi_xmax[i]);
@@ -261,7 +260,8 @@ void softmax_stable(data_T data[CONFIG_T::n_slice], res_T res[CONFIG_T::n_slice]
     // Explicitly sum the results with an adder tree.
     // Rounding & Saturation mode, which improve accuracy, prevent Vivado from expression balancing
     Op_add<typename CONFIG_T::accum_t> op_add;
-    exp_sum = reduce<typename CONFIG_T::accum_t, CONFIG_T::n_slice, Op_add<typename CONFIG_T::accum_t>>(exp_res, op_add);
+    typename CONFIG_T::inv_inp_t exp_sum =
+        reduce<typename CONFIG_T::accum_t, CONFIG_T::n_slice, Op_add<typename CONFIG_T::accum_t>>(exp_res, op_add);
 
     typename CONFIG_T::inv_table_t inv_exp_sum =
         invert_table[softmax_idx_from_real_val<typename CONFIG_T::inv_inp_t, CONFIG_T::inv_table_size>(exp_sum)];
@@ -271,18 +271,18 @@ void softmax_stable(data_T data[CONFIG_T::n_slice], res_T res[CONFIG_T::n_slice]
     }
 }
 
-template <typename CONFIG_T, int N_TABLE> void init_exp_table_legacy(typename CONFIG_T::table_t table_out[N_TABLE]) {
+template <typename CONFIG_T, int N_TABLE> void init_exp_table_legacy(typename CONFIG_T::exp_table_t table_out[N_TABLE]) {
     for (int ii = 0; ii < N_TABLE; ii++) {
         // First, convert from table index to X-value (signed 8-bit, range -8 to +8)
         float in_val = 2 * 8.0 * (ii - float(N_TABLE) / 2.0) / float(N_TABLE);
         // Next, compute lookup table function
-        typename CONFIG_T::table_t real_val = exp_fcn_float(in_val);
+        typename CONFIG_T::exp_table_t real_val = exp_fcn_float(in_val);
         // std::cout << "Lookup table In Value: " << in_val << " Result: " << real_val << std::endl;
         table_out[ii] = real_val;
     }
 }
 
-template <typename CONFIG_T, int N_TABLE> void init_invert_table_legacy(typename CONFIG_T::table_t table_out[N_TABLE]) {
+template <typename CONFIG_T, int N_TABLE> void init_invert_table_legacy(typename CONFIG_T::inv_table_t table_out[N_TABLE]) {
     // Inversion function:
     //   result = 1/x
     for (int ii = 0; ii < N_TABLE; ii++) {
@@ -301,12 +301,12 @@ void softmax_legacy(data_T data[CONFIG_T::n_slice], res_T res[CONFIG_T::n_slice]
     // Initialize the lookup table
 #ifdef __HLS_SYN__
     bool initialized = false;
-    typename CONFIG_T::table_t exp_table[CONFIG_T::exp_table_size];
-    typename CONFIG_T::table_t invert_table[CONFIG_T::inv_table_size];
+    typename CONFIG_T::exp_table_t exp_table[CONFIG_T::exp_table_size];
+    typename CONFIG_T::inv_table_t invert_table[CONFIG_T::inv_table_size];
 #else
     static bool initialized = false;
-    static typename CONFIG_T::table_t exp_table[CONFIG_T::exp_table_size];
-    static typename CONFIG_T::table_t invert_table[CONFIG_T::inv_table_size];
+    static typename CONFIG_T::exp_table_t exp_table[CONFIG_T::exp_table_size];
+    static typename CONFIG_T::inv_table_t invert_table[CONFIG_T::inv_table_size];
 #endif
     if (!initialized) {
         init_exp_table_legacy<CONFIG_T, CONFIG_T::exp_table_size>(exp_table);
@@ -317,22 +317,23 @@ void softmax_legacy(data_T data[CONFIG_T::n_slice], res_T res[CONFIG_T::n_slice]
     #pragma HLS PIPELINE
 
     // Index into the lookup table based on data for exponentials
-    typename CONFIG_T::table_t exp_res[CONFIG_T::n_slice]; // different, independent, fixed point precision
-    typename CONFIG_T::table_t exp_diff_res;               // different, independent, fixed point precision
+    typename CONFIG_T::accum_t exp_res[CONFIG_T::n_slice]; // different, independent, fixed point precision
+    typename CONFIG_T::exp_table_t exp_diff_res;           // different, independent, fixed point precision
     data_T data_cache[CONFIG_T::n_slice];
-    int data_round;
     int index;
+
     for (int ii = 0; ii < CONFIG_T::n_slice; ii++) {
         data_cache[ii] = data[ii];
         exp_res[ii] = 0;
     }
 
+    // first calculate 1/softmax as a sum over fractions.
     for (int ii = 0; ii < CONFIG_T::n_slice; ii++) {
         for (int jj = 0; jj < CONFIG_T::n_slice; jj++) {
             if (ii == jj)
                 exp_diff_res = 1;
             else {
-                data_round = (data_cache[jj] - data_cache[ii]) * CONFIG_T::exp_table_size / 16;
+                auto data_round = (data_cache[jj] - data_cache[ii]) * CONFIG_T::exp_table_size / 16;
                 index = data_round + 8 * CONFIG_T::exp_table_size / 16;
                 if (index < 0)
                     index = 0;
@@ -352,7 +353,7 @@ void softmax_legacy(data_T data[CONFIG_T::n_slice], res_T res[CONFIG_T::n_slice]
         if (exp_res_index > CONFIG_T::inv_table_size - 1)
             exp_res_index = CONFIG_T::inv_table_size - 1;
         // typename CONFIG_T::table_t exp_res_invert = invert_table[exp_res_index];
-        res[ii] = (res_T)invert_table[exp_res_index];
+        res[ii] = static_cast<res_T>(invert_table[exp_res_index]);
     }
 }