Support gather with negative indices (onnx#678)

Signed-off-by: Kevin Chen <kevinch@nvidia.com>

Author: kevinch-nv

builtin_op_importers.cpp

@@ -1190,17 +1190,22 @@ DEFINE_BUILTIN_OP_IMPORTER(Floor)
 
 DEFINE_BUILTIN_OP_IMPORTER(Gather)
 {
-    nvinfer1::ITensor& data = convertToTensor(inputs.at(0), ctx);
+    nvinfer1::ITensor* data = &convertToTensor(inputs.at(0), ctx);
     // TRT does not support BOOL input types for this node
-    ASSERT(data.getType() != nvinfer1::DataType::kBOOL, ErrorCode::kUNSUPPORTED_NODE);
-    nvinfer1::ITensor& indices = convertToTensor(inputs.at(1), ctx);
+    ASSERT( (data->getType() != nvinfer1::DataType::kBOOL) && "This version of TensorRT does not support BOOL input type for the Gather operator.", ErrorCode::kUNSUPPORTED_NODE);
+
+    nvinfer1::ITensor* indices = &convertToTensor(inputs.at(1), ctx);
     OnnxAttrs attrs(node, ctx);
-    int axis = attrs.get<int>("axis", 0);
-    int nbDims = inputs.at(0).shape().nbDims;
+    int32_t axis = attrs.get<int32_t>("axis", 0);
+    int32_t nbDims = inputs.at(0).shape().nbDims;
     TRT_CHECK(convertAxis(axis, nbDims));
     LOG_VERBOSE("Using Gather axis: " << axis);
-    auto* layer = ctx->network()->addGather(data, indices, axis);
-    ctx->registerLayer(layer, node.name());
+
+    // Convert any negative indices to positive ones
+    indices = convertGatherIndices(ctx, data, indices, axis);
+
+    auto* layer = ctx->network()->addGather(*data, *indices, axis);
+    ctx->registerLayer(layer, getNodeName(node));
     RETURN_FIRST_OUTPUT(layer);
 }
 
@@ -1231,11 +1236,11 @@ DEFINE_BUILTIN_OP_IMPORTER(GatherElements)
     */
     // clang-format on
 
-    nvinfer1::ITensor& data = convertToTensor(inputs.at(0), ctx);
-    nvinfer1::ITensor& index = convertToTensor(inputs.at(1), ctx);
+    nvinfer1::ITensor* data = &convertToTensor(inputs.at(0), ctx);
+    nvinfer1::ITensor* index = &convertToTensor(inputs.at(1), ctx);
 
-    const nvinfer1::Dims& idxDims = index.getDimensions();
-    const nvinfer1::Dims& daDims = data.getDimensions();
+    const nvinfer1::Dims& idxDims = index->getDimensions();
+    const nvinfer1::Dims& daDims = data->getDimensions();
 
     // Note the above tranformation requires dimensions to be known at parse time, so check for dynamic shapes
     ASSERT(!isDynamic(daDims) && !isDynamic(idxDims)
@@ -1246,6 +1251,9 @@ DEFINE_BUILTIN_OP_IMPORTER(GatherElements)
     int32_t axis = attrs.get<int32_t>("axis", 0);
     int32_t dataNbDims = daDims.nbDims;
 
+    // Convert any negative indices to positive ones
+    index = convertGatherIndices(ctx, data, index, axis);
+
     TRT_CHECK(convertAxis(axis, dataNbDims));
     LOG_VERBOSE("Using Gather axis: " << axis);
 
@@ -1262,12 +1270,12 @@ DEFINE_BUILTIN_OP_IMPORTER(GatherElements)
     auto* biasTensor = addConstant(ctx, biasVector, ::ONNX_NAMESPACE::TensorProto::INT32, idxDims)->getOutput(0);
 
     auto* mul
-        = ctx->network()->addElementWise(index, *axisPitchTensor, nvinfer1::ElementWiseOperation::kPROD)->getOutput(0);
+        = ctx->network()->addElementWise(*index, *axisPitchTensor, nvinfer1::ElementWiseOperation::kPROD)->getOutput(0);
     auto* newIndices
         = ctx->network()->addElementWise(*mul, *biasTensor, nvinfer1::ElementWiseOperation::kSUM)->getOutput(0);
 
     nvinfer1::Dims flattenDataDims{1, {static_cast<int32_t>(volume(daDims))}};
-    auto* reshape = ctx->network()->addShuffle(data);
+    auto* reshape = ctx->network()->addShuffle(*data);
     reshape->setReshapeDimensions(flattenDataDims);
     reshape->setZeroIsPlaceholder(false);
 
@@ -1277,7 +1285,6 @@ DEFINE_BUILTIN_OP_IMPORTER(GatherElements)
     RETURN_FIRST_OUTPUT(layer);
 }
 
-
 DEFINE_BUILTIN_OP_IMPORTER(Gemm)
 {
     OnnxAttrs attrs(node, ctx);

onnx2trt_utils.cpp

@@ -405,6 +405,37 @@ onnx2trt::ShapedWeights createZeroShifts(const onnx2trt::ShapedWeights& shiftInt
     return shift;
 }
 
+nvinfer1::ITensor* createZeroTensor(IImporterContext* ctx, nvinfer1::ITensor* data)
+{
+    nvinfer1::ITensor* zero;
+    if (data->getType() == nvinfer1::DataType::kFLOAT)
+    {
+        zero
+            = addConstant(ctx, std::vector<float>{0.f}, ::ONNX_NAMESPACE::TensorProto::FLOAT, {0, {1}})->getOutput(0);
+    }
+    else
+    {
+        zero
+            = addConstant(ctx, std::vector<int>{0}, ::ONNX_NAMESPACE::TensorProto::INT32, {0, {1}})->getOutput(0);
+    }
+    broadcastTensors(ctx, zero, data);
+    zero = ctx->network()->addElementWise(*data, *zero, nvinfer1::ElementWiseOperation::kPROD)->getOutput(0);
+    return zero;
+}
+
+nvinfer1::ITensor* convertGatherIndices(IImporterContext* ctx, nvinfer1::ITensor* data, nvinfer1::ITensor* indices, int32_t axis)
+{
+    // Create a condition tensor that is 1 for the elements in indices that are < 0 or 0 otherwise
+    auto condition = ctx->network()->addElementWise(*indices, *createZeroTensor(ctx, indices), nvinfer1::ElementWiseOperation::kLESS)->getOutput(0);
+    auto axisLength = getAxisLength(ctx, data, axis);
+    broadcastTensors(ctx, axisLength, indices);
+    // Create a shifted tensor that is indices + axisLength
+    auto shifted = ctx->network()->addElementWise(*indices, *axisLength, nvinfer1::ElementWiseOperation::kSUM)->getOutput(0);
+    // Select between the shifted and original data based on condition
+    auto select = ctx->network()->addSelect(*condition, *shifted, *indices);
+    return select->getOutput(0);
+}
+
 template <typename DataType>
 DataType* convertINT32Data(const int32_t* weightValues, nvinfer1::Dims shape, int32_t onnxdtype, IImporterContext* ctx)
 {

onnx2trt_utils.hpp

@@ -181,13 +181,19 @@ bool convertDtype(int32_t onnx_dtype, nvinfer1::DataType* trt_dtype);
 // Helper function to convert INT64 weight values into INT32
 int32_t* convertINT64(const int64_t* weightValues, nvinfer1::Dims shape, IImporterContext* ctx);
 
+// Helper function to convert negative gather indices into positive ones
+nvinfer1::ITensor* convertGatherIndices(IImporterContext* ctx, nvinfer1::ITensor* data, nvinfer1::ITensor* indices, int32_t axis);
+
 // Helper function to convert ONNX padding into TRT padding
 bool convertOnnxPadding(
     const std::vector<int64_t>& onnxPadding, nvinfer1::Dims2* begPadding, nvinfer1::Dims2* endPadding);
 
 // Helper function to create zero shifts for QuantizeLinear/DequantizeLinear ops
 onnx2trt::ShapedWeights createZeroShifts(const onnx2trt::ShapedWeights& shiftInt8, int32_t type, IImporterContext* ctx);
 
+// Helper function to create a tensor of all zeros with the same shape as a data tensor
+nvinfer1::ITensor* createZeroTensor(IImporterContext* ctx, nvinfer1::ITensor* data);
+
 // Helper function to convert an ONNX weight into a ShapedWeights object
 bool convertOnnxWeights(
     const ::ONNX_NAMESPACE::TensorProto& onnxTensor, onnx2trt::ShapedWeights* weights, IImporterContext* ctx);