add kernel for max with axis (#698)

* add kernel for max with axis

Author: Rubtsowa

dpnp/backend/kernels/dpnp_krnl_statistics.cpp

@@ -146,36 +146,188 @@ class dpnp_max_c_kernel;
 template <typename _DataType>
 void dpnp_max_c(void* array1_in, void* result1, const size_t* shape, size_t ndim, const size_t* axis, size_t naxis)
 {
-    __attribute__((unused)) void* tmp = (void*)(axis + naxis);
+    if (naxis == 0)
+    {
+        __attribute__((unused)) void* tmp = (void*)(axis + naxis);
 
-    _DataType* array_1 = reinterpret_cast<_DataType*>(array1_in);
-    _DataType* result = reinterpret_cast<_DataType*>(result1);
+        _DataType* array_1 = reinterpret_cast<_DataType*>(array1_in);
+        _DataType* result = reinterpret_cast<_DataType*>(result1);
 
-    size_t size = 1;
-    for (size_t i = 0; i < ndim; ++i)
-    {
-        size *= shape[i];
-    }
+        size_t size = 1;
+        for (size_t i = 0; i < ndim; ++i)
+        {
+            size *= shape[i];
+        }
 
-    if constexpr (std::is_same<_DataType, double>::value || std::is_same<_DataType, float>::value)
-    {
-        // Required initializing the result before call the function
-        result[0] = array_1[0];
+        if constexpr (std::is_same<_DataType, double>::value || std::is_same<_DataType, float>::value)
+        {
+            // Required initializing the result before call the function
+            result[0] = array_1[0];
 
-        auto dataset = mkl_stats::make_dataset<mkl_stats::layout::row_major>(1, size, array_1);
+            auto dataset = mkl_stats::make_dataset<mkl_stats::layout::row_major>(1, size, array_1);
 
-        cl::sycl::event event = mkl_stats::max(DPNP_QUEUE, dataset, result);
+            cl::sycl::event event = mkl_stats::max(DPNP_QUEUE, dataset, result);
 
-        event.wait();
+            event.wait();
+        }
+        else
+        {
+            auto policy = oneapi::dpl::execution::make_device_policy<class dpnp_max_c_kernel<_DataType>>(DPNP_QUEUE);
+
+            _DataType* res = std::max_element(policy, array_1, array_1 + size);
+            policy.queue().wait();
+
+            result[0] = *res;
+        }
     }
     else
     {
-        auto policy = oneapi::dpl::execution::make_device_policy<class dpnp_max_c_kernel<_DataType>>(DPNP_QUEUE);
+        _DataType* array_1 = reinterpret_cast<_DataType*>(array1_in);
+        _DataType* result = reinterpret_cast<_DataType*>(result1);
+
+        size_t res_ndim = ndim - naxis;
+        size_t res_shape[res_ndim];
+        int ind = 0;
+        for (size_t i = 0; i < ndim; i++)
+        {
+            bool found = false;
+            for (size_t j = 0; j < naxis; j++)
+            {
+                if (axis[j] == i)
+                {
+                    found = true;
+                    break;
+                }
+            }
+            if (!found)
+            {
+                res_shape[ind] = shape[i];
+                ind++;
+            }
+        }
+
+        size_t size_input = 1;
+        for (size_t i = 0; i < ndim; ++i)
+        {
+            size_input *= shape[i];
+        }
+
+        size_t input_shape_offsets[ndim];
+        size_t acc = 1;
+        for (size_t i = ndim - 1; i > 0; --i)
+        {
+            input_shape_offsets[i] = acc;
+            acc *= shape[i];
+        }
+        input_shape_offsets[0] = acc;
+
+        size_t output_shape_offsets[res_ndim];
+        acc = 1;
+        if (res_ndim > 0)
+        {
+            for (size_t i = res_ndim - 1; i > 0; --i)
+            {
+                output_shape_offsets[i] = acc;
+                acc *= res_shape[i];
+            }
+        }
+        output_shape_offsets[0] = acc;
+
+        size_t size_result = 1;
+        for (size_t i = 0; i < res_ndim; ++i)
+        {
+            size_result *= res_shape[i];
+        }
 
-        _DataType* res = std::max_element(policy, array_1, array_1 + size);
-        policy.queue().wait();
+        //init result array
+        for (size_t result_idx = 0; result_idx < size_result; ++result_idx)
+        {
+            size_t xyz[res_ndim];
+            size_t remainder = result_idx;
+            for (size_t i = 0; i < res_ndim; ++i)
+            {
+                xyz[i] = remainder / output_shape_offsets[i];
+                remainder = remainder - xyz[i] * output_shape_offsets[i];
+            }
+
+            size_t source_axis[ndim];
+            size_t result_axis_idx = 0;
+            for (size_t idx = 0; idx < ndim; ++idx)
+            {
+                bool found = false;
+                for (size_t i = 0; i < naxis; ++i)
+                {
+                    if (axis[i] == idx)
+                    {
+                        found = true;
+                        break;
+                    }
+                }
+                if (found)
+                {
+                    source_axis[idx] = 0;
+                }
+                else
+                {
+                    source_axis[idx] = xyz[result_axis_idx];
+                    result_axis_idx++;
+                }
+            }
+
+            size_t source_idx = 0;
+            for (size_t i = 0; i < ndim; ++i)
+            {
+                source_idx += input_shape_offsets[i] * source_axis[i];
+            }
+
+            result[result_idx] = array_1[source_idx];
+        }
 
-        result[0] = *res;
+        for (size_t source_idx = 0; source_idx < size_input; ++source_idx)
+        {
+            // reconstruct x,y,z from linear source_idx
+            size_t xyz[ndim];
+            size_t remainder = source_idx;
+            for (size_t i = 0; i < ndim; ++i)
+            {
+                xyz[i] = remainder / input_shape_offsets[i];
+                remainder = remainder - xyz[i] * input_shape_offsets[i];
+            }
+
+            // extract result axis
+            size_t result_axis[res_ndim];
+            size_t result_idx = 0;
+            for (size_t idx = 0; idx < ndim; ++idx)
+            {
+                // try to find current idx in axis array
+                bool found = false;
+                for (size_t i = 0; i < naxis; ++i)
+                {
+                    if (axis[i] == idx)
+                    {
+                        found = true;
+                        break;
+                    }
+                }
+                if (!found)
+                {
+                    result_axis[result_idx] = xyz[idx];
+                    result_idx++;
+                }
+            }
+
+            // Construct result offset
+            size_t result_offset = 0;
+            for (size_t i = 0; i < res_ndim; ++i)
+            {
+                result_offset += output_shape_offsets[i] * result_axis[i];
+            }
+
+            if (result[result_offset] < array_1[source_idx])
+            {
+                result[result_offset] = array_1[source_idx];
+            }
+        }
     }
 
     return;

dpnp/dpnp_algo/dpnp_algo_statistics.pyx

@@ -144,89 +144,60 @@ cpdef dparray dpnp_cov(dparray array1):
     return result
 
 
-cpdef dparray _dpnp_max(dparray input):
+cpdef dparray _dpnp_max(dparray input, _axis_, output_shape):
     cdef DPNPFuncType param1_type = dpnp_dtype_to_DPNPFuncType(input.dtype)
 
     cdef DPNPFuncData kernel_data = get_dpnp_function_ptr(DPNP_FN_MAX, param1_type, param1_type)
 
-    result_type = dpnp_DPNPFuncType_to_dtype(< size_t > kernel_data.return_type)
-    cdef dparray result = dparray((1,), dtype=result_type)
+    result_type = dpnp_DPNPFuncType_to_dtype( < size_t > kernel_data.return_type)
+    cdef dparray result = dparray(output_shape, dtype=result_type)
 
     cdef custom_statistic_1in_1out_func_ptr_t func = <custom_statistic_1in_1out_func_ptr_t > kernel_data.ptr
-
-    # stub for interface support
     cdef dparray_shape_type axis
     cdef Py_ssize_t axis_size = 0
+    cdef dparray_shape_type axis_ = axis
 
-    func(input.get_data(), result.get_data(), < size_t * > input._dparray_shape.data(), input.ndim, < size_t * > axis.data(), axis_size)
+    if _axis_ is not None:
+        axis = _axis_
+        axis_.reserve(len(axis))
+        for shape_it in axis:
+            axis_.push_back(shape_it)
+        axis_size = len(axis)
 
-    return result
+    func(input.get_data(), result.get_data(), < size_t * > input._dparray_shape.data(), input.ndim, < size_t * > axis_.data(), axis_size)
 
+    dpnp_array = dpnp.array(result, dtype=input.dtype)
+    dpnp_result_array = dpnp_array.reshape(output_shape)
+    return dpnp_result_array
 
-cpdef dparray dpnp_max(dparray input, axis):
-    if axis is None:
-        return _dpnp_max(input)
 
+cpdef dparray dpnp_max(dparray input, axis):
     cdef dparray_shape_type shape_input = input.shape
-    cdef long size_input = input.size
-    if isinstance(axis, int):
-        axis_ = tuple([axis])
-    else:
+    if axis is None:
         axis_ = axis
-
-    output_shape = dparray(len(shape_input) - len(axis_), dtype=numpy.int64)
-    ind = 0
-    for id, shape_axis in enumerate(shape_input):
-        if id not in axis_:
-            output_shape[ind] = shape_axis
-            ind += 1
-    cdef long prod = 1
-    for i in range(len(output_shape)):
-        if output_shape[i] != 0:
-            prod *= output_shape[i]
-    result_array = [None] * prod
-    input_shape_offsets = [None] * len(shape_input)
-    acc = 1
-    for i in range(len(shape_input)):
-        ind = len(shape_input) - 1 - i
-        input_shape_offsets[ind] = acc
-        acc *= shape_input[ind]
-    output_shape_offsets = [None] * len(shape_input)
-    acc = 1
-    if len(output_shape) > 0:
-        for i in range(len(output_shape)):
-            ind = len(output_shape) - 1 - i
-            output_shape_offsets[ind] = acc
-            acc *= output_shape[ind]
-
-    for source_idx in range(size_input):
-
-        # reconstruct x,y,z from linear source_idx
-        xyz = []
-        remainder = source_idx
-        for i in input_shape_offsets:
-            quotient, remainder = divmod(remainder, i)
-            xyz.append(quotient)
-
-        # extract result axis
-        result_axis = []
-        for idx, offset in enumerate(xyz):
-            if idx not in axis_:
-                result_axis.append(offset)
-
-        # Construct result offset
-        result_offset = 0
-        for i, result_axis_val in enumerate(result_axis):
-            result_offset += (output_shape_offsets[i] * result_axis_val)
-
-        input_elem = input.item(source_idx)
-        if result_array[result_offset] is None:
-            result_array[result_offset] = input_elem
+        output_shape = 1
+    else:
+        if isinstance(axis, int):
+            if axis < 0:
+                axis_ = tuple([input.ndim - axis])
+            else:
+                axis_ = tuple([axis])
         else:
-            result_array[result_offset] = max(result_array[result_offset], input_elem)
-    dpnp_array = dpnp.array(result_array, dtype=input.dtype)
-    dpnp_result_array = dpnp_array.reshape(output_shape)
-    return dpnp_result_array
+            _axis_ = []
+            for i in range(len(axis)):
+                if axis[i] < 0:
+                    _axis_.append(input.ndim - axis[i])
+                else:
+                    _axis_.append(axis[i])
+            axis_ = tuple(_axis_)
+
+        output_shape = dparray(len(shape_input) - len(axis_), dtype=numpy.int64)
+        ind = 0
+        for id, shape_axis in enumerate(shape_input):
+            if id not in axis_:
+                output_shape[ind] = shape_axis
+                ind += 1
+    return _dpnp_max(input, axis_, output_shape)
 
 
 cpdef dparray _dpnp_mean(dparray input):

dpnp/dpnp_iface_statistics.py

@@ -302,8 +302,30 @@ def max(input, axis=None, out=None, keepdims=numpy._NoValue, initial=numpy._NoVa
     """
 
     if not use_origin_backend(input):
+        # Negative values in 'shape' are not allowed in dparray
+        # 306-322 check on negative and duplicate axis
+        isaxis = True
+        if axis is not None:
+            if dpnp.isscalar(axis):
+                if axis < 0:
+                    isaxis = False
+            else:
+                for val in axis:
+                    if val < 0:
+                        isaxis = False
+                        break
+                if isaxis:
+                    for i in range(len(axis)):
+                        for j in range(len(axis)):
+                            if i != j:
+                                if axis[i] == axis[j]:
+                                    isaxis = False
+                                    break
+
         if not isinstance(input, dparray):
             pass
+        elif not isaxis:
+            pass
         elif out is not None:
             pass
         elif keepdims is not numpy._NoValue:

tests/test_statistics.py

@@ -20,6 +20,17 @@ def test_median(type, size):
     numpy.testing.assert_allclose(dpnp_res, np_res)
 
 
+@pytest.mark.parametrize("axis",
+                         [0, 1, -1, 2, -2, (1, 2), (0, -2)])
+def test_max(axis):
+    a = numpy.arange(768, dtype=numpy.float64).reshape((4, 4, 6, 8))
+    ia = dpnp.array(a)
+
+    np_res = numpy.max(a, axis=axis)
+    dpnp_res = dpnp.max(ia, axis=axis)
+
+    numpy.testing.assert_allclose(dpnp_res, np_res)
+
 @pytest.mark.parametrize("array",
                          [[2, 0, 6, 2],
                           [2, 0, 6, 2, 5, 6, 7, 8],