speedup matmul of CuMatrix and OneHotMatrix

2020-04-20 19:11:57 +08:00 · 2020-04-20 19:11:57 +08:00 · 427c55af92
parent 32e2435729
commit 427c55af92
1 changed files with 4 additions and 0 deletions
--- a/src/onehot.jl
+++ b/src/onehot.jl
@ -41,6 +41,10 @@ import .CuArrays: CuArray, CuArrayStyle, cudaconvert
 import Base.Broadcast: BroadcastStyle, ArrayStyle
 BroadcastStyle(::Type{<:OneHotMatrix{<:CuArray}}) = CuArrayStyle{2}()
 cudaconvert(x::OneHotMatrix{<:CuArray}) = OneHotMatrix(x.height, cudaconvert(x.data))
+function Base.:(*)(A::CuArrays.CuMatrix, B::OneHotMatrix{CuArrays.CuArray{OneHotVector,1}})
+  I = CuArrays.CuArray{UInt32, 1}(B.data.buf, 2 .* B.data.dims, offset = B.data.offset)[1:2:end]
+  A[:, Array(I)]
+end

 """
    onehot(l, labels[, unk])