diff --git a/lib/mkl/utils.jl b/lib/mkl/utils.jl index ba8ddbcb..033f328d 100644 --- a/lib/mkl/utils.jl +++ b/lib/mkl/utils.jl @@ -108,6 +108,18 @@ function Base.convert(::Type{onemklRangev}, range::Char) end end +# Allocate a one-element buffer for a scalar oneMKL output (nrm2, dot, asum, iamax, +# iamin). These are written by the device but consumed on the host, so we place them in +# host USM: oneMKL writes straight into host-visible memory and we read the value with a +# plain load, instead of allocating a device buffer and copying it back. The C wrappers +# synchronize the queue before returning, so the result is readable as soon as the +# `ccall` does. +function scalar_result(::Type{T}) where {T} + res = oneArray{T, 1, oneL0.HostBuffer}(undef, 1) + res[1] = zero(T) + return res +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) diff --git a/lib/mkl/wrappers_blas.jl b/lib/mkl/wrappers_blas.jl index 1071040e..657d177b 100644 --- a/lib/mkl/wrappers_blas.jl +++ b/lib/mkl/wrappers_blas.jl @@ -587,10 +587,9 @@ for (fname, elty, ret_type) in @eval begin function nrm2(n::Integer, x::oneStridedArray{$elty}) queue = global_queue(context(x), device(x)) - result = oneArray{$ret_type}([0]); + result = scalar_result($ret_type) $fname(sycl_queue(queue), n, x, stride(x,1), result) - res = Array(result) - return res[1] + return result[1] end end end @@ -617,10 +616,9 @@ for (jname, fname, elty) in x::oneStridedArray{$elty}, y::oneStridedArray{$elty}) queue = global_queue(context(x), device(x)) - result = oneArray{$elty}([0]); + result = scalar_result($elty) $fname(sycl_queue(queue), n, x, stride(x,1), y, stride(y,1), result) - res = Array(result) - return res[1] + return result[1] end end end @@ -806,11 +804,10 @@ for (fname, elty, ret_type) in @eval begin function asum(n::Integer, x::oneStridedArray{$elty}) - result = oneArray{$ret_type}([0]) + result = scalar_result($ret_type) queue = global_queue(context(x), device(x)) $fname(sycl_queue(queue), n, x, stride(x, 1), result) - res = Array(result) - return res[1] + return result[1] end end end @@ -825,9 +822,9 @@ for (fname, elty) in function iamax(x::oneStridedArray{$elty}) n = length(x) queue = global_queue(context(x), device(x)) - result = oneArray{Int64}([0]); + result = scalar_result(Int64) $fname(sycl_queue(queue), n, x, stride(x, 1), result, 'O') - return Array(result)[1] + return result[1] end end end @@ -841,10 +838,10 @@ for (fname, elty) in @eval begin function iamin(x::StridedArray{$elty}) n = length(x) - result = oneArray{Int64}([0]); + result = scalar_result(Int64) queue = global_queue(context(x), device(x)) $fname(sycl_queue(queue),n, x, stride(x, 1), result, 'O') - return Array(result)[1] + return result[1] end end end