[incubator-mxnet] branch v1.x updated: [1.x] Enable CUDA Graphs for TRT (#19184)

[sa...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

samskalicky pushed a commit to branch v1.x
in repository https://gitbox.apache.org/repos/asf/incubator-mxnet.git


The following commit(s) were added to refs/heads/v1.x by this push:
     new a35d568  [1.x] Enable CUDA Graphs for TRT (#19184)
a35d568 is described below

commit a35d568a09c7e54338644130d01f0697e3b19a57
Author: Serge Panev <sp...@nvidia.com>
AuthorDate: Sat Sep 19 15:17:12 2020 -0700

    [1.x] Enable CUDA Graphs for TRT (#19184)
    
    Signed-off-by: Serge Panev <sp...@nvidia.com>
---
 src/operator/subgraph/tensorrt/tensorrt.cc | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/src/operator/subgraph/tensorrt/tensorrt.cc b/src/operator/subgraph/tensorrt/tensorrt.cc
index 1c60d02..71d096e 100644
--- a/src/operator/subgraph/tensorrt/tensorrt.cc
+++ b/src/operator/subgraph/tensorrt/tensorrt.cc
@@ -370,6 +370,11 @@ NNVM_REGISTER_OP(_TensorRT)
     .set_attr<nnvm::FListInputNames>("FListInputNames", TRTListInputNames)
     .set_attr<nnvm::FListOutputNames>("FListOutputNames", DefaultSubgraphOpListOutputs)
     .set_attr<FCreateOpState>("FCreateOpState", TRTCreateState)
+    .set_attr<FIsCUDAGraphsCompatible>("FIsCUDAGraphsCompatible",
+        [](const NodeAttrs& attrs, const bool) {
+          const TRTParam& param = nnvm::get<TRTParam>(attrs.parsed);
+          return !param.int8_mode;
+        })
     .set_attr<FInferStorageType>("FInferStorageType", TRTInferStorageType);
 
 MXNET_REGISTER_SUBGRAPH_BACKEND(TensorRT);