Skip to content
New issue

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

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

Already on GitHub? # to your account

Support scatter for CUDA gradient #13

Merged
merged 1 commit into from
Jul 2, 2021
Merged

Conversation

yuehhua
Copy link
Member

@yuehhua yuehhua commented Jun 16, 2021

No description provided.

@yuehhua yuehhua requested a review from CarloLucibello June 16, 2021 07:50
@CarloLucibello
Copy link
Member

@CarloLucibello
Copy link
Member

what's the blocker here?

@yuehhua
Copy link
Member Author

yuehhua commented Jun 21, 2021

The gradient of scatter * and / for cuda need a cuda kernel. I wrote the cuda kernel and I am debugging it. It's almost there.

add count_indices for cuarray

add CUDA kernel for divide_by_counts!

add NNlib.∇scatter_src for cuda gradient

support scatter mean AD for CUDA

support scatter *,/ AD for CUDA
@yuehhua
Copy link
Member Author

yuehhua commented Jul 1, 2021

Finally!!

# multiply all values to be aggregated but not itself
x = one(T)
for k in inds
jk = Base._to_linear_index(src, Tuple(cart_j)..., Tuple(k)...)
Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Base._to_linear_index is introduced here to transform index of any form into integer index. Integer index is required to index a cuarray.

@CarloLucibello CarloLucibello merged commit c3e1331 into FluxML:master Jul 2, 2021
@yuehhua yuehhua deleted the scatter branch July 2, 2021 14:40
# for free to join this conversation on GitHub. Already have an account? # to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

2 participants