Add UnaryLUT support

calad0i · calad0i · commit 7f19b03a761c · 2024-07-21T23:24:05.000-07:00
Support UnaryLUT w/ vivado&io_parallel hotfix format support vivado_stream fix syn issue fix syn issue2 Revert unnecessary "fix syn issue1/2" Revert "fix syn issue2" This reverts commit af3c347. Revert "fix syn issue" This reverts commit 532cb9d. rm redundant pipeline pragma unary lut vitis fix unary lut vitis fix - leftover
diff --git a/hls4ml/backends/fpga/passes/hgq_proxy_model.py b/hls4ml/backends/fpga/passes/hgq_proxy_model.py
@@ -4,7 +4,7 @@
 from hls4ml.backends.template import FunctionCallTemplate
 from hls4ml.model.layers import Layer
 from hls4ml.model.optimizer import OptimizerPass
-from hls4ml.model.optimizer.passes.hgq_proxy_model import FixedPointQuantizer
+from hls4ml.model.optimizer.passes.hgq_proxy_model import FixedPointQuantizer, UnaryLUT
 from hls4ml.model.types import Source
 
 
@@ -40,7 +40,6 @@ def generate_mask_fn(
 template<typename input_t, typename output_t>
 void {name}(input_t *inp, output_t *out) {{
     #pragma HLS INLINE
-    #pragma HLS PIPELINE
 
 {body}
 }}
@@ -84,6 +83,25 @@ def format(self, node):
         return self.template.format(**params)
 
 
+class ProcessUnaryLUTCall(FunctionCallTemplate):
+    def __init__(self):
+        super().__init__(UnaryLUT, include_header=[])
+        self.template = 'nnet::unary_lut<{input_t}, {output_t}, {config}>({input}, {output}, {table});'
+        self.include_header = [
+            'nnet_utils/nnet_activation.h',
+            'nnet_utils/nnet_activation_stream.h',
+        ]
+
+    def format(self, node):
+        params = self._default_function_params(node)
+        node.attributes['result_t'].precision = node.attributes['table_t'].precision
+        params['config'] = f'unary_lut_config{node.index}'
+        params['table'] = node.get_weights('table').name
+
+        return self.template.format(**params)
+
+
 def register_hgq_proxy_model(backend: Backend):
     backend.register_pass('process_fixed_point_quantizer_layer', ProcessFixedPointQuantizerLayer)
     backend.register_template(ProcessFixedPointQuantizerCall)
+    backend.register_template(ProcessUnaryLUTCall)
diff --git a/hls4ml/backends/quartus/passes/core_templates.py b/hls4ml/backends/quartus/passes/core_templates.py
@@ -1,6 +1,7 @@
 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.optimizer.passes.hgq_proxy_model import UnaryLUT
 
 # Dense templates
 
@@ -152,7 +153,7 @@ def format(self, node):
 
 class ActivationConfigTemplate(LayerConfigTemplate):
     def __init__(self):
-        super().__init__((Activation, ParametrizedActivation, PReLU))
+        super().__init__((Activation, ParametrizedActivation, PReLU, UnaryLUT))
         self.template = activ_config_template
 
     def format(self, node):
diff --git a/hls4ml/backends/vivado/passes/core_templates.py b/hls4ml/backends/vivado/passes/core_templates.py
@@ -1,6 +1,7 @@
 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.optimizer.passes.hgq_proxy_model import UnaryLUT
 
 # Dense templates
 
@@ -144,7 +145,7 @@ def format(self, node):
 
 class ActivationConfigTemplate(LayerConfigTemplate):
     def __init__(self):
-        super().__init__((Activation, ParametrizedActivation, PReLU))
+        super().__init__((Activation, ParametrizedActivation, PReLU, UnaryLUT))
         self.template = activ_config_template
 
     def format(self, node):
diff --git a/hls4ml/converters/keras/hgq_proxy_model.py b/hls4ml/converters/keras/hgq_proxy_model.py
@@ -1,8 +1,8 @@
-from hls4ml.converters.keras_to_hls import keras_handler, parse_default_keras_layer
+from hls4ml.converters.keras_to_hls import KerasReader, keras_handler, parse_default_keras_layer
 
 
 @keras_handler('FixedPointQuantizer')
-def fixedpoint_quantizer_handler(keras_layer, input_names, input_shapes, data_reader):
+def fixedpoint_quantizer_handler(keras_layer, input_names, input_shapes, data_reader: KerasReader):
     config = parse_default_keras_layer(keras_layer, input_names)
 
     name = config['name']
@@ -19,3 +19,19 @@ def fixedpoint_quantizer_handler(keras_layer, input_names, input_shapes, data_re
 
     layer = config
     return layer, input_shapes[0]
+
+
+@keras_handler('UnaryLUT')
+def unary_lut_keras_handler(keras_layer, input_names, input_shapes, data_reader: KerasReader):
+    config = parse_default_keras_layer(keras_layer, input_names)
+
+    table = data_reader.get_weights_data(config['name'], 'table')
+    k, i, f = keras_layer['config']['kif_out']
+    k, b, i = k, k + i + f, k + i
+    config['table_t'] = f'{"" if k else "u"}fixed<{b},{i}>'
+    config['table'] = table
+    config['table_size'] = len(table)
+    config['activation'] = 'unary_lut'
+
+    layer = config
+    return layer, input_shapes[0]
diff --git a/hls4ml/converters/keras_to_hls.py b/hls4ml/converters/keras_to_hls.py
@@ -205,6 +205,7 @@ def parse_keras_model(model_arch, reader):
         'Softmax',
         'TernaryTanh',
         'HardActivation',
+        'UnaryLUT',
     ]
     # Recurrent layers
     recurrent_layers = ['SimpleRNN', 'LSTM', 'GRU']
diff --git a/hls4ml/model/optimizer/passes/hgq_proxy_model.py b/hls4ml/model/optimizer/passes/hgq_proxy_model.py
@@ -23,6 +23,20 @@ def initialize(self):
         self.mask_kbi = self.attributes.get('mask_kbi', None)
 
 
+class UnaryLUT(Layer):
+    def initialize(self):
+        inp = self.get_input_variable()
+        shape = inp.shape
+        dims = inp.dim_names
+        self.add_output_variable(shape, dims)
+        self.set_attr('n_in', inp.size())
+        self.table = self.attributes['table']
+        self.table_size = self.attributes['table_size']
+
+        table_t = to_hls4ml_fixed(self.attributes['table_t'])
+        self.add_weights_variable(name='table', var_name='table{index}', precision=table_t, data=self.table)
+
+
 def to_hls4ml_fixed(fixed: str):
     matched = re_parse_fixed.match(re_purge_prefix.sub('', fixed))
     assert matched is not None, f'Cannot parse {fixed}'
@@ -125,4 +139,5 @@ def transform(self, model, node: FixedPointQuantizer):
 
 def register_hgq_proxy_model():
     register_layer('FixedPointQuantizer', FixedPointQuantizer)
+    register_layer('UnaryLUT', UnaryLUT)
     register_pass('enforce_proxy_model_embedded_config', EnforceProxyModelEmbeddedConfig)
diff --git a/hls4ml/templates/vivado/nnet_utils/nnet_activation.h b/hls4ml/templates/vivado/nnet_utils/nnet_activation.h
@@ -441,6 +441,28 @@ template <class data_T, class res_T, typename CONFIG_T> void tanh(data_T data[CO
     }
 }
 
+// *************************************************
+//       UnaryLUT Activation
+// *************************************************
+template <int table_size, class data_T> inline unsigned get_index_unary_lut(data_T x) {
+    // Slice the top N bits to get an index into the table
+    static constexpr int N = ceillog2(table_size);
+    return (unsigned)(x(x.width - 1, 0));
+}
+
+template <class data_T, class res_T, typename CONFIG_T>
+void unary_lut(data_T data[CONFIG_T::n_in], res_T res[CONFIG_T::n_in],
+               typename CONFIG_T::table_t table[CONFIG_T::table_size]) {
+    #pragma HLS function_instantiate variable=table
+    #pragma HLS ARRAY_PARTITION variable=table
+
+    for (int ii = 0; ii < CONFIG_T::n_in; ii++) {
+        #pragma HLS UNROLL
+        unsigned index = get_index_unary_lut<CONFIG_T::table_size>(data[ii]);
+        res[ii] = (res_T)table[index];
+    }
+}
+
 // *************************************************
 //       Hard sigmoid Activation
 // *************************************************
diff --git a/hls4ml/templates/vivado/nnet_utils/nnet_activation_stream.h b/hls4ml/templates/vivado/nnet_utils/nnet_activation_stream.h
@@ -412,6 +412,34 @@ template <class data_T, class res_T, typename CONFIG_T> void tanh(hls::stream<da
     }
 }
 
+// *************************************************
+//       UnaryLUT Activation
+// *************************************************
+
+template <class data_T, class res_T, typename CONFIG_T>
+void unary_lut(hls::stream<data_T> &data, hls::stream<res_T> &res, typename CONFIG_T::table_t table[CONFIG_T::table_size]) {
+    #pragma HLS function_instantiate variable=table
+    #pragma HLS ARRAY_PARTITION variable=table complete
+
+UnaryLUTActLoop:
+    for (int i = 0; i < CONFIG_T::n_in / res_T::size; i++) {
+        #pragma HLS PIPELINE II=CONFIG_T::reuse_factor rewind
+
+        data_T in_data = data.read();
+        res_T out_data;
+        PRAGMA_DATA_PACK(out_data)
+
+    UnaryLUTPackLoop:
+        for (int j = 0; j < res_T::size; j++) {
+            #pragma HLS UNROLL
+            unsigned index = get_index_unary_lut<CONFIG_T::table_size>(in_data[j].V);
+            out_data[j] = table[index];
+        }
+
+        res.write(out_data);
+    }
+}
+
 // *************************************************
 //       Hard sigmoid Activation
 // *************************************************

Original file line number	Diff line number	Diff line change
`@@ -205,6 +205,7 @@ def parse_keras_model(model_arch, reader):`
`205`	`205`	`'Softmax',`
`206`	`206`	`'TernaryTanh',`
`207`	`207`	`'HardActivation',`
	`208`	`+ 'UnaryLUT',`
`208`	`209`	`]`
`209`	`210`	`# Recurrent layers`
`210`	`211`	`recurrent_layers = ['SimpleRNN', 'LSTM', 'GRU']`