From 39313454e8f19905db5e11311a9c5e1d58ad250d Mon Sep 17 00:00:00 2001
From: Alexander Pivovarov <pivovaa@amazon.com>
Date: Mon, 2 May 2022 10:04:36 -0700
Subject: [PATCH] [TRT] Add check to support split op with TRT 5.1.5+ (#11154)

---
 docs/how_to/deploy/tensorrt.rst              | 4 ++++
 src/runtime/contrib/tensorrt/tensorrt_ops.cc | 4 +++-
 2 files changed, 7 insertions(+), 1 deletion(-)

diff --git a/docs/how_to/deploy/tensorrt.rst b/docs/how_to/deploy/tensorrt.rst
index 7950fcfbdbc9..b91f7b1a874a 100644
--- 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
index 2c5f293bc431..e7e83bf9840a 100644
--- 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)