Softmax update

hls4ml/backends/fpga/fpga_backend.py

@@ -129,22 +129,33 @@ 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=FixedPrecisionType(
+                    18, 8, signed=False, rounding_mode=RoundingMode.RND, saturation_mode=SaturationMode.SAT
+                ),
                 description=descriptions.table_type,
             ),
             TypeAttribute(
                 'inv_table',
-                default=FixedPrecisionType(18, 8, rounding_mode=RoundingMode.RND, saturation_mode=SaturationMode.SAT),
+                default=FixedPrecisionType(
+                    18, 8, signed=False, rounding_mode=RoundingMode.RND, saturation_mode=SaturationMode.SAT
+                ),
                 description=descriptions.table_type,
             ),
             TypeAttribute(
                 'inv_inp',
-                default=FixedPrecisionType(18, 8, rounding_mode=RoundingMode.RND, saturation_mode=SaturationMode.SAT),
+                default=FixedPrecisionType(
+                    18, 8, signed=False, rounding_mode=RoundingMode.RND, saturation_mode=SaturationMode.SAT
+                ),
+                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=FixedPrecisionType(
+                    18, 8, signed=False, rounding_mode=RoundingMode.RND, saturation_mode=SaturationMode.SAT
+                ),
+                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/oneapi/oneapi_backend.py

@@ -19,7 +19,6 @@
     Embedding,
     Layer,
     SimpleRNN,
-    Softmax,
 )
 from hls4ml.model.optimizer import get_backend_passes, layer_optimizer
 from hls4ml.model.types import FixedPrecisionType, IntegerPrecisionType, NamedType
@@ -210,7 +209,11 @@ def build(self, model, build_type='fpga_emu', run=False):
         try:
             subprocess.run('which icpx', shell=True, cwd=builddir, check=True)
         except subprocess.CalledProcessError:
-            raise RuntimeError('Could not find icpx. Please configure oneAPI appropriately')
+            print('icpx not found. Trying ahls')
+            try:
+                subprocess.run('which ahls', shell=True, cwd=builddir, check=True)
+            except subprocess.CalledProcessError:
+                raise RuntimeError('Could not find icpx or ahls. Please configure oneAPI appropriately')
         subprocess.run('cmake ..', shell=True, cwd=builddir, check=True)
         subprocess.run(f'make {build_type}', shell=True, cwd=builddir, check=True)
 
@@ -257,13 +260,6 @@ def init_activation(self, layer):
         if layer.get_attr('recurrent_activation') == 'tanh':
             layer.set_attr('recurrent_activation', 'dense_tanh')
 
-    @layer_optimizer(Softmax)
-    def init_softmax(self, layer):
-        if layer.model.config.get_config_value('IOType') == 'io_parallel':
-            assert len(layer.get_input_variable().shape) == 1, (
-                'Softmax with io_parallel strategy cannot be used on multidimensional tensors.'
-            )
-
     @layer_optimizer(Embedding)
     def init_embed(self, layer):
         if layer.attributes['n_in'] is None:

hls4ml/backends/oneapi/passes/core_templates.py

@@ -194,12 +194,42 @@ def format(self, node):
 
 softmax_config_template = """struct {type}_config{index} : nnet::activ_config {{
     static constexpr unsigned n_in = {n_in};
+
+    // For multi-dim softmax
+    static const unsigned n_slice = {n_slice};
+    static const unsigned n_outer = {n_outer};
+    static const unsigned n_inner = {n_inner};
+
+    // For legacy softmax
+    typedef {table_t.name} table_t;
     static constexpr unsigned table_size = {table_size};
+
+    static constexpr unsigned exp_table_size = {exp_table_size};
+    static constexpr unsigned inv_table_size = {inv_table_size};
     static constexpr unsigned io_type = nnet::{iotype};
     static constexpr unsigned reuse_factor = {reuse};
+
     static constexpr nnet::softmax_implementation implementation = nnet::softmax_implementation::{implementation};
+    typedef {smax_accum_t} accum_t;
     typedef {exp_table_t.name} exp_table_t;
-    typedef {inv_table_t.name} inv_table_t;
+    typedef {inv_table_t.name} inv_table_t;"""
+
+softmax_config_table_template = """
+
+    using {exp_table_name}_arr_t = nnet::array<exp_table_t, exp_table_size>;
+    using {inv_table_name}_arr_t = nnet::array<inv_table_t, inv_table_size>;
+    static constexpr const {exp_table_name}_arr_t exp_table = {exp_table_name};
+    static constexpr const {inv_table_name}_arr_t invert_table = {inv_table_name};
+}};\n"""
+
+softmax_config_table_template_stable = """
+    typedef {inv_inp_t.name} inv_inp_t;
+    typedef {inp_norm_t.name} inp_norm_t;
+
+    using {exp_table_name}_arr_t = nnet::array<exp_table_t, exp_table_size>;
+    using {inv_table_name}_arr_t = nnet::array<inv_table_t, inv_table_size>;
+    static constexpr const {exp_table_name}_arr_t exp_table = {exp_table_name};
+    static constexpr const {inv_table_name}_arr_t invert_table = {inv_table_name};
 }};\n"""
 
 activ_function_template = 'nnet::{activation}<{input_t}, {output_t}, {config}>({input}, {output});'
@@ -221,6 +251,36 @@ def format(self, node):
         params = self._default_config_params(node)
         params['type'] = node.get_attr('activation')
 
+        if (params['type'] == 'softmax') or (params['type'] == 'softmax_multidim'):
+            params.setdefault('n_inner', 1)
+            params.setdefault('n_outer', 1)
+            n_slice = params['n_in'] // params['n_inner'] // params['n_outer']
+            params['n_slice'] = n_slice
+
+            params['exp_table_name'] = node.name + '_exp_table'
+            params['inv_table_name'] = node.name + '_inv_table'
+
+            if params['implementation'] == 'stable':
+                self.template = softmax_config_template + softmax_config_table_template_stable
+
+                if 'exp_table_size' not in params:
+                    params['exp_table_size'] = 2 ** params['inp_norm_t'].precision.width
+                    node.set_attr('exp_table_size', params['exp_table_size'])
+
+                if 'inv_table_size' not in params:
+                    params['inv_table_size'] = 2 ** params['inv_inp_t'].precision.width
+                    node.set_attr('inv_table_size', params['inv_table_size'])
+
+                params['smax_accum_t'] = params['accum_t'].name
+                # params.setdefault('table_size', params['exp_table_size'])  # Not sure if necessary
+
+            else:
+                # TODO: For latency check the table sizes correctly, match them
+                if 'exp_table_size' not in params:
+                    params['exp_table_size'] = params['table_size']
+                if 'inv_table_size' not in params:
+                    params['inv_table_size'] = params['table_size']
+
         return self.template.format(**params)
 
 
@@ -251,7 +311,7 @@ def format(self, node):
 class SoftmaxConfigTemplate(ActivationConfigTemplate):
     def __init__(self):
         super(ActivationConfigTemplate, self).__init__(Softmax)  # Skip ActivationConfigTemplate's __init__
-        self.template = softmax_config_template
+        self.template = softmax_config_template + softmax_config_table_template
 
 
 class ActivationFunctionTemplate(FunctionCallTemplate):
@@ -304,6 +364,7 @@ def format(self, node):
         params = self._default_function_params(node)
         params['activation'] = node.get_attr('activation').lower()
         params['config'] = f'{node.get_attr("activation")}_config{node.index}'
+
         return self.template.format(**params)
 
 

hls4ml/converters/keras_v3/hgq2/softmax.py

@@ -12,10 +12,16 @@
     from keras import KerasTensor
 
 
-def fixed_quantizer_to_hls4ml_t(q: 'FixedPointQuantizerBase', take_max=False):
+def fixed_quantizer_to_hls4ml_t(q: 'FixedPointQuantizerBase', take_max=False, force_unsigned=False):
     from keras import ops
 
     k, i, f = q.kif
+
+    if force_unsigned:
+        k = 0.0
+        i += 1.0
+        f += 1.0
+
     k = ops.convert_to_numpy(k)
     i = ops.convert_to_numpy(i)
     f = ops.convert_to_numpy(f)
@@ -78,10 +84,12 @@ def handle(
         exp_oq = layer.exp_table.oq.quantizer
         inv_oq = layer.inv_table.oq.quantizer
         inv_iq = layer.inv_table.iq.quantizer
+
         assert isinstance(exp_oq, FixedPointQuantizerBase), 'Only fixed-point quantizer is supported for exp_table'
-        exp_table_t = fixed_quantizer_to_hls4ml_t(exp_oq)
-        inv_table_t = fixed_quantizer_to_hls4ml_t(inv_oq)
-        inv_inp_t = fixed_quantizer_to_hls4ml_t(inv_iq)
+        # Enforce unsigned quantisers on all three types
+        exp_table_t = fixed_quantizer_to_hls4ml_t(exp_oq, force_unsigned=True)
+        inv_table_t = fixed_quantizer_to_hls4ml_t(inv_oq, force_unsigned=True)
+        inv_inp_t = fixed_quantizer_to_hls4ml_t(inv_iq, force_unsigned=True)
         exp_scale = layer.input_scaler
 
         inv_table_size = 2**inv_inp_t.width
@@ -99,13 +107,19 @@ def handle(
         else:
             raise ValueError(f'Too many inputs for softmax layer {layer.name}: expected 1 or 2, got {len(in_tensors)}')
 
+        # For masked implementation assume first input is the tensor we are operating on
+        activation = 'softmax'
+        if len(in_tensors[0].shape[1:]) > 1:
+            if (1 not in in_tensors[0].shape[1:]) or (len(in_tensors[0].shape[1:]) > 2):
+                activation = 'softmax_multidim'
+
         config = {}
         config.update(self.default_config)
         config.update(
             {
                 'axis': ax,
                 'n_in': n_in,
-                'activation': 'softmax',
+                'activation': activation,
                 'n_outer': n_outer,
                 'n_inner': n_inner,
                 'implementation': impl,
@@ -122,7 +136,8 @@ def handle(
         )
 
         if layer.stable:
-            inp_norm_t = fixed_quantizer_to_hls4ml_t(layer.exp_table.iq.quantizer)
+            # Force unsigned since norm >= 0
+            inp_norm_t = fixed_quantizer_to_hls4ml_t(layer.exp_table.iq.quantizer, force_unsigned=True)
             config['inp_norm_t'] = inp_norm_t
 
         return (config,)

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,26 @@ def _infer_prelu_act_precision(self, node, types_to_infer):
 
         return inferred_types
 
+    def _infer_softmax_precision(self, node, types_to_infer):
+        inferred_types = []
+
+        # for softmax, the table parameters have a default setting, so they don't need to be inferred
+        # here. We never expect them to be of type auto.
+
+        # For result, we leave it to be set externally (model default if not set). We expect it to
+        # likely be the output value, in which case the output format would determine it's precision.
+        # Therefore, only the accum is configured here
+
+        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
+            ceillog = math.ceil(np.log2(node.get_attr('n_in')))
+            node.types['accum_t'].precision = FixedPrecisionType(exp_w + ceillog, exp_i + ceillog, signed=exp_s)
+            inferred_types.append('accum_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