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

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

Merged
merged 8 commits into from
Jul 29, 2021

Conversation

jcf94
Copy link
Contributor

@jcf94 jcf94 commented Jul 22, 2021

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

@jcf94 jcf94 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 matmul Jul 22, 2021
@jcf94 jcf94 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 matmul Jul 22, 2021
@jcf94 jcf94 marked this pull request as draft July 22, 2021 12:45
@jcf94
Copy link
Contributor Author

jcf94 commented Jul 22, 2021

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. 😆

Copy link
Contributor

@comaniac comaniac left a comment

Choose a reason for hiding this comment

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

Overall LGTM. Thanks.

python/tvm/relay/op/nn/_nn.py Outdated Show resolved Hide resolved
python/tvm/relay/op/nn/nn.py Outdated Show resolved Hide resolved
python/tvm/topi/nn/batch_matmul.py Outdated Show resolved Hide resolved
python/tvm/topi/nn/batch_matmul.py Outdated Show resolved Hide resolved
python/tvm/topi/nn/batch_matmul.py Outdated Show resolved Hide resolved
src/relay/op/nn/nn.cc Outdated Show resolved Hide resolved
src/relay/op/nn/nn.cc Show resolved Hide resolved
src/relay/qnn/op/batch_matmul.cc Outdated Show resolved Hide resolved
src/relay/transforms/combine_parallel_batch_matmul.cc Outdated Show resolved Hide resolved
@jcf94 jcf94 marked this pull request as ready for review July 24, 2021 03:05
Re-triggle CI

Bug fix

Re-triggle CI

Re-triggle CI

Re-triggle CI
@jcf94
Copy link
Contributor Author

jcf94 commented Jul 29, 2021

@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.

Copy link
Contributor

@comaniac comaniac left a comment

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:main Jul 29, 2021
@comaniac
Copy link
Contributor

Thanks @jcf94 @elvin-n

@jcf94 jcf94 deleted the transpose_batch_matmul branch July 30, 2021 02:05
ylc pushed a commit to ylc/tvm that referenced this pull request Sep 29, 2021
…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
ylc pushed a commit to ylc/tvm that referenced this pull request Jan 13, 2022
…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 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.

3 participants