Adding dynamic split support (#326)

Author: kevinch-nv

builtin_op_importers.cpp

@@ -1839,48 +1839,59 @@ DEFINE_BUILTIN_OP_IMPORTER(SpatialBN)
 
 DEFINE_BUILTIN_OP_IMPORTER(Split)
 {
-    ASSERT(inputs.size() == 1, ErrorCode::kUNSUPPORTED_NODE);
-    nvinfer1::ITensor* tensor_ptr = &convertToTensor(inputs.at(0), ctx);
-    nvinfer1::Dims dims = tensor_ptr->getDimensions();
-    int nbDims = dims.nbDims;
+    const int numOutputs = node.output().size();
+
+    nvinfer1::ITensor* tensorPtr = &convertToTensor(inputs.at(0), ctx);
+    const int rank = tensorPtr->getDimensions().nbDims;
+    nvinfer1::ITensor* shape = ctx->network()->addShape(*tensorPtr)->getOutput(0);
+
     OnnxAttrs attrs(node);
     int axis = attrs.get<int>("axis", 0);
-    TRT_CHECK(convert_axis(axis, nbDims));
-    std::vector<int> output_lengths;
-    int noutput = node.output().size();
-    std::vector<int> start_index (noutput, 0);
+    TRT_CHECK(convert_axis(axis, rank));
+
+    std::vector<int> outputLengths;
     if (attrs.count("split"))
     {
-        output_lengths = attrs.get<std::vector<int>>("split");
-        ASSERT(static_cast<int>(output_lengths.size()) == noutput, ErrorCode::kINVALID_NODE);
-    }
-    else
-    {
-        ASSERT(dims.d[axis] == -1 || dims.d[axis] % noutput == 0, ErrorCode::kINVALID_NODE);
-        output_lengths.assign(noutput, dims.d[axis] / noutput);
-    }
-    for (size_t i = 1; i < output_lengths.size(); i++)
-    {
-        start_index[i] = start_index[i - 1] + output_lengths[i - 1];
+        outputLengths = attrs.get<std::vector<int>>("split");
+        ASSERT(static_cast<int>(outputLengths.size()) == numOutputs, ErrorCode::kINVALID_NODE);
     }
 
-    nvinfer1::Dims sliceStart = makeDims(nbDims, 0);
-    nvinfer1::Dims sliceSize = dims;
-    nvinfer1::Dims sliceStride = makeDims(nbDims, 1);
-    std::vector<TensorOrWeights> outputs;
-    for (int i = 0; i < noutput; ++i)
-    {
-        sliceStart.d[axis] = start_index[i];
-        sliceSize.d[axis] = output_lengths[i];
-        auto const layer = ctx->network()->addSlice(*tensor_ptr, sliceStart, sliceSize, sliceStride);
-        if (std::any_of(sliceSize.d, sliceSize.d + sliceSize.nbDims, [](int i){return i == -1;})){
-            layer->setInput(1, dimension_to_tensor(ctx, sliceStart));
-            layer->setInput(2, dimension_to_tensor(ctx, sliceSize));
-            layer->setInput(3, dimension_to_tensor(ctx, sliceStride));
+    nvinfer1::ITensor* startSliceAxis{addConstantScalar<int32_t>(ctx, 0, ::ONNX_NAMESPACE::TensorProto::INT32, nvinfer1::Dims{1, 1})->getOutput(0)};
+    // sizeSliceAxis = axisLength / numOutputs
+    nvinfer1::ITensor* sizeSliceAxis{ctx->network()->addElementWise(
+        *gatherDimension(ctx, shape, axis, nvinfer1::Dims{1, 1}),
+        *addConstantScalar(ctx, numOutputs, ::ONNX_NAMESPACE::TensorProto::INT32, nvinfer1::Dims{1, 1})->getOutput(0),
+        nvinfer1::ElementWiseOperation::kDIV
+    )->getOutput(0)};
+
+    nvinfer1::Dims zeroStartsDims{rank};
+    std::fill(zeroStartsDims.d, zeroStartsDims.d + zeroStartsDims.nbDims, 0);
+    nvinfer1::ITensor* zeroStarts = &makeShapeTensor(ctx, zeroStartsDims);
+
+    nvinfer1::Dims strides{rank};
+    std::fill(strides.d, strides.d + strides.nbDims, 1);
+
+    std::vector<TensorOrWeights> outputs{};
+    for (int i = 0; i < numOutputs; ++i)
+    {
+        if (!outputLengths.empty())
+        {
+            sizeSliceAxis = addConstantScalar(ctx, outputLengths.at(i), ::ONNX_NAMESPACE::TensorProto::INT32, nvinfer1::Dims{1, 1})->getOutput(0);
         }
-        outputs.push_back(layer->getOutput(0));
+
+        nvinfer1::ITensor* starts{overwriteDim(ctx, zeroStarts, startSliceAxis, axis)};
+        nvinfer1::ITensor* sizes{overwriteDim(ctx, shape, sizeSliceAxis, axis)};
+
+        nvinfer1::ISliceLayer* slice = ctx->network()->addSlice(*tensorPtr, nvinfer1::Dims{rank}, nvinfer1::Dims{rank}, strides);
+        slice->setInput(1, *starts);
+        slice->setInput(2, *sizes);
+        outputs.emplace_back(slice->getOutput(0));
+
+        startSliceAxis = ctx->network()->addElementWise(*startSliceAxis, *sizeSliceAxis, nvinfer1::ElementWiseOperation::kSUM)->getOutput(0);
     }
+
     return outputs;
+
 }
 
 DEFINE_BUILTIN_OP_IMPORTER(Sqrt)

onnx2trt_utils.cpp

@@ -709,6 +709,12 @@ nvinfer1::ITensor* flattenTensor(IImporterContext* ctx, nvinfer1::ITensor& tenso
     return flattenTensorStatic(ctx, tensor, axis);
 }
 
+nvinfer1::ITensor* gatherDimension(IImporterContext* ctx, nvinfer1::ITensor* shapeTensor, int dim, nvinfer1::Dims shape)
+{
+    auto& axisValue = *addConstantScalar(ctx, dim, ::ONNX_NAMESPACE::TensorProto_DataType_INT32, shape)->getOutput(0);
+    return ctx->network()->addGather(*shapeTensor, axisValue, 0)->getOutput(0);
+}
+
 bool isDynamic(nvinfer1::Dims const& dims)
 {
     return std::any_of(dims.d, dims.d + dims.nbDims, [](int dim) {return dim == -1;});
@@ -954,6 +960,26 @@ nvinfer1::ITensor& makeShapeTensor(IImporterContext* ctx, nvinfer1::Dims dims)
     return convertToTensor(valueWeights, ctx);
 }
 
+nvinfer1::ITensor* overwriteDim(IImporterContext* ctx, nvinfer1::ITensor* shape, nvinfer1::ITensor* dim, int axis)
+{
+    const int shapeLength = shape->getDimensions().d[0];
+
+    std::vector<nvinfer1::ITensor*> dims{};
+
+    for (int i = 0; i < shapeLength; ++i)
+    {
+        if (i == axis)
+        {
+            dims.emplace_back(dim);
+        }
+        else
+        {
+            dims.emplace_back(gatherDimension(ctx, shape, i, nvinfer1::Dims{1, 1}));
+        }
+    }
+    return ctx->network()->addConcatenation(dims.data(), dims.size())->getOutput(0);
+}
+
 NodeImportResult poolingHelper(IImporterContext* ctx, ::ONNX_NAMESPACE::NodeProto const& node, std::vector<TensorOrWeights>& inputs, nvinfer1::PoolingType type)
 {
     nvinfer1::ITensor* tensorPtr = &convertToTensor(inputs.at(0), ctx);

onnx2trt_utils.hpp

@@ -167,6 +167,10 @@ NodeImportResult elementwiseHelper(IImporterContext* ctx, ::ONNX_NAMESPACE::Node
 // Helper functino to flatten a tensor on a specified axis
 nvinfer1::ITensor* flattenTensor(IImporterContext* ctx, nvinfer1::ITensor& tensor, int axis);
 
+// Gathers the specified dimension from a shape tensor. e.g. gatherDimension(shape=(7, 6, 5), dim=2) would return 5.
+// shape specifies the shape of the returned Tensor. Must have a volume of 1.
+nvinfer1::ITensor* gatherDimension(IImporterContext* ctx, nvinfer1::ITensor* shapeTensor, int dim, nvinfer1::Dims shape);
+
 // Helper function to check if any input dimensions are dynamic
 bool isDynamic (nvinfer1::Dims const& dims);
 
@@ -214,6 +218,10 @@ nvinfer1::Dims makeDims(int nbDims, int val);
 // Helper function to create a shape tensor from a Dims object for dynamic reshape
 nvinfer1::ITensor& makeShapeTensor(IImporterContext* ctx, nvinfer1::Dims dims);
 
+
+// Helper function to overwrite the value of a single axis in a shape tensor
+nvinfer1::ITensor* overwriteDim(IImporterContext* ctx, nvinfer1::ITensor* shape, nvinfer1::ITensor* dim, int axis);
+
 // Helper function to map various ONNX pooling ops into TensorRT.
 NodeImportResult poolingHelper(IImporterContext* ctx, ::ONNX_NAMESPACE::NodeProto const& node, std::vector<TensorOrWeights>& inputs, nvinfer1::PoolingType type);