general transpose for oneAPI

hls4ml/backends/oneapi/passes/reshaping_templates.py

@@ -161,29 +161,64 @@ 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']
 
 
+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=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 TransposeConfigTemplate(LayerConfigTemplate):
     def __init__(self):
         super().__init__(Transpose)
         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}'
+        return permute_config_gen(name, shape, perm)
 
 
 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