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

Scatter on Cuda #6533

Merged
merged 4 commits into from
Oct 28, 2020
Merged

Conversation

mbrookhart
Copy link
Contributor

I was unable to get scatter working with te schedules on cuda due to the two loops that update values in place, so I resorted to using ir_builder directly.

Attempts to better parallelize the algorithm have created some strange behavior. If I can get a correct and faster implementation to work, I'll submit another PR.

Thanks to @tkonolige for very helpful discussions. @notoraptor, this may be useful for your scatter_add work.

cc @zhiics

@zhiics
Copy link
Member

zhiics commented Sep 22, 2020

cc @Laurawly @vinx13 @icemelon9


with ib.new_scope():
bx = te.thread_axis("blockIdx.x")
ib.scope_attr(bx, "thread_extent", 1)
Copy link
Contributor

Choose a reason for hiding this comment

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

if bx is not used in this context, we can remove the above two lines.

Copy link
Contributor Author

@mbrookhart mbrookhart Oct 5, 2020

Choose a reason for hiding this comment

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

If I don't define at least one block in the scope, the generated cuda code fails to compile.

@mbrookhart mbrookhart force-pushed the mbrookhart/scatter_cuda_schedule branch from 963147a to 488ad3d Compare October 8, 2020 15:31
@tqchen tqchen changed the base branch from master to main October 11, 2020 18:19
Copy link
Member

@zhiics zhiics left a comment

Choose a reason for hiding this comment

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

LGTM. just one minor nitpick. @Laurawly could you please take another look?

python/tvm/topi/cuda/scatter.py Outdated Show resolved Hide resolved
@masahi
Copy link
Member

masahi commented Oct 27, 2020

@zhiics Can we merge this? I want to send a PR to enable scatter_add on GPU, building off of this PR

@zhiics
Copy link
Member

zhiics commented Oct 27, 2020

yeah, let's merge this. before that, we probably want to rebase again just in case bad things would happen

@mbrookhart mbrookhart force-pushed the mbrookhart/scatter_cuda_schedule branch from 7881b2e to 03de673 Compare October 28, 2020 00:03
@mbrookhart
Copy link
Contributor Author

I rebased, I'll keep an eye on it this evening to make sure it still passes CI. Thanks!

@zhiics zhiics merged commit 3d624ec into apache:main Oct 28, 2020
@zhiics
Copy link
Member

zhiics commented Oct 28, 2020

Thanks @mbrookhart @Laurawly @tkonolige @masahi

@mbrookhart mbrookhart deleted the mbrookhart/scatter_cuda_schedule branch October 28, 2020 14:47
trevor-m pushed a commit to trevor-m/tvm that referenced this pull request Oct 29, 2020
* working cuda scatter

fix lint

fix pylint again

* cuda scatter with threading

* add dynamic shape tests

* remove unused variable
trevor-m pushed a commit to trevor-m/tvm that referenced this pull request Dec 2, 2020
* working cuda scatter

fix lint

fix pylint again

* cuda scatter with threading

* add dynamic shape tests

* remove unused variable
trevor-m pushed a commit to trevor-m/tvm that referenced this pull request Dec 4, 2020
* working cuda scatter

fix lint

fix pylint again

* cuda scatter with threading

* add dynamic shape tests

* remove unused variable
trevor-m pushed a commit to neo-ai/tvm that referenced this pull request Dec 4, 2020
* working cuda scatter

fix lint

fix pylint again

* cuda scatter with threading

* add dynamic shape tests

* remove unused variable
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Projects
None yet
Development

Successfully merging this pull request may close these issues.

6 participants