Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Move strided batch pointer conversion to GPU #2608

Merged
merged 4 commits into from
Jan 8, 2025
Merged

Move strided batch pointer conversion to GPU #2608

Changes from 3 commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
727dff7
Move strided batch pointer conversion to GPU
THargreaves Dec 20, 2024
18a4652
Remove redundant/incorrect nblock check
THargreaves Dec 20, 2024
6d9b85d
Generalise address logic
THargreaves Jan 7, 2025
8bac8cc
Clean-up.
maleadt Jan 8, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
16 changes: 14 additions & 2 deletions lib/cublas/wrappers.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1215,8 +1215,20 @@ end
@inline function unsafe_strided_batch(strided::DenseCuArray{T}) where {T}
batchsize = last(size(strided))
stride = prod(size(strided)[1:end-1])
ptrs = [pointer(strided, (i-1)*stride + 1) for i in 1:batchsize]
return CuArray(ptrs)

ptrs = CuArray{CuPtr{T}}(undef, batchsize)
nblocks = cld(batchsize, 256)
@cuda threads = 256 blocks = nblocks create_ptrs_kernel!(ptrs, strided, stride)
return ptrs
end

function create_ptrs_kernel!(ptrs::CuDeviceArray{T}, A, batch_stride) where {T}
index = (blockIdx().x - 1i32) * blockDim().x + threadIdx().x
stride = gridDim().x * blockDim().x
for i in index:stride:length(ptrs)
ptrs[i] = reinterpret(CuPtr{T}, pointer(A, (i - 1i32) * batch_stride + 1i32))
end
return nothing
end

## (GE) general matrix-matrix multiplication grouped batched
Expand Down
Loading