[CUDNN] Add cuDNN as a Relay partitioning target (BYOC)#10871
[CUDNN] Add cuDNN as a Relay partitioning target (BYOC)#10871masahi merged 3 commits intoapache:mainfrom
Conversation
mbs-octoml
left a comment
mbs-octoml
left a comment
There was a problem hiding this comment.
LGTM modulo two nits.
python/tvm/relay/op/contrib/cudnn.py
Outdated
| assert isinstance(partition.body.op, relay.Function) | ||
|
|
||
| global_name = str(partition.attrs.global_symbol) | ||
| target = tvm.target.cuda() |
There was a problem hiding this comment.
Just notice this, I think Target.current() is better so that cuda params are not lost.
There was a problem hiding this comment.
I had a go using this based on your suggestion, but it seems the Target only ends up in the context if you use the 'with' way of specifying targets (not just directly passing them to relay.build). In an ideal world, this somehow gets plumbed directly through the BYOC/partitioning mechanism - but I think that exceeds the immediate scope of this PR. Perhaps we could leave as a future improvement for now?
python/tvm/relay/op/contrib/cudnn.py
Outdated
| return _register | ||
|
|
||
|
|
||
| @tvm._ffi.register_func("relay.ext.cudnn") |
There was a problem hiding this comment.
Now would be a good time to hoist this boilerplate into a library_byoc.py helper or something similar?
There was a problem hiding this comment.
Done - let me know what you think.
mikepapadim
left a comment
mikepapadim
left a comment
There was a problem hiding this comment.
LGTM, just trigger again CI. It looks like a random fail
mbaret
force-pushed
the
cudnn-relay
branch
2 times, most recently
from
5c2e300 to
0e06391
Compare
|
it looks like some flaky test in aarch64 |
This adds infrastructure to support offloading of Relay patterns to cuDNN. In this initial commit, only softmax is supported.
|
cc @masahi @mbrookhart @junrushao1994 PTAL and merge if you're happy :) |
* [CUDNN] Add cuDNN as a Relay partitioning target (BYOC) This adds infrastructure to support offloading of Relay patterns to cuDNN. In this initial commit, only softmax is supported. * Refactor common TE BYOC code into separate file * Add test guard
* [CUDNN] Add cuDNN as a Relay partitioning target (BYOC) This adds infrastructure to support offloading of Relay patterns to cuDNN. In this initial commit, only softmax is supported. * Refactor common TE BYOC code into separate file * Add test guard
* [CUDNN] Add cuDNN as a Relay partitioning target (BYOC) This adds infrastructure to support offloading of Relay patterns to cuDNN. In this initial commit, only softmax is supported. * Refactor common TE BYOC code into separate file * Add test guard
5 participants
This adds infrastructure to support offloading of Relay patterns to cuDNN. In this initial commit, only softmax is supported. Later PRs will include support for more operators, including some limited fused patterns.