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.

Sign up for GitHub

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

Jump to bottom

[CUDNN] Support gradient kernels #9986

Merged
merged 21 commits into from
Jan 22, 2022
Merged

[CUDNN] Support gradient kernels #9986

merged 21 commits into from
Jan 22, 2022

Conversation

masahi
Copy link
Member

@masahi masahi commented Jan 19, 2022 •
edited
Loading

This adds support for computing conv2d gradient wrt data and filter by cudnn. See the change in test_cudnn.py. The diff is large but majority of them are boilerplate.

Also enabled offloading the relay conv2d_backward_weight op introduced in #9954 directly to the cudnn equivalent without legalization. The op strategy for that op is only populated when cudnn is in enabled in the target, otherwise emit an error and suggest running Legalize pass. This feature can go into a separate PR if desired (the last part of commits).

After opening this PR, I realized that I haven't enabled offloading relay conv2d_transpose op to the cudnn dgrad kernel. I'll follow up later. (UPDATE: Done, can also go into a separate PR)

cc @vinx13 @comaniac @tkonolige @Laurawly @Hzfengsy @YuchenJin

@driazati driazati mentioned this pull request Jan 19, 2022
Closed
@masahi masahi force-pushed the cudnn-grad-kernels branch 4 times, most recently from 63a5548 to 6a2f900 Compare January 21, 2022 07:23
Hzfengsy
Hzfengsy approved these changes Jan 22, 2022
View reviewed changes
Copy link
Member

@Hzfengsy Hzfengsy left a comment

Choose a reason for hiding this comment

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

LGTM

@masahi
Dgrad nchw, nhwc, fp16 working
e084caa 
commit 426e5dc
Author: Masahiro Masuda <masahi129@gmail.com>
Date:   Tue Jan 18 11:48:53 2022 +0900

    black

commit 211a58b
Author: Masahiro Masuda <masahi129@gmail.com>
Date:   Tue Jan 18 11:43:52 2022 +0900

    fp16 also works

commit c2a34d4
Author: Masahiro Masuda <masahi129@gmail.com>
Date:   Tue Jan 18 11:36:36 2022 +0900

    nhwc test also worked

commit c0609ab
Author: Masahiro Masuda <masahi129@gmail.com>
Date:   Tue Jan 18 11:21:23 2022 +0900

    nchw test worked

commit 2bf68c7
Author: Masahiro Masuda <masahi129@gmail.com>
Date:   Tue Jan 18 10:41:35 2022 +0900

    add test stub

commit c86b128
Author: Masahiro Masuda <masahi129@gmail.com>
Date:   Tue Jan 18 10:32:09 2022 +0900

    add python definition stub

commit 3166952
Author: Masahiro Masuda <masahi129@gmail.com>
Date:   Tue Jan 18 06:57:18 2022 +0900

    bwd filter compiled

commit e311ba3
Author: Masahiro Masuda <masahi129@gmail.com>
Date:   Tue Jan 18 06:27:55 2022 +0900

    dgrad compiled

commit 47f35be
Author: Masahiro Masuda <masahi129@gmail.com>
Date:   Tue Jan 18 06:16:43 2022 +0900

    add dgrad stub

commit ebed032
Author: Masahiro Masuda <masahi129@gmail.com>
Date:   Mon Jan 17 17:01:56 2022 +0900

    cpplint

commit 834f54a
Author: Masahiro Masuda <masahi129@gmail.com>
Date:   Mon Jan 17 16:55:58 2022 +0900

    remove cudnn get output

commit dcbd9c9
Author: Masahiro Masuda <masahi129@gmail.com>
Date:   Mon Jan 17 16:28:07 2022 +0900

    more refactor

commit 146464e
Author: Masahiro Masuda <masahi129@gmail.com>
Date:   Mon Jan 17 15:57:35 2022 +0900

    Introduce SetConvdescriptors to refactor cudnn/conv_forward.cc
@masahi
add python function for cudnn wgrad
204eb13
@masahi
adding wgrad test
03bb809
@masahi
black
61060bb
@masahi
wgrad nchw and nhwc worked
5f13439
@masahi
remove bwd algo name stuff
64f76c8
@masahi
compute output shape properly
0ede216
@masahi
swap arg order in wgrad
0771b5f
@masahi
add kernel size arg in test
282dd05
@masahi
black
ce8fde0
@masahi
cleanup
cebc8d1
@masahi
more fix
62c96c9
@masahi
fix dgrad test
42bc825
@masahi
support running relay conv2d_backward_weight directly with cudnn
367ad62
@masahi
black
09b32ba
@masahi
refactor reference function to support nhwc
64b4279
@masahi
removed unused function
06870e4
@masahi
lint
615b90d
@masahi
enable offloading conv2d_transpose to cudnn dgrad
5b5992c
@masahi
relax tol
c595676
@masahi
name fix, remove print
7ed295c
YuchenJin
YuchenJin approved these changes Jan 22, 2022
View reviewed changes
Copy link
Contributor

@YuchenJin YuchenJin left a comment

Choose a reason for hiding this comment

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

LGTM! Thanks @masahi!

@masahi masahi merged commit d35b858 into apache:main Jan 22, 2022
@masahi
Copy link
Member Author

masahi commented Jan 22, 2022

Thanks @comaniac @Hzfengsy @YuchenJin

yuanfz98 pushed a commit to yuanfz98/tvm that referenced this pull request Jan 24, 2022
@masahi @yuanfz98
[CUDNN] Support gradient kernels (apache#9986)
7b76524 
* Dgrad nchw, nhwc, fp16 working

commit 426e5dc
Author: Masahiro Masuda <masahi129@gmail.com>
Date:   Tue Jan 18 11:48:53 2022 +0900

    black

commit 211a58b
Author: Masahiro Masuda <masahi129@gmail.com>
Date:   Tue Jan 18 11:43:52 2022 +0900

    fp16 also works

commit c2a34d4
Author: Masahiro Masuda <masahi129@gmail.com>
Date:   Tue Jan 18 11:36:36 2022 +0900

    nhwc test also worked

commit c0609ab
Author: Masahiro Masuda <masahi129@gmail.com>
Date:   Tue Jan 18 11:21:23 2022 +0900

    nchw test worked

commit 2bf68c7
Author: Masahiro Masuda <masahi129@gmail.com>
Date:   Tue Jan 18 10:41:35 2022 +0900

    add test stub

commit c86b128
Author: Masahiro Masuda <masahi129@gmail.com>
Date:   Tue Jan 18 10:32:09 2022 +0900

    add python definition stub

commit 3166952
Author: Masahiro Masuda <masahi129@gmail.com>
Date:   Tue Jan 18 06:57:18 2022 +0900

    bwd filter compiled

commit e311ba3
Author: Masahiro Masuda <masahi129@gmail.com>
Date:   Tue Jan 18 06:27:55 2022 +0900

    dgrad compiled

commit 47f35be
Author: Masahiro Masuda <masahi129@gmail.com>
Date:   Tue Jan 18 06:16:43 2022 +0900

    add dgrad stub

commit ebed032
Author: Masahiro Masuda <masahi129@gmail.com>
Date:   Mon Jan 17 17:01:56 2022 +0900

    cpplint

commit 834f54a
Author: Masahiro Masuda <masahi129@gmail.com>
Date:   Mon Jan 17 16:55:58 2022 +0900

    remove cudnn get output

commit dcbd9c9
Author: Masahiro Masuda <masahi129@gmail.com>
Date:   Mon Jan 17 16:28:07 2022 +0900

    more refactor

commit 146464e
Author: Masahiro Masuda <masahi129@gmail.com>
Date:   Mon Jan 17 15:57:35 2022 +0900

    Introduce SetConvdescriptors to refactor cudnn/conv_forward.cc

* add python function for cudnn wgrad

* adding wgrad test

* black

* wgrad nchw and nhwc worked

* remove bwd algo name stuff

* compute output shape properly

* swap arg order in wgrad

* add kernel size arg in test

* black

* cleanup

* more fix

* fix dgrad test

* support running relay conv2d_backward_weight directly with cudnn

* black

* refactor reference function to support nhwc

* removed unused function

* lint

* enable offloading conv2d_transpose to cudnn dgrad

* relax tol

* name fix, remove print
ylc pushed a commit to ylc/tvm that referenced this pull request Feb 16, 2022
@masahi @ylc
[CUDNN] Support gradient kernels (apache#9986)
03b4444 
* Dgrad nchw, nhwc, fp16 working

commit 426e5dc
Author: Masahiro Masuda <masahi129@gmail.com>
Date:   Tue Jan 18 11:48:53 2022 +0900

    black

commit 211a58b
Author: Masahiro Masuda <masahi129@gmail.com>
Date:   Tue Jan 18 11:43:52 2022 +0900

    fp16 also works

commit c2a34d4
Author: Masahiro Masuda <masahi129@gmail.com>
Date:   Tue Jan 18 11:36:36 2022 +0900

    nhwc test also worked

commit c0609ab
Author: Masahiro Masuda <masahi129@gmail.com>
Date:   Tue Jan 18 11:21:23 2022 +0900

    nchw test worked

commit 2bf68c7
Author: Masahiro Masuda <masahi129@gmail.com>
Date:   Tue Jan 18 10:41:35 2022 +0900

    add test stub

commit c86b128
Author: Masahiro Masuda <masahi129@gmail.com>
Date:   Tue Jan 18 10:32:09 2022 +0900

    add python definition stub

commit 3166952
Author: Masahiro Masuda <masahi129@gmail.com>
Date:   Tue Jan 18 06:57:18 2022 +0900

    bwd filter compiled

commit e311ba3
Author: Masahiro Masuda <masahi129@gmail.com>
Date:   Tue Jan 18 06:27:55 2022 +0900

    dgrad compiled

commit 47f35be
Author: Masahiro Masuda <masahi129@gmail.com>
Date:   Tue Jan 18 06:16:43 2022 +0900

    add dgrad stub

commit ebed032
Author: Masahiro Masuda <masahi129@gmail.com>
Date:   Mon Jan 17 17:01:56 2022 +0900

    cpplint

commit 834f54a
Author: Masahiro Masuda <masahi129@gmail.com>
Date:   Mon Jan 17 16:55:58 2022 +0900

    remove cudnn get output

commit dcbd9c9
Author: Masahiro Masuda <masahi129@gmail.com>
Date:   Mon Jan 17 16:28:07 2022 +0900

    more refactor

commit 146464e
Author: Masahiro Masuda <masahi129@gmail.com>
Date:   Mon Jan 17 15:57:35 2022 +0900

    Introduce SetConvdescriptors to refactor cudnn/conv_forward.cc

* add python function for cudnn wgrad

* adding wgrad test

* black

* wgrad nchw and nhwc worked

* remove bwd algo name stuff

* compute output shape properly

* swap arg order in wgrad

* add kernel size arg in test

* black

* cleanup

* more fix

* fix dgrad test

* support running relay conv2d_backward_weight directly with cudnn

* black

* refactor reference function to support nhwc

* removed unused function

* lint

* enable offloading conv2d_transpose to cudnn dgrad

* relax tol

* name fix, remove print
@driazati driazati mentioned this pull request Jul 14, 2022
Closed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers

@YuchenJin YuchenJin YuchenJin approved these changes

@anijain2305 anijain2305 Awaiting requested review from anijain2305

@areusch areusch Awaiting requested review from areusch

@Huyuwei Huyuwei Awaiting requested review from Huyuwei

@icemelon icemelon Awaiting requested review from icemelon

@jcf94 jcf94 Awaiting requested review from jcf94

@jroesch jroesch Awaiting requested review from jroesch

@junrushao junrushao Awaiting requested review from junrushao

@jwfromm jwfromm Awaiting requested review from jwfromm

@kazum kazum Awaiting requested review from kazum

@kevinthesun kevinthesun Awaiting requested review from kevinthesun

@Laurawly Laurawly Awaiting requested review from Laurawly

@liangfu liangfu Awaiting requested review from liangfu

@MarisaKirisame MarisaKirisame Awaiting requested review from MarisaKirisame

@mbrookhart mbrookhart Awaiting requested review from mbrookhart

@merrymercy merrymercy Awaiting requested review from merrymercy

@slyubomirsky slyubomirsky Awaiting requested review from slyubomirsky

@tmoreau89 tmoreau89 Awaiting requested review from tmoreau89

@tqchen tqchen Awaiting requested review from tqchen

@vinx13 vinx13 Awaiting requested review from vinx13

@wweic wweic Awaiting requested review from wweic

@yzhliu yzhliu Awaiting requested review from yzhliu

@zhiics zhiics Awaiting requested review from zhiics

@ZihengJiang ZihengJiang Awaiting requested review from ZihengJiang

@comaniac comaniac Awaiting requested review from comaniac

@Hzfengsy Hzfengsy Awaiting requested review from Hzfengsy

Assignees
No one assigned
Labels
None yet
Projects
None yet
Milestone
No milestone
Development

Successfully merging this pull request may close these issues.
4 participants
@masahi @comaniac @YuchenJin @Hzfengsy