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

[Relax] Support input_axis_separator to allow 2D to 1D conversion #17115

Merged
merged 3 commits into from
Jul 1, 2024

Conversation

abhikran-quic
Copy link
Contributor

@abhikran-quic abhikran-quic commented Jun 25, 2024

  • Introduce input_axis_separator in relax.transform_layout op for 2D input buffer to allow conversion of 2D buffers to 1D buffers. The conversion from 2D->1D is handled while lowering of transform_layout operator.
  • This is needed because once the tensor is converted into 2D via axis_separator, the input_axis_separator helps in identifying the buffers with 2 dimensions especially when the conversion is needed to 1D operation.
  • Also introducing support for input_axis_separator in AlterOpImpl pass.

Introduce input_axis_separator in relax.transform_layout op to allow conversion of 2D buffers to 1D buffers.
The conversion from 2D->1D is handled while lowering of transform_layout operator.
Also introducing support for input_axis_separator in AlterOpImpl pass.
@abhikran-quic abhikran-quic force-pushed the input_axis_sep branch 2 times, most recently from a7018ad to 778557e Compare June 26, 2024 03:04
@abhikran-quic abhikran-quic changed the title [Relax] Support input axis_separator to allow 2D to 1D conversion [Relax] Support input_axis_separator to allow 2D to 1D conversion Jun 26, 2024
@abhikran-quic
Copy link
Contributor Author

cc: @Lunderberg , @tqchen

python/tvm/relax/transform/legalize_ops/manipulate.py Outdated Show resolved Hide resolved
python/tvm/relax/transform/legalize_ops/manipulate.py Outdated Show resolved Hide resolved
src/relax/op/tensor/manipulate.h Outdated Show resolved Hide resolved
src/relax/transform/alter_op_impl.cc Outdated Show resolved Hide resolved
src/relax/transform/alter_op_impl.cc Outdated Show resolved Hide resolved
tests/python/relax/test_transform_alter_op_impl.py Outdated Show resolved Hide resolved
Copy link
Contributor Author

@abhikran-quic abhikran-quic left a comment

Choose a reason for hiding this comment

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

Thank you @Lunderberg for your comments. I've addressed them in a new commit.

python/tvm/relax/transform/legalize_ops/manipulate.py Outdated Show resolved Hide resolved
python/tvm/relax/transform/legalize_ops/manipulate.py Outdated Show resolved Hide resolved
src/relax/op/tensor/manipulate.h Outdated Show resolved Hide resolved
src/relax/transform/alter_op_impl.cc Outdated Show resolved Hide resolved
src/relax/transform/alter_op_impl.cc Outdated Show resolved Hide resolved
tests/python/relax/test_transform_alter_op_impl.py Outdated Show resolved Hide resolved
Copy link
Contributor

@Lunderberg Lunderberg left a comment

Choose a reason for hiding this comment

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

Thank you for making the changes, and it looks good!

@abhikran-quic
Copy link
Contributor Author

@tvm-bot rerun

@quic-sanirudh quic-sanirudh merged commit ab7c1a9 into apache:main Jul 1, 2024
20 checks passed
@quic-sanirudh
Copy link
Contributor

Thanks @abhikran-quic @Lunderberg

@abhikran-quic abhikran-quic deleted the input_axis_sep branch July 1, 2024 08:10
@abhikran-quic
Copy link
Contributor Author

Thank you @Lunderberg, @quic-sanirudh !

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