Skip to content

Rework memory pinning and speed up async ops on unpinned memory - #760

Merged
maleadt merged 5 commits into
masterfrom
tb/async_pin
Mar 12, 2021
Merged

maleadt merged 5 commits into
masterfrom
tb/async_pin

Conversation

@maleadt

@maleadt maleadt commented Mar 11, 2021

Copy link
Copy Markdown
Member

Closes #242, closes #735.

With the recent stream-per-task execution model, we can start executing a whole bunch of operations asynchronously (because synchronization is cheaper now that it can yield back to the Julia scheduler). There was one exception though: memory copies. Basically, when transferring memory to or from the CPU, that memory should not be paged out. You can avoid that from happening by pinning or page-locking memory. If you don't, the operation will still execute synchronously, and worse, it will do so without yielding back to the Julia scheduler.

Demo for the rest of this PR:

using CUDA, LinearAlgebra

# an expensive computation
function compute(a,b,c)
    mul!(c, a, b)         # library call
    broadcast!(sin, c, c) # Julia kernel
    c
end

# one "iteration", computing in two tasks
# and comparing the output.
function iteration(a,b,c)
    results = Vector{Any}(undef, 2)
    NVTX.@range "computation" @sync begin
        @async begin
            results[1] = Array(compute(a,b,c))
        end
        @async begin
            results[2] = Array(compute(a,b,c))
        end
    end
    NVTX.@range "comparison" results[1] == results[2]
end

function main(N=1024)
    a = CUDA.rand(N,N)
    b = CUDA.rand(N,N)
    c = CUDA.rand(N,N)

    # make sure this data can be used by other tasks!
    synchronize()

    # warm-up
    iteration(a,b,c)
    GC.gc(true)

    CUDA.@profile begin
        NVTX.@range "run" iteration(a,b,c)
    end
end

isinteractive() || main()

Similar example to #662, but with one crucial difference: a memory copy is performed within the computational task (Array(compute(a,b,c))). That's bad, and results in the following profile:

image

The copies behave synchronously, which results in the first task blocking outside of the scheduler so that the second task can't get scheduled. One solution, as sneakily used in #662, is to synchronize() in each task such that we can start work on the other task before we perform a memory operation. That results in the following trace:

image

Much better. The void before each memory copy is because of allocating memory. A better solution is to pre-allocate that memory and pin it (so results[...] = Mem.pin(Array{eltype(c)}(undef, size(c))) in the main task, and a copyto! on the results instead of just calling the Array constructor), not only avoiding that 'blank space' but also making sure the memory operations can happen asynchronously. That gives a best-case profile trace:

image

Note that the copies themselves are also a lot faster. So you might think, why not eagerly pin every CPU buffer before a memory operation? Alas, pinning isn't free:

image

So, that brings me to the final design of this PR, where we try and figure out whether a CPU buffer is pinned before a memory operation, and if it isn't, eagerly synchronize execution. By doing so, we can at least finish the current work on the GPU concurrently to work queued up by other Julia tasks, even though we'll still end up synchronizing in the driver afterwards:

image

The overhead of these operations are minimal, couple of 10s of ns is barely visible on the profiler trace.

cc @staticfloat

@maleadt maleadt added enhancement New feature or request cuda array Stuff about CuArray. labels Mar 11, 2021
@maleadt
maleadt requested a review from vchuravy March 11, 2021 09:52
@maleadt

maleadt commented Mar 11, 2021

Copy link
Copy Markdown
Member Author

@vchuravy (or @kshyatt, not sure who added this): I removed the flags support for pinning, since it's not clear what to do if you request to pin multiple arrays with different flags. Thoughts? Is this feature required by any of your use cases?

@codecov

codecov Bot commented Mar 11, 2021 •

Copy link
Copy Markdown

Codecov Report

Merging #760 (062ef3d) into master (29fa836) will increase coverage by 0.06%.
The diff coverage is 91.22%.

Impacted file tree graph

@@            Coverage Diff             @@
##           master     #760      +/-   ##
==========================================
+ Coverage   78.20%   78.26%   +0.06%     
==========================================
  Files         121      122       +1     
  Lines        7245     7280      +35     
==========================================
+ Hits         5666     5698      +32     
- Misses       1579     1582       +3     
Impacted Files Coverage Δ
src/pool/binned.jl 0.00% <ø> (ø)
src/pool/none.jl 100.00% <ø> (ø)
src/pool/simple.jl 0.00% <ø> (ø)
src/pool/split.jl 97.12% <ø> (ø)
src/pool/utils.jl 33.33% <ø> (-28.21%) ⬇️
lib/utils/threading.jl 85.71% <85.71%> (ø)
lib/cudadrv/memory.jl 85.03% <91.89%> (+0.96%) ⬆️
src/array.jl 89.82% <100.00%> (+0.32%) ⬆️
src/pool.jl 83.52% <0.00%> (-0.38%) ⬇️
... and 1 more

Continue to review full report at Codecov.

Legend - Click here to learn more
Δ = absolute <relative> (impact), ø = not affected, ? = missing data
Powered by Codecov. Last update 29fa836...062ef3d. Read the comment docs.

@maleadt

maleadt commented Mar 12, 2021

Copy link
Copy Markdown
Member Author

Let's go ahead and merge this, 3.0 is some weeks away anyway so please let me know if there's missing features.

@maleadt
maleadt merged commit a19275f into master Mar 12, 2021
@maleadt
maleadt deleted the tb/async_pin branch March 12, 2021 06:56
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

cuda array Stuff about CuArray. enhancement New feature or request

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Pin CPU buffers when doing memory copies Memory pinning needs more features

1 participant