-
Notifications
You must be signed in to change notification settings - Fork 3.4k
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
[TVMC][TensorRT] Enable TensorRT target through TVMC #7658
Conversation
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Thanks for the PR. I suggest we coordinate this with #7577, as you're gonna see on my comment. Once that comes in, with a little refactoring we can have this support done.
python/tvm/driver/tvmc/compiler.py
Outdated
@@ -196,9 +196,9 @@ def compile_model( | |||
for codegen_from_cli in extra_targets: | |||
codegen = composite_target.get_codegen_by_target(codegen_from_cli["name"]) | |||
partition_function = codegen["pass_pipeline"] | |||
mod = partition_function(mod, params) | |||
mod, codegen_config = partition_function(mod, params) |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
In the interest of standardizing all the partition_for_*
functions, I think it would be good to do a little refactoring in the partition_for_tensorrt
so that it only returns mod
. This is important not only for TVMC, but I think it would be a good move in terms of future partitioning support in TVM.
There is some new functionality (yet to be merged) in #7577 to have an init
function per target in tvmc, in which we can then populate the configs using the existing get_tensorrt_version
.
I suggest we coordinate to get #7577 landed and then adjust this to comply with the suggestions above, so that we can get a uniform support for all targets. What do you think?
(cc @comaniac)
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I agree that it would be easier to ask all pritition functions return a single module. However, the init function introduced by #7577 seems not applicable to this case to me. The config required by TensorRT is used when building the module; while the init function used by Vitis-AI imports the required package. I have no idea how could we make it general enough to support both Vitis-AI and TensorRT.
On the other hand, let a partition function return a tuple of module and build config seems more straightforward at the first glance, as I cannot think of a better idea for now. First, partition functions won't be a Relay pass in a pass sequence (at least for now). Second, it is reasonable that the build config may need to be constructed during the partition. Separating them to two APIs may introduce overheads and be more confusing to non TVMC users, which are still majority in this community.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I agree with @comaniac, I don't see a good way to separate this yet. The parameters used for parition_for_tensorrt
affect both partitioning and building the module.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Fair enough. I thought a bit more about it and as with #7577, we will be already sending a config dictionary for the partition function. so I think it is ok to get the same, or a modified version of that dictionary as a result, which fits with the existing support on TensorRT.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I guess an alternative could be for the partitioning to be done within the pass context, and have the partition function set the config in the passcontext.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
That's our ultimate goal of the composite target, which integrates partition and other codegen specific functions. Once we have it, the partition logic could be completely hidden from not only users but also frontend interface.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
For now, I tried to receive and apply the config returned partiton_for_tensrort.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Yeah, I agree this is simpler and effective for now. Happy to approve once we get the linting issues sorted out and CI passing here.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
cc @trevor-m who is the main contributor of TensorRT integration.
python/tvm/driver/tvmc/compiler.py
Outdated
if codegen["config_key"] is not None: | ||
config[codegen["config_key"]] = codegen_from_cli["opts"] | ||
config[codegen["config_key"]] = codegen_from_cli["opts"] or codegen_config |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
It means the generated codegen_config will be completely overwritten by the user config, but we should make it like
codegen_config.update(codegen_from_cli["opts"] if codegen_from_cli["opts"] is not None else {})
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
The config returned by partiton_for_tensorrt
reflects the user config by thecodegen_from_cli
argument.
And, the other composite_targets don't return config.
So, config with a extra_codegen
can be overwritten codegen_config if any.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
The config returned by partiton_for_tensorrt reflects the user config by the codegen_from_cli argument.
This is not guaranteed. What if partition_for_xxx
only uses a part of the configs in codegen_from_cli
?
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Sorry, I don't know it means.
The implementation before this change use all of codegen_from_cli
as the config.
If should consider only uses a part of partiton_for_xxx
, it's due to not this change but original implementation.
Or, should it apply also in this change?
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Sorry I misinterpreted the meaning of codegen_from_cli
(their names are sort of confusing...).
Since codegen_from_cli
is given by each composite target, this way should be fine.
python/tvm/driver/tvmc/compiler.py
Outdated
@@ -196,9 +196,9 @@ def compile_model( | |||
for codegen_from_cli in extra_targets: | |||
codegen = composite_target.get_codegen_by_target(codegen_from_cli["name"]) | |||
partition_function = codegen["pass_pipeline"] | |||
mod = partition_function(mod, params) | |||
mod, codegen_config = partition_function(mod, params) |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I agree that it would be easier to ask all pritition functions return a single module. However, the init function introduced by #7577 seems not applicable to this case to me. The config required by TensorRT is used when building the module; while the init function used by Vitis-AI imports the required package. I have no idea how could we make it general enough to support both Vitis-AI and TensorRT.
On the other hand, let a partition function return a tuple of module and build config seems more straightforward at the first glance, as I cannot think of a better idea for now. First, partition functions won't be a Relay pass in a pass sequence (at least for now). Second, it is reasonable that the build config may need to be constructed during the partition. Separating them to two APIs may introduce overheads and be more confusing to non TVMC users, which are still majority in this community.
0b9945f
to
e5b1a92
Compare
Sorry for my late fix, please review again! |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
LGTM. Please fix the CI.
@akmaru Any updates on this? |
Gentle ping @akmaru please fix the CI |
@akmaru can you update? |
Adds TensorRT integration as a composite target into TVMC.