ibsidorenko commented on code in PR #16895:
URL: https://github.com/apache/tvm/pull/16895#discussion_r1568429222
##########
python/tvm/relax/backend/contrib/cublas.py:
##########
@@ -68,11 +69,30 @@ def _check_matmul(context: PatternCheckContext) -> bool:
# Rows number must be multiples of 4 for IGEMM
return False
elif lhs_dtype == "e4m3_float8" and rhs_dtype == "e4m3_float8":
- # Matrix dimensions must be multiples of 16. This requirement is
missing from the cuBLAS
- # docs, but it was observed during testing.
- if not isinstance(rhs_shape[-1], (tvm.tir.expr.IntImm, int)) or
rhs_shape[-1] % 16 != 0:
+ matmul_rhs_var = matmul_call.args[1]
+ rhs_transposed = False
+ if matmul_rhs_var in context.matched_bindings:
+ matmul_rhs_call = context.matched_bindings[matmul_rhs_var]
+ assert (
+ isinstance(matmul_rhs_call, tvm.relax.Call)
+ and matmul_rhs_call.op.name == "relax.permute_dims"
+ )
Review Comment:
I am Ok, thank you! Just a nit question: do we need here assert for the case
when rhs_call is something but not `permute_dims`? Just to leave rhs_transposed
== False and return False in the next IF (without crash):
```
if not rhs_transposed:
# cuBLAS FP8 operations require rhs being transposed
return False
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
