From 6badafa9f19b73e090fa29bade31e5dd75d024ff Mon Sep 17 00:00:00 2001
From: Michel Schanen <michel.schanen@gmail.com>
Date: Tue, 4 Aug 2026 11:16:41 -0500
Subject: [PATCH] Use host USM for scalar oneMKL results

The scalar-returning BLAS wrappers (nrm2, dot/dotc/dotu, asum, iamax,
iamin) allocated their one-element result in device memory and copied it
back to the host, so every call paid for a device allocation, an H2D
transfer to zero-initialize it, and a D2H transfer to read it.

Allocate the result in host USM instead. oneMKL writes straight into
host-visible memory and we read the value with a plain load; the C
wrappers already wait_and_throw() before returning, so the result is
ready as soon as the ccall is.

nrm2 over a 4096-element Float32 vector on an Arc A750: 3079.5 -> 331.5
us/call. The win is call overhead, so it is independent of length.
---
 lib/mkl/utils.jl         | 12 ++++++++++++
 lib/mkl/wrappers_blas.jl | 23 ++++++++++-------------
 2 files changed, 22 insertions(+), 13 deletions(-)

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