Add LayerNorm support for Vivado (#1110)

* paser_mht

* change parser and modify keras_to_hls

* IR_mutihead_attention

* IR done

* create mha file in template

* mha .h file dummy algo

* config of mha

* update mha config

* dummy mha

* add transpose into mha

* projection_of_qkv_in_mha

* mha_first_draft

* able to predict model correct

* delete some unnassary comments

* delete comments

* resource strategy of transformer

* change sm lagacy

* update MHA, optimized

* support resource

* update

* dense_muti_dim_support

* parallel execute dense

* updates

* add_layerNorm_support

* MHA updated

* LayerNorm_bug_fix

* update bit precision

* config update

* add some comment

* run pre-commit

* Added support on QMultiHeadAttention, QLayerNormalization, and quantized_softmax

* updated on hls4ml transformer

* trying to clean the diff

* trying to clean the diff

* trying to clean the diff

* trying to clean the diff

* trying to clean the diff

* undo vhdl -> verilog change

* halfway working layernorm + test

* layernorm is now pretty functional

* layernorm on pytorch also

* minor cleanup

* more cleanup, pre-commit

* test for mha which kinda works maybe if you squint

* multihead attention working on keras and pytorch

* fiddly precision / accuracy changes for layernorm

* fix lookup table and label loops

* remove dense_seq

* undo qkeras changes

* fix merge conflict residue

* remove non-layernorm changes

* change to uniform LUT and fix precision

* [pre-commit.ci] auto fixes from pre-commit hooks

* fix encodings issue with dos2unix

* add Vitis as another tested backend

* Address PR feedback

* [pre-commit.ci] auto fixes from pre-commit hooks

* fix too-long lines

* fix merge issue

* trigger pre-commit

* re-add missing math import

* [pre-commit.ci] auto fixes from pre-commit hooks

* addressing Vladimir's latest comments

* change also pytorch test for layernorm and revert change to build command

* sideport changes to channels-last converter from 1352

---------

Co-authored-by: Ethan <[email protected]>
Co-authored-by: Jan-Frederik Schulte <[email protected]>
Co-authored-by: LostEcho365 <[email protected]>
Co-authored-by: Rian Flynn <[email protected]>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: 6 people

hls4ml/backends/fpga/fpga_backend.py

@@ -25,6 +25,7 @@
     GarNetStack,
     GlobalPooling1D,
     GlobalPooling2D,
+    LayerNormalization,
     MatMul,
     Merge,
     Pooling1D,
@@ -73,6 +74,7 @@ def __init__(self, name):
             Dot,
             Conv,
             MatMul,
+            LayerNormalization,
         ]
 
         for layer in accum_layers:

hls4ml/backends/vivado/passes/core_templates.py

@@ -2,7 +2,16 @@
 
 from hls4ml.backends.backend import get_backend
 from hls4ml.backends.template import FunctionCallTemplate, LayerConfigTemplate
-from hls4ml.model.layers import Activation, BatchNormalization, Dense, HardActivation, ParametrizedActivation, PReLU, Softmax
+from hls4ml.model.layers import (
+    Activation,
+    BatchNormalization,
+    Dense,
+    HardActivation,
+    LayerNormalization,
+    ParametrizedActivation,
+    PReLU,
+    Softmax,
+)
 from hls4ml.model.optimizer.passes.hgq_proxy_model import UnaryLUT
 
 # Dense templates
@@ -136,6 +145,58 @@ def format(self, node):
         return self.template.format(**params)
 
 
+# LayerNormalization templates
+
+layernorm_config_template = """struct config{index} : nnet::layernorm_config {{
+    static const unsigned n_in = {n_in};
+    static const unsigned seq_len = {seq_len};
+    static const unsigned axis = {axis};
+    static const unsigned epsilon_power_of_10 = {epsilon_power_of_10};
+    static const unsigned table_range_power2 = {table_range_power2};
+    static const unsigned table_size = {table_size};
+    typedef {accum_t.name} accum_t;
+    typedef {bias_t.name} bias_t;
+    typedef {scale_t.name} scale_t;
+    typedef {table_t.name} table_t;
+    static const unsigned io_type = nnet::{iotype};
+    static const unsigned reuse_factor = {reuse};
+    template<class x_T, class y_T>
+    using product = nnet::product::{product_type}<x_T, y_T>;
+}};\n"""
+
+layernorm_function_template = 'nnet::layernormalize<{input_t}, {output_t}, {config}>({input}, {output}, {scale}, {bias});'
+
+layernorm_include_list = ['nnet_utils/nnet_layernorm.h']
+
+
+class LayerNormalizationConfigTemplate(LayerConfigTemplate):
+    def __init__(self):
+        super().__init__(LayerNormalization)
+        self.template = layernorm_config_template
+
+    def format(self, node):
+        params = self._default_config_params(node)
+        params['n_in'] = node.get_input_variable().size_cpp()
+        params['product_type'] = get_backend('vivado').product_type(
+            node.get_input_variable().type.precision, node.get_weights('scale').type.precision
+        )
+
+        return self.template.format(**params)
+
+
+class LayerNormalizationFunctionTemplate(FunctionCallTemplate):
+    def __init__(self):
+        super().__init__(LayerNormalization, include_header=layernorm_include_list)
+        self.template = layernorm_function_template
+
+    def format(self, node):
+        params = self._default_function_params(node)
+        params['scale'] = node.get_weights('scale').name
+        params['bias'] = node.get_weights('bias').name
+
+        return self.template.format(**params)
+
+
 # Activation templates
 
 activ_config_template = """struct {type}_config{index} : nnet::activ_config {{

hls4ml/backends/vivado/vivado_backend.py

@@ -24,6 +24,7 @@
     GarNet,
     GarNetStack,
     Layer,
+    LayerNormalization,
     Pooling1D,
     Pooling2D,
     SeparableConv1D,
@@ -32,7 +33,7 @@
     TimeDistributed,
 )
 from hls4ml.model.optimizer import get_backend_passes, layer_optimizer
-from hls4ml.model.types import FixedPrecisionType, IntegerPrecisionType, NamedType, PackedType
+from hls4ml.model.types import FixedPrecisionType, IntegerPrecisionType, NamedType, PackedType, RoundingMode, SaturationMode
 from hls4ml.report import parse_vivado_report
 from hls4ml.utils import attribute_descriptions as descriptions
 from hls4ml.utils.einsum_utils import parse_einsum
@@ -101,6 +102,32 @@ def _register_layer_attributes(self):
             )
             self.attribute_map[layer] = attrs
 
+        # Add LayerNorm attributes
+        ln_layers = [LayerNormalization]
+        for layer in ln_layers:
+            attrs = self.attribute_map.get(layer, [])
+            attrs.append(ConfigurableAttribute('table_range_power2', default=0, description=descriptions.table_range_power2))
+            attrs.append(ConfigurableAttribute('table_size', default=4096, description=descriptions.table_size))
+            attrs.append(
+                TypeAttribute(
+                    'table',
+                    default=FixedPrecisionType(
+                        8, 5, signed=False, rounding_mode=RoundingMode.RND_CONV, saturation_mode=SaturationMode.SAT
+                    ),
+                    description=descriptions.table_type,
+                )
+            )
+            attrs.append(
+                TypeAttribute(
+                    'accum',
+                    default=FixedPrecisionType(
+                        14, 4, signed=True, rounding_mode=RoundingMode.RND_CONV, saturation_mode=SaturationMode.SAT
+                    ),
+                    description=descriptions.accum_type,
+                )
+            )
+            self.attribute_map[layer] = attrs
+
         # Add TimeStepLoopParallelism to TimeDistributed
         attrs = self.attribute_map.get(TimeDistributed, [])
         attrs.append(

hls4ml/converters/keras/core.py

@@ -1,3 +1,5 @@
+import math
+
 from hls4ml.converters.keras_v2_to_hls import get_weights_data, keras_handler, parse_default_keras_layer
 from hls4ml.model.quantizers import BinaryQuantizer, TernaryQuantizer
 from hls4ml.model.types import IntegerPrecisionType
@@ -131,6 +133,40 @@ def parse_batchnorm_layer(keras_layer, input_names, input_shapes, data_reader):
     return layer, [shape for shape in input_shapes[0]]
 
 
+@keras_handler('LayerNormalization')
+def parse_layernorm_layer(keras_layer, input_names, input_shapes, data_reader):
+    assert 'LayerNormalization' in keras_layer['class_name']
+
+    layer = parse_default_keras_layer(keras_layer, input_names)
+
+    in_size = 1
+    for dim in input_shapes[0][1:]:
+        in_size *= dim
+    layer['n_in'] = layer['n_out'] = in_size
+
+    if not ((len(input_shapes[0])) == 3):
+        raise Exception(
+            'input size is not currently supported by hls4ml; '
+            'only three-dimensional input (including batch dimension) is supported'
+        )
+    layer['seq_len'] = input_shapes[0][-2]
+
+    if not (keras_layer['config']['axis'][0] == 2):
+        raise Exception('assigning the axis is not currently supported by hls4ml; only axis 2 is supported')
+    layer['axis'] = keras_layer['config']['axis'][0]
+
+    layer['gamma_data'] = get_weights_data(data_reader, layer['name'], 'gamma')
+    layer['beta_data'] = get_weights_data(data_reader, layer['name'], 'beta')
+
+    if keras_layer['config']['epsilon'] <= 0:
+        raise Exception('epsilon must be positive')
+    layer['epsilon_power_of_10'] = -round(math.log10(keras_layer['config']['epsilon']))
+    if layer['epsilon_power_of_10'] <= 0:
+        raise Exception('epsilon must be less than 1e-1')
+
+    return layer, [shape for shape in input_shapes[0]]
+
+
 @keras_handler('Embedding')
 def parse_embedding_layer(keras_layer, input_names, input_shapes, data_reader):
     assert 'Embedding' in keras_layer['class_name']

hls4ml/converters/pytorch/core.py

@@ -1,3 +1,5 @@
+import math
+
 import numpy as np
 
 from hls4ml.converters.pytorch_to_hls import pytorch_handler
@@ -160,6 +162,42 @@ def parse_batchnorm_layer(operation, layer_name, input_names, input_shapes, node
     return layer, [shape for shape in input_shapes[0]]
 
 
+@pytorch_handler('LayerNorm')
+def parse_layernorm_layer(operation, layer_name, input_names, input_shapes, node, class_object, data_reader, config):
+    assert 'LayerNorm' in operation
+
+    layer = {}
+
+    layer['class_name'] = 'LayerNormalization'
+    layer['name'] = layer_name
+    layer['inputs'] = input_names
+
+    in_size = 1
+    for dim in input_shapes[0][1:]:
+        in_size *= dim
+    layer['n_in'] = layer['n_out'] = in_size
+
+    if not ((len(input_shapes[0])) == 3):
+        raise Exception(
+            f'Input shape {input_shapes[0]} is not currently supported for LayerNorm; '
+            'only three-dimensional inputs (including batch dimension) are supported'
+        )
+    layer['seq_len'] = input_shapes[0][-2]
+
+    layer['axis'] = 2
+
+    layer['gamma_data'] = class_object.weight.data.numpy()
+    layer['beta_data'] = class_object.bias.data.numpy()
+
+    if class_object.eps <= 0:
+        raise Exception('epsilon must be positive')
+    layer['epsilon_power_of_10'] = -round(math.log10(class_object.eps))
+    if layer['epsilon_power_of_10'] <= 0:
+        raise Exception('epsilon must be less than 1e-1')
+
+    return layer, [shape for shape in input_shapes[0]]
+
+
 @pytorch_handler('einsum')
 def parse_einsum_layer(operation, layer_name, input_names, input_shapes, node, class_object, data_reader, config):
     assert 'einsum' in operation

hls4ml/model/layers.py

@@ -1132,6 +1132,31 @@ def add_bias(self, bias, quantizer=None, precision=None):
         self.add_weights_variable(name='bias', var_name='b{index}', data=bias, quantizer=quantizer, precision=precision)
 
 
+class LayerNormalization(Layer):
+    _expected_attributes = [
+        Attribute('n_in'),
+        Attribute('seq_len'),
+        Attribute('axis', value_type=int, default=2),
+        Attribute('epsilon_power_of_10', value_type=int, default=3),
+        WeightAttribute('scale'),
+        WeightAttribute('bias'),
+        TypeAttribute('scale'),
+        TypeAttribute('bias'),
+    ]
+
+    def initialize(self):
+        inp = self.get_input_variable()
+        shape = inp.shape
+        dims = inp.dim_names
+        self.add_output_variable(shape, dims)
+
+        scale = self.get_attr('gamma_data')
+        bias = self.get_attr('beta_data')
+
+        self.add_weights_variable(name='scale', var_name='s{index}', data=scale)
+        self.add_weights_variable(name='bias', var_name='b{index}', data=bias)
+
+
 class Merge(Layer):
     def initialize(self):
         assert len(self.inputs) == 2
@@ -1902,6 +1927,7 @@ def initialize(self):
     'BatchNormOnnx': BatchNormOnnx,
     'LayerGroup': LayerGroup,
     'SymbolicExpression': SymbolicExpression,
+    'LayerNormalization': LayerNormalization,
     'EinsumDense': EinsumDense,
     'Einsum': Einsum,
     # TensorFlow-specific layers:

hls4ml/model/optimizer/passes/convert_to_channels_last.py

@@ -2,7 +2,7 @@
 # Based on https://github.com/fastmachinelearning/qonnx/blob/
 # 12c96a3ded06beacab08e0f554e4ed014476c0aa/src/qonnx/transformation/channels_last.py
 
-from hls4ml.model.layers import Concatenate, Dense, Input, Reshape, Transpose
+from hls4ml.model.layers import Concatenate, Dense, Input, LayerNormalization, Reshape, Transpose
 from hls4ml.model.optimizer import OptimizerPass
 from hls4ml.model.types import WeightVariable
 
@@ -13,8 +13,9 @@ class ChannelsLastConverter(OptimizerPass):
 
     def match(self, node):
         # If this parameter has not been set, this model does not need to be converted
-        if 'ChannelsLastConversion' not in node.model.config.config['HLSConfig']['Model']:
-            return False  # No littering of unused property
+        do_convert = node.model.config.config['HLSConfig']['Model'].get('ChannelsLastConversion', 'off')
+        if do_convert == 'off':
+            return False
         if not hasattr(node, 'channels_last_converted'):
             return True
 
@@ -44,6 +45,22 @@ def transform(self, model, node):
                 node.get_output_variable().shape = input_shape
                 dim_names = [f'N_INPUT_{i}_{node.index}' for i in range(1, len(input_shape) + 1)]
                 node.get_output_variable().dim_names = dim_names
+        elif isinstance(node, LayerNormalization):
+            # LayerNorm only works on the last dimension in PyTorch
+            perm = [1, 0]
+            pre_transpose = model.make_node(
+                'Transpose', f'pre_transpose_for_{node.get_attr("name")}', {'perm': perm}, [node.get_input_node().name]
+            )
+            pre_transpose.channels_last_converted = True
+            model.insert_node(pre_transpose)
+
+            # If not the output layer, transpose again
+            if not node.get_attr('name') in model.outputs or model.config.config['HLSConfig']['Model']['TransposeOutputs']:
+                post_transpose = model.make_node(
+                    'Transpose', f'post_transpose_for_{node.get_attr("name")}', {'perm': perm}, [node.name]
+                )
+                post_transpose.channels_last_converted = True
+                model.insert_node(post_transpose)
         else:
             # Transpose weight tensors
             tensors = ['weight', 'depthwise', 'pointwise', 'zero_bias', 'scale', 'recurrent_weight']

hls4ml/model/optimizer/passes/infer_precision.py

@@ -51,7 +51,7 @@ def _infer_precision(self, node, types_to_infer):
         if node_class in ['Dense']:
             return self._infer_dense_precision(node, types_to_infer)
 
-        if node_class in ['BatchNormalization', 'ApplyAlpha']:
+        if node_class in ['BatchNormalization', 'ApplyAlpha', 'LayerNormalization']:
             return self._infer_bn_precision(node, types_to_infer)
 
         if node_class in ['Conv1D', 'Conv2D', 'PointwiseConv1D', 'PointwiseConv2D', 'Conv2DBatchnorm']:

hls4ml/model/profiling.py

@@ -293,6 +293,18 @@ def _keras_layer(layer):
     return layer.get_weights(), ['w', 'b']
 
 
+def _keras_layernorm(layer):
+    weights = layer.get_weights()
+
+    gamma = weights[0]
+    beta = weights[1]
+
+    scale = gamma
+    bias = beta
+
+    return [scale, bias], ['s', 'b']
+
+
 def _keras_lstm(layer):
     return layer.get_weights(), ['w', 'u', 'b']
 
@@ -302,6 +314,7 @@ def _keras_lstm(layer):
     {
         'BatchNormalization': _keras_batchnorm,
         'QBatchNormalization': _keras_batchnorm,
+        'LayerNormalization': _keras_layernorm,
         'LSTM': _keras_lstm,
         'QLSTM': _keras_lstm,
     },