fix: clamp start and end to [0, input_dim] in aten::slice

Signed-off-by: inocsin <[email protected]>

Author: inocsin

core/conversion/converters/converter_util.cpp

@@ -219,6 +219,24 @@ nvinfer1::ITensor* clamp(ConversionCtx* ctx, const torch::jit::Node* n, nvinfer1
   return min_itensor;
 }
 
+// clamp x to [0, input_dim]
+nvinfer1::ITensor* clamp_to_input_dim(ConversionCtx* ctx, const torch::jit::Node* n, nvinfer1::ITensor* x,
+                        nvinfer1::ITensor* input_dim) {
+  auto nbdims = input_dim->getDimensions().d[0];
+  auto zero = torch::zeros({nbdims}).to(torch::kI32);
+  auto zero_itensor = toITensor(ctx, n, &zero);
+  auto one = torch::ones({nbdims}).to(torch::kI32);
+  auto one_itensor = toITensor(ctx, n, &one);
+  auto upper_bound_layer = ctx->net->addElementWise(*input_dim, *one_itensor, nvinfer1::ElementWiseOperation::kSUB);
+  auto upper_bound = upper_bound_layer->getOutput(0);
+  auto max_layer = ctx->net->addElementWise(*x, *zero_itensor, nvinfer1::ElementWiseOperation::kMAX);
+  auto max_itensor = max_layer->getOutput(0);
+  auto min_layer = ctx->net->addElementWise(*max_itensor, *upper_bound, nvinfer1::ElementWiseOperation::kMIN);
+  auto min_itensor = min_layer->getOutput(0);
+  return min_itensor;
+}
+
+
 // return indices < 0 ? inputDims + indices : indices
 nvinfer1::ITensor* bump_if_negtive(ConversionCtx* ctx, const torch::jit::Node* n, nvinfer1::ITensor* input_dim,
                                    nvinfer1::ITensor* indices) {
@@ -238,8 +256,10 @@ nvinfer1::ITensor* bump_if_negtive(ConversionCtx* ctx, const torch::jit::Node* n
 void update_start_and_end(ConversionCtx* ctx, const torch::jit::Node* n, nvinfer1::ITensor* in_shape, 
                         nvinfer1::ITensor* in_start, nvinfer1::ITensor* in_end,
                         nvinfer1::ITensor** out_start, nvinfer1::ITensor** out_end) {
-    *out_start = bump_if_negtive(ctx, n, in_shape, in_start);
-    *out_end = bump_if_negtive(ctx, n, in_shape, in_end);
+    auto start = bump_if_negtive(ctx, n, in_shape, in_start);
+    *out_start = clamp_to_input_dim(ctx, n, start, in_shape);
+    auto end = bump_if_negtive(ctx, n, in_shape, in_end);
+    *out_end = clamp_to_input_dim(ctx, n, end, in_shape);
 }
 
 bool is_dynamic_shape(nvinfer1::ITensor* tensor) {

tests/core/conversion/converters/test_select.cpp

@@ -394,6 +394,35 @@ TEST(Converters, ATenSliceDynamicBatchConvertsCorrectly) {
   ASSERT_TRUE(torch_tensorrt::tests::util::almostEqual(jit_results[0], trt, 2e-6));
 }
 
+TEST(Converters, ATenSliceDynamicBatchLargeEndConvertsCorrectly) {
+  const auto graph = R"IR(
+        graph(%x.1 : Tensor):
+              %2 : None = prim::Constant()
+              %dim : int = prim::Constant[value=0]()
+              %start : int = prim::Constant[value=1]()
+              %end : int = prim::Constant[value=99999]()
+              %step : int = prim::Constant[value=2]()
+              %9 : Tensor = aten::slice(%x.1, %dim, %start, %end, %step)
+              return (%9))IR";
+
+  auto g = std::make_shared<torch::jit::Graph>();
+
+  torch::jit::parseIR(graph, g.get());
+
+  auto in = at::randint(1, 10, {16, 32}, {at::kCUDA});
+
+  auto jit_in = at::clone(in);
+  auto params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {});
+  auto jit_results = torch_tensorrt::tests::util::RunGraph(g, params, {jit_in});
+
+  auto trt_in = at::clone(in);
+  // dynamic shape in batch
+  auto trt_results = torch_tensorrt::tests::util::RunGraphEngineDynamic(g, params, {trt_in}, true);
+  auto trt = trt_results[0].reshape(jit_results[0].sizes());
+
+  ASSERT_TRUE(torch_tensorrt::tests::util::almostEqual(jit_results[0], trt, 2e-6));
+}
+
 TEST(Converters, ATenSliceDynamicNegStartBatchConvertsCorrectly) {
   const auto graph = R"IR(
         graph(%x.1 : Tensor):