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

[Unity][BYOC] Use arith.Analyzer to check batch equality of matmul in cublas #16982

Merged
merged 3 commits into from
May 9, 2024
Merged

[Unity][BYOC] Use arith.Analyzer to check batch equality of matmul in cublas #16982

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
ef6919f
[Unity][BYOC] Use arith.Analyzer to check batch equality of matmul in…
rickzx May 8, 2024
b735e4e
Fix lint, add unit tests
rickzx May 8, 2024
96a1dd5
Trigger Build
rickzx May 9, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Jump to
Jump to file
Failed to load files.
Diff view
Diff view
5 changes: 4 additions & 1 deletion python/tvm/relax/backend/contrib/cublas.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -21,6 +21,7 @@

import tvm
from tvm import DataType
from tvm.arith import Analyzer
from tvm.relax import transform
from tvm.relax.transform import PatternCheckContext

Expand Down Expand Up @@ -119,14 +120,16 @@ def _check_matmul(context: PatternCheckContext) -> bool:
# cuBLAS only supports bias vector
return False

analyzer = Analyzer()

# cuBLASLt does not seem to support batched GEMM with one of matrices having
# one batch (with batch_stride 0). So for batched GEMM, the two batch counts
# must be equal. If lhs is batched but rhs is not, we can use the regular GEMM by
# flattening all batch axes into the M axis.
return (
isinstance(lhs_batches, tvm.tir.Var)
or isinstance(rhs_batches, tvm.tir.Var)
or (int(lhs_batches) == int(rhs_batches))
or (analyzer.can_prove_equal(lhs_batches, rhs_batches))
or (lhs_batches >= 1 and rhs_batches == 1)
)

Expand Down
2 changes: 2 additions & 0 deletions tests/python/relax/test_codegen_cublas.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -149,6 +149,8 @@ def get_relax_matmul_dequantize_module(
((_vars["a"], 32, 8), (_vars["a"], 8, 10), True, "gelu"),
# ND x ND
((5, 3, 32, 8), (5, 3, 8, 10), True, "relu"),
((_vars["a"], 3, 32, 8), (_vars["a"], 3, 8, 10), True, "relu"),
((_vars["a"], _vars["b"], 32, 8), (_vars["a"], _vars["b"], 8, 10), True, "relu"),
# ND x 2D
((5, 3, 32, 8), (8, 10), False, "none"),
],
Expand Down