[TRT] Add check to support split op with TRT 5.1.5+ (apache#11154)

SebastianBoblest · May 27, 2022 · 3931345 · 3931345
1 parent 13f866c
commit 3931345
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 1 deletion.
diff --git a/docs/how_to/deploy/tensorrt.rst b/docs/how_to/deploy/tensorrt.rst
@@ -193,6 +193,8 @@ Operator support
 +------------------------+------------------------------------+
 | nn.softmax             |                                    |
 +------------------------+------------------------------------+
+| nn.conv1d              |                                    |
++------------------------+------------------------------------+
 | nn.conv2d              |                                    |
 +------------------------+------------------------------------+
 | nn.dense               |                                    |
@@ -279,6 +281,8 @@ Operator support
 +------------------------+------------------------------------+
 | floor                  | Requires TensorRT 5.1.5 or greater |
 +------------------------+------------------------------------+
+| split                  | Requires TensorRT 5.1.5 or greater |
++------------------------+------------------------------------+
 | strided_slice          | Requires TensorRT 5.1.5 or greater |
 +------------------------+------------------------------------+
 | nn.conv3d              | Requires TensorRT 6.0.1 or greater |

diff --git a/src/runtime/contrib/tensorrt/tensorrt_ops.cc b/src/runtime/contrib/tensorrt/tensorrt_ops.cc
@@ -857,6 +857,7 @@ class ConcatOpConverter : public TensorRTOpConverter {
   }
 };
 
+#if TRT_VERSION_GE(5, 1, 5)
 class SplitOpConverter : public TensorRTOpConverter {
  public:
   SplitOpConverter() : TensorRTOpConverter({kTensor}) {}
@@ -903,6 +904,7 @@ class SplitOpConverter : public TensorRTOpConverter {
     }
   }
 };
+#endif
 
 class BiasAddOpConverter : public TensorRTOpConverter {
  public:
@@ -1285,7 +1287,6 @@ GetOpConverters() {
   map->emplace("expand_dims", std::make_shared<ExpandDimsOpConverter>());
   map->emplace("squeeze", std::make_shared<SqueezeOpConverter>());
   map->emplace("concatenate", std::make_shared<ConcatOpConverter>());
-  map->emplace("split", std::make_shared<SplitOpConverter>());
   map->emplace("nn.conv2d_transpose", std::make_shared<Conv2DTransposeOpConverter>());
   map->emplace("transpose", std::make_shared<TransposeOpConverter>());
   map->emplace("layout_transform", std::make_shared<LayoutTransformOpConverter>());
@@ -1307,6 +1308,7 @@ GetOpConverters() {
   map->emplace("atan", std::make_shared<UnaryOpConverter>());
   map->emplace("ceil", std::make_shared<UnaryOpConverter>());
   map->emplace("floor", std::make_shared<UnaryOpConverter>());
+  map->emplace("split", std::make_shared<SplitOpConverter>());
   map->emplace("strided_slice", std::make_shared<StridedSliceOpConverter>());
 #endif  // TRT_VERSION_GE(5, 1, 5)
 #if TRT_VERSION_GE(6, 0, 1)