-
Notifications
You must be signed in to change notification settings - Fork 3.5k
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
[Topi] Allow relay.nn.dense support arbitrary number dimensions #8412
Comments
See https://discuss.tvm.apache.org/t/relay-nn-does-relay-nn-dense-supports-multi-dimensional-input/10343/7. I realized that supporting more than 2D input for It would also improve TVM numbers for @AndrewZhaoLuo Are you going to work on this one? If not, I need to do it anyway before or after #9261 lands. |
I think a good way to add N-D input support to |
Hey @masahi you can go ahead and tackle this one. I'll have time probably next week if this is super high priority. |
I have a concern that we have more than one backends and we need to make sure all backend implementations have the same semantic/spec. For example, previously I encountered the issue (#7730) that CuBLAS batch_matmul implementation doesn't supports implicit broadcasting, which results in errors when lowering batch_matmul to CuBLAS. If we allow arbitrary dimension in |
I'm happy to take this, but I also want to continue working on cutlass byoc #9261. Moreover, things look more complicated than I thought yesterday, as pointed out by @comaniac and discussed below. Making sure that N-D input dense is supported by external backends is a really good point. It is kind of odd I didn't think about it yesterday because I was trying to fuse 3-D dense + activation for cutlass byoc. Now I realized that I don't even know if cutlass supports such N-D dense operation. @comaniac N-D dense input is common in all frameworks so we should definitely support it. My hope is that since those extra batch dimensions can be conceptually fused into 1D batch dim, external libs can directly work on N-D input, i.e. whether or not we explicitly do reshape doesn't change underlying memory addresses. I guess PyTorch simply uses Other related points I realized now:
|
Another workaround that allows dense + activation without supporting N-D dense is to swap the order of reshape and activations. So right now we are doing |
Thanks for the analysis and I deeply agree with it as we (with @icemelon) are also experiencing the same issue. Making the pattern like |
I found something interesting wrt the reshape swapping approach. If there are back-to-back
Note that there are two |
Right now the schedule assumes two dimensions for the input tensor though the documentation suggests otherwise:
https://tvm.apache.org/docs/api/python/relay/nn.html#tvm.relay.nn.dense
Alternatively we can just adjust the documentation to stipulate relay.nn.dense has an input tensor of only two dimensions. However, other frameworks allow their dense layers to be arbitrary number dimensions. E.g. https://pytorch.org/docs/stable/generated/torch.nn.Linear.html#torch.nn.Linear
For forum discussion see: https://discuss.tvm.apache.org/t/relay-nn-does-relay-nn-dense-supports-multi-dimensional-input/10343/5
The text was updated successfully, but these errors were encountered: