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.

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

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
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

2 participants