diff --git a/hls4ml/backends/fpga/passes/hgq_proxy_model.py b/hls4ml/backends/fpga/passes/hgq_proxy_model.py index a8135f655a..0e79df1c1a 100644 --- a/hls4ml/backends/fpga/passes/hgq_proxy_model.py +++ b/hls4ml/backends/fpga/passes/hgq_proxy_model.py @@ -34,7 +34,8 @@ def generate_mask_fn( masks = [] to_fixed = to_acfixed if backend.lower() in ['oneapi', 'quartus'] else to_apfixed for idx, (k, b, i) in enumerate(zip(Ks, Bs, Is)): - if b == 0: + if b <= 0: + # b <= 0 means nothing is representable: the channel is a constant zero fn = f'out[{idx}] = 0;' else: fn = f'out[{idx}] = {to_fixed(k, b, i, RND, SAT)}(inp[{idx}]);' diff --git a/hls4ml/backends/vivado/passes/core_templates.py b/hls4ml/backends/vivado/passes/core_templates.py index 777a435ab0..e275d752d0 100644 --- a/hls4ml/backends/vivado/passes/core_templates.py +++ b/hls4ml/backends/vivado/passes/core_templates.py @@ -250,6 +250,9 @@ def format(self, node): param_activ_function_template = ( 'nnet::{activation}<{input_t}, {param_t.name}, {output_t}, {config}>({input}, {param}, {output});' ) +softmax_lut_function_template = ( + 'nnet::{activation}<{input_t}, {output_t}, {config}>({input}, {output}, {exp_table}, {inv_table});' +) activ_include_list = ['nnet_utils/nnet_activation.h', 'nnet_utils/nnet_activation_stream.h'] @@ -349,6 +352,14 @@ def format(self, node): params['activation'] = 'softmax' if not use_multidim else 'softmax_multidim' params['config'] = '{}_config{}'.format(node.get_attr('activation'), node.index) + # HGQ2 softmax carries its trained tables as weights; every other softmax lets the + # HLS code build them at runtime. + if 'exp_table' in node.weights and 'inv_table' in node.weights: + params['activation'] += '_lut' + params['exp_table'] = node.get_weights('exp_table').name + params['inv_table'] = node.get_weights('inv_table').name + return softmax_lut_function_template.format(**params) + return self.template.format(**params) diff --git a/hls4ml/backends/vivado/passes/hgq_softmax_tables.py b/hls4ml/backends/vivado/passes/hgq_softmax_tables.py new file mode 100644 index 0000000000..d5cf193f2e --- /dev/null +++ b/hls4ml/backends/vivado/passes/hgq_softmax_tables.py @@ -0,0 +1,68 @@ +from warnings import warn + +from hls4ml.model.layers import Layer, Softmax +from hls4ml.model.optimizer import OptimizerPass +from hls4ml.model.types import FixedPrecisionType + + +class MaterializeSoftmaxTables(OptimizerPass): + """Turn HGQ's trained softmax exp/reciprocal lookup tables into weight arrays. + + Without this the generated C++ rebuilds both tables at runtime from std::exp and 1/x + (nnet::init_exp_table / nnet::init_invert_table), which does not reproduce HGQ's own + quantized values. QSoftmaxHandler stashes the builders as _exp_table_fn / _inv_table_fn; + both take the (k, i, f) of the type the C++ addresses the table with, which is known only + here: for the latency implementation it is the softmax input precision, decided by + bit_exact. This pass therefore runs after bit_exact and before transform_types. + """ + + def match(self, node: Layer): + # Keyed on the stash rather than on the layer type: only the HGQ frontend leaves it, + # and it must always be removed again, being a callable no saved model can hold. + return isinstance(node, Softmax) and '_exp_table_fn' in node.attributes + + def transform(self, model, node: Layer): + exp_table_fn = node.attributes.pop('_exp_table_fn') + inv_table_fn = node.attributes.pop('_inv_table_fn') + + impl = node.get_attr('implementation') + if impl not in ('latency', 'stable'): + return False # argmax and legacy use no lookup table + + # The exp table is indexed by the normalized input (x_max - x) for the stable + # implementation and by the layer input itself for the latency one. Reading the + # element precision off the variable is only unambiguous because this runs before + # transform_types wraps io_stream variables. + if impl == 'stable': + exp_inp_t: FixedPrecisionType = node.attributes['inp_norm_t'].precision + else: + exp_inp_t = node.get_input_variable().type.precision + inv_inp_t: FixedPrecisionType = node.attributes['inv_inp_t'].precision + + for name, addr_t, fn in (('exp_table', exp_inp_t, exp_table_fn), ('inv_table', inv_inp_t, inv_table_fn)): + # softmax_idx_from_real_val slices the top ceillog2(table_size) bits of the address + # word. Sizing the table as 2**width makes that the whole word, so no low bits are + # dropped and the slice cannot run off the end of a narrower word. + size = 1 << int(addr_t.width) + if int(node.get_attr(f'{name}_size') or 0) != size: + warn( + f'{node.name}: {name}_size {node.get_attr(f"{name}_size")} does not match the {addr_t.width}-bit ' + f'address type {addr_t}; overriding to {size} to keep table indexing bit-exact.', + stacklevel=1, + ) + node.set_attr(f'{name}_size', size) + k = int(bool(addr_t.signed)) + data = fn((k, addr_t.integer - k, int(addr_t.width) - addr_t.integer)) + + # Capture the type bit_exact derived first: storing a WeightVariable under `name` + # makes AttributeDict overwrite `{name}_t` with the variable's own type. + named_t = node.attributes[f'{name}_t'] + node.set_attr(f'{name}_data', data) + node.add_weights_variable( + name=name, var_name=name + '{index}', precision=named_t.precision, type_name=named_t.name + ) + node.get_weights(name).type = named_t + node.attributes[f'{name}_t'] = named_t + model.config.layer_name_precision[f'{node.name}_{name}'] = str(named_t.precision) + + return False diff --git a/hls4ml/backends/vivado/vivado_backend.py b/hls4ml/backends/vivado/vivado_backend.py index 5014f6836f..cc80d29b79 100644 --- a/hls4ml/backends/vivado/vivado_backend.py +++ b/hls4ml/backends/vivado/vivado_backend.py @@ -171,6 +171,7 @@ def _register_flows(self): 'vivado:inplace_stream_flatten', 'vivado:skip_softmax', 'vivado:fix_softmax_table_size', + 'vivado:materialize_softmax_tables', 'infer_precision_types', 'vivado:distributed_arithmetic_codegen', 'vivado:distributed_arithmetic_einsum_codegen', diff --git a/hls4ml/converters/keras_v3/hgq2/_base.py b/hls4ml/converters/keras_v3/hgq2/_base.py index b6f705e27f..de1f20b5af 100644 --- a/hls4ml/converters/keras_v3/hgq2/_base.py +++ b/hls4ml/converters/keras_v3/hgq2/_base.py @@ -1,6 +1,7 @@ from collections.abc import Sequence from math import prod from typing import TYPE_CHECKING, Any +from warnings import warn import numpy as np @@ -39,6 +40,17 @@ def extract_fixed_quantizer_config(q, tensor: 'KerasTensor', is_input: bool) -> B = np.ravel(B).astype(np.int16) I = np.ravel(I).astype(np.int16) # noqa: E741 + # A channel can be trained down to a negative total width. Clamp to 0, not 1: B == 0 is the + # constant-zero encoding generate_mask_fn understands, while B == 1 would bring the channel + # back as a live 1-bit one. I is legitimately negative and already zeroed for these above. + if np.any(B < 0): + warn( + f'Quantizer {q.name} has {int(np.sum(B < 0))} channel(s) with total bitwidth < 0; ' + 'treating them as constant zero.', + stacklevel=2, + ) + B = np.maximum(B, 0) + overflow_mode: str = internal_q.overflow_mode round_mode: str = internal_q.round_mode if round_mode.startswith('S_'): diff --git a/hls4ml/converters/keras_v3/hgq2/softmax.py b/hls4ml/converters/keras_v3/hgq2/softmax.py index 136662beba..60b4cae49a 100644 --- a/hls4ml/converters/keras_v3/hgq2/softmax.py +++ b/hls4ml/converters/keras_v3/hgq2/softmax.py @@ -1,10 +1,12 @@ import typing from collections.abc import Sequence +from functools import partial from math import prod from hls4ml.model.types import FixedPrecisionType, RoundingMode, SaturationMode from ._base import QLayerHandler +from .unary_lut import extract_lut_table if typing.TYPE_CHECKING: import hgq @@ -118,6 +120,16 @@ def handle( 'parallelization_factor': parallelization_factor, 'class_name': class_name, '_bit_exact': True, + # Not materialized here: for the latency implementation the domain is the + # softmax input precision, only final after bit_exact. materialize_softmax_tables + # calls these with the (k, i, f) of the address type, and pops them again -- they + # are not JSON-serializable, so they must not reach a saved model. + '_exp_table_fn': partial( + extract_lut_table, layer.exp_table.activation, layer.exp_table.oq if layer.exp_table.enable_oq else None + ), + '_inv_table_fn': partial( + extract_lut_table, layer.inv_table.activation, layer.inv_table.oq if layer.inv_table.enable_oq else None + ), } ) diff --git a/hls4ml/converters/keras_v3/hgq2/unary_lut.py b/hls4ml/converters/keras_v3/hgq2/unary_lut.py index e1b28c42e5..64f6cff122 100644 --- a/hls4ml/converters/keras_v3/hgq2/unary_lut.py +++ b/hls4ml/converters/keras_v3/hgq2/unary_lut.py @@ -1,5 +1,5 @@ import typing -from collections.abc import Sequence +from collections.abc import Callable, Sequence import numpy as np from quantizers import get_fixed_quantizer_np @@ -10,11 +10,76 @@ if typing.TYPE_CHECKING: import hgq + from hgq.quantizer import Quantizer + from hgq.quantizer.internal import FixedPointQuantizerBase from keras import KerasTensor from decimal import Decimal +def fixed_grid_by_bit_pattern(k: int, i: int, f: int) -> np.ndarray: + """All values a fixed-point type can hold, indexed by their two's-complement bit pattern. + + This is the addressing convention of ``nnet::get_index_unary_lut`` and + ``nnet::softmax_real_val_from_idx`` in the generated C++. + """ + K, I, F = Decimal(int(k)), Decimal(int(i)), Decimal(int(f)) # noqa: E741 + _eps = Decimal(2) ** -F + _min = -K * Decimal(2) ** I + _max = Decimal(2) ** I - _eps + N = (_max - _min) / _eps + 1 + assert float(N).is_integer(), 'Invalid quantizer range' + N = int(N) + assert N <= 1e6, 'Too large quantizer range' + assert np.log2(N).is_integer(), f'Invalid quantizer range: N must be power of 2, got {N}' + + grid = np.linspace(float(_min), float(_max), N, dtype=np.float32) + if k: + # idx by binary repr, move the positive part to the front + grid = np.concatenate([grid[N // 2 :], grid[: N // 2]]) + return grid + + +def kif_of(q: 'FixedPointQuantizerBase') -> tuple[int, int, int]: + """(k, i, f) of a homogeneous quantizer, with entries representing nothing masked out.""" + from keras import ops + + k, i, f = q.kif + mask = k + i + f > 0 + i, f = np.where(mask, i, -32), np.where(mask, f, -32) # type: ignore + return int(ops.max(k)), int(ops.max(i)), int(ops.max(f)) # type: ignore + + +def extract_lut_table(activation: Callable, oq: 'Quantizer|None', kif: tuple[int, int, int]) -> np.ndarray: + """Tabulate activation over the fixed-point type kif describes, quantized by oq. + + kif is the type the generated C++ addresses the table with, passed in rather than read off + a layer: QSoftmax.exp_table has no input quantizer when stable=False, and for the latency + softmax the domain is the softmax input precision, only final after the bit_exact pass. + """ + from hgq.quantizer.internal import FixedPointQuantizerBase + from keras import ops + + grid = fixed_grid_by_bit_pattern(*kif) + table = activation(grid) + + if oq is not None: + internal_q = oq.quantizer + if not isinstance(internal_q, FixedPointQuantizerBase): + raise NotImplementedError('FloatPointQuantizer is not supported yet') + + # Not oq(table): the Quantizer layer broadcasts against the shape it was built for, + # which is not the rank-1 grid here. Homogeneous, so this is the same operation. + round_mode = internal_q.round_mode + if round_mode.startswith('S_'): + round_mode = round_mode[2:] + fixed_q = get_fixed_quantizer_np(round_mode, internal_q.overflow_mode) + k, i, f = (ops.convert_to_numpy(x).ravel().item() for x in internal_q.kif) + table = fixed_q(table, k, i, f) # type: ignore + + return np.asarray(ops.convert_to_numpy(table)) + + class QUnaryLUTHandler(QLayerHandler, KerasV3LayerHandler): handles = ('hgq.layers.activation.QUnaryFunctionLUT',) @@ -24,7 +89,7 @@ def handle( in_tensors: Sequence['KerasTensor'], out_tensors: Sequence['KerasTensor'], ): - from hgq.quantizer.internal import FixedPointQuantizerBase, FloatPointQuantizer + from hgq.quantizer.internal import FixedPointQuantizerBase from keras import ops if not layer.enable_iq and not layer.enable_oq: @@ -32,53 +97,20 @@ def handle( assert not layer._allow_heterogeneous_table, 'Heterogeneous table is not supported in QUnaryFunctionLUT layer' iq = layer.iq.quantizer - if isinstance(iq, FixedPointQuantizerBase): - k, i, f = iq.kif - mask = k + i + f > 0 - i, f = np.where(mask, i, -32), np.where(mask, f, -32) # type: ignore - k, i, f = (Decimal(int(ops.max(x))) for x in (k, i, f)) # type: ignore - _min = -k * 2**i - _eps = 2**-f - _max = 2**i - _eps - N = (_max - _min) / _eps + 1 - assert float(N).is_integer(), 'Invalid quantizer range' - N = int(N) - assert N <= 1e6, 'Too large quantizer range' - assert np.log2(N).is_integer(), f'Invalid quantizer range: N must be power of 2, got {N}' - - all_inputs = np.linspace(float(_min), float(_max), N, dtype=np.float32) - - config = {} - config.update(self.default_config) - table = layer.activation(all_inputs) - if layer.enable_oq: - table = layer.oq(table[None, ...])[0] - table = ops.convert_to_numpy(table) - if k: - # idx by binary repr, move the positive part to the front - table_pos, table_neg = table[N // 2 :], table[: N // 2] - table = np.concatenate([table_pos, table_neg]) - else: + if not isinstance(iq, FixedPointQuantizerBase): raise NotImplementedError('FloatPointQuantizer is not supported yet') + table = extract_lut_table(layer.activation, layer.oq if layer.enable_oq else None, kif_of(iq)) + oq = layer.oq.quantizer - if isinstance(oq, FixedPointQuantizerBase): - round_mode = oq.round_mode - if round_mode.startswith('S_'): - round_mode = round_mode[2:] - overflow_mode = oq.overflow_mode - fixed_q = get_fixed_quantizer_np(round_mode, overflow_mode) - k, i, f = (ops.convert_to_numpy(x).ravel().item() for x in oq.kif) - table = fixed_q(table, k, i, f) # type: ignore - - k, b, I = bool(k), k + i + f, k + i # noqa: E741 - table_t = FixedPrecisionType(b, I, k) - else: - assert isinstance(oq, FloatPointQuantizer) + if not isinstance(oq, FixedPointQuantizerBase): raise NotImplementedError('FloatPointQuantizer is not supported yet') + k, i, f = (ops.convert_to_numpy(x).ravel().item() for x in oq.kif) + k, b, I = bool(k), k + i + f, k + i # noqa: E741 + table_t = FixedPrecisionType(b, I, k) - table = ops.convert_to_numpy(table) - + config = {} + config.update(self.default_config) config.update( { 'class_name': 'UnaryLUT', diff --git a/hls4ml/templates/vivado/nnet_utils/nnet_activation.h b/hls4ml/templates/vivado/nnet_utils/nnet_activation.h index ac85e0b2cc..f36a8a8a20 100644 --- a/hls4ml/templates/vivado/nnet_utils/nnet_activation.h +++ b/hls4ml/templates/vivado/nnet_utils/nnet_activation.h @@ -172,27 +172,11 @@ void init_invert_table(typename CONFIG_T::inv_table_t table_out[CONFIG_T::inv_ta } template -void softmax_latency(data_T data[CONFIG_T::n_slice], res_T res[CONFIG_T::n_slice]) { - #pragma HLS pipeline - // Initialize the lookup tables -#ifdef __HLS_SYN__ - bool initialized = false; - 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::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) { - // Note we are exponentiating the inputs, which have type data_T - init_exp_table(exp_table); - // Note we are inverting the exponentials, which have type exp_table_t - init_invert_table(invert_table); - initialized = true; - } - +void softmax_latency_impl(data_T data[CONFIG_T::n_slice], res_T res[CONFIG_T::n_slice], + 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]) { + // Inlined so that the caller's pipeline pragma covers this body + #pragma HLS inline // 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 @@ -217,27 +201,11 @@ void softmax_latency(data_T data[CONFIG_T::n_slice], res_T res[CONFIG_T::n_slice } template -void softmax_stable(data_T data[CONFIG_T::n_slice], res_T res[CONFIG_T::n_slice]) { - #pragma HLS pipeline - // Initialize the lookup tables -#ifdef __HLS_SYN__ - bool initialized = false; - 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::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) { - // Note we are exponentiating the inputs, which have type data_T - init_exp_table(exp_table, true); - // Note we are inverting the exponentials, which have type exp_table_t - init_invert_table(invert_table); - initialized = true; - } - +void softmax_stable_impl(data_T data[CONFIG_T::n_slice], res_T res[CONFIG_T::n_slice], + 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]) { + // Inlined so that the caller's pipeline pragma covers this body + #pragma HLS inline // Find the max and compute all delta(x_i, x_max) Op_max op_max; data_T x_max = reduce>(data, op_max); @@ -271,6 +239,81 @@ void softmax_stable(data_T data[CONFIG_T::n_slice], res_T res[CONFIG_T::n_slice] } } +// Entry points building the tables at runtime. Used by every softmax whose frontend does +// not supply pre-computed tables (plain Keras, QKeras, ...). +template +void softmax_latency(data_T data[CONFIG_T::n_slice], res_T res[CONFIG_T::n_slice]) { + #pragma HLS pipeline + // Initialize the lookup tables +#ifdef __HLS_SYN__ + bool initialized = false; + 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::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) { + // Note we are exponentiating the inputs, which have type data_T + init_exp_table(exp_table); + // Note we are inverting the exponentials, which have type exp_table_t + init_invert_table(invert_table); + initialized = true; + } + + softmax_latency_impl(data, res, exp_table, invert_table); +} + +template +void softmax_stable(data_T data[CONFIG_T::n_slice], res_T res[CONFIG_T::n_slice]) { + #pragma HLS pipeline + // Initialize the lookup tables +#ifdef __HLS_SYN__ + bool initialized = false; + 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::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) { + // Note we are exponentiating the inputs, which have type data_T + init_exp_table(exp_table, true); + // Note we are inverting the exponentials, which have type exp_table_t + init_invert_table(invert_table); + initialized = true; + } + + softmax_stable_impl(data, res, exp_table, invert_table); +} + +// Entry points taking pre-computed tables. Used when the frontend knows the exact table +// contents (e.g. HGQ2's trained QSoftmax), so that CONFIG_T::exp_scale and the generic +// exp()/1-over-x reconstruction are bypassed entirely. +template +void softmax_latency(data_T data[CONFIG_T::n_slice], res_T res[CONFIG_T::n_slice], + 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]) { + #pragma HLS pipeline + #pragma HLS function_instantiate variable=exp_table + #pragma HLS function_instantiate variable=invert_table + softmax_latency_impl(data, res, exp_table, invert_table); +} + +template +void softmax_stable(data_T data[CONFIG_T::n_slice], res_T res[CONFIG_T::n_slice], + 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]) { + #pragma HLS pipeline + #pragma HLS function_instantiate variable=exp_table + #pragma HLS function_instantiate variable=invert_table + softmax_stable_impl(data, res, exp_table, invert_table); +} + template void init_exp_table_legacy(typename CONFIG_T::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) @@ -419,6 +462,56 @@ void softmax_multidim(data_T data[CONFIG_T::n_in], res_T res[CONFIG_T::n_in]) { } } +// Dispatchers for pre-computed tables. Separate names rather than overloads of softmax / +// softmax_multidim, so that the allocation pragma below still names a single function. +template +void softmax_lut(data_T data[CONFIG_T::n_slice], res_T res[CONFIG_T::n_slice], + 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]) { + #pragma HLS inline + switch (CONFIG_T::implementation) { + case softmax_implementation::latency: + softmax_latency(data, res, exp_table, invert_table); + break; + case softmax_implementation::stable: + softmax_stable(data, res, exp_table, invert_table); + break; + case softmax_implementation::legacy: + // legacy addresses CONFIG_T::table_t tables of its own; supplied tables are unused + softmax_legacy(data, res); + break; + case softmax_implementation::argmax: + // argmax needs no table at all + softmax_argmax(data, res); + break; + } +} + +template +void softmax_multidim_lut(data_T data[CONFIG_T::n_in], res_T res[CONFIG_T::n_in], + 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]) { + #pragma HLS inline + #pragma HLS allocation instances = softmax_lut limit = CONFIG_T::parallelization_factor function + data_T buffer_in[CONFIG_T::n_slice]; + res_T buffer_out[CONFIG_T::n_slice]; + for (signed i = 0; i < CONFIG_T::n_outer; i++) { + #pragma HLS UNROLL + for (signed k = 0; k < CONFIG_T::n_inner; k++) { + #pragma HLS UNROLL + for (signed j = 0; j < CONFIG_T::n_slice; j++) { + #pragma HLS UNROLL + buffer_in[j] = data[i * CONFIG_T::n_slice * CONFIG_T::n_inner + j * CONFIG_T::n_inner + k]; + } + softmax_lut(buffer_in, buffer_out, exp_table, invert_table); + for (signed j = 0; j < CONFIG_T::n_slice; j++) { + #pragma HLS UNROLL + res[i * CONFIG_T::n_slice * CONFIG_T::n_inner + j * CONFIG_T::n_inner + k] = buffer_out[j]; + } + } + } +} + // ************************************************* // TanH Activation // ************************************************* diff --git a/hls4ml/templates/vivado/nnet_utils/nnet_activation_stream.h b/hls4ml/templates/vivado/nnet_utils/nnet_activation_stream.h index 50c6c4068c..87b2177ae4 100644 --- a/hls4ml/templates/vivado/nnet_utils/nnet_activation_stream.h +++ b/hls4ml/templates/vivado/nnet_utils/nnet_activation_stream.h @@ -105,26 +105,9 @@ template void sigmoid(hls::stream // ************************************************* template -void softmax_latency(hls::stream &data, hls::stream &res) { - // Initialize the lookup tables -#ifdef __HLS_SYN__ - bool initialized = false; - 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::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) { - // Note we are exponentiating the inputs, which have type data_T - init_exp_table(exp_table); - // Note we are inverting the exponentials, which have type exp_table_t - init_invert_table(invert_table); - initialized = true; - } - +void softmax_latency_impl(hls::stream &data, hls::stream &res, + 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]) { constexpr unsigned multiplier_limit = DIV_ROUNDUP(data_T::size, CONFIG_T::reuse_factor); constexpr unsigned ii = data_T::size / multiplier_limit; @@ -166,26 +149,9 @@ void softmax_latency(hls::stream &data, hls::stream &res) { } template -void softmax_stable(hls::stream &data, hls::stream &res) { - // Initialize the lookup tables -#ifdef __HLS_SYN__ - bool initialized = false; - 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::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) { - // Note we are exponentiating the inputs, which have type data_T - init_exp_table(exp_table, true); - // Note we are inverting the exponentials, which have type exp_table_t - init_invert_table(invert_table); - initialized = true; - } - +void softmax_stable_impl(hls::stream &data, hls::stream &res, + 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]) { constexpr unsigned multiplier_limit = DIV_ROUNDUP(data_T::size, CONFIG_T::reuse_factor); constexpr unsigned ii = data_T::size / multiplier_limit; @@ -244,6 +210,75 @@ void softmax_stable(hls::stream &data, hls::stream &res) { } } +// Entry points building the tables at runtime. Used by every softmax whose frontend does +// not supply pre-computed tables (plain Keras, QKeras, ...). +template +void softmax_latency(hls::stream &data, hls::stream &res) { + // Initialize the lookup tables +#ifdef __HLS_SYN__ + bool initialized = false; + 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::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) { + // Note we are exponentiating the inputs, which have type data_T + init_exp_table(exp_table); + // Note we are inverting the exponentials, which have type exp_table_t + init_invert_table(invert_table); + initialized = true; + } + + softmax_latency_impl(data, res, exp_table, invert_table); +} + +template +void softmax_stable(hls::stream &data, hls::stream &res) { + // Initialize the lookup tables +#ifdef __HLS_SYN__ + bool initialized = false; + 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::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) { + // Note we are exponentiating the inputs, which have type data_T + init_exp_table(exp_table, true); + // Note we are inverting the exponentials, which have type exp_table_t + init_invert_table(invert_table); + initialized = true; + } + + softmax_stable_impl(data, res, exp_table, invert_table); +} + +// Entry points taking pre-computed tables (e.g. HGQ2's trained QSoftmax). +template +void softmax_latency(hls::stream &data, hls::stream &res, + 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]) { + #pragma HLS function_instantiate variable=exp_table + #pragma HLS function_instantiate variable=invert_table + softmax_latency_impl(data, res, exp_table, invert_table); +} + +template +void softmax_stable(hls::stream &data, hls::stream &res, + 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]) { + #pragma HLS function_instantiate variable=exp_table + #pragma HLS function_instantiate variable=invert_table + softmax_stable_impl(data, res, exp_table, invert_table); +} + template void softmax_legacy(hls::stream &data, hls::stream &res) { // Initialize the lookup table @@ -367,6 +402,30 @@ template void softmax(hls::stream } } +template +void softmax_lut(hls::stream &data, hls::stream &res, + 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]) { + assert(CONFIG_T::axis == -1); + + switch (CONFIG_T::implementation) { + case softmax_implementation::latency: + softmax_latency(data, res, exp_table, invert_table); + break; + case softmax_implementation::stable: + softmax_stable(data, res, exp_table, invert_table); + break; + case softmax_implementation::legacy: + // legacy addresses CONFIG_T::table_t tables of its own; supplied tables are unused + softmax_legacy(data, res); + break; + case softmax_implementation::argmax: + // argmax needs no table at all + softmax_argmax(data, res); + break; + } +} + // ************************************************* // TanH Activation // ************************************************* diff --git a/test/pytest/generate_ci_yaml.py b/test/pytest/generate_ci_yaml.py index 91c1c730a0..c17bc42e2b 100644 --- a/test/pytest/generate_ci_yaml.py +++ b/test/pytest/generate_ci_yaml.py @@ -38,6 +38,7 @@ KERAS3_LIST = { 'test_keras_v3_api', 'test_hgq2_mha', + 'test_hgq2_softmax', 'test_einsum_dense', 'test_qeinsum', 'test_multiout_onnx', diff --git a/test/pytest/test_hgq2_softmax.py b/test/pytest/test_hgq2_softmax.py new file mode 100644 index 0000000000..1f438dffc9 --- /dev/null +++ b/test/pytest/test_hgq2_softmax.py @@ -0,0 +1,157 @@ +"""Regression tests for fastmachinelearning/hls4ml#1523. + +``hgq.layers.QSoftmax`` trains two lookup tables (``exp_table`` and ``inv_table``). hls4ml +used to propagate only their sizes and types, leaving the generated C++ to rebuild the +contents from ``std::exp`` / ``1/x`` at runtime. These tests pin down that the trained +tables now reach the generated code, and that plain Keras softmax is unaffected. +""" + +from contextlib import nullcontext +from pathlib import Path + +import keras +import numpy as np +import pytest + +hgq = pytest.importorskip('hgq') + +from hgq.config import QuantizerConfigScope # noqa: E402 +from hgq.layers import QDense, QSoftmax # noqa: E402 +from hgq.utils import trace_minmax # noqa: E402 + +import hls4ml # noqa: E402 +from hls4ml.backends.fpga.passes.hgq_proxy_model import generate_mask_fn # noqa: E402 +from hls4ml.converters.keras_v3.hgq2._base import extract_fixed_quantizer_config # noqa: E402 +from hls4ml.converters.keras_v3.hgq2.unary_lut import extract_lut_table, fixed_grid_by_bit_pattern # noqa: E402 + +test_root_path = Path(__file__).parent + + +def _build_model(stable, shape=(16,), io_type='io_parallel', n_out=8, seed=42): + keras.utils.set_random_seed(seed) + # Heterogeneous activation quantization is io_parallel-only in hls4ml, so the datalane + # quantizers have to be pinned to a single bitwidth to reach the io_stream kernels. + scope = QuantizerConfigScope(place='datalane', heterogeneous_axis=()) if io_type == 'io_stream' else nullcontext() + with scope: + inp = keras.Input(shape) + x = QDense(n_out)(inp) + out = QSoftmax(axis=-1, stable=stable, name='sm')(x) + model = keras.Model(inp, out) + + rng = np.random.default_rng(seed) + X = rng.uniform(-2.0, 2.0, size=(200,) + shape).astype(np.float32) + trace_minmax(model, X) + return model, X + + +@pytest.mark.parametrize('backend', ['Vivado', 'Vitis']) +@pytest.mark.parametrize('io_type', ['io_parallel', 'io_stream']) +@pytest.mark.parametrize('stable', [True, False]) +@pytest.mark.parametrize('shape', [(16,), (4, 16)], ids=['1d', 'multidim']) +def test_hgq2_softmax_uses_trained_tables(test_case_id, backend, io_type, stable, shape): + model, X = _build_model(stable, shape=shape, io_type=io_type) + r_keras = np.asarray(model(X)) + + impl = 'stable' if stable else 'latency' + odir = str(test_root_path / test_case_id) + hls_model = hls4ml.converters.convert_from_keras_model( + model, backend=backend, io_type=io_type, output_dir=odir, part='xcvu13p-flga2577-2-e' + ) + + node = hls_model.graph['sm'] + assert node.get_attr('implementation') == impl + + # The trained tables are carried as weights ... + assert 'exp_table' in node.weights and 'inv_table' in node.weights + exp_var, inv_var = node.get_weights('exp_table'), node.get_weights('inv_table') + + # ... sized so that the C++ table index is a plain reinterpretation of the address word + if stable: + exp_addr_t = node.attributes['inp_norm_t'].precision + else: + exp_addr_t = node.get_input_variable().type.precision + inv_addr_t = node.attributes['inv_inp_t'].precision + assert len(exp_var.data) == node.get_attr('exp_table_size') == 2**exp_addr_t.width + assert len(inv_var.data) == node.get_attr('inv_table_size') == 2**inv_addr_t.width + + # ... typed with what bit_exact derived, not with a type re-inferred from the data + assert exp_var.type is node.attributes['exp_table_t'] + assert inv_var.type is node.attributes['inv_table_t'] + + # The non-serializable builders stashed by the converter must not survive + assert '_exp_table_fn' not in node.attributes and '_inv_table_fn' not in node.attributes + + # Table contents match an independent recomputation straight from hgq ... + sm_layer = model.get_layer('sm') + k = int(bool(exp_addr_t.signed)) + kif = (k, exp_addr_t.integer - k, exp_addr_t.width - exp_addr_t.integer) + expected = extract_lut_table(sm_layer.exp_table.activation, sm_layer.exp_table.oq, kif) + np.testing.assert_array_equal(np.asarray(exp_var.data).ravel(), np.asarray(expected).ravel()) + + # ... and are not simply what the generic C++ reconstruction would produce unquantized + naive = np.exp(fixed_grid_by_bit_pattern(*kif) * (-1.0 if stable else 1.0) * node.get_attr('exp_scale')) + assert not np.allclose(np.asarray(exp_var.data).ravel(), naive) + + # The generated code passes the tables in + hls_model.write() + call = [ln for ln in open(f'{odir}/firmware/myproject.cpp') if 'nnet::softmax' in ln] + assert len(call) == 1 + assert exp_var.name in call[0] and inv_var.name in call[0] + assert Path(f'{odir}/firmware/weights/{exp_var.name}.h').exists() + assert f'#include "weights/{exp_var.name}.h"' in open(f'{odir}/firmware/parameters.h').read() + + # End to end: HGQ2 + bit_exact promises exact agreement with Keras + hls_model.compile() + r_hls = np.asarray(hls_model.predict(X)).reshape(r_keras.shape) + assert np.std(r_hls) > 0 + np.testing.assert_array_equal(r_hls, r_keras) + + +@pytest.mark.parametrize('backend', ['Vivado', 'Vitis']) +def test_plain_keras_softmax_keeps_runtime_tables(test_case_id, backend): + """Negative control: non-HGQ softmax must keep the two-argument call.""" + keras.utils.set_random_seed(0) + inp = keras.Input((8,)) + out = keras.layers.Softmax(name='sm')(inp) + model = keras.Model(inp, out) + + odir = str(test_root_path / test_case_id) + cfg = hls4ml.utils.config_from_keras_model(model, granularity='name', backend=backend) + hls_model = hls4ml.converters.convert_from_keras_model( + model, hls_config=cfg, backend=backend, output_dir=odir, part='xcvu13p-flga2577-2-e' + ) + + node = hls_model.graph['sm'] + assert 'exp_table' not in node.weights and 'inv_table' not in node.weights + + hls_model.write() + (call,) = (ln for ln in open(f'{odir}/firmware/myproject.cpp') if 'nnet::softmax' in ln) + args = call.split('>(')[1].split(')')[0] + assert len(args.split(',')) == 2, call + + +def test_dead_channels_stay_zero(): + """A heterogeneous quantizer channel trained to a negative total width is a dead + channel and must be rendered as a constant zero, not resurrected as a 1-bit channel. + """ + from hgq.quantizer import Quantizer, QuantizerConfig + + q = Quantizer(QuantizerConfig('kif', 'weight', heterogeneous_axis=(0,))) + x = keras.ops.arange(8.0, dtype='float32') + q.build(x.shape) + q(x) + + i_var, f_var, _ = q.quantizer.weights + i_var.assign(np.array([1, 1, 1, -2, 1, 1, 1, 1], dtype=np.float32)) + f_var.assign(np.array([3, 3, 3, -4, 3, 3, 3, 3], dtype=np.float32)) + + inp = keras.Input(shape=(8,), name='x') + conf = extract_fixed_quantizer_config(q, inp, is_input=True) + k_arr, b_arr, i_arr = (np.ravel(a) for a in conf['mask_kbi']) + + assert b_arr[3] == 0, f'dead channel clamped to {b_arr[3]}, expected 0' + assert i_arr[3] == 0 + assert (b_arr >= 0).all() + + mask_fn = generate_mask_fn('mask', (8,), *(a.reshape(1, 8) for a in (k_arr, b_arr, i_arr)), 'RND', 'SAT', 'vivado') + assert 'out[3] = 0;' in mask_fn