From 271567047ce26e8a38ba4f59d6a4b87949f83a1a Mon Sep 17 00:00:00 2001
From: Alexander Kalistratov <alexander.kalistratov@intel.com>
Date: Mon, 17 Jul 2023 18:00:22 +0200
Subject: [PATCH 1/3] Initial commit

---
 dpnp/CMakeLists.txt                           |   1 +
 .../extensions/sycl_ext/CMakeLists.txt        |  72 +++
 .../extensions/sycl_ext/dispatcher_utils.hpp  | 558 ++++++++++++++++++
 dpnp/backend/extensions/sycl_ext/sum_mean.cpp | 113 ++++
 dpnp/backend/extensions/sycl_ext/sum_mean.hpp | 209 +++++++
 dpnp/dpnp_iface_mathematical.py               |  25 +
 6 files changed, 978 insertions(+)
 create mode 100644 dpnp/backend/extensions/sycl_ext/CMakeLists.txt
 create mode 100644 dpnp/backend/extensions/sycl_ext/dispatcher_utils.hpp
 create mode 100644 dpnp/backend/extensions/sycl_ext/sum_mean.cpp
 create mode 100644 dpnp/backend/extensions/sycl_ext/sum_mean.hpp

diff --git a/dpnp/CMakeLists.txt b/dpnp/CMakeLists.txt
index 89524ab1c58..5ee84b28153 100644
--- a/dpnp/CMakeLists.txt
+++ b/dpnp/CMakeLists.txt
@@ -49,6 +49,7 @@ build_dpnp_cython_ext_with_backend(dparray ${CMAKE_CURRENT_SOURCE_DIR}/dparray.p
 add_subdirectory(backend)
 add_subdirectory(backend/extensions/lapack)
 add_subdirectory(backend/extensions/vm)
+add_subdirectory(backend/extensions/sycl_ext)
 
 add_subdirectory(dpnp_algo)
 add_subdirectory(dpnp_utils)
diff --git a/dpnp/backend/extensions/sycl_ext/CMakeLists.txt b/dpnp/backend/extensions/sycl_ext/CMakeLists.txt
new file mode 100644
index 00000000000..9a18630eb3b
--- /dev/null
+++ b/dpnp/backend/extensions/sycl_ext/CMakeLists.txt
@@ -0,0 +1,72 @@
+# *****************************************************************************
+# Copyright (c) 2016-2023, Intel Corporation
+# All rights reserved.
+#
+# Redistribution and use in source and binary forms, with or without
+# modification, are permitted provided that the following conditions are met:
+# - Redistributions of source code must retain the above copyright notice,
+#   this list of conditions and the following disclaimer.
+# - Redistributions in binary form must reproduce the above copyright notice,
+#   this list of conditions and the following disclaimer in the documentation
+#   and/or other materials provided with the distribution.
+#
+# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
+# ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
+# LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
+# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
+# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
+# INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
+# CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
+# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF
+# THE POSSIBILITY OF SUCH DAMAGE.
+# *****************************************************************************
+
+
+set(python_module_name _sycl_ext_impl)
+pybind11_add_module(${python_module_name} MODULE
+    sum_mean.cpp
+)
+
+if (WIN32)
+    if (${CMAKE_VERSION} VERSION_LESS "3.27")
+        # this is a work-around for target_link_options inserting option after -link option, cause
+        # linker to ignore it.
+        set(CMAKE_CXX_LINK_FLAGS "${CMAKE_CXX_LINK_FLAGS} -fsycl-device-code-split=per_kernel")
+    endif()
+endif()
+
+set_target_properties(${python_module_name} PROPERTIES CMAKE_POSITION_INDEPENDENT_CODE ON)
+
+target_include_directories(${python_module_name} PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/../../include)
+target_include_directories(${python_module_name} PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/../../src)
+
+target_include_directories(${python_module_name} PUBLIC ${Dpctl_INCLUDE_DIRS})
+target_include_directories(${python_module_name} PUBLIC ${Dpctl_TENSOR_INCLUDE_DIR})
+
+if (WIN32)
+  target_compile_options(${python_module_name} PRIVATE
+    /clang:-fno-approx-func
+    /clang:-fno-finite-math-only
+    )
+else()
+  target_compile_options(${python_module_name} PRIVATE
+    -fno-approx-func
+    -fno-finite-math-only
+    )
+endif()
+
+target_link_options(${python_module_name} PUBLIC -fsycl-device-code-split=per_kernel)
+if (UNIX)
+    # this option is support on Linux only
+    target_link_options(${python_module_name} PUBLIC -fsycl-link-huge-device-code)
+endif()
+
+if (DPNP_GENERATE_COVERAGE)
+    target_link_options(${python_module_name} PRIVATE -fprofile-instr-generate -fcoverage-mapping)
+endif()
+
+install(TARGETS ${python_module_name}
+  DESTINATION "dpnp/backend/extensions/sycl_ext"
+)
diff --git a/dpnp/backend/extensions/sycl_ext/dispatcher_utils.hpp b/dpnp/backend/extensions/sycl_ext/dispatcher_utils.hpp
new file mode 100644
index 00000000000..4bbb7a3c367
--- /dev/null
+++ b/dpnp/backend/extensions/sycl_ext/dispatcher_utils.hpp
@@ -0,0 +1,558 @@
+#include <array>
+#include <tuple>
+
+namespace dpnp
+{
+namespace backend
+{
+namespace ext
+{
+namespace sycl_ext
+{
+
+template <template <typename...> class C, typename Tuple>
+struct resolve_type
+{
+};
+
+template <template <typename...> class C, typename... Args>
+struct resolve_type<C, std::tuple<Args...>>
+{
+    using type = C<Args...>;
+};
+
+template <typename Tuple>
+struct pop_tuple_type
+{
+};
+
+template <typename T, typename... Args>
+struct pop_tuple_type<std::tuple<T, Args...>>
+{
+    using element = T;
+    using rest = std::tuple<Args...>;
+};
+
+template <typename FirstTuple, typename SecondTuple>
+struct merge_two_tuple_types
+{
+};
+
+template <typename... FirstTupleArgs, typename... SecondTupleArgs>
+struct merge_two_tuple_types<std::tuple<FirstTupleArgs...>,
+                             std::tuple<SecondTupleArgs...>>
+{
+    using type = std::tuple<FirstTupleArgs..., SecondTupleArgs...>;
+};
+
+template <typename FirstTuple, typename SecondTuple>
+struct merge_two_inner_tuple_types
+{
+};
+
+template <typename... FirstTupleArgs, typename... SecondTupleArgs>
+struct merge_two_inner_tuple_types<std::tuple<std::tuple<FirstTupleArgs...>>,
+                                   std::tuple<std::tuple<SecondTupleArgs...>>>
+{
+    using type = std::tuple<std::tuple<FirstTupleArgs..., SecondTupleArgs...>>;
+};
+
+template <typename... FirstTupleArgs, typename... SecondTupleArgs>
+struct merge_two_inner_tuple_types<std::tuple<FirstTupleArgs...>,
+                                   std::tuple<SecondTupleArgs...>>
+{
+    using T1 = typename pop_tuple_type<std::tuple<FirstTupleArgs...>>::element;
+    using rest1 = typename pop_tuple_type<std::tuple<FirstTupleArgs...>>::rest;
+
+    using T2 = typename pop_tuple_type<std::tuple<SecondTupleArgs...>>::element;
+    using rest2 = typename pop_tuple_type<std::tuple<SecondTupleArgs...>>::rest;
+
+    using merged_first = typename merge_two_tuple_types<T1, T2>::type;
+    using merged_rest =
+        typename merge_two_inner_tuple_types<rest1, rest2>::type;
+
+    using type = typename merge_two_tuple_types<std::tuple<merged_first>,
+                                                merged_rest>::type;
+};
+
+template <typename Tuple, typename T>
+struct extend_tuple_type
+{
+};
+
+template <typename... Args, typename T>
+struct extend_tuple_type<std::tuple<Args...>, T>
+{
+    using type = std::tuple<Args..., T>;
+};
+
+template <typename... Args>
+struct tuple_of_tuples
+{
+};
+
+template <typename T>
+struct tuple_of_tuples<std::tuple<T>>
+{
+    using type = std::tuple<std::tuple<T>>;
+};
+
+template <typename... Args>
+struct tuple_of_tuples<std::tuple<Args...>>
+{
+    using T = typename pop_tuple_type<std::tuple<Args...>>::element;
+    using other = typename pop_tuple_type<std::tuple<Args...>>::rest;
+
+    using tuple_t = std::tuple<T>;
+    using tuple_other = typename tuple_of_tuples<other>::type;
+    using type =
+        typename merge_two_tuple_types<std::tuple<tuple_t>, tuple_other>::type;
+};
+
+template <typename T, int N>
+struct tuple_of_n_elements
+{
+    using type =
+        typename extend_tuple_type<typename tuple_of_n_elements<T, N - 1>::type,
+                                   T>::type;
+};
+
+template <typename T>
+struct tuple_of_n_elements<T, 1>
+{
+    using type = std::tuple<T>;
+};
+
+template <typename T>
+struct tuple_of_n_elements<T, 0>
+{
+    using type = std::tuple<>;
+};
+
+template <typename... Args>
+struct cartesian_product_impl
+{
+};
+
+template <typename T, typename... SecondTupleArgs>
+struct cartesian_product_impl<std::tuple<std::tuple<T>>,
+                              std::tuple<SecondTupleArgs...>>
+{
+    using extended_T_tuple =
+        typename tuple_of_n_elements<T, sizeof...(SecondTupleArgs)>::type;
+    using extended_T_tuple_of_tuples =
+        typename tuple_of_tuples<extended_T_tuple>::type;
+    using type = typename merge_two_inner_tuple_types<
+        extended_T_tuple_of_tuples,
+        std::tuple<SecondTupleArgs...>>::type;
+};
+
+template <typename T, typename... FirstTupleArgs, typename... SecondTupleArgs>
+struct cartesian_product_impl<std::tuple<std::tuple<T>, FirstTupleArgs...>,
+                              std::tuple<SecondTupleArgs...>>
+{
+    using extended_T_tuple =
+        typename tuple_of_n_elements<T, sizeof...(SecondTupleArgs)>::type;
+    using extended_T_tuple_of_tuples =
+        typename tuple_of_tuples<extended_T_tuple>::type;
+    using cartesian_with_first = typename merge_two_inner_tuple_types<
+        extended_T_tuple_of_tuples,
+        std::tuple<SecondTupleArgs...>>::type;
+    using cartesian_with_other =
+        typename cartesian_product_impl<std::tuple<FirstTupleArgs...>,
+                                        std::tuple<SecondTupleArgs...>>::type;
+    using type = typename merge_two_tuple_types<cartesian_with_first,
+                                                cartesian_with_other>::type;
+};
+
+template <typename... Args>
+struct cartesian_product
+{
+};
+
+template <typename... FirstTupleArgs>
+struct cartesian_product<std::tuple<std::tuple<FirstTupleArgs...>>>
+{
+    using type = typename tuple_of_tuples<std::tuple<FirstTupleArgs...>>::type;
+};
+
+template <typename... Tuples>
+struct cartesian_product<std::tuple<Tuples...>>
+{
+    using first_tuple = typename pop_tuple_type<std::tuple<Tuples...>>::element;
+    using other = typename pop_tuple_type<std::tuple<Tuples...>>::rest;
+    using cartesian_other = typename cartesian_product<other>::type;
+    using first_tuple_of_tuples = typename tuple_of_tuples<first_tuple>::type;
+    using type = typename cartesian_product_impl<first_tuple_of_tuples,
+                                                 cartesian_other>::type;
+};
+
+template <template <typename...> class C, typename FnT, typename Tuple>
+struct populate_table
+{
+};
+
+template <template <typename...> class C, typename FnT, typename... Args>
+struct populate_table<C, FnT, std::tuple<std::tuple<Args...>>>
+{
+    static void populate(FnT *data)
+    {
+        data[0] = resolve_type<C, std::tuple<Args...>>::type::call;
+    }
+};
+
+template <template <typename...> class C, typename FnT, typename... Args>
+struct populate_table<C, FnT, std::tuple<Args...>>
+{
+    static void populate(FnT *data)
+    {
+        using T = typename pop_tuple_type<std::tuple<Args...>>::element;
+        using rest = typename pop_tuple_type<std::tuple<Args...>>::rest;
+
+        data[0] = resolve_type<C, T>::type::call;
+        populate_table<C, FnT, rest>::populate(data + 1);
+    }
+};
+
+template <size_t idx, typename DispatchT, typename Matcher, typename Tuple>
+struct index_on_axis
+{
+};
+
+template <size_t idx, typename DispatchT, typename Matcher, typename T>
+struct index_on_axis<idx, DispatchT, Matcher, std::tuple<T>>
+{
+    static int get_index(const DispatchT &val, int *result)
+    {
+        if (Matcher::template match<T>(val)) {
+            result[0] = idx;
+            return true;
+        }
+
+        result[0] = idx + 1;
+        return false;
+    }
+};
+
+template <size_t idx,
+          typename DispatchT,
+          typename Matcher,
+          typename... AxisData>
+struct index_on_axis<idx, DispatchT, Matcher, std::tuple<AxisData...>>
+{
+    using T = typename pop_tuple_type<std::tuple<AxisData...>>::element;
+    using rest = typename pop_tuple_type<std::tuple<AxisData...>>::rest;
+
+    static int get_index(const DispatchT &val, int *result)
+    {
+        if (Matcher::template match<T>(val)) {
+            result[0] = idx;
+            return true;
+        }
+
+        return index_on_axis<idx + 1, DispatchT, Matcher, rest>::get_index(
+            val, result);
+    }
+};
+
+template <typename DispatchT, typename Matcher, typename... AxisData>
+struct coord_in_space
+{
+};
+
+template <typename DispatchT, typename Matcher, typename CurrentAxis>
+struct coord_in_space<DispatchT, Matcher, std::tuple<CurrentAxis>>
+{
+    static bool get(const DispatchT *data, int *result)
+    {
+        auto found =
+            index_on_axis<0, DispatchT, Matcher, CurrentAxis>::get_index(
+                data[0], result);
+
+        return found;
+    }
+};
+
+template <typename DispatchT, typename Matcher, typename... Axises>
+struct coord_in_space<DispatchT, Matcher, std::tuple<Axises...>>
+{
+    using CurrentAxis = typename pop_tuple_type<std::tuple<Axises...>>::element;
+    using OtherAxises = typename pop_tuple_type<std::tuple<Axises...>>::rest;
+
+    static bool get(const DispatchT *data, int *result)
+    {
+        auto found =
+            index_on_axis<0, DispatchT, Matcher, CurrentAxis>::get_index(
+                data[0], result);
+        if (found) {
+            found = coord_in_space<DispatchT, Matcher, OtherAxises>::get(
+                data + 1, result + 1);
+
+            return found;
+        }
+
+        return found;
+    }
+};
+
+template <typename Tuple>
+struct elements_count
+{
+};
+
+template <typename T>
+struct elements_count<std::tuple<T>>
+{
+    static constexpr int count()
+    {
+        return std::tuple_size_v<T>;
+    }
+};
+
+template <typename T, typename... Axis>
+struct elements_count<std::tuple<T, Axis...>>
+{
+    static constexpr int count()
+    {
+        return std::tuple_size_v<T> *
+               elements_count<std::tuple<Axis...>>::count();
+    }
+};
+
+template <typename Tuple>
+struct get_linear_id
+{
+};
+
+template <typename T>
+struct get_linear_id<std::tuple<T>>
+{
+    static int get(const int *data)
+    {
+        return data[0];
+    }
+};
+
+template <typename T, typename... Axis>
+struct get_linear_id<std::tuple<T, Axis...>>
+{
+    static int get(const int *data)
+    {
+        return data[0] * elements_count<std::tuple<Axis...>>::count() +
+               get_linear_id<std::tuple<Axis...>>::get(data + 1);
+    }
+};
+
+template <template <typename...> class C,
+          typename FnT,
+          typename DispatchT,
+          typename Matcher,
+          typename... Axis>
+struct CartesianDispatcher
+{
+    using space = typename cartesian_product<std::tuple<Axis...>>::type;
+
+    std::array<FnT, std::tuple_size_v<space>> dispatch_table;
+
+    CartesianDispatcher()
+    {
+        populate_table<C, FnT, space>::populate(dispatch_table.data());
+    }
+
+    FnT operator()(std::initializer_list<DispatchT> values)
+    {
+        // though std::initializer_list::size is constexpr it doesnt work for
+        // some reson constexpr auto const list_size = values.size();
+        std::array<int, sizeof...(Axis)> idx;
+        bool found =
+            coord_in_space<DispatchT, Matcher, std::tuple<Axis...>>::get(
+                std::data(values), idx.data());
+
+        if (found) {
+            int linear_idx =
+                get_linear_id<std::tuple<Axis...>>::get(idx.data());
+            return dispatch_table[linear_idx];
+        }
+
+        return nullptr;
+    }
+};
+
+struct UsmArrayMatcher
+{
+    template <class T>
+    static bool match(const dpctl::tensor::usm_ndarray &)
+    {
+        return false;
+    }
+};
+
+// int typenum_to_lookup_id(int typenum) const
+// {
+//     using typenum_t = ::dpctl::tensor::type_dispatch::typenum_t;
+//     auto const &api = ::dpctl::detail::dpctl_capi::get();
+
+//     if (typenum == api.UAR_DOUBLE_) {
+//         return static_cast<int>(typenum_t::DOUBLE);
+//     }
+//     else if (typenum == api.UAR_INT64_) {
+//         return static_cast<int>(typenum_t::INT64);
+//     }
+//     else if (typenum == api.UAR_INT32_) {
+//         return static_cast<int>(typenum_t::INT32);
+//     }
+//     else if (typenum == api.UAR_BOOL_) {
+//         return static_cast<int>(typenum_t::BOOL);
+//     }
+//     else if (typenum == api.UAR_CDOUBLE_) {
+//         return static_cast<int>(typenum_t::CDOUBLE);
+//     }
+//     else if (typenum == api.UAR_FLOAT_) {
+//         return static_cast<int>(typenum_t::FLOAT);
+//     }
+//     else if (typenum == api.UAR_INT16_) {
+//         return static_cast<int>(typenum_t::INT16);
+//     }
+//     else if (typenum == api.UAR_INT8_) {
+//         return static_cast<int>(typenum_t::INT8);
+//     }
+//     else if (typenum == api.UAR_UINT64_) {
+//         return static_cast<int>(typenum_t::UINT64);
+//     }
+//     else if (typenum == api.UAR_UINT32_) {
+//         return static_cast<int>(typenum_t::UINT32);
+//     }
+//     else if (typenum == api.UAR_UINT16_) {
+//         return static_cast<int>(typenum_t::UINT16);
+//     }
+//     else if (typenum == api.UAR_UINT8_) {
+//         return static_cast<int>(typenum_t::UINT8);
+//     }
+//     else if (typenum == api.UAR_CFLOAT_) {
+//         return static_cast<int>(typenum_t::CFLOAT);
+//     }
+//     else if (typenum == api.UAR_HALF_) {
+//         return static_cast<int>(typenum_t::HALF);
+//     }
+//     else if (typenum == api.UAR_INT_ || typenum == api.UAR_UINT_) {
+//         switch (sizeof(int)) {
+//         case sizeof(std::int32_t):
+//             return ((typenum == api.UAR_INT_)
+//                         ? static_cast<int>(typenum_t::INT32)
+//                         : static_cast<int>(typenum_t::UINT32));
+//         case sizeof(std::int64_t):
+//             return ((typenum == api.UAR_INT_)
+//                         ? static_cast<int>(typenum_t::INT64)
+//                         : static_cast<int>(typenum_t::UINT64));
+//         default:
+//             throw_unrecognized_typenum_error(typenum);
+//         }
+//     }
+//     else if (typenum == api.UAR_LONGLONG_ || typenum == api.UAR_ULONGLONG_)
+//     {
+//         switch (sizeof(long long)) {
+//         case sizeof(std::int64_t):
+//             return ((typenum == api.UAR_LONGLONG_)
+//                         ? static_cast<int>(typenum_t::INT64)
+//                         : static_cast<int>(typenum_t::UINT64));
+//         default:
+//             throw_unrecognized_typenum_error(typenum);
+//         }
+//     }
+//     else {
+//         throw_unrecognized_typenum_error(typenum);
+//     }
+//     // return code signalling error, should never be reached
+//     assert(false);
+//     return -1;
+// }
+
+template <>
+bool UsmArrayMatcher::match<int8_t>(const dpctl::tensor::usm_ndarray &arr)
+{
+    return arr.get_typenum() == UAR_BYTE;
+}
+
+template <>
+bool UsmArrayMatcher::match<uint8_t>(const dpctl::tensor::usm_ndarray &arr)
+{
+    return arr.get_typenum() == UAR_UBYTE;
+}
+
+template <>
+bool UsmArrayMatcher::match<int16_t>(const dpctl::tensor::usm_ndarray &arr)
+{
+    return arr.get_typenum() == UAR_SHORT;
+}
+
+template <>
+bool UsmArrayMatcher::match<uint16_t>(const dpctl::tensor::usm_ndarray &arr)
+{
+    return arr.get_typenum() == UAR_USHORT;
+}
+
+template <>
+bool UsmArrayMatcher::match<int32_t>(const dpctl::tensor::usm_ndarray &arr)
+{
+    return (arr.get_typenum() == UAR_INT and sizeof(int) == sizeof(int32_t)) or
+           (arr.get_typenum() == UAR_LONG and sizeof(long) == sizeof(int32_t));
+}
+
+template <>
+bool UsmArrayMatcher::match<uint32_t>(const dpctl::tensor::usm_ndarray &arr)
+{
+    return (arr.get_typenum() == UAR_UINT and
+            sizeof(unsigned int) == sizeof(uint32_t)) or
+           (arr.get_typenum() == UAR_ULONG and
+            sizeof(unsigned long) == sizeof(uint32_t));
+}
+
+template <>
+bool UsmArrayMatcher::match<int64_t>(const dpctl::tensor::usm_ndarray &arr)
+{
+    return (arr.get_typenum() == UAR_INT and sizeof(int) == sizeof(int64_t)) or
+           (arr.get_typenum() == UAR_LONG and
+            sizeof(long) == sizeof(int64_t)) or
+           (arr.get_typenum() == UAR_LONGLONG and
+            sizeof(long long) == sizeof(int64_t));
+}
+
+template <>
+bool UsmArrayMatcher::match<uint64_t>(const dpctl::tensor::usm_ndarray &arr)
+{
+    return (arr.get_typenum() == UAR_UINT and
+            sizeof(unsigned int) == sizeof(uint64_t)) or
+           (arr.get_typenum() == UAR_ULONG and
+            sizeof(unsigned long) == sizeof(uint64_t)) or
+           (arr.get_typenum() == UAR_ULONGLONG and
+            sizeof(unsigned long long) == sizeof(uint64_t));
+}
+
+template <>
+bool UsmArrayMatcher::match<float>(const dpctl::tensor::usm_ndarray &arr)
+{
+    return arr.get_typenum() == UAR_FLOAT;
+}
+
+template <>
+bool UsmArrayMatcher::match<double>(const dpctl::tensor::usm_ndarray &arr)
+{
+    return arr.get_typenum() == UAR_DOUBLE;
+}
+
+template <class NumT, class DenT>
+NumT DivUp(NumT numerator, DenT denominator)
+{
+    return (numerator + denominator - 1) / denominator;
+}
+
+template <class VT, class BT>
+VT Align(VT value, BT base)
+{
+    return base * DivUp(value, base);
+}
+
+} // namespace sycl_ext
+} // namespace ext
+} // namespace backend
+} // namespace dpnp
diff --git a/dpnp/backend/extensions/sycl_ext/sum_mean.cpp b/dpnp/backend/extensions/sycl_ext/sum_mean.cpp
new file mode 100644
index 00000000000..f8cc57d9ef4
--- /dev/null
+++ b/dpnp/backend/extensions/sycl_ext/sum_mean.cpp
@@ -0,0 +1,113 @@
+//*****************************************************************************
+// Copyright (c) 2023, Intel Corporation
+// All rights reserved.
+//
+// Redistribution and use in source and binary forms, with or without
+// modification, are permitted provided that the following conditions are met:
+// - Redistributions of source code must retain the above copyright notice,
+//   this list of conditions and the following disclaimer.
+// - Redistributions in binary form must reproduce the above copyright notice,
+//   this list of conditions and the following disclaimer in the documentation
+//   and/or other materials provided with the distribution.
+//
+// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+// AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+// IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
+// ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
+// LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
+// CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
+// SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
+// INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
+// CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
+// ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF
+// THE POSSIBILITY OF SUCH DAMAGE.
+//*****************************************************************************
+//
+// This file defines functions of dpnp.backend._sycl_ext_impl extensions
+//
+//*****************************************************************************
+
+#include <memory>
+
+#include <pybind11/functional.h>
+#include <pybind11/pybind11.h>
+#include <pybind11/stl.h>
+
+#include "dpctl4pybind11.hpp"
+#include "sum_mean.hpp"
+
+namespace sycl_ext = dpnp::backend::ext::sycl_ext;
+namespace py = pybind11;
+
+using dpctl::tensor::usm_ndarray;
+
+std::unique_ptr<sycl_ext::SumOverAxisContigDispatcher> sum_dispatcher;
+std::unique_ptr<sycl_ext::MeanOverAxisContigDispatcher> mean_dispatcher;
+
+void init_dispatchers(void)
+{
+    sum_dispatcher.reset(new sycl_ext::SumOverAxisContigDispatcher());
+    mean_dispatcher.reset(new sycl_ext::MeanOverAxisContigDispatcher());
+}
+
+using SumMeanFnT = sycl::event (*)(const usm_ndarray &,
+                                   usm_ndarray &,
+                                   const std::vector<sycl::event> &);
+
+sycl::event sum_over_axis(usm_ndarray input,
+                          usm_ndarray output,
+                          std::vector<sycl::event> depends)
+{
+    auto sum_fn = (*sum_dispatcher)({input, output});
+
+    if (sum_fn == nullptr)
+        throw py::value_error("No suitable implementation found");
+
+    return sum_fn(input, output, depends);
+}
+
+sycl::event mean_over_axis(usm_ndarray input,
+                           usm_ndarray output,
+                           std::vector<sycl::event> depends)
+{
+    auto mean_fn = (*mean_dispatcher)({input, output});
+
+    if (mean_fn == nullptr)
+        throw py::value_error("No suitable implementation found");
+
+    return mean_fn(input, output, depends);
+}
+
+py::cpp_function get_sum_over_axis(usm_ndarray input, usm_ndarray output)
+{
+    if (not sycl_ext::check_limitations(input, output))
+        return nullptr;
+
+    return (*sum_dispatcher)({input, output});
+}
+
+py::cpp_function get_mean_over_axis(usm_ndarray input, usm_ndarray output)
+{
+    if (not sycl_ext::check_limitations(input, output))
+        return nullptr;
+
+    return (*mean_dispatcher)({input, output});
+}
+
+PYBIND11_MODULE(_sycl_ext_impl, m)
+{
+    import_dpctl();
+    init_dispatchers();
+
+    m.def("_sum_over_axis_0", &sum_over_axis, "Sum over axis 0",
+          py::arg("input"), py::arg("output"), py::arg("depends") = py::list());
+
+    m.def("_mean_over_axis_0", &mean_over_axis, "Mean over axis 0",
+          py::arg("input"), py::arg("output"), py::arg("depends") = py::list());
+
+    m.def("_get_sum_over_axis_0", &get_sum_over_axis, "Sum over axis 0",
+          py::arg("input"), py::arg("output"));
+
+    m.def("_get_mean_over_axis_0", &get_mean_over_axis, "Mean over axis 0",
+          py::arg("input"), py::arg("output"));
+}
diff --git a/dpnp/backend/extensions/sycl_ext/sum_mean.hpp b/dpnp/backend/extensions/sycl_ext/sum_mean.hpp
new file mode 100644
index 00000000000..6437c783c85
--- /dev/null
+++ b/dpnp/backend/extensions/sycl_ext/sum_mean.hpp
@@ -0,0 +1,209 @@
+//*****************************************************************************
+// Copyright (c) 2023, Intel Corporation
+// All rights reserved.
+//
+// Redistribution and use in source and binary forms, with or without
+// modification, are permitted provided that the following conditions are met:
+// - Redistributions of source code must retain the above copyright notice,
+//   this list of conditions and the following disclaimer.
+// - Redistributions in binary form must reproduce the above copyright notice,
+//   this list of conditions and the following disclaimer in the documentation
+//   and/or other materials provided with the distribution.
+//
+// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+// AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+// IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
+// ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
+// LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
+// CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
+// SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
+// INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
+// CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
+// ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF
+// THE POSSIBILITY OF SUCH DAMAGE.
+//*****************************************************************************
+
+#pragma once
+
+#include "dispatcher_utils.hpp"
+#include <CL/sycl.hpp>
+#include <tuple>
+
+#include "utils/memory_overlap.hpp"
+#include "utils/type_utils.hpp"
+
+namespace dpnp
+{
+namespace backend
+{
+namespace ext
+{
+namespace sycl_ext
+{
+
+namespace py = pybind11;
+namespace type_utils = dpctl::tensor::type_utils;
+
+bool check_limitations(const dpctl::tensor::usm_ndarray &in,
+                       const dpctl::tensor::usm_ndarray &out,
+                       bool throw_on_fail = false)
+{
+    if (not in.is_c_contiguous() or not out.is_c_contiguous()) {
+        if (throw_on_fail)
+            throw py::value_error(
+                "Input and output arrays must be c-contiguos");
+
+        return false;
+    }
+
+    // auto local_mem_size = in.get_sycl_device().local_mem_size();
+    // auto out_full_size = out.get_size()*out.get_elemsize();
+    // if (out_full_size > local_mem_size)
+    // {
+    //     if (throw_on_fail)
+    //         throw py::value_error("Resulting array exceeds local memroy size
+    //         " + std::to_string(local_mem_size));
+
+    //     return false;
+    // }
+
+    return true;
+}
+
+using SumMeanFuncT = sycl::event (*)(const dpctl::tensor::usm_ndarray &,
+                                     dpctl::tensor::usm_ndarray &,
+                                     const std::vector<sycl::event> &);
+
+template <typename InT, typename OutT, bool Mean>
+struct sum_mean
+{
+    static sycl::event call(sycl::queue exec_q,
+                            const InT *in,
+                            OutT *out,
+                            const size_t height,
+                            const size_t width,
+                            const std::vector<sycl::event> &depends)
+    {
+        int local_size = std::min(
+            1024,
+            int(exec_q.get_device()
+                    .get_info<sycl::info::device::max_work_group_size>()));
+
+        int WorkPI = 32; // empirically found number
+        auto _total_size = height * width;
+        auto global_size = Align(DivUp(_total_size, WorkPI), local_size);
+
+        auto fill_e = exec_q.fill(out, OutT(0), width, depends);
+
+        return exec_q.submit([&](sycl::handler &cgh) {
+            cgh.depends_on(fill_e);
+            auto local_result =
+                sycl::local_accessor<OutT>(sycl::range<1>(width), cgh);
+            cgh.parallel_for<sum_mean<InT, OutT, Mean>>(
+                sycl::nd_range<1>(global_size, local_size),
+                [=](sycl::nd_item<1> item) {
+                    auto grid = item.get_group_linear_id();
+                    auto group = item.get_group();
+                    auto llid = item.get_local_linear_id();
+
+                    for (size_t i = llid; i < width; i += local_size) {
+                        local_result[i] = 0;
+                    }
+
+                    sycl::group_barrier(group);
+
+                    for (int i = 0; i < WorkPI; ++i) {
+                        auto id =
+                            grid * WorkPI * local_size + i * local_size + llid;
+                        if (id < _total_size) {
+                            sycl::atomic_ref<OutT, sycl::memory_order::relaxed,
+                                             sycl::memory_scope::work_group>
+                                aresult(local_result[id % width]);
+                            aresult += in[id];
+                        }
+                    }
+
+                    sycl::group_barrier(group);
+
+                    for (size_t i = llid; i < width; i += local_size) {
+                        sycl::atomic_ref<OutT, sycl::memory_order::relaxed,
+                                         sycl::memory_scope::device>
+                            autput(out[i]);
+
+                        auto r = local_result[i];
+                        autput += Mean ? r / height : r;
+                    }
+                });
+        });
+    }
+
+    static sycl::event call(const dpctl::tensor::usm_ndarray &in,
+                            dpctl::tensor::usm_ndarray &out,
+                            const std::vector<sycl::event> &depends)
+    {
+        validate_params(in, out);
+        return call(in.get_queue(), in.get_data<InT>(), out.get_data<OutT>(),
+                    in.get_shape(0), in.get_shape(1), depends);
+    }
+
+    static void validate_params(const dpctl::tensor::usm_ndarray &in,
+                                dpctl::tensor::usm_ndarray &out)
+    {
+        auto exec_q = in.get_queue();
+
+        type_utils::validate_type_for_device<InT>(exec_q);
+        type_utils::validate_type_for_device<OutT>(exec_q);
+
+        // check compatibility of execution queue and allocation queue
+        if (!dpctl::utils::queues_are_compatible(exec_q, {in, out})) {
+            throw py::value_error(
+                "Execution queue is not compatible with allocation queues");
+        }
+
+        if (in.get_shape(1) != out.get_shape(0))
+            throw py::value_error(
+                "Input array axis 1 size must match output array size");
+
+        check_limitations(in, out, true);
+    }
+};
+
+template <typename InT, typename OutT>
+using SumOverAxisContig = sum_mean<InT, OutT, false>;
+using SumInputTypes = std::tuple<int8_t,
+                                 uint8_t,
+                                 int16_t,
+                                 uint16_t,
+                                 int32_t,
+                                 uint32_t,
+                                 int64_t,
+                                 uint64_t,
+                                 float,
+                                 double>;
+using SumOutputTypes =
+    std::tuple<int32_t, uint32_t, int64_t, uint64_t, float, double>;
+
+template <typename InT, typename OutT>
+using MeanOverAxisContig = sum_mean<InT, OutT, true>;
+using MeanInputTypes = SumInputTypes;
+using MeanOutputTypes = std::tuple<float, double>;
+
+using SumOverAxisContigDispatcher =
+    CartesianDispatcher<SumOverAxisContig,
+                        SumMeanFuncT,
+                        dpctl::tensor::usm_ndarray,
+                        UsmArrayMatcher,
+                        SumInputTypes,
+                        SumOutputTypes>;
+
+using MeanOverAxisContigDispatcher =
+    CartesianDispatcher<MeanOverAxisContig,
+                        SumMeanFuncT,
+                        dpctl::tensor::usm_ndarray,
+                        UsmArrayMatcher,
+                        SumInputTypes,
+                        SumOutputTypes>;
+} // namespace sycl_ext
+} // namespace ext
+} // namespace backend
+} // namespace dpnp
diff --git a/dpnp/dpnp_iface_mathematical.py b/dpnp/dpnp_iface_mathematical.py
index 26c033032ea..23089444ffb 100644
--- a/dpnp/dpnp_iface_mathematical.py
+++ b/dpnp/dpnp_iface_mathematical.py
@@ -1814,6 +1814,13 @@ def sum(
 
     """
 
+    if not isinstance(axis, (tuple, list)):
+        axis = (axis,)
+
+    from numpy.core.numeric import normalize_axis_tuple
+
+    axis = normalize_axis_tuple(axis, x.ndim, "axis")
+
     if out is not None:
         pass
     elif initial != 0:
@@ -1821,6 +1828,24 @@ def sum(
     elif where is not True:
         pass
     else:
+        if (axis == (0,) and len(x.shape) == 2):
+            from dpctl.tensor._reduction import _default_reduction_dtype
+
+            from dpnp.backend.extensions.sycl_ext import _sycl_ext_impl
+
+            input = dpnp.get_usm_ndarray(x)
+
+            queue = input.sycl_queue
+            out_dtype = _default_reduction_dtype(input.dtype, queue)
+            output = dpt.empty(input.shape[1], dtype=out_dtype, sycl_queue=queue)
+
+            get_sum = _sycl_ext_impl._get_sum_over_axis_0
+            sum = get_sum(input, output)
+
+            if sum:
+                sum(input, output, []).wait()
+                return dpnp_array._create_from_usm_ndarray(output)
+
         y = dpt.sum(
             dpnp.get_usm_ndarray(x), axis=axis, dtype=dtype, keepdims=keepdims
         )

From fdec1ee26eef5fa4a96c61c063783a27ba75cdd1 Mon Sep 17 00:00:00 2001
From: Alexander Kalistratov <alexander.kalistratov@intel.com>
Date: Mon, 17 Jul 2023 20:57:31 +0200
Subject: [PATCH 2/3] Cleanups and fixes

---
 .../extensions/sycl_ext/dispatcher_utils.hpp  | 107 +++++-------------
 dpnp/backend/extensions/sycl_ext/sum_mean.hpp |  27 +++--
 dpnp/dpnp_iface_mathematical.py               |  16 +--
 3 files changed, 53 insertions(+), 97 deletions(-)

diff --git a/dpnp/backend/extensions/sycl_ext/dispatcher_utils.hpp b/dpnp/backend/extensions/sycl_ext/dispatcher_utils.hpp
index 4bbb7a3c367..57f1a4e43a0 100644
--- a/dpnp/backend/extensions/sycl_ext/dispatcher_utils.hpp
+++ b/dpnp/backend/extensions/sycl_ext/dispatcher_utils.hpp
@@ -1,3 +1,30 @@
+//*****************************************************************************
+// Copyright (c) 2023, Intel Corporation
+// All rights reserved.
+//
+// Redistribution and use in source and binary forms, with or without
+// modification, are permitted provided that the following conditions are met:
+// - Redistributions of source code must retain the above copyright notice,
+//   this list of conditions and the following disclaimer.
+// - Redistributions in binary form must reproduce the above copyright notice,
+//   this list of conditions and the following disclaimer in the documentation
+//   and/or other materials provided with the distribution.
+//
+// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+// AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+// IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
+// ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
+// LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
+// CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
+// SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
+// INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
+// CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
+// ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF
+// THE POSSIBILITY OF SUCH DAMAGE.
+//*****************************************************************************
+
+#pragma once
+
 #include <array>
 #include <tuple>
 
@@ -387,86 +414,6 @@ struct UsmArrayMatcher
     }
 };
 
-// int typenum_to_lookup_id(int typenum) const
-// {
-//     using typenum_t = ::dpctl::tensor::type_dispatch::typenum_t;
-//     auto const &api = ::dpctl::detail::dpctl_capi::get();
-
-//     if (typenum == api.UAR_DOUBLE_) {
-//         return static_cast<int>(typenum_t::DOUBLE);
-//     }
-//     else if (typenum == api.UAR_INT64_) {
-//         return static_cast<int>(typenum_t::INT64);
-//     }
-//     else if (typenum == api.UAR_INT32_) {
-//         return static_cast<int>(typenum_t::INT32);
-//     }
-//     else if (typenum == api.UAR_BOOL_) {
-//         return static_cast<int>(typenum_t::BOOL);
-//     }
-//     else if (typenum == api.UAR_CDOUBLE_) {
-//         return static_cast<int>(typenum_t::CDOUBLE);
-//     }
-//     else if (typenum == api.UAR_FLOAT_) {
-//         return static_cast<int>(typenum_t::FLOAT);
-//     }
-//     else if (typenum == api.UAR_INT16_) {
-//         return static_cast<int>(typenum_t::INT16);
-//     }
-//     else if (typenum == api.UAR_INT8_) {
-//         return static_cast<int>(typenum_t::INT8);
-//     }
-//     else if (typenum == api.UAR_UINT64_) {
-//         return static_cast<int>(typenum_t::UINT64);
-//     }
-//     else if (typenum == api.UAR_UINT32_) {
-//         return static_cast<int>(typenum_t::UINT32);
-//     }
-//     else if (typenum == api.UAR_UINT16_) {
-//         return static_cast<int>(typenum_t::UINT16);
-//     }
-//     else if (typenum == api.UAR_UINT8_) {
-//         return static_cast<int>(typenum_t::UINT8);
-//     }
-//     else if (typenum == api.UAR_CFLOAT_) {
-//         return static_cast<int>(typenum_t::CFLOAT);
-//     }
-//     else if (typenum == api.UAR_HALF_) {
-//         return static_cast<int>(typenum_t::HALF);
-//     }
-//     else if (typenum == api.UAR_INT_ || typenum == api.UAR_UINT_) {
-//         switch (sizeof(int)) {
-//         case sizeof(std::int32_t):
-//             return ((typenum == api.UAR_INT_)
-//                         ? static_cast<int>(typenum_t::INT32)
-//                         : static_cast<int>(typenum_t::UINT32));
-//         case sizeof(std::int64_t):
-//             return ((typenum == api.UAR_INT_)
-//                         ? static_cast<int>(typenum_t::INT64)
-//                         : static_cast<int>(typenum_t::UINT64));
-//         default:
-//             throw_unrecognized_typenum_error(typenum);
-//         }
-//     }
-//     else if (typenum == api.UAR_LONGLONG_ || typenum == api.UAR_ULONGLONG_)
-//     {
-//         switch (sizeof(long long)) {
-//         case sizeof(std::int64_t):
-//             return ((typenum == api.UAR_LONGLONG_)
-//                         ? static_cast<int>(typenum_t::INT64)
-//                         : static_cast<int>(typenum_t::UINT64));
-//         default:
-//             throw_unrecognized_typenum_error(typenum);
-//         }
-//     }
-//     else {
-//         throw_unrecognized_typenum_error(typenum);
-//     }
-//     // return code signalling error, should never be reached
-//     assert(false);
-//     return -1;
-// }
-
 template <>
 bool UsmArrayMatcher::match<int8_t>(const dpctl::tensor::usm_ndarray &arr)
 {
diff --git a/dpnp/backend/extensions/sycl_ext/sum_mean.hpp b/dpnp/backend/extensions/sycl_ext/sum_mean.hpp
index 6437c783c85..9b7c12db6f3 100644
--- a/dpnp/backend/extensions/sycl_ext/sum_mean.hpp
+++ b/dpnp/backend/extensions/sycl_ext/sum_mean.hpp
@@ -56,16 +56,23 @@ bool check_limitations(const dpctl::tensor::usm_ndarray &in,
         return false;
     }
 
-    // auto local_mem_size = in.get_sycl_device().local_mem_size();
-    // auto out_full_size = out.get_size()*out.get_elemsize();
-    // if (out_full_size > local_mem_size)
-    // {
-    //     if (throw_on_fail)
-    //         throw py::value_error("Resulting array exceeds local memroy size
-    //         " + std::to_string(local_mem_size));
-
-    //     return false;
-    // }
+    auto device = in.get_queue().get_device();
+    auto local_mem_size = device.get_info<sycl::info::device::local_mem_size>();
+    size_t out_full_size = out.get_size() * out.get_elemsize();
+    if (out_full_size > local_mem_size) {
+        if (throw_on_fail)
+            throw py::value_error("Resulting array exceeds local memroy size" +
+                                  std::to_string(local_mem_size));
+
+        return false;
+    }
+
+    if (out.get_elemsize() == 64 and not device.has(sycl::aspect::atomic64)) {
+        if (throw_on_fail)
+            throw py::value_error("64-bit atomics are not supported");
+
+        return false;
+    }
 
     return true;
 }
diff --git a/dpnp/dpnp_iface_mathematical.py b/dpnp/dpnp_iface_mathematical.py
index 23089444ffb..f4bcce40882 100644
--- a/dpnp/dpnp_iface_mathematical.py
+++ b/dpnp/dpnp_iface_mathematical.py
@@ -42,6 +42,7 @@
 
 import dpctl.tensor as dpt
 import numpy
+from numpy.core.numeric import normalize_axis_tuple
 
 import dpnp
 from dpnp.dpnp_array import dpnp_array
@@ -1814,12 +1815,11 @@ def sum(
 
     """
 
-    if not isinstance(axis, (tuple, list)):
-        axis = (axis,)
+    if axis is not None:
+        if not isinstance(axis, (tuple, list)):
+            axis = (axis,)
 
-    from numpy.core.numeric import normalize_axis_tuple
-
-    axis = normalize_axis_tuple(axis, x.ndim, "axis")
+        axis = normalize_axis_tuple(axis, x.ndim, "axis")
 
     if out is not None:
         pass
@@ -1828,7 +1828,7 @@ def sum(
     elif where is not True:
         pass
     else:
-        if (axis == (0,) and len(x.shape) == 2):
+        if axis == (0,) and len(x.shape) == 2:
             from dpctl.tensor._reduction import _default_reduction_dtype
 
             from dpnp.backend.extensions.sycl_ext import _sycl_ext_impl
@@ -1837,7 +1837,9 @@ def sum(
 
             queue = input.sycl_queue
             out_dtype = _default_reduction_dtype(input.dtype, queue)
-            output = dpt.empty(input.shape[1], dtype=out_dtype, sycl_queue=queue)
+            output = dpt.empty(
+                input.shape[1], dtype=out_dtype, sycl_queue=queue
+            )
 
             get_sum = _sycl_ext_impl._get_sum_over_axis_0
             sum = get_sum(input, output)

From 24df081be94ae7c14def087b21814abf95fdcb33 Mon Sep 17 00:00:00 2001
From: Alexander Kalistratov <alexander.kalistratov@intel.com>
Date: Tue, 18 Jul 2023 00:23:03 +0200
Subject: [PATCH 3/3] Review comments, bug fixing and tests

---
 dpnp/backend/extensions/sycl_ext/sum_mean.cpp |  20 +-
 dpnp/backend/extensions/sycl_ext/sum_mean.hpp |   6 +-
 dpnp/dpnp_iface_mathematical.py               |   8 +-
 tests/test_extensions.py                      | 238 ++++++++++++++++++
 4 files changed, 263 insertions(+), 9 deletions(-)
 create mode 100644 tests/test_extensions.py

diff --git a/dpnp/backend/extensions/sycl_ext/sum_mean.cpp b/dpnp/backend/extensions/sycl_ext/sum_mean.cpp
index f8cc57d9ef4..e5097cd060a 100644
--- a/dpnp/backend/extensions/sycl_ext/sum_mean.cpp
+++ b/dpnp/backend/extensions/sycl_ext/sum_mean.cpp
@@ -78,20 +78,28 @@ sycl::event mean_over_axis(usm_ndarray input,
     return mean_fn(input, output, depends);
 }
 
-py::cpp_function get_sum_over_axis(usm_ndarray input, usm_ndarray output)
+py::object get_sum_over_axis(usm_ndarray input, usm_ndarray output)
 {
     if (not sycl_ext::check_limitations(input, output))
-        return nullptr;
+        return py::none();
 
-    return (*sum_dispatcher)({input, output});
+    auto sum = (*sum_dispatcher)({input, output});
+    if (sum == nullptr)
+        return py::none();
+
+    return py::cpp_function(sum);
 }
 
-py::cpp_function get_mean_over_axis(usm_ndarray input, usm_ndarray output)
+py::object get_mean_over_axis(usm_ndarray input, usm_ndarray output)
 {
     if (not sycl_ext::check_limitations(input, output))
-        return nullptr;
+        return py::none();
+
+    auto mean = (*mean_dispatcher)({input, output});
+    if (mean == nullptr)
+        return py::none();
 
-    return (*mean_dispatcher)({input, output});
+    return py::cpp_function(mean);
 }
 
 PYBIND11_MODULE(_sycl_ext_impl, m)
diff --git a/dpnp/backend/extensions/sycl_ext/sum_mean.hpp b/dpnp/backend/extensions/sycl_ext/sum_mean.hpp
index 9b7c12db6f3..a3779901501 100644
--- a/dpnp/backend/extensions/sycl_ext/sum_mean.hpp
+++ b/dpnp/backend/extensions/sycl_ext/sum_mean.hpp
@@ -171,6 +171,10 @@ struct sum_mean
             throw py::value_error(
                 "Input array axis 1 size must match output array size");
 
+        auto const &overlap = dpctl::tensor::overlap::MemoryOverlap();
+        if (overlap(in, out))
+            throw py::value_error("Input and output array are overlapped");
+
         check_limitations(in, out, true);
     }
 };
@@ -209,7 +213,7 @@ using MeanOverAxisContigDispatcher =
                         dpctl::tensor::usm_ndarray,
                         UsmArrayMatcher,
                         SumInputTypes,
-                        SumOutputTypes>;
+                        MeanOutputTypes>;
 } // namespace sycl_ext
 } // namespace ext
 } // namespace backend
diff --git a/dpnp/dpnp_iface_mathematical.py b/dpnp/dpnp_iface_mathematical.py
index f4bcce40882..2c45f05e2b9 100644
--- a/dpnp/dpnp_iface_mathematical.py
+++ b/dpnp/dpnp_iface_mathematical.py
@@ -1828,7 +1828,7 @@ def sum(
     elif where is not True:
         pass
     else:
-        if axis == (0,) and len(x.shape) == 2:
+        if axis == (0,) and len(x.shape) == 2 and not keepdims:
             from dpctl.tensor._reduction import _default_reduction_dtype
 
             from dpnp.backend.extensions.sycl_ext import _sycl_ext_impl
@@ -1836,7 +1836,11 @@ def sum(
             input = dpnp.get_usm_ndarray(x)
 
             queue = input.sycl_queue
-            out_dtype = _default_reduction_dtype(input.dtype, queue)
+            out_dtype = (
+                _default_reduction_dtype(input.dtype, queue)
+                if dtype is None
+                else dtype
+            )
             output = dpt.empty(
                 input.shape[1], dtype=out_dtype, sycl_queue=queue
             )
diff --git a/tests/test_extensions.py b/tests/test_extensions.py
new file mode 100644
index 00000000000..a020be77637
--- /dev/null
+++ b/tests/test_extensions.py
@@ -0,0 +1,238 @@
+from itertools import product
+
+import dpctl
+import dpctl.tensor as dpt
+import pytest
+
+import dpnp
+from dpnp.backend.extensions.sycl_ext import _sycl_ext_impl
+
+all_devices = [
+    device for device in dpctl.get_devices() if device.is_cpu or device.is_gpu
+]
+sum_supported_input_dtypes = [
+    dpt.dtype("i1"),
+    dpt.dtype("u1"),
+    dpt.dtype("i2"),
+    dpt.dtype("u2"),
+    dpt.dtype("i4"),
+    dpt.dtype("u4"),
+    dpt.dtype("i8"),
+    dpt.dtype("u8"),
+    dpt.float32,
+    dpt.float64,
+]
+sum_supported_output_dtypes = [
+    dpt.dtype("i4"),
+    dpt.dtype("u4"),
+    dpt.dtype("i8"),
+    dpt.dtype("u8"),
+    dpt.float32,
+    dpt.float64,
+]
+
+mean_supported_output_dtypes = [dpt.float32, dpt.float64]
+sum_without_mean_supported_dtypes = list(
+    set(sum_supported_output_dtypes) - set(mean_supported_output_dtypes)
+)
+
+sum_unsupported_input_dtypes = [
+    dpt.bool,
+    dpt.float16,
+    dpt.complex64,
+    dpt.complex128,
+]
+sum_unsupported_output_dtypes = [
+    dpt.bool,
+    dpt.dtype("i1"),
+    dpt.dtype("u1"),
+    dpt.dtype("i2"),
+    dpt.dtype("u2"),
+    dpt.float16,
+    dpt.complex64,
+    dpt.complex128,
+]
+
+mean_unsupported_output_dtypes = [
+    dpt.bool,
+    dpt.dtype("i1"),
+    dpt.dtype("u1"),
+    dpt.dtype("i2"),
+    dpt.dtype("u2"),
+    dpt.dtype("i4"),
+    dpt.dtype("u4"),
+    dpt.dtype("i8"),
+    dpt.dtype("u8"),
+    dpt.float16,
+    dpt.complex64,
+    dpt.complex128,
+]
+
+
+sum_only = [_sycl_ext_impl._get_sum_over_axis_0]
+mean_only = [_sycl_ext_impl._get_mean_over_axis_0]
+mean_sum = sum_only + mean_only
+
+
+def supported_by_device(device, typ):
+    if typ == dpt.float64 or typ == dpt.complex128:
+        return device.has_aspect_fp64
+
+    if typ == dpt.float16:
+        return device.has_aspect_fp16
+
+    return True
+
+
+def skip_unsupported(device, typ):
+    if not supported_by_device(device, typ):
+        pytest.skip(f"{typ} type is not supported by {device}")
+
+
+@pytest.mark.parametrize(
+    "func, device, input_type, output_type",
+    product(
+        mean_sum,
+        all_devices,
+        sum_supported_input_dtypes,
+        mean_supported_output_dtypes,
+    ),
+)
+def test_mean_sum_over_axis_0_supported_types(
+    func, device, input_type, output_type
+):
+    skip_unsupported(device, input_type)
+    skip_unsupported(device, output_type)
+
+    height = 20
+    width = 10
+
+    input = dpt.empty((height, width), dtype=input_type, device=device)
+    output = dpt.empty(width, dtype=output_type, device=device)
+
+    assert func(input, output) is not None
+
+
+@pytest.mark.parametrize(
+    "func, device, input_type, output_type",
+    product(
+        sum_only,
+        all_devices,
+        sum_supported_input_dtypes,
+        sum_without_mean_supported_dtypes,
+    ),
+)
+def test_sum_over_axis_0_supported_types(func, device, input_type, output_type):
+    skip_unsupported(device, input_type)
+    skip_unsupported(device, output_type)
+
+    height = 20
+    width = 10
+
+    input = dpt.empty((height, width), dtype=input_type, device=device)
+    output = dpt.empty(width, dtype=output_type, device=device)
+
+    assert func(input, output) is not None
+
+
+@pytest.mark.parametrize(
+    "func, device, input_type, output_type",
+    product(mean_sum, all_devices, sum_unsupported_input_dtypes, [dpt.float32]),
+)
+def test_mean_sum_over_axis_0_unsupported_in_types(
+    func, device, input_type, output_type
+):
+    skip_unsupported(device, input_type)
+    skip_unsupported(device, output_type)
+
+    height = 1
+    width = 1
+
+    input = dpt.empty((height, width), dtype=input_type, device=device)
+    output = dpt.empty(width, dtype=output_type, device=device)
+
+    assert func(input, output) is None
+
+
+@pytest.mark.parametrize(
+    "func, device, input_type, output_type",
+    product(
+        sum_only, all_devices, [dpt.float32], sum_unsupported_output_dtypes
+    ),
+)
+def test_sum_over_axis_0_unsupported_out_types(
+    func, device, input_type, output_type
+):
+    skip_unsupported(device, input_type)
+    skip_unsupported(device, output_type)
+
+    height = 1
+    width = 1
+
+    input = dpt.empty((height, width), dtype=input_type, device=device)
+    output = dpt.empty(width, dtype=output_type, device=device)
+
+    assert func(input, output) is None
+
+
+@pytest.mark.parametrize(
+    "func, device, input_type, output_type",
+    product(
+        mean_only, all_devices, [dpt.float32], mean_unsupported_output_dtypes
+    ),
+)
+def test_mean_over_axis_0_unsupported_out_types(
+    func, device, input_type, output_type
+):
+    skip_unsupported(device, input_type)
+    skip_unsupported(device, output_type)
+
+    height = 1
+    width = 1
+
+    input = dpt.empty((height, width), dtype=input_type, device=device)
+    output = dpt.empty(width, dtype=output_type, device=device)
+
+    if func(input, output):
+        print(output_type)
+    assert func(input, output) is None
+
+
+@pytest.mark.parametrize(
+    "func, device, input_type, output_type",
+    product(mean_sum, all_devices, [dpt.float32], [dpt.float32]),
+)
+def test_mean_over_axis_0_f_contig_input(func, device, input_type, output_type):
+    skip_unsupported(device, input_type)
+    skip_unsupported(device, output_type)
+
+    height = 20
+    width = 10
+
+    input = dpt.empty((height, width), dtype=input_type, device=device).T
+    output = dpt.empty(width, dtype=output_type, device=device)
+
+    if func(input, output):
+        print(output_type)
+    assert func(input, output) is None
+
+
+@pytest.mark.parametrize(
+    "func, device, input_type, output_type",
+    product(mean_sum, all_devices, [dpt.float32], [dpt.float32]),
+)
+def test_mean_over_axis_0_f_contig_output(
+    func, device, input_type, output_type
+):
+    skip_unsupported(device, input_type)
+    skip_unsupported(device, output_type)
+
+    height = 1
+    width = 10
+
+    input = dpt.empty((height, 10), dtype=input_type, device=device)
+    output = dpt.empty(20, dtype=output_type, device=device)[::2]
+
+    if func(input, output):
+        print(output_type)
+    assert func(input, output) is None