Fix cholesky (#538)

* fix linalg.cholesky

Author: Rubtsowa

dpnp/backend/include/dpnp_iface.hpp

@@ -240,12 +240,13 @@ INP_DLLEXPORT void dpnp_sort_c(void* array, void* result, size_t size);
  * @ingroup BACKEND_API
  * @brief math library implementation of cholesky function
  *
- * @param [in]  array   Input array with data.
- * @param [out] result  Output array.
- * @param [in]  shape   Shape of input array.
+ * @param [in]  array       Input array with data.
+ * @param [out] result      Output array.
+ * @param [in]  size        Number of elements in input arrays.
+ * @param [in]  data_size   Last element of shape arrays.
  */
 template <typename _DataType>
-INP_DLLEXPORT void dpnp_cholesky_c(void* array1_in, void* result1, size_t* shape);
+INP_DLLEXPORT void dpnp_cholesky_c(void* array1_in, void* result1, const size_t size, const size_t data_size);
 
 /**
  * @ingroup BACKEND_API

dpnp/backend/kernels/dpnp_krnl_linalg.cpp

@@ -34,47 +34,78 @@
 namespace mkl_blas = oneapi::mkl::blas::row_major;
 namespace mkl_lapack = oneapi::mkl::lapack;
 
-template <typename _DataType>
-class dpnp_cholesky_c_kernel;
 
 template <typename _DataType>
-void dpnp_cholesky_c(void* array1_in, void* result1, size_t* shape)
+void dpnp_cholesky_c(void* array1_in, void* result1, const size_t size, const size_t data_size)
 {
-    _DataType* array_1 = reinterpret_cast<_DataType*>(array1_in);
-    _DataType* l_result = reinterpret_cast<_DataType*>(result1);
+    cl::sycl::event event;
 
-    size_t n = shape[0];
+    _DataType* in_array = reinterpret_cast<_DataType*>(array1_in);
+    _DataType* result = reinterpret_cast<_DataType*>(result1);
 
-    l_result[0] = sqrt(array_1[0]);
+    size_t iters = size / (data_size * data_size);
 
-    for (size_t j = 1; j < n; j++)
+    for (size_t k = 0; k < iters; ++k)
     {
-        l_result[j * n] = array_1[j * n] / l_result[0];
-    }
+        _DataType matrix[data_size * data_size];
+        _DataType result_[data_size * data_size];
 
-    for (size_t i = 1; i < n; i++)
-    {
-        _DataType sum_val = 0;
-        for (size_t p = 0; p < i - 1; p++)
+        for (size_t t = 0; t < data_size * data_size; ++t)
         {
-            sum_val += l_result[i * n + p - 1] * l_result[i * n + p - 1];
+            matrix[t] = in_array[k * (data_size * data_size) + t];
+
         }
-        l_result[i * n + i - 1] = sqrt(array_1[i * n + i - 1] - sum_val);
-    }
 
-    for (size_t i = 1; i < n - 1; i++)
-    {
-        for (size_t j = i; j < n; j++)
+        for (size_t it = 0; it < data_size * data_size; ++it)
         {
-            _DataType sum_val = 0;
-            for (size_t p = 0; p < i - 1; p++)
+            result_[it] = matrix[it];
+        }
+
+        const std::int64_t n = data_size;
+
+        const std::int64_t lda = std::max<size_t>(1UL, n);
+
+        const std::int64_t scratchpad_size = mkl_lapack::potrf_scratchpad_size<_DataType>(
+            DPNP_QUEUE, oneapi::mkl::uplo::upper, n, lda);
+
+        _DataType* scratchpad = reinterpret_cast<_DataType*>(dpnp_memory_alloc_c(scratchpad_size * sizeof(_DataType)));
+
+        event = mkl_lapack::potrf(DPNP_QUEUE,
+                                  oneapi::mkl::uplo::upper,
+                                  n,
+                                  result_,
+                                  lda,
+                                  scratchpad,
+                                  scratchpad_size);
+
+        event.wait();
+
+        for (size_t i = 0; i < data_size; i++)
+        {
+            bool arg = false;
+            for (size_t j = 0; j < data_size; j++)
             {
-                sum_val += l_result[i * n + p - 1] * l_result[j * n + p - 1];
+                if (i == j - 1)
+                {
+                    arg = true;
+                }
+                if (arg)
+                {
+                    result_[i * data_size + j] = 0;
+                }
             }
-            l_result[j * n + i - 1] = (1 / l_result[i * n + i - 1]) * (array_1[j * n + i - 1] - sum_val);
         }
+
+        dpnp_memory_free_c(scratchpad);
+
+        for (size_t t = 0; t < data_size * data_size; ++t)
+        {
+            result[k * (data_size * data_size) + t] = result_[t];
+
+        }
+
     }
-    return;
+
 }
 
 template <typename _DataType>
@@ -431,8 +462,6 @@ void dpnp_svd_c(void* array1_in, void* result1, void* result2, void* result3, si
 
 void func_map_init_linalg_func(func_map_t& fmap)
 {
-    fmap[DPNPFuncName::DPNP_FN_CHOLESKY][eft_INT][eft_INT] = {eft_INT, (void*)dpnp_cholesky_c<int>};
-    fmap[DPNPFuncName::DPNP_FN_CHOLESKY][eft_LNG][eft_LNG] = {eft_LNG, (void*)dpnp_cholesky_c<long>};
     fmap[DPNPFuncName::DPNP_FN_CHOLESKY][eft_FLT][eft_FLT] = {eft_FLT, (void*)dpnp_cholesky_c<float>};
     fmap[DPNPFuncName::DPNP_FN_CHOLESKY][eft_DBL][eft_DBL] = {eft_DBL, (void*)dpnp_cholesky_c<double>};
 

dpnp/linalg/dpnp_algo_linalg.pyx

@@ -57,22 +57,30 @@ __all__ = [
 ctypedef void(*custom_linalg_1in_1out_func_ptr_t)(void *, void * , size_t * , size_t)
 ctypedef void(*custom_linalg_1in_1out_func_ptr_t_)(void * , void * , size_t * )
 ctypedef void(*custom_linalg_1in_1out_with_size_func_ptr_t_)(void *, void * , size_t)
+ctypedef void(*custom_linalg_1in_1out_with_2size_func_ptr_t_)(void *, void * , size_t, size_t)
 ctypedef void(*custom_linalg_1in_3out_shape_t)(void *, void * , void * , void * , size_t , size_t )
 
 
 cpdef dparray dpnp_cholesky(dparray input):
-    cdef DPNPFuncType param1_type = dpnp_dtype_to_DPNPFuncType(input.dtype)
+    if input.dtype == dpnp.int32 or input.dtype == dpnp.int64:
+        input_ = input.astype(dpnp.float64)
+    else:
+        input_ = input
+
+    size_ = input_.shape[-1]
+
+    cdef DPNPFuncType param1_type = dpnp_dtype_to_DPNPFuncType(input_.dtype)
 
     cdef DPNPFuncData kernel_data = get_dpnp_function_ptr(DPNP_FN_CHOLESKY, param1_type, param1_type)
 
     result_type = dpnp_DPNPFuncType_to_dtype(< size_t > kernel_data.return_type)
-    cdef dparray result = dparray(input.size, dtype=result_type)
+    cdef dparray result = dparray(input_.shape, dtype=result_type)
+
+    cdef custom_linalg_1in_1out_with_2size_func_ptr_t_ func = <custom_linalg_1in_1out_with_2size_func_ptr_t_ > kernel_data.ptr
 
-    cdef custom_linalg_1in_1out_func_ptr_t_ func = <custom_linalg_1in_1out_func_ptr_t_ > kernel_data.ptr
+    func(input_.get_data(), result.get_data(), input.size, size_)
 
-    func(input.get_data(), result.get_data(), < size_t * > input._dparray_shape.data())
-    l_result = result.reshape(input.shape)
-    return l_result
+    return result
 
 
 cpdef dparray dpnp_cond(dparray input, p):

dpnp/linalg/dpnp_iface_linalg.py

@@ -86,13 +86,16 @@ def cholesky(input):
         Upper or lower-triangular Cholesky factor of `input`.  Returns a
         matrix object if `input` is a matrix object.
     """
-    is_input_dparray = isinstance(input, dparray)
-
-    if not use_origin_backend(input) and is_input_dparray and input.ndim == 2 and \
-            input.shape[0] == input.shape[1] and input.shape[0] > 0:
-        result = dpnp_cholesky(input)
 
-        return result
+    if not use_origin_backend(input):
+        if not isinstance(input, dparray):
+            pass
+        elif input.shape[-1] != input.shape[-2]:
+            pass
+        elif input.ndim < 3:
+            pass
+        else:
+            return dpnp_cholesky(input)
 
     return call_origin(numpy.linalg.cholesky, input)
 

tests/test_linalg.py

@@ -22,8 +22,15 @@ def vvsort(val, vec, size):
             vec[k, imax] = temp
 
 
-def test_cholesky():
-    a = numpy.array([[[1, -2], [2, 5]]])
+@pytest.mark.parametrize("array",
+                         [[[[1, -2], [2, 5]]],
+                          [[[1., -2.], [2., 5.]]],
+                          [[[1., -2.], [2., 5.]], [[1., -2.], [2., 5.]]]],
+                         ids=['[[[1, -2], [2, 5]]]',
+                              '[[[1., -2.], [2., 5.]]]',
+                              '[[[1., -2.], [2., 5.]], [[1., -2.], [2., 5.]]]'])
+def test_cholesky(array):
+    a = numpy.array(array)
     ia = inp.array(a)
     result = inp.linalg.cholesky(ia)
     expected = numpy.linalg.cholesky(a)