Avoid extra allocation of memory (#679)

* Avoid extra allocation of memory in USM SYCL iterator

Author: densmirn

dpnp/backend/kernels/dpnp_krnl_fft.cpp

@@ -67,7 +67,6 @@ void dpnp_fft_fft_c(const void* array1_in,
     long* output_shape_offsets = reinterpret_cast<long*>(dpnp_memory_alloc_c(shape_size * sizeof(long)));
     long* input_shape_offsets = reinterpret_cast<long*>(dpnp_memory_alloc_c(shape_size * sizeof(long)));
     // must be a thread local storage.
-    long* xyz = reinterpret_cast<long*>(dpnp_memory_alloc_c(result_size * shape_size * sizeof(long)));
     long* axis_iterator = reinterpret_cast<long*>(dpnp_memory_alloc_c(result_size * shape_size * sizeof(long)));
 
     get_shape_offsets_inkernel<long>(output_shape, shape_size, output_shape_offsets);
@@ -79,14 +78,14 @@ void dpnp_fft_fft_c(const void* array1_in,
 
         double sum_real = 0.0;
         double sum_imag = 0.0;
-        // need to replace these arrays by thread local storage
-        long* xyz_thread = xyz + (output_id * shape_size);
+        // need to replace this array by thread local storage
         long* axis_iterator_thread = axis_iterator + (output_id * shape_size);
 
-        get_xyz_by_id_inkernel(output_id, output_shape_offsets, shape_size, xyz_thread);
+        size_t xyz_id;
         for (size_t i = 0; i < shape_size; ++i)
         {
-            axis_iterator_thread[i] = xyz_thread[i];
+            xyz_id = get_xyz_id_by_id_inkernel(output_id, output_shape_offsets, shape_size, i);
+            axis_iterator_thread[i] = xyz_id;
         }
 
         const long axis_length = input_boundarie;
@@ -114,7 +113,8 @@ void dpnp_fft_fft_c(const void* array1_in,
                 }
             }
 
-            const size_t output_local_id = xyz_thread[axis];
+            xyz_id = get_xyz_id_by_id_inkernel(output_id, output_shape_offsets, shape_size, axis);
+            const size_t output_local_id = xyz_id;
             const double angle = 2.0 * kernel_pi * it * output_local_id / axis_length;
 
             const double angle_cos = cl::sycl::cos(angle);
@@ -153,7 +153,6 @@ void dpnp_fft_fft_c(const void* array1_in,
     dpnp_memory_free_c(input_shape_offsets);
     dpnp_memory_free_c(output_shape_offsets);
     dpnp_memory_free_c(axis_iterator);
-    dpnp_memory_free_c(xyz);
 
     return;
 }

dpnp/backend/src/dpnp_iterator.hpp

@@ -290,10 +290,6 @@ class DPNPC_id final
             get_shape_offsets_inkernel<size_type>(output_shape, output_shape_size, output_shape_strides);
 
             iteration_size = 1;
-
-            // make thread private storage for each shape by multiplying memory
-            sycl_output_xyz =
-                reinterpret_cast<size_type*>(dpnp_memory_alloc_c(output_size * output_shape_size_in_bytes));
         }
     }
 
@@ -400,10 +396,6 @@ class DPNPC_id final
             {
                 axes_shape_strides[i] = input_shape_strides[axes[i]];
             }
-
-            // make thread private storage for each shape by multiplying memory
-            sycl_output_xyz =
-                reinterpret_cast<size_type*>(dpnp_memory_alloc_c(output_size * output_shape_size_in_bytes));
         }
     }
 
@@ -485,35 +477,30 @@ class DPNPC_id final
         {
             assert(output_global_id < output_size);
 
-            // use thread private storage
-            size_type* sycl_output_xyz_thread = sycl_output_xyz + (output_global_id * output_shape_size);
-
-            get_xyz_by_id_inkernel(output_global_id, output_shape_strides, output_shape_size, sycl_output_xyz_thread);
-
             for (size_t iit = 0, oit = 0; iit < input_shape_size; ++iit)
             {
                 if (std::find(axes.begin(), axes.end(), iit) == axes.end())
                 {
-                    input_global_id += (sycl_output_xyz_thread[oit] * input_shape_strides[iit]);
+                    const size_type output_xyz_id = get_xyz_id_by_id_inkernel(output_global_id, output_shape_strides,
+                                                                              output_shape_size, oit);
+                    input_global_id += (output_xyz_id * input_shape_strides[iit]);
                     ++oit;
                 }
             }
         }
         else if (broadcast_use)
         {
             assert(output_global_id < output_size);
-
-            // use thread private storage
-            size_type* sycl_output_xyz_thread = sycl_output_xyz + (output_global_id * output_shape_size);
-
-            get_xyz_by_id_inkernel(output_global_id, output_shape_strides, output_shape_size, sycl_output_xyz_thread);
+            assert(input_shape_size <= output_shape_size);
 
             for (int irit = input_shape_size - 1, orit = output_shape_size - 1; irit >= 0; --irit, --orit)
             {
                 size_type* broadcast_axes_end = broadcast_axes + broadcast_axes_size;
                 if (std::find(broadcast_axes, broadcast_axes_end, orit) == broadcast_axes_end)
                 {
-                    input_global_id += (sycl_output_xyz_thread[orit] * input_shape_strides[irit]);
+                    const size_type output_xyz_id = get_xyz_id_by_id_inkernel(output_global_id, output_shape_strides,
+                                                                              output_shape_size, orit);
+                    input_global_id += (output_xyz_id * input_shape_strides[irit]);
                 }
             }
         }
@@ -565,10 +552,8 @@ class DPNPC_id final
         output_shape_size = size_type{};
         dpnp_memory_free_c(output_shape);
         dpnp_memory_free_c(output_shape_strides);
-        dpnp_memory_free_c(sycl_output_xyz);
         output_shape = nullptr;
         output_shape_strides = nullptr;
-        sycl_output_xyz = nullptr;
     }
 
     void free_memory()
@@ -602,9 +587,6 @@ class DPNPC_id final
     size_type iteration_shape_size = size_type{};
     size_type* iteration_shape_strides = nullptr;
     size_type* axes_shape_strides = nullptr;
-
-    // data allocated to use inside SYCL kernels
-    size_type* sycl_output_xyz = nullptr;
 };
 
 #endif // DPNP_ITERATOR_H

dpnp/backend/src/dpnp_utils.hpp

@@ -27,6 +27,7 @@
 #ifndef BACKEND_UTILS_H // Cython compatibility
 #define BACKEND_UTILS_H
 
+#include <cassert>
 #include <algorithm>
 #include <iostream>
 #include <iterator>
@@ -69,33 +70,38 @@ void get_shape_offsets_inkernel(const _DataType* shape, size_t shape_size, _Data
 
 /**
  * @ingroup BACKEND_UTILS
- * @brief Calculate ids for all given axes from linear index
+ * @brief Calculate xyz id for given axis from linear index
  *
- * Calculates ids of the array with given shape. This is reverse operation of @ref get_id_by_xyz_inkernel
+ * Calculates xyz id of the array with given shape.
  * for example:
  *   input_array_shape_offsets[20, 5, 1]
  *   global_id == 5
- *   xyz array ids should be [0, 1, 0]
+ *   axis == 1
+ *   xyz_id should be 1
  *
- * @param [in]  global_id     linear index id of the element in multy-D array.
+ * @param [in]  global_id     linear index of the element in multy-D array.
  * @param [in]  offsets       array with input offsets.
  * @param [in]  offsets_size  array size for @ref offsets parameter.
- * @param [out] xyz           Result array with @ref offsets_size size.
+ * @param [in]  axis          axis.
  */
 template <typename _DataType>
-void get_xyz_by_id_inkernel(size_t global_id, const _DataType* offsets, size_t offsets_size, _DataType* xyz)
+_DataType get_xyz_id_by_id_inkernel(size_t global_id, const _DataType* offsets, size_t offsets_size, size_t axis)
 {
+    /* avoid warning unused variable*/
+    (void)offsets_size;
+
+    assert(axis < offsets_size);
+
+    _DataType xyz_id = 0;
     long reminder = global_id;
-    for (size_t axis = 0; axis < offsets_size; ++axis)
+    for (size_t i = 0; i < axis + 1; ++i)
     {
-        /* reconstruct [x][y][z] from given linear idx */
-        const _DataType axis_val = offsets[axis];
-        _DataType xyz_id = reminder / axis_val;
+        const _DataType axis_val = offsets[i];
+        xyz_id = reminder / axis_val;
         reminder = reminder % axis_val;
-        xyz[axis] = xyz_id;
     }
 
-    return;
+    return xyz_id;
 }
 
 /**