Inplace version of batched adjoint/transpose
- Dominant language
- Julia
- Stars
- 258
- Forks
- 138
- Avg merge
- 2d 1h
- Merged PRs (30d)
- 3
Description
We are missing an the inplace version of batched adjoint/transpose. They are required to avoid gpu scalar indexing with `Base.copy` like `copy(batched_adjoint(CUDA.randn(3,5,2)))`. They can be implemented as:
```julia
# Inplace
function batched_transpose_f!(f, B::AbstractArray{T, 3}, A::AbstractArray{T, 3}) where T
axes(B,1) == axes(A,2) && axes(B,2) == axes(A,1) && axes(A,3) == axes(B,3) || throw(DimensionMismatch(string(f)))
Threads.@threads for i in axes(A,3)
Bi = @view B[:, :, i]
Ai = @view A[:, :, i]
LinearAlgebra.transpose_f!(f, Bi, Ai)
end
return B
end
using GPUArrays
function batched_transpose_f!(f, B::AnyGPUArray{T, 3}, A::AnyGPUArray{T, 3}) where T
axes(B,1) == axes(A,2) && axes(B,2) == axes(A,1) && axes(A,3) == axes(B,3) || throw(DimensionMismatch(string(f)))
GPUArrays.gpu_call(B, A) do ctx, B, A
idx = GPUArrays.@cartesianidx A
@inbounds B[idx[2], idx[1], idx[3]] = f(A[idx[1], idx[2], idx[3]])
return
end
return B
end
batched_adjoint!(B, A) = batched_transpose_f!(adjoint, B, A)
batched_transpose!(B, A) = batched_transpose_f!(transpose, B, A)
# copy
function Base.copy(x::BatchedAdjoint)
p = parent(x)
a1, a2, a3 = axes(p)
return batched_adjoint!(similar(p, (a2, a1, a3)), p)
end
function Base.copy(x::BatchedTranspose)
p = parent(x)
a1, a2, a3 = axes(p)
return batched_transpose!(similar(p, (a2, a1, a3)), p)
end
```
which require an extra dependency of `GPUArrays`. I have no idea where should we put these code under the `ext`.
Contributor guide
No contributing guide indexed for this repository
Research direction
Start by locating the batched adjoint/transpose implementations and the package's ext layout. Review how optional GPU backends and GPUArrays dependencies are declared, then determine where the proposed CPU and GPU methods belong. Done means inplace batched transpose and adjoint work for three-dimensional arrays and Base.copy avoids GPU scalar indexing.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- julia
- Domain
- backend-api-design, machine-learning
- Issue type
- Feature
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Stale
- Clarity
- Mostly clear
- Newbie friendliness
- 35/100