Skip to content

[TOPI] Add transpose_a/b & dynamic shape support for batch matmul - #8527

Merged
comaniac merged 8 commits into
apache:mainfrom
jcf94:transpose_batch_matmul
Jul 29, 2021
Merged

[TOPI] Add transpose_a/b & dynamic shape support for batch matmul#8527
comaniac merged 8 commits into
apache:mainfrom
jcf94:transpose_batch_matmul

Conversation

@jcf94

@jcf94jcf94 commented Jul 22, 2021•
edited
Loading

Copy link
Copy Markdown
Member

In #8234, a new op nn.matmul was added to extend nn.dense with transpose/non-transpose inputs support.

This PR is going to also extend nn.batch_matmul to support inputs to be in transpose or non-transposed format.

Since the original topi schedule is still for NT format, I set the default format of nn.batch_matmul to be NT. And for the qnn.batch_matmul and some pass like CombineParallelBatchMatmul, I just left them to only support NT currently. Guys who are interested in these may continue to fix :)

Things TODO in this PR:

  • Add tf frontend support for NN batch_matmul(NN matmul supported has already been added)

cc @comaniac@altanh@tkonolige

@jcf94jcf94 changed the title [TOPI] Add transpose_a/b support for batch matmul & Bug fix for dynamic shape with batch matmul[TOPI] Add transpose_a/b support for batch matmul & dynamic shape for batch matmulJul 22, 2021
@jcf94jcf94 changed the title [TOPI] Add transpose_a/b support for batch matmul & dynamic shape for batch matmul[TOPI] Add transpose_a/b & dynamic shape support for batch matmulJul 22, 2021
@jcf94
jcf94 marked this pull request as draft July 22, 2021 12:45
@jcf94

jcf94 commented Jul 22, 2021•
edited
Loading

Copy link
Copy Markdown
MemberAuthor

Also cc @wyc-ruiker , I remember you mentioned the requirement of this in #8402? Seems you need to add topi schedules yourself after merging this. 😆

comaniac
comaniac requested changes Jul 22, 2021
View reviewed changes

@comaniaccomaniac left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

Overall LGTM. Thanks.

Comment threadpython/tvm/relay/op/nn/_nn.py Outdated
Comment threadpython/tvm/relay/op/nn/nn.py Outdated
Comment threadpython/tvm/topi/nn/batch_matmul.py Outdated
Comment threadpython/tvm/topi/nn/batch_matmul.py Outdated
Comment threadpython/tvm/topi/nn/batch_matmul.py Outdated
Comment threadsrc/relay/op/nn/nn.cc Outdated
Comment threadsrc/relay/op/nn/nn.cc
Comment threadsrc/relay/qnn/op/batch_matmul.cc Outdated
Comment threadsrc/relay/transforms/combine_parallel_batch_matmul.cc
Comment threadsrc/relay/transforms/combine_parallel_batch_matmul.cc Outdated
@jcf94
jcf94force-pushed the transpose_batch_matmul branch from f79d12b to 76e2b46CompareJuly 23, 2021 07:05
@jcf94
jcf94 marked this pull request as ready for review July 24, 2021 03:05
@jcf94
jcf94force-pushed the transpose_batch_matmul branch from 492ee2f to 8dac20aCompareJuly 24, 2021 07:21
@jcf94
Lint fix
0649ec8
Re-triggle CI
Bug fix
Re-triggle CI
Re-triggle CI
Re-triggle CI
@jcf94
jcf94force-pushed the transpose_batch_matmul branch from 295777b to 0649ec8CompareJuly 28, 2021 06:38
@jcf94

jcf94 commented Jul 29, 2021•
edited
Loading

Copy link
Copy Markdown
MemberAuthor

@comaniac Thanks! I think I've addressed all of the comments. And thankfully the CI finally get passed.
Please take another look when you have time~

I have tested that the default topi schedule works well for all 4 input formats, at least won't cause any error.

comaniac
comaniac approved these changes Jul 29, 2021
View reviewed changes

@comaniaccomaniac left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

LGTM

@comaniac
comaniac merged commit 850abb0 into apache:mainJul 29, 2021
@comaniac

comaniac commented Jul 29, 2021

Copy link
Copy Markdown
Contributor

Thanks @jcf94@elvin-n

@jcf94
jcf94 deleted the transpose_batch_matmul branch July 30, 2021 02:05
@jcf94jcf94 mentioned this pull request Aug 18, 2021
Merged
ylc pushed a commit to ylc/tvm that referenced this pull request Sep 29, 2021
@jcf94@ylc
[TOPI] Add transpose_a/b & dynamic shape support for batch matmul (ap…
27d88b7
…ache#8527)
* Add basic support for batch matmul transpose
* Update
* Lint fix & add tf convert support
* Update
Lint fix
* Bug fix for qnn.batch_matmul
* Bug fix for tensorflow test
* Add grad support for batch_matmul
* Lint fix
Re-triggle CI
Bug fix
Re-triggle CI
Re-triggle CI
Re-triggle CI
@junrushaojunrushao mentioned this pull request Nov 1, 2021
Closed
ylc pushed a commit to ylc/tvm that referenced this pull request Jan 13, 2022
@jcf94@ylc
[TOPI] Add transpose_a/b & dynamic shape support for batch matmul (ap…
f4698af
…ache#8527)
* Add basic support for batch matmul transpose
* Update
* Lint fix & add tf convert support
* Update
Lint fix
* Bug fix for qnn.batch_matmul
* Bug fix for tensorflow test
* Add grad support for batch_matmul
* Lint fix
Re-triggle CI
Bug fix
Re-triggle CI
Re-triggle CI
Re-triggle CI
Sign up for freeto join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@comaniaccomaniaccomaniac approved these changes
@anijain2305anijain2305Awaiting requested review from anijain2305
@areuschareuschAwaiting requested review from areusch
@HuyuweiHuyuweiAwaiting requested review from Huyuwei
@jroeschjroeschAwaiting requested review from jroesch
@junrushaojunrushaoAwaiting requested review from junrushao
@jwfrommjwfrommAwaiting requested review from jwfromm
@kevinthesunkevinthesunAwaiting requested review from kevinthesun
@LaurawlyLaurawlyAwaiting requested review from Laurawly
@MarisaKirisameMarisaKirisameAwaiting requested review from MarisaKirisame
@masahimasahiAwaiting requested review from masahi
@mbrookhartmbrookhartAwaiting requested review from mbrookhart
@merrymercymerrymercyAwaiting requested review from merrymercy
@slyubomirskyslyubomirskyAwaiting requested review from slyubomirsky
@tqchentqchenAwaiting requested review from tqchen
@vinx13vinx13Awaiting requested review from vinx13
@wweicwweicAwaiting requested review from wweic
@yzhliuyzhliuAwaiting requested review from yzhliu
@zhiicszhiicsAwaiting requested review from zhiics
@ZihengJiangZihengJiangAwaiting requested review from ZihengJiang
@kazumkazumAwaiting requested review from kazum
@siju-samuelsiju-samuelAwaiting requested review from siju-samuel
@srkreddy1238srkreddy1238Awaiting requested review from srkreddy1238
+1 more reviewer
@elvin-nelvin-nelvin-n left review comments
Reviewers whose approvals may not affect merge requirements

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

3 participants

@jcf94@comaniac@elvin-n