This is an automated email from the ASF dual-hosted git repository.
tqchen pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/main by this push:
new 37a6200185 [Relax] Fix cublas dispatch for corner cases (#17139)
37a6200185 is described below
commit 37a62001857c812afed1f6f7df3b49ff01bd2988
Author: Siyuan Feng <[email protected]>
AuthorDate: Thu Jul 11 00:24:06 2024 +0800
[Relax] Fix cublas dispatch for corner cases (#17139)
Fix case when `lhs_batches` and `rhs_batches` are symbolic expressions,
but not standalone variables.
---
python/tvm/relax/backend/contrib/cublas.py | 2 +-
1 file changed, 1 insertion(+), 1 deletion(-)
diff --git a/python/tvm/relax/backend/contrib/cublas.py
b/python/tvm/relax/backend/contrib/cublas.py
index febb401bc0..287b18b440 100644
--- a/python/tvm/relax/backend/contrib/cublas.py
+++ b/python/tvm/relax/backend/contrib/cublas.py
@@ -134,7 +134,7 @@ def _check_matmul(context: PatternCheckContext) -> bool:
isinstance(lhs_batches, tvm.tir.Var)
or isinstance(rhs_batches, tvm.tir.Var)
or (analyzer.can_prove_equal(lhs_batches, rhs_batches))
- or (lhs_batches >= 1 and rhs_batches == 1)
+ or (analyzer.can_prove(lhs_batches >= 1) and
analyzer.can_prove(rhs_batches == 1))
)
