Skip to content
Snippets Groups Projects
Commit b071317e authored by Wissam Boussella's avatar Wissam Boussella
Browse files

Add transpose function for 4D tensors and related templates

parent 8cb3a2c1
No related branches found
No related tags found
2 merge requests!39Update 0.2.1 -> 0.3.0,!30Add transpose function for 4D tensors and related templates
#ifndef __AIDGE_EXPORT_CPP_KERNELS_TRANSPOSE__
#define __AIDGE_EXPORT_CPP_KERNELS_TRANSPOSE__
#include "network/typedefs.hpp"
/**
* @brief Transposes a 4-dimensional tensor based on the specified permutation.
*
* This function rearranges the dimensions of a 4D tensor according to the
* permutation array provided. The input tensor is expected to have dimensions
* [in_dim1][in_dim2][in_dim3][in_dim4], and the output tensor will have
* dimensions reordered as specified by the permute array.
*
* @tparam T Data type of the tensor elements (e.g., float, double).
* @param[in] inputs Pointer to the input tensor data stored in contiguous memory.
* @param[out] outputs Pointer to the pre-allocated memory for the transposed tensor.
* Ensure this memory is appropriately sized to hold the transposed data.
* @param[in] in_dim1 Size of the first dimension of the input tensor.
* @param[in] in_dim2 Size of the second dimension of the input tensor.
* @param[in] in_dim3 Size of the third dimension of the input tensor.
* @param[in] in_dim4 Size of the fourth dimension of the input tensor.
* @param[in] permute Array of four unsigned integers specifying the desired permutation
* of dimensions. Each value should be in the range [0, 3], defining
* the new order of dimensions for the output tensor.
*/
template <typename T>
void transpose_4D_forward(const T* inputs,
T* outputs,
unsigned int in_dim1,
unsigned int in_dim2,
unsigned int in_dim3,
unsigned int in_dim4,
const unsigned int* permute)
{
unsigned int in_dims[4] = {in_dim1, in_dim2, in_dim3, in_dim4};
unsigned int out_dims[4];
for (unsigned int i = 0; i < 4; ++i) {
out_dims[i] = in_dims[permute[i]];
}
unsigned int in_strides[4] = {
in_dim2 * in_dim3 * in_dim4,
in_dim3 * in_dim4,
in_dim4,
1
};
unsigned int out_strides[4] = {
out_dims[1] * out_dims[2] * out_dims[3],
out_dims[2] * out_dims[3],
out_dims[3],
1
};
for (unsigned int i = 0; i < in_dim1; ++i) {
for (unsigned int j = 0; j < in_dim2; ++j) {
for (unsigned int k = 0; k < in_dim3; ++k) {
for (unsigned int l = 0; l < in_dim4; ++l) {
// Compute the linear index in the input tensor
unsigned int input_index = i * in_strides[0] +
j * in_strides[1] +
k * in_strides[2] +
l * in_strides[3];
unsigned int in_idx[4] = {i, j, k, l};
unsigned int output_index = 0;
for (unsigned int m = 0; m < 4; ++m) {
output_index += in_idx[permute[m]] * out_strides[m];
}
outputs[output_index] = inputs[input_index];
}
}
}
}
}
#endif // __AIDGE_EXPORT_CPP_KERNELS_TRANSPOSE__
\ No newline at end of file
......@@ -56,8 +56,8 @@ class ProducerCPP(ExportNode):
super().__init__(node, mem_info)
self.values = np.array(self.operator.get_output(0))
if len(self.values.shape) == 4: # Note: export in HWC
self.values = np.transpose(self.values, (0, 2, 3, 1))
# if len(self.values.shape) == 4: # Note: export in HWC
# self.values = np.transpose(self.values, (0, 2, 3, 1))
def export(self, export_folder: Path):
header_path = f"include/parameters/{self.attributes['name']}.h"
......@@ -295,3 +295,16 @@ class FcCPP(ExportNodeCpp):
str(ROOT / "kernels" / "activation.hpp"),
str(ROOT / "kernels" / "rescaling.hpp")
]
@ExportLibCpp.register("Transpose", aidge_core.ImplSpec(aidge_core.IOSpec(aidge_core.dtype.float32)))
class TransposeCPP(ExportNodeCpp):
def __init__(self, node, mem_info):
super().__init__(node, mem_info)
self.config_template = str(
ROOT / "templates" / "configuration" / "transpose_4D_config.jinja")
self.forward_template = str(
ROOT / "templates" / "kernel_forward" / "transpose_4D_forward.jinja")
self.include_list = []
self.kernels_to_copy = [
str(ROOT / "kernels" / "transpose.hpp")
]
\ No newline at end of file
{#- For name header -#}
#ifndef {{ name|upper }}_LAYER_H
#define {{ name|upper }}_LAYER_H
{# For layer configuration -#}
{% include "./_def_io.jinja" %}
{% include "./_meminfo.jinja" %}
{# 4D kernels are not supported yet that's why define NAME_DIM1 #}
#define {{ name|upper }}_DIM1 1
static const unsigned int {{ name|upper }}_PERMUTE[] = { {{ output_dims_order | join(', ') }} };
#endif /* {{ name|upper }}_LAYER_H */
transpose_4D_forward({{in_name[0]}},
{{out_name[0]}},
{{name|upper}}_DIM1,
{{ in_name[0]|upper }}_NB_CHANNELS,
{{ in_name[0]|upper }}_IN_HEIGHT,
{{ in_name[0]|upper }}_IN_WIDTH,
{{name|upper}}_PERMUTE);
\ No newline at end of file
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment