Skip to content

Commit

Permalink
general transpose for oneAPI (fastmachinelearning#1165)
Browse files Browse the repository at this point in the history 
* general transpose for oneAPI

* move the permute_config_gen to a central location

---------

Co-authored-by: Jan-Frederik Schulte <[email protected]>
  • Loading branch information
@jmitrevs @JanFSchulte
jmitrevs and JanFSchulte authored Feb 7, 2025
1 parent 3b7e595 commit 4d23e9f
Show file tree
Hide file tree
Showing 6 changed files with 86 additions and 86 deletions.
27 changes: 27 additions & 0 deletions hls4ml/backends/fpga/fpga_backend.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down
35 changes: 23 additions & 12 deletions hls4ml/backends/oneapi/passes/reshaping_templates.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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']


Expand All @@ -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):
Expand Down
47 changes: 10 additions & 37 deletions hls4ml/backends/vivado/passes/reshaping_templates.py
View file
Original file line number Diff line number Diff line change
@@ -1,5 +1,3 @@
from math import prod

import numpy as np

from hls4ml.backends.template import FunctionCallTemplate, LayerConfigTemplate
Expand Down Expand Up @@ -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)
Expand All @@ -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):
Expand Down
51 changes: 20 additions & 31 deletions hls4ml/templates/oneapi/firmware/nnet_utils/nnet_transpose.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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];
}
}

Expand Down
10 changes: 5 additions & 5 deletions hls4ml/templates/oneapi/firmware/nnet_utils/nnet_transpose_stream.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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);
Expand Down
2 changes: 1 addition & 1 deletion test/pytest/test_transpose_concat.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down

0 comments on commit 4d23e9f

Please sign in to comment.