Implement a node module replacements for GNN support. #117
Conversation
…and SDFG with some dummy logic is generated.
…inference has access to module
…with exapansions.
…and SDFG with some dummy logic is generated.
…inference has access to module
…with exapansions.
# Conflicts: # daceml/onnx/nodes/replacement.py # daceml/onnx/nodes/replacement_entries.py # daceml/onnx/onnx_importer.py # daceml/onnx/op_implementations/replacement_implementations.py # tests/module_replacements/test_gcnconv.py
Codecov Report
@@ Coverage Diff @@
## master #117 +/- ##
===========================================
- Coverage 62.69% 10.85% -51.84%
===========================================
Files 65 68 +3
Lines 7239 7561 +322
===========================================
- Hits 4538 820 -3718
- Misses 2701 6741 +4040
... and 58 files with indirect coverage changes 📣 We’re building smart automated test selection to slash your CI/CD build times. Learn more |
orausch
left a comment
orausch
left a comment
There was a problem hiding this comment.
Thanks for the PR, looks very cool! I will need some time to understand this and try it out.
Two things I noticed on a first glance:
- You should add
pytorch_geometricto thesetup.pyto allow CI to run. - I see that you've made some changes to
symbolic_shape_infer.py. In general, it would be good to avoid touching this file since it is just vendored from https://github.com/microsoft/onnxruntime/blob/main/onnxruntime/python/tools/symbolic_shape_infer.py. At a first glance it looks like we should be able to get your functionality if we just add your replacements toself.dispatcher_after instantiation. I also see that the MSFT folks have been doing something similar to you with theiratenreplacements for theirORTModulewrapper. Maybe the new version ofsymbolic_shape_inferhas something similar to what you need?
…x bias bug, simplify GAT tests, add edge weights to GCN conv.
setup.py
Outdated
| 'coverage', 'pytest', 'yapf==0.31', 'pytest-cov', 'transformers', | ||
| 'pytest-xdist', 'torchvision', 'tabulate', 'efficientnet_pytorch', | ||
| 'pytest-timeout' | ||
| 'pytest-timeout', 'pytorch_geometric' |
There was a problem hiding this comment.
That's not the right line. https://pytorch-geometric.readthedocs.io/en/latest/notes/installation.html
There was a problem hiding this comment.
Changed it to the correct dependency but I'm not sure how to test it correctly.
…rings in replacement entries.
…nference class gets edited on use. (very ugly :( )
tbennun
left a comment
tbennun
left a comment
There was a problem hiding this comment.
Looks much better! some comments remain
examples/gnn_benchmark/util.py
Outdated
| print("Changing schdeule to TB dynamic: ", node[0].map) | ||
| node[0].schedule = dace.dtypes.ScheduleType.GPU_ThreadBlock_Dynamic | ||
| and len(node[0].map.params): | ||
| if node[0].label in [ |
There was a problem hiding this comment.
Still not a big fan of this. Probably needs some explanation comment.
There was a problem hiding this comment.
I restructured it a bit: now you provide an external exclude list that's specified in the benchmark.py file. All exclusions have comments describing why they're excluded. This way if the loop naming changes due to implementation changes:
- the code will not work because erroneous loops will be tb-scheduled,
- there will be a warning saying that some loops you wanted to exclude from tb scheduling do not exist in the graph.
I think that's way better than previous behavior (which was silently not changing the schedule of the loops).
| 'coverage', 'pytest', 'yapf==0.31', 'pytest-cov', 'transformers', | ||
| 'pytest-xdist', 'torchvision', 'tabulate', 'efficientnet_pytorch', | ||
| 'pytest-timeout', 'pytorch_geometric' | ||
| 'pytest-timeout', 'torch-geometric' |
There was a problem hiding this comment.
this is not the right way to install pyg: https://github.com/pyg-team/pytorch_geometric#pytorch-113
There was a problem hiding this comment.
@lamyiowce if there is something you can't put in the setup.py, you can always add it to the makefile, which is called by CI: https://github.com/spcl/daceml/blob/master/Makefile#L36
Co-authored-by: Tal Ben-Nun <tbennun@users.noreply.github.com>
3 participants
Implementing module replacements:
daceml/onnx/nodes/replacement.pyanddaceml/pytorch/module_replacement.pydaceml/onnx/nodes/replacement_entries.py, along with appropriate tests comparing to Pytorch Geometric implementationsexamples