Use host USM for scalar oneMKL results - #604
Merged
Merged
Conversation
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.
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## main #604 +/- ##
==========================================
+ Coverage 80.24% 80.33% +0.09%
==========================================
Files 50 50
Lines 3487 3488 +1
==========================================
+ Hits 2798 2802 +4
+ Misses 689 686 -3 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Closes #234.
The scalar-returning oneMKL BLAS wrappers hand the library a one-element device
oneArrayand copy the value back:The two transfers dominate the call. Since we now support host-USM-backed arrays (
oneArray{T,N,oneL0.HostBuffer}with host-sidegetindex/setindex!), the result can live in host USM instead: oneMKL writes straight into host-visible memory and we read it with a plain load. The C wrapperswait_and_throw()before returning, so the value is ready as soon as theccallis.New helper in
lib/mkl/utils.jl:Applied to all five sites in
lib/mkl/wrappers_blas.jl:nrm2,dot/dotc/dotu,asum,iamax,iamin.This also makes #234 moot as originally written — a dedicated
ZeHostReftype isn't needed, since host-USMoneArrays cover the use case andoneL0.host_allochandles the raw-pointer case.Performance
nrm2over a 4096-elementFloat32vector, Arc A750:~9.5x. It's pure call overhead, so the win is independent of vector length.
Testing
Arc A750, NEO 26.18.38308, oneMKL 2025.3.0, Julia 1.12.6:
test/runtests.jl onemkl— 1112 pass, 48 broken, SUCCESStest/runtests.jl gpuarrays/linalg array— 2492 pass, SUCCESSLinearAlgebraforFloat32/Float64/ComplexF32/ComplexF64Notes
Two pre-existing things I left alone, happy to fold in if you'd rather:
iaminis typedx::StridedArray{$elty}where every sibling usesoneStridedArray(wrappers_blas.jl:838).intstatus the C functions return, so a oneMKL exception surfaces as a silent zero rather than an error. That's whyscalar_resultzero-initializes rather than leaving the bufferundef— it preserves the current behaviour on that path instead of returning garbage.