[TOPI] Enhance topi.nn.matmul
#16052
Conversation
|
We have a separate |
|
I agree, in this case, moving to use batch_matmul might be easier |
@tqchen Thanks for the comment, but what I need is a matmul where the shape of A is [b, m, k] and B is [n, k], and A or/and B could be transposed. But batch_matmul only supports 3-dim A and B. And I guess having a matmul operator supporting all input dimensions may be helpful. As described in data apis, the additional dimensions will be broadcasted. |
|
OK, in this case i think we can allow the extended version, provided that it does not break original usage and we have testcases to cover the extended usecase in form of TIR |
|
@tvm-bot rerun |
|
|
||
| if auto_scheduler_rewritten_layout: | ||
| # Infer shape for the rewritten layout | ||
| out_dim, red_dim = auto_scheduler.get_shape_from_rewritten_layout( | ||
| assert len(tensor_b).shape == 2, "only support 2-dim matmul when using auto-scheduler" |
There was a problem hiding this comment.
This should be len(tensor_b.shape). (nobody probably tested this)
This PR enhances
topi.nn.matmulto support batch matmuls (i.e. inputs with dims larger than 2).This is useful in certain cases where we need to generate a batch matmul kernel.