fix: fix aten::slice when index is LONG_MAX

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

Author: inocsin

core/conversion/converters/converter_util.cpp

@@ -218,8 +218,8 @@ nvinfer1::ITensor* clamp(
 }
 
 // clamp x to [0, input_dim]
-nvinfer1::ITensor* clamp_to_input_dim(ConversionCtx* ctx, nvinfer1::ITensor* x, nvinfer1::ITensor* input_dim) {
-  auto nbdims = input_dim->getDimensions().d[0];
+nvinfer1::ITensor* clamp_to_input_dim(ConversionCtx* ctx, nvinfer1::ITensor* x, nvinfer1::ITensor* input_dim, int nbdims) {
+  // auto nbdims = input_dim->getDimensions().d[0];
   auto zero = torch::zeros({nbdims}).to(torch::kI32);
   auto zero_itensor = tensor_to_const(ctx, zero);
   auto one = torch::ones({nbdims}).to(torch::kI32);
@@ -243,8 +243,7 @@ nvinfer1::ITensor* clamp_to_input_dim(ConversionCtx* ctx, nvinfer1::ITensor* x,
 }
 
 // return indices < 0 ? inputDims + indices : indices
-nvinfer1::ITensor* bump_if_negtive(ConversionCtx* ctx, nvinfer1::ITensor* input_dim, nvinfer1::ITensor* indices) {
-  auto nbdims = input_dim->getDimensions().d[0];
+nvinfer1::ITensor* bump_if_negtive(ConversionCtx* ctx, nvinfer1::ITensor* input_dim, nvinfer1::ITensor* indices, int nbdims) {
   auto zero = torch::zeros({nbdims}).to(torch::kI32);
   auto neg = -torch::ones({nbdims}).to(torch::kI32);
   auto zero_itensor = tensor_to_const(ctx, zero);
@@ -270,11 +269,12 @@ std::vector<nvinfer1::ITensor*> update_start_and_end(
     ConversionCtx* ctx,
     nvinfer1::ITensor* in_shape,
     nvinfer1::ITensor* in_start,
-    nvinfer1::ITensor* in_end) {
-  auto start = bump_if_negtive(ctx, in_shape, in_start);
-  auto out_start = clamp_to_input_dim(ctx, start, in_shape);
-  auto end = bump_if_negtive(ctx, in_shape, in_end);
-  auto out_end = clamp_to_input_dim(ctx, end, in_shape);
+    nvinfer1::ITensor* in_end,
+    int nbdims) {
+  auto start = bump_if_negtive(ctx, in_shape, in_start, nbdims);
+  auto out_start = clamp_to_input_dim(ctx, start, in_shape, nbdims);
+  auto end = bump_if_negtive(ctx, in_shape, in_end, nbdims);
+  auto out_end = clamp_to_input_dim(ctx, end, in_shape, nbdims);
   std::vector<nvinfer1::ITensor*> outputs;
   outputs.push_back(out_start);
   outputs.push_back(out_end);

core/conversion/converters/converter_util.h

@@ -63,7 +63,8 @@ std::vector<nvinfer1::ITensor*> update_start_and_end(
     ConversionCtx* ctx,
     nvinfer1::ITensor* in_shape,
     nvinfer1::ITensor* in_start,
-    nvinfer1::ITensor* in_end);
+    nvinfer1::ITensor* in_end,
+    int nbdims);
 
 nvinfer1::ITensor* calculate_output_size(
     ConversionCtx* ctx,

core/conversion/converters/impl/select.cpp

@@ -318,7 +318,8 @@ auto select_registrations TORCHTRT_UNUSED =
                int startIdx = 0;
                auto startIdxIVal = args[2].IValue();
                if (!startIdxIVal->isNone()) {
-                 startIdx = std::min((int64_t)std::numeric_limits<int32_t>::max(), startIdxIVal->toInt());
+                 startIdx = startIdxIVal->toInt() > std::numeric_limits<int32_t>::max() ? maxDim : startIdxIVal->toInt();
+                 startIdx = maxDim == -1 ? startIdx : std::min(startIdx, maxDim);
                }
                // Handle case when given tensor index is negative
                if (maxDim > 0) { // only for static shape
@@ -329,7 +330,7 @@ auto select_registrations TORCHTRT_UNUSED =
                int endIdx = maxDim; // -1 for dynamic shape
                auto endIdxIVal = args[3].IValue();
                if (!endIdxIVal->isNone()) {
-                 int truncate_value = std::min((int64_t)std::numeric_limits<int32_t>::max(), endIdxIVal->toInt());
+                 int truncate_value = endIdxIVal->toInt() > std::numeric_limits<int32_t>::max() ? maxDim : endIdxIVal->toInt();
                  endIdx = maxDim == -1 ? truncate_value : std::min(truncate_value, maxDim);
                }
                if (maxDim > 0) {
@@ -373,7 +374,7 @@ auto select_registrations TORCHTRT_UNUSED =
                  at::Tensor end_tensor = torch::zeros({nbdims}).to(torch::kI32);
                  for (int i = 0; i < nbdims; i++) {
                    if (i == axis) {
-                     end_tensor[i] = endIdxIVal->isNone() ? -1 : endIdx - 1;
+                     end_tensor[i] = endIdx == -1 ? -1 : endIdx - 1;
                    } else {
                      end_tensor[i] = input_dim.d[i] == -1 ? -1 : input_dim.d[i] - 1;
                    }
@@ -383,7 +384,7 @@ auto select_registrations TORCHTRT_UNUSED =
                  // update start and end
                  nvinfer1::ITensor* out_start;
                  nvinfer1::ITensor* out_end;
-                 auto start_end = update_start_and_end(ctx, ishape_tensor, start_itensor, end_itensor);
+                 auto start_end = update_start_and_end(ctx, ishape_tensor, start_itensor, end_itensor, nbdims);
                  out_start = start_end[0];
                  out_end = start_end[1];