[oneMKL] Interface sparse routines

Author: amontoison

Project.toml

@@ -19,6 +19,7 @@ Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 SPIRV_LLVM_Translator_unified_jll = "85f0d8ed-5b39-5caa-b1ae-7472de402361"
 SPIRV_Tools_jll = "6ac6d60f-d740-5983-97d7-a4482c0689f4"
+SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 oneAPI_Level_Zero_Headers_jll = "f4bc562b-d309-54f8-9efb-476e56f0410d"
@@ -40,7 +41,7 @@ SpecialFunctions = "1.3, 2"
 StaticArrays = "1"
 julia = "1.8"
 oneAPI_Level_Zero_Loader_jll = "1.9"
-oneAPI_Support_jll = "~0.3.1"
+oneAPI_Support_jll = "~0.3.2"
 
 [extras]
 libigc_jll = "94295238-5935-5bd7-bb0f-b00942e9bdd5"

deps/generate_interfaces.jl

@@ -232,7 +232,7 @@ function generate_headers(library::String, filename::String, output::String)
           occursin("int64_t", header) && (suffix = "_64")
         end
         header = replace(header, "$(name_routine)(" => "onemkl$(version)$(name_routine)$(suffix)(")
-        if name_routine ∉ ("init_matrix_handle", "init_matmat_descr", "release_matmat_descr", "set_matmat_data")
+        if name_routine ∉ ("init_matrix_handle", "init_matmat_descr", "release_matmat_descr", "set_matmat_data", "set_csr_data")
           header = replace(header, "void onemkl" => "int onemkl")
         end
         if library == "sparse"
@@ -369,7 +369,7 @@ function generate_cpp(library::String, filename::String, output::String)
       !occursin("scratchpad_size", name) && write(oneapi_cpp, "   auto status = oneapi::mkl::$library::$variant$name<$type>($parameters);\n")
       occursin("scratchpad_size", name)  && write(oneapi_cpp, "   int64_t scratchpad_size = oneapi::mkl::$library::$variant$name<$type>($parameters);\n")
     else
-      if name ∉ ("init_matrix_handle", "init_matmat_descr", "release_matmat_descr", "set_matmat_data")
+      if name ∉ ("init_matrix_handle", "init_matmat_descr", "release_matmat_descr", "set_matmat_data", "set_csr_data")
         write(oneapi_cpp, "   auto status = oneapi::mkl::$library::$variant$name($parameters);\n")
       else
         write(oneapi_cpp, "   oneapi::mkl::$library::$variant$name($parameters);\n")
@@ -378,8 +378,8 @@ function generate_cpp(library::String, filename::String, output::String)
     if occursin("scratchpad_size", name)
       write(oneapi_cpp, "   return scratchpad_size;\n")
     else
-      (name ∉ ("init_matrix_handle", "init_matmat_descr", "release_matmat_descr", "set_matmat_data")) && write(oneapi_cpp, "   __FORCE_MKL_FLUSH__(status);\n")
-      (name ∉ ("init_matrix_handle", "init_matmat_descr", "release_matmat_descr", "set_matmat_data")) && write(oneapi_cpp, "   return 0;\n")
+      (name ∉ ("init_matrix_handle", "init_matmat_descr", "release_matmat_descr", "set_matmat_data", "set_csr_data")) && write(oneapi_cpp, "   __FORCE_MKL_FLUSH__(status);\n")
+      (name ∉ ("init_matrix_handle", "init_matmat_descr", "release_matmat_descr", "set_matmat_data", "set_csr_data")) && write(oneapi_cpp, "   return 0;\n")
     end
     write(oneapi_cpp, "}")
     write(oneapi_cpp, "\n\n")

deps/src/onemkl.cpp

@@ -3465,52 +3465,36 @@ extern "C" void onemklXsparse_init_matrix_handle(matrix_handle_t *handle) {
    oneapi::mkl::sparse::init_matrix_handle((oneapi::mkl::sparse::matrix_handle_t*) handle);
 }
 
-extern "C" int onemklSsparse_set_csr_data(syclQueue_t device_queue, matrix_handle_t handle, int32_t num_rows, int32_t num_cols, onemklIndex index, int32_t *row_ptr, int32_t *col_ind, float *val) {
-   auto status = oneapi::mkl::sparse::set_csr_data(device_queue->val, (oneapi::mkl::sparse::matrix_handle_t) handle, num_rows, num_cols, convert(index), row_ptr, col_ind, val);
-   __FORCE_MKL_FLUSH__(status);
-   return 0;
+extern "C" void onemklSsparse_set_csr_data(syclQueue_t device_queue, matrix_handle_t handle, int32_t num_rows, int32_t num_cols, onemklIndex index, int32_t *row_ptr, int32_t *col_ind, float *val) {
+   oneapi::mkl::sparse::set_csr_data(device_queue->val, (oneapi::mkl::sparse::matrix_handle_t) handle, num_rows, num_cols, convert(index), row_ptr, col_ind, val);
 }
 
-extern "C" int onemklSsparse_set_csr_data_64(syclQueue_t device_queue, matrix_handle_t handle, int64_t num_rows, int64_t num_cols, onemklIndex index, int64_t *row_ptr, int64_t *col_ind, float *val) {
-   auto status = oneapi::mkl::sparse::set_csr_data(device_queue->val, (oneapi::mkl::sparse::matrix_handle_t) handle, num_rows, num_cols, convert(index), row_ptr, col_ind, val);
-   __FORCE_MKL_FLUSH__(status);
-   return 0;
+extern "C" void onemklSsparse_set_csr_data_64(syclQueue_t device_queue, matrix_handle_t handle, int64_t num_rows, int64_t num_cols, onemklIndex index, int64_t *row_ptr, int64_t *col_ind, float *val) {
+   oneapi::mkl::sparse::set_csr_data(device_queue->val, (oneapi::mkl::sparse::matrix_handle_t) handle, num_rows, num_cols, convert(index), row_ptr, col_ind, val);
 }
 
-extern "C" int onemklDsparse_set_csr_data(syclQueue_t device_queue, matrix_handle_t handle, int32_t num_rows, int32_t num_cols, onemklIndex index, int32_t *row_ptr, int32_t *col_ind, double *val) {
-   auto status = oneapi::mkl::sparse::set_csr_data(device_queue->val, (oneapi::mkl::sparse::matrix_handle_t) handle, num_rows, num_cols, convert(index), row_ptr, col_ind, val);
-   __FORCE_MKL_FLUSH__(status);
-   return 0;
+extern "C" void onemklDsparse_set_csr_data(syclQueue_t device_queue, matrix_handle_t handle, int32_t num_rows, int32_t num_cols, onemklIndex index, int32_t *row_ptr, int32_t *col_ind, double *val) {
+   oneapi::mkl::sparse::set_csr_data(device_queue->val, (oneapi::mkl::sparse::matrix_handle_t) handle, num_rows, num_cols, convert(index), row_ptr, col_ind, val);
 }
 
-extern "C" int onemklDsparse_set_csr_data_64(syclQueue_t device_queue, matrix_handle_t handle, int64_t num_rows, int64_t num_cols, onemklIndex index, int64_t *row_ptr, int64_t *col_ind, double *val) {
-   auto status = oneapi::mkl::sparse::set_csr_data(device_queue->val, (oneapi::mkl::sparse::matrix_handle_t) handle, num_rows, num_cols, convert(index), row_ptr, col_ind, val);
-   __FORCE_MKL_FLUSH__(status);
-   return 0;
+extern "C" void onemklDsparse_set_csr_data_64(syclQueue_t device_queue, matrix_handle_t handle, int64_t num_rows, int64_t num_cols, onemklIndex index, int64_t *row_ptr, int64_t *col_ind, double *val) {
+   oneapi::mkl::sparse::set_csr_data(device_queue->val, (oneapi::mkl::sparse::matrix_handle_t) handle, num_rows, num_cols, convert(index), row_ptr, col_ind, val);
 }
 
-extern "C" int onemklCsparse_set_csr_data(syclQueue_t device_queue, matrix_handle_t handle, int32_t num_rows, int32_t num_cols, onemklIndex index, int32_t *row_ptr, int32_t *col_ind, float _Complex *val) {
-   auto status = oneapi::mkl::sparse::set_csr_data(device_queue->val, (oneapi::mkl::sparse::matrix_handle_t) handle, num_rows, num_cols, convert(index), row_ptr, col_ind, reinterpret_cast<std::complex<float>*>(val));
-   __FORCE_MKL_FLUSH__(status);
-   return 0;
+extern "C" void onemklCsparse_set_csr_data(syclQueue_t device_queue, matrix_handle_t handle, int32_t num_rows, int32_t num_cols, onemklIndex index, int32_t *row_ptr, int32_t *col_ind, float _Complex *val) {
+   oneapi::mkl::sparse::set_csr_data(device_queue->val, (oneapi::mkl::sparse::matrix_handle_t) handle, num_rows, num_cols, convert(index), row_ptr, col_ind, reinterpret_cast<std::complex<float>*>(val));
 }
 
-extern "C" int onemklCsparse_set_csr_data_64(syclQueue_t device_queue, matrix_handle_t handle, int64_t num_rows, int64_t num_cols, onemklIndex index, int64_t *row_ptr, int64_t *col_ind, float _Complex *val) {
-   auto status = oneapi::mkl::sparse::set_csr_data(device_queue->val, (oneapi::mkl::sparse::matrix_handle_t) handle, num_rows, num_cols, convert(index), row_ptr, col_ind, reinterpret_cast<std::complex<float>*>(val));
-   __FORCE_MKL_FLUSH__(status);
-   return 0;
+extern "C" void onemklCsparse_set_csr_data_64(syclQueue_t device_queue, matrix_handle_t handle, int64_t num_rows, int64_t num_cols, onemklIndex index, int64_t *row_ptr, int64_t *col_ind, float _Complex *val) {
+   oneapi::mkl::sparse::set_csr_data(device_queue->val, (oneapi::mkl::sparse::matrix_handle_t) handle, num_rows, num_cols, convert(index), row_ptr, col_ind, reinterpret_cast<std::complex<float>*>(val));
 }
 
-extern "C" int onemklZsparse_set_csr_data(syclQueue_t device_queue, matrix_handle_t handle, int32_t num_rows, int32_t num_cols, onemklIndex index, int32_t *row_ptr, int32_t *col_ind, double _Complex *val) {
-   auto status = oneapi::mkl::sparse::set_csr_data(device_queue->val, (oneapi::mkl::sparse::matrix_handle_t) handle, num_rows, num_cols, convert(index), row_ptr, col_ind, reinterpret_cast<std::complex<double>*>(val));
-   __FORCE_MKL_FLUSH__(status);
-   return 0;
+extern "C" void onemklZsparse_set_csr_data(syclQueue_t device_queue, matrix_handle_t handle, int32_t num_rows, int32_t num_cols, onemklIndex index, int32_t *row_ptr, int32_t *col_ind, double _Complex *val) {
+   oneapi::mkl::sparse::set_csr_data(device_queue->val, (oneapi::mkl::sparse::matrix_handle_t) handle, num_rows, num_cols, convert(index), row_ptr, col_ind, reinterpret_cast<std::complex<double>*>(val));
 }
 
-extern "C" int onemklZsparse_set_csr_data_64(syclQueue_t device_queue, matrix_handle_t handle, int64_t num_rows, int64_t num_cols, onemklIndex index, int64_t *row_ptr, int64_t *col_ind, double _Complex *val) {
-   auto status = oneapi::mkl::sparse::set_csr_data(device_queue->val, (oneapi::mkl::sparse::matrix_handle_t) handle, num_rows, num_cols, convert(index), row_ptr, col_ind, reinterpret_cast<std::complex<double>*>(val));
-   __FORCE_MKL_FLUSH__(status);
-   return 0;
+extern "C" void onemklZsparse_set_csr_data_64(syclQueue_t device_queue, matrix_handle_t handle, int64_t num_rows, int64_t num_cols, onemklIndex index, int64_t *row_ptr, int64_t *col_ind, double _Complex *val) {
+   oneapi::mkl::sparse::set_csr_data(device_queue->val, (oneapi::mkl::sparse::matrix_handle_t) handle, num_rows, num_cols, convert(index), row_ptr, col_ind, reinterpret_cast<std::complex<double>*>(val));
 }
 
 extern "C" int onemklSsparse_gemv(syclQueue_t device_queue, onemklTranspose transpose_flag, float alpha, matrix_handle_t handle, float *x, float beta, float *y) {

deps/src/onemkl.h

@@ -2009,37 +2009,37 @@ int64_t onemklZgels_batch_scratchpad_size(syclQueue_t device_queue, onemklTransp
 // SPARSE
 void onemklXsparse_init_matrix_handle(matrix_handle_t *handle);
 
-int onemklSsparse_set_csr_data(syclQueue_t device_queue, matrix_handle_t handle, int32_t
-                               num_rows, int32_t num_cols, onemklIndex index, int32_t *row_ptr,
-                               int32_t *col_ind, float *val);
+void onemklSsparse_set_csr_data(syclQueue_t device_queue, matrix_handle_t handle, int32_t
+                                num_rows, int32_t num_cols, onemklIndex index, int32_t *row_ptr,
+                                int32_t *col_ind, float *val);
 
-int onemklSsparse_set_csr_data_64(syclQueue_t device_queue, matrix_handle_t handle, int64_t
-                                  num_rows, int64_t num_cols, onemklIndex index, int64_t
-                                  *row_ptr, int64_t *col_ind, float *val);
+void onemklSsparse_set_csr_data_64(syclQueue_t device_queue, matrix_handle_t handle, int64_t
+                                   num_rows, int64_t num_cols, onemklIndex index, int64_t
+                                   *row_ptr, int64_t *col_ind, float *val);
 
-int onemklDsparse_set_csr_data(syclQueue_t device_queue, matrix_handle_t handle, int32_t
-                               num_rows, int32_t num_cols, onemklIndex index, int32_t *row_ptr,
-                               int32_t *col_ind, double *val);
+void onemklDsparse_set_csr_data(syclQueue_t device_queue, matrix_handle_t handle, int32_t
+                                num_rows, int32_t num_cols, onemklIndex index, int32_t *row_ptr,
+                                int32_t *col_ind, double *val);
 
-int onemklDsparse_set_csr_data_64(syclQueue_t device_queue, matrix_handle_t handle, int64_t
-                                  num_rows, int64_t num_cols, onemklIndex index, int64_t
-                                  *row_ptr, int64_t *col_ind, double *val);
+void onemklDsparse_set_csr_data_64(syclQueue_t device_queue, matrix_handle_t handle, int64_t
+                                   num_rows, int64_t num_cols, onemklIndex index, int64_t
+                                   *row_ptr, int64_t *col_ind, double *val);
 
-int onemklCsparse_set_csr_data(syclQueue_t device_queue, matrix_handle_t handle, int32_t
-                               num_rows, int32_t num_cols, onemklIndex index, int32_t *row_ptr,
-                               int32_t *col_ind, float _Complex *val);
+void onemklCsparse_set_csr_data(syclQueue_t device_queue, matrix_handle_t handle, int32_t
+                                num_rows, int32_t num_cols, onemklIndex index, int32_t *row_ptr,
+                                int32_t *col_ind, float _Complex *val);
 
-int onemklCsparse_set_csr_data_64(syclQueue_t device_queue, matrix_handle_t handle, int64_t
-                                  num_rows, int64_t num_cols, onemklIndex index, int64_t
-                                  *row_ptr, int64_t *col_ind, float _Complex *val);
+void onemklCsparse_set_csr_data_64(syclQueue_t device_queue, matrix_handle_t handle, int64_t
+                                   num_rows, int64_t num_cols, onemklIndex index, int64_t
+                                   *row_ptr, int64_t *col_ind, float _Complex *val);
 
-int onemklZsparse_set_csr_data(syclQueue_t device_queue, matrix_handle_t handle, int32_t
-                               num_rows, int32_t num_cols, onemklIndex index, int32_t *row_ptr,
-                               int32_t *col_ind, double _Complex *val);
+void onemklZsparse_set_csr_data(syclQueue_t device_queue, matrix_handle_t handle, int32_t
+                                num_rows, int32_t num_cols, onemklIndex index, int32_t *row_ptr,
+                                int32_t *col_ind, double _Complex *val);
 
-int onemklZsparse_set_csr_data_64(syclQueue_t device_queue, matrix_handle_t handle, int64_t
-                                  num_rows, int64_t num_cols, onemklIndex index, int64_t
-                                  *row_ptr, int64_t *col_ind, double _Complex *val);
+void onemklZsparse_set_csr_data_64(syclQueue_t device_queue, matrix_handle_t handle, int64_t
+                                   num_rows, int64_t num_cols, onemklIndex index, int64_t
+                                   *row_ptr, int64_t *col_ind, double _Complex *val);
 
 int onemklSsparse_gemv(syclQueue_t device_queue, onemklTranspose transpose_flag, float alpha,
                        matrix_handle_t handle, float *x, float beta, float *y);

lib/mkl/oneMKL.jl

@@ -11,11 +11,17 @@ using ..SYCL: syclQueue_t
 
 using GPUArrays
 
+using LinearAlgebra
+using SparseArrays
+
 # Exclude Float16 for now, since many oneMKL functions - copy, scal, do not take Float16
 const onemklFloat = Union{Float64,Float32,ComplexF64,ComplexF32}
 const onemklComplex = Union{ComplexF32,ComplexF64}
 const onemklHalf = Union{Float16,ComplexF16}
-include("wrappers.jl")
+
+include("utils.jl")
+include("wrappers_blas.jl")
+include("wrappers_sparse.jl")
 include("linalg.jl")
 
 function band(A::StridedArray, kl, ku)

lib/mkl/utils.jl

@@ -0,0 +1,71 @@
+#
+# Auxiliary
+#
+
+function Base.convert(::Type{onemklSide}, side::Char)
+    if side == 'L'
+        return ONEMKL_SIDE_LEFT
+    elseif side == 'R'
+        return ONEMKL_SIDE_RIGHT
+    else
+        throw(ArgumentError("Unknown transpose $side"))
+    end
+end
+
+function Base.convert(::Type{onemklTranspose}, trans::Char)
+    if trans == 'N'
+        return ONEMKL_TRANSPOSE_NONTRANS
+    elseif trans == 'T'
+        return ONEMKL_TRANSPOSE_TRANS
+    elseif trans == 'C'
+        return ONEMLK_TRANSPOSE_CONJTRANS
+    else
+        throw(ArgumentError("Unknown transpose $trans"))
+    end
+end
+
+function Base.convert(::Type{onemklUplo}, uplo::Char)
+    if uplo == 'U'
+        return ONEMKL_UPLO_UPPER
+    elseif uplo == 'L'
+        return ONEMKL_UPLO_LOWER
+    else
+        throw(ArgumentError("Unknown transpose $uplo"))
+    end
+end
+
+function Base.convert(::Type{onemklDiag}, diag::Char)
+    if diag == 'N'
+        return ONEMKL_DIAG_NONUNIT
+    elseif diag == 'U'
+        return ONEMKL_DIAG_UNIT
+    else
+        throw(ArgumentError("Unknown transpose $diag"))
+    end
+end
+
+function Base.convert(::Type{onemklIndex}, index::Char)
+    if index == 'O'
+        return ONEMKL_INDEX_ONE
+    elseif index == 'Z'
+        return ONEMKL_INDEX_ZERO
+    else
+        throw(ArgumentError("Unknown index $index"))
+    end
+end
+
+function Base.convert(::Type{onemklLayout}, index::Char)
+    if index == 'R'
+        return ONEMKL_LAYOUT_ROW
+    elseif index == 'C'
+        return ONEMKL_LAYOUT_COL
+    else
+        throw(ArgumentError("Unknown layout $layout"))
+    end
+end
+
+# create a batch of pointers in device memory from a batch of device arrays
+@inline function unsafe_batch(batch::Vector{<:oneArray{T}}) where {T}
+    ptrs = pointer.(batch)
+    return oneArray(ptrs)
+end

lib/mkl/wrappers.jl renamed to lib/mkl/wrappers_blas.jl

@@ -1,65 +1,3 @@
-#
-# Auxiliary
-#
-
-function Base.convert(::Type{onemklSide}, side::Char)
-    if side == 'L'
-        return ONEMKL_SIDE_LEFT
-    elseif side == 'R'
-        return ONEMKL_SIDE_RIGHT
-    else
-        throw(ArgumentError("Unknown transpose $side"))
-    end
-end
-
-function Base.convert(::Type{onemklTranspose}, trans::Char)
-    if trans == 'N'
-        return ONEMKL_TRANSPOSE_NONTRANS
-    elseif trans == 'T'
-        return ONEMKL_TRANSPOSE_TRANS
-    elseif trans == 'C'
-        return ONEMLK_TRANSPOSE_CONJTRANS
-    else
-        throw(ArgumentError("Unknown transpose $trans"))
-    end
-end
-
-function Base.convert(::Type{onemklUplo}, uplo::Char)
-    if uplo == 'U'
-        return ONEMKL_UPLO_UPPER
-    elseif uplo == 'L'
-        return ONEMKL_UPLO_LOWER
-    else
-        throw(ArgumentError("Unknown transpose $uplo"))
-    end
-end
-
-function Base.convert(::Type{onemklDiag}, diag::Char)
-    if diag == 'N'
-        return ONEMKL_DIAG_NONUNIT
-    elseif diag == 'U'
-        return ONEMKL_DIAG_UNIT
-    else
-        throw(ArgumentError("Unknown transpose $diag"))
-    end
-end
-
-function Base.convert(::Type{onemklIndex}, index::Char)
-    if index == 'O'
-        return ONEMKL_INDEX_ONE
-    elseif index == 'Z'
-        return ONEMKL_INDEX_ZERO
-    else
-        throw(ArgumentError("Unknown index $index"))
-    end
-end
-
-# create a batch of pointers in device memory from a batch of device arrays
-@inline function unsafe_batch(batch::Vector{<:oneArray{T}}) where {T}
-    ptrs = pointer.(batch)
-    return oneArray(ptrs)
-end
-
 ## (GE) general matrix-matrix multiplication batched
 for (fname, elty) in
         ((:onemklDgemmBatched,:Float64),
@@ -1263,4 +1201,4 @@ end
 function gemm_strided_batched(transA::Char, transB::Char, A::AbstractArray{T, 3},
                               B::AbstractArray{T,3}) where T
     gemm_strided_batched(transA, transB, one(T), A, B)
-end
+end