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.

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

Add ResMLP test #14

Open
wants to merge 4 commits into
base: 3la-rebase-complete
Choose a base branch
from

Conversation

slyubomirsky
Copy link
Collaborator

This PR adds an end-to-end run of ResMLP (from PyTorch import to running) and compares the TVM result against the PyTorch result.

There is an additional rewrite here of linear layers. The TVM PT importer produces linear layers that look like this

reshape(nn.dense(*, *), (1, 256, 512)) + *

whereas the FlexNLP codegen matches patterns of the form nn.bias_add(nn.dense(*, *), *).

Fortunately, the above is equivalent to reshape(nn.bias_add(nn.dense(*, *), *), (1, 256, 512)), which does match the pattern. The rewrite is currently implemented manually, but this would be a useful test case for flexible matching in the future.

Please review @AD1024 -- you will have to add in the calls to our custom codegens.

@AD1024
Copy link
Member

AD1024 commented Jun 25, 2021

Thank you for implementing this. I tried to annotate the module with flexnlp and ilavta codegen, but it seems that the rewrite does not work properly (linear pattern is still not matching for flexnlp). Maybe I missed anything?

Also, Conv1D can be offloaded to ILA simulator but the size is too large to terminate.

@slyubomirsky
Copy link
Collaborator Author

slyubomirsky commented Jun 25, 2021

I think you're right, the rewrite isn't firing. I'll dig into it some more

For Conv1D, we could try to use a smaller model size

@slyubomirsky
Copy link
Collaborator Author

It was a simple bug (I was not checking for the operator's name when I should have)

@slyubomirsky
Copy link
Collaborator Author

Note: For now the rewrite I have written here will be correct in the ResMLP model provided. There need to be more conditions on the dimensions for it to work in general. That might also be an interesting use case for flexible matching, since the rewrite system would probably have to reason about the shapes of all the terms.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

2 participants