Merge branch 'main' into initialRecurr

Author: JanFSchulte

hls4ml/backends/fpga/fpga_backend.py

@@ -913,6 +913,33 @@ def generate_conv2d_line_buffer_fn(
 
         return generated_code
 
+    @staticmethod
+    def permute_config_gen(name: str, shape: tuple[int, ...], perm: tuple[int, ...]):
+        """
+        Generate new shape and perm_strides for a permute operation. Operates by mapping the output index
+        to input input index by:
+        - unravel the output index
+        - map each dimension to the corresponding stride in the input tensor, sum
+        The operation can be expressed as:
+
+        new_shape = tuple(shape[i] for i in perm)
+        strides = np.cumprod((shapes[1:] + (1,))[::-1])[::-1]
+        perm_strides = [strides[i] for i in perm]
+        out[index] = inp[np.dot(np.unravel_index(index, new_shape), perm_strides)]
+
+        Args:
+            name (str): The name of the configuration.
+            shape (tuple[int, ...]): The shape of the input tensor.
+            perm (tuple[int, ...]): The permutation of the dimensions.
+
+        Returns:
+            (new_shape, perm_strides) (tuple, tuple):  the output shape and permutation strides.
+        """
+        new_shape = tuple(shape[i] for i in perm)
+        strides = np.cumprod((shape[1:] + (1,))[::-1])[::-1]
+        perm_strides = tuple(int(strides[i]) for i in perm)
+        return (new_shape, perm_strides)
+
     @model_optimizer()
     def write_hls(self, model):
         self.writer.write_hls(model)

hls4ml/backends/oneapi/passes/reshaping_templates.py

@@ -161,17 +161,17 @@ def format(self, node):
 
 # Transpose templates
 
-transpose_config_template = """struct config{index} : nnet::transpose_config {{
-    static const unsigned depth = {depth};
-    static const unsigned height = {height};
-    static const unsigned width = {width};
-    static constexpr unsigned perm[3] = {{{perm_str}}};
+transpose_config_template = """struct {config_name} : nnet::transpose_config {{
+    static constexpr unsigned dims = {dims};
+    static constexpr unsigned N = {N};
+    static constexpr std::array<unsigned, dims> from_shape = {{{from_shape}}};
+    static constexpr std::array<unsigned, dims> to_shape = {{{to_shape}}};
+    static constexpr std::array<unsigned, dims> perm = {{{perm}}};
+    static constexpr std::array<unsigned, dims> perm_strides = {{{perm_strides}}};
 }};\n"""
 
-transpose_function_template = 'nnet::transpose_{dim}<{input_t}, {output_t}, {config}>({input}, {output});'
-transpose_task_sequence_template = (
-    'task_sequence<nnet::transpose_{dim}_stream<{input_pipe}, {output_pipe}, {config}>> {name};'
-)
+transpose_function_template = 'nnet::transpose<{input_t}, {output_t}, {config}>({input}, {output});'
+transpose_task_sequence_template = 'task_sequence<nnet::transpose_stream<{input_pipe}, {output_pipe}, {config}>> {name};'
 transpose_include_list = ['nnet_utils/nnet_transpose.h', 'nnet_utils/nnet_transpose_stream.h']
 
 
@@ -181,9 +181,20 @@ def __init__(self):
         self.template = transpose_config_template
 
     def format(self, node):
-        params = self._default_config_params(node)
-
-        return self.template.format(**params)
+        shape = tuple(node.get_input_variable().shape)
+        perm = tuple(node.get_attr('perm'))
+        name = f'config{node.index}'
+
+        new_shape, perm_strides = node.model.config.backend.permute_config_gen(name, shape, perm)
+        return transpose_config_template.format(
+            dims=len(shape),
+            N=np.prod(shape),
+            from_shape=', '.join(str(x) for x in shape),
+            perm=', '.join(str(x) for x in perm),
+            perm_strides=', '.join(str(x) for x in perm_strides),
+            to_shape=', '.join(str(x) for x in new_shape),
+            config_name=name,
+        )
 
 
 class TransposeFunctionTemplate(FunctionCallTemplate):

hls4ml/backends/vivado/passes/reshaping_templates.py

@@ -1,5 +1,3 @@
-from math import prod
-
 import numpy as np
 
 from hls4ml.backends.template import FunctionCallTemplate, LayerConfigTemplate
@@ -127,40 +125,6 @@ def format(self, node):
 transpose_function_template = 'nnet::transpose<{input_t}, {output_t}, {config_name}>({input}, {output});'
 
 
-def permute_config_gen(name: str, shape: tuple[int, ...], perm: tuple[int, ...]):
-    """
-    Generate a configuration string for a permute operation. Operates by mapping the output index to input input index by:
-     - unravel the output index
-     - map each dimension to the corresponding stride in the input tensor, sum
-    The operation can be expressed as:
-
-    new_shape = tuple(shape[i] for i in perm)
-    strides = np.cumprod((shapes[1:] + (1,))[::-1])[::-1]
-    perm_strides = [strides[i] for i in perm]
-    out[index] = inp[np.dot(np.unravel_index(index, new_shape), perm_strides)]
-
-    Args:
-        name (str): The name of the configuration.
-        shape (tuple[int, ...]): The shape of the input tensor.
-        perm (tuple[int, ...]): The permutation of the dimensions.
-
-    Returns:
-        str: The formatted configuration string for the permute operation.
-    """
-    new_shape = tuple(shape[i] for i in perm)
-    strides = np.cumprod((shape[1:] + (1,))[::-1])[::-1]
-    perm_strides = tuple(int(strides[i]) for i in perm)
-    return transpose_config_template.format(
-        dims=len(shape),
-        N=prod(shape),
-        from_shape=', '.join(str(x) for x in shape),
-        perm=', '.join(str(x) for x in perm),
-        perm_strides=', '.join(str(x) for x in perm_strides),
-        to_shape=', '.join(str(x) for x in new_shape),
-        config_name=name,
-    )
-
-
 class TransposeConfigTemplate(LayerConfigTemplate):
     def __init__(self):
         super().__init__(Transpose)
@@ -170,7 +134,16 @@ def format(self, node):
         shape = tuple(node.get_input_variable().shape)
         perm = tuple(node.get_attr('perm'))
         name = f'config{node.index}'
-        return permute_config_gen(name, shape, perm)
+        new_shape, perm_strides = node.model.config.backend.permute_config_gen(name, shape, perm)
+        return transpose_config_template.format(
+            dims=len(shape),
+            N=np.prod(shape),
+            from_shape=', '.join(str(x) for x in shape),
+            perm=', '.join(str(x) for x in perm),
+            perm_strides=', '.join(str(x) for x in perm_strides),
+            to_shape=', '.join(str(x) for x in new_shape),
+            config_name=name,
+        )
 
 
 class TransposeFunctionTemplate(FunctionCallTemplate):

hls4ml/templates/oneapi/firmware/nnet_utils/nnet_transpose.h

@@ -4,42 +4,31 @@
 namespace nnet {
 
 struct transpose_config {
-    static const unsigned height = 10;
-    static const unsigned width = 10;
-    static const unsigned depth = 10;
-    static constexpr unsigned perm[3] = {2, 0, 1};
+    static constexpr unsigned dims = 0;
+    static constexpr unsigned N = 0;
+
+    // Inherited struct should define these
+    // static constexpr std::array<unsigned, dims> from_shape;
+    // static constexpr std::array<unsigned, dims> to_shape;
+    // static constexpr std::array<unsigned, dims> perm;
+    // static constexpr std::array<unsigned, dims> perm_strides;
 };
 
-template <class data_T, class res_T, typename CONFIG_T> void transpose_2d(const data_T &data, res_T &res) {
-    for (int i = 0; i < CONFIG_T::height; i++) {
-        #pragma unroll
-        for (int j = 0; j < CONFIG_T::width; j++) {
-            res[j * CONFIG_T::height + i] = static_cast<typename res_T::value_type>(data[i * CONFIG_T::width + j]);
-        }
+template <typename CONFIG_T> unsigned transfer_idx(int index) {
+    // Given output idx in c-order flat array, return input idx
+    int idx = 0;
+    for (int i = CONFIG_T::dims - 1; i >= 0; i--) {
+        idx += (index % CONFIG_T::to_shape[i]) * CONFIG_T::perm_strides[i];
+        index /= CONFIG_T::to_shape[i];
     }
+    return idx;
 }
 
-template <class data_T, class res_T, typename CONFIG_T> void transpose_3d(const data_T &data, res_T &res) {
-    static constexpr unsigned dim_data[3] = {CONFIG_T::depth, CONFIG_T::height, CONFIG_T::width};
-    static constexpr unsigned dim_res[3] = {dim_data[CONFIG_T::perm[0]], dim_data[CONFIG_T::perm[1]],
-                                            dim_data[CONFIG_T::perm[2]]};
-
-    int index_data[3] = {0}, index_res[3] = {0};
-
-    for (index_data[0] = 0; index_data[0] < dim_data[0]; index_data[0]++) {
-        #pragma unroll
-        for (index_data[1] = 0; index_data[1] < dim_data[1]; index_data[1]++) {
-            #pragma unroll
-            for (index_data[2] = 0; index_data[2] < dim_data[2]; index_data[2]++) {
-                index_res[0] = index_data[CONFIG_T::perm[0]];
-                index_res[1] = index_data[CONFIG_T::perm[1]];
-                index_res[2] = index_data[CONFIG_T::perm[2]];
-
-                res[index_res[0] * dim_res[1] * dim_res[2] + index_res[1] * dim_res[2] + index_res[2]] =
-                    static_cast<typename res_T::value_type>(
-                        data[index_data[0] * dim_data[1] * dim_data[2] + index_data[1] * dim_data[2] + index_data[2]]);
-            }
-        }
+template <class data_T, class res_T, typename CONFIG_T> void transpose(const data_T &data, res_T &res) {
+    #pragma unroll
+    for (int i = 0; i < CONFIG_T::N; i++) {
+        int idx = transfer_idx<CONFIG_T>(i);
+        res[i] = data[idx];
     }
 }
 

hls4ml/templates/oneapi/firmware/nnet_utils/nnet_transpose_stream.h

@@ -3,17 +3,17 @@
 
 namespace nnet {
 
-template <class data_pipe, class res_pipe, typename CONFIG_T> void transpose_2d_stream() {
+template <class data_pipe, class res_pipe, typename CONFIG_T> void transpose_stream() {
 
     using data_T = typename ExtractPipeType<data_pipe>::value_type;
     using res_T = typename ExtractPipeType<res_pipe>::value_type;
 
     constexpr auto data_size = std::tuple_size<typename ExtractPipeType<data_pipe>::value_type>{};
     constexpr auto res_size = std::tuple_size<typename ExtractPipeType<res_pipe>::value_type>{};
 
-    [[intel::fpga_register]] typename data_T::value_type data_array[CONFIG_T::height * CONFIG_T::width];
+    [[intel::fpga_register]] typename data_T::value_type data_array[CONFIG_T::N];
 
-    for (int i = 0; i < CONFIG_T::height * CONFIG_T::width / data_size; i++) {
+    for (int i = 0; i < CONFIG_T::N / data_size; i++) {
         [[intel::fpga_register]] data_T in_data = data_pipe::read();
 
         #pragma unroll
@@ -22,12 +22,12 @@ template <class data_pipe, class res_pipe, typename CONFIG_T> void transpose_2d_
         }
     }
 
-    for (int i = 0; i < CONFIG_T::height * CONFIG_T::width / res_size; i++) {
+    for (int i = 0; i < CONFIG_T::N / res_size; i++) {
         [[intel::fpga_register]] res_T out_data;
 
         #pragma unroll
         for (int j = 0; j < res_size; j++) {
-            out_data[j] = typename res_T::value_type(data_array[j * data_size + i]);
+            out_data[j] = typename res_T::value_type(data_array[transfer_idx<CONFIG_T>(i * res_size + j)]);
         }
 
         res_pipe::write(out_data);

test/pytest/test_transpose_concat.py

@@ -72,7 +72,7 @@ def data_highdim():
 
 
 @pytest.mark.parametrize('io_type', ['io_stream', 'io_parallel'])
-@pytest.mark.parametrize('backend', ['Vivado', 'Vitis'])
+@pytest.mark.parametrize('backend', ['Vivado', 'Vitis', 'oneAPI'])
 def test_highdim_permute(data_highdim, keras_model_highdim, io_type, backend):
     X = data_highdim
     model = keras_model_highdim