fix!(aten::adaptive_avg_pool_2d, aten::interpolate): Moves interpolate

align_corners cases to TensorRT and also creates a work around for
dynamic adaptive_avg_pool_2d that runs the kernel on CPU

HACK: Runs adaptive_avg_pool_2d on CPU, this should be replaced when it
can be determined why creating a Tensor on GPU causes segfault. Current
idea is there is a cuDNN version mismatch.

Signed-off-by: Naren Dasan <[email protected]>
Signed-off-by: Naren Dasan <[email protected]>

Author: narendasan

core/conversion/converters/impl/interpolate.cpp

@@ -16,42 +16,22 @@ namespace {
  * Helper functions
  */
 
-void create_plugin(ConversionCtx* ctx, const torch::jit::Node* n, nvinfer1::ITensor* in, const char* name, 
-                                                                                         std::vector<int64_t> in_shape, 
-                                                                                         std::vector<int64_t> out_shape, 
-                                                                                         std::vector<int64_t> out_size, 
-                                                                                         std::string mode) {
-    LOG_WARNING("Interpolation layer will be run through ATen, not TensorRT. Performance may differ.");
-    
-    auto creator = new plugins::InterpolatePluginCreator();
-    auto plugin = creator->createPlugin(name, in_shape, out_shape, out_size, mode, false);
-
-    auto resize_layer = ctx->net->addPluginV2(reinterpret_cast<nvinfer1::ITensor* const*>(&in), 1, *plugin);
-    TRTORCH_CHECK(resize_layer, "Unable to create interpolation plugin from node" << *n);
-
-    resize_layer->setName(util::node_info(n).c_str());
-
-    auto layer_output = ctx->AssociateValueAndTensor(n->outputs()[0], resize_layer->getOutput(0));
-
-    LOG_DEBUG("Output tensor shape: " << layer_output->getDimensions());
-}
-
-void resize_layer_size(ConversionCtx* ctx, const torch::jit::Node* n, nvinfer1::ITensor* in, std::vector<int64_t> out_shape, 
-                                                                                             nvinfer1::ResizeMode mode) {
+void resize_layer_size(ConversionCtx* ctx, const torch::jit::Node* n, nvinfer1::ITensor* in, std::vector<int64_t> out_shape,
+                                                                                             nvinfer1::ResizeMode mode, bool align_corners=false) {
     auto resize_layer = ctx->net->addResize(*in);
     TRTORCH_CHECK(resize_layer, "Unable to create interpolation (resizing) layer from node" << *n);
 
     resize_layer->setOutputDimensions(util::toDims(out_shape));
     resize_layer->setResizeMode(mode);
     resize_layer->setName(util::node_info(n).c_str());
-    
+
     // if interpolation mode is linear, align corners must have been set to true. else, don't use align corners.
     if (mode == nvinfer1::ResizeMode::kLINEAR) {
-        resize_layer->setAlignCorners(true);
+        resize_layer->setAlignCorners(align_corners);
     }
 
     auto layer_output = ctx->AssociateValueAndTensor(n->outputs()[0], resize_layer->getOutput(0));
-    
+
     LOG_DEBUG("Output tensor shape: " << layer_output->getDimensions());
 }
 
@@ -72,7 +52,7 @@ auto interpolate_registrations TRTORCH_UNUSED = RegisterNodeConversionPatterns()
                 auto out_size = util::toVec(util::toDims(args[1].unwrapToIntList()));
 
                 TRTORCH_ASSERT(out_size.size() == 1, "aten::upsample_nearest1d input Tensor and output size dimension mismatch");
-                
+
                 auto out_shape = in_shape;
                 std::copy(out_size.begin(), out_size.end(), out_shape.begin() + (in_shape.size() - out_size.size()));
 
@@ -94,10 +74,10 @@ auto interpolate_registrations TRTORCH_UNUSED = RegisterNodeConversionPatterns()
                 auto out_size = util::toVec(util::toDims(args[1].unwrapToIntList()));
 
                 TRTORCH_ASSERT(out_size.size() == 2, "aten::upsample_nearest2d input Tensor and output size dimension mismatch");
-                
+
                 auto out_shape = in_shape;
                 std::copy(out_size.begin(), out_size.end(), out_shape.begin() + (in_shape.size() - out_size.size()));
-                
+
                 resize_layer_size(ctx, n, in, out_shape, nvinfer1::ResizeMode::kNEAREST);
             } else {
                 TRTORCH_THROW_ERROR("Unable to convert node: " << util::node_info(n) << "\nScale factor parameter for upsample_nearest2d not supported yet.");
@@ -116,7 +96,7 @@ auto interpolate_registrations TRTORCH_UNUSED = RegisterNodeConversionPatterns()
                 auto out_size = util::toVec(util::toDims(args[1].unwrapToIntList()));
 
                 TRTORCH_ASSERT(out_size.size() == 3, "aten::upsample_nearest3d input Tensor and output size dimension mismatch");
-                
+
                 auto out_shape = in_shape;
                 std::copy(out_size.begin(), out_size.end(), out_shape.begin() + (in_shape.size() - out_size.size()));
 
@@ -139,16 +119,11 @@ auto interpolate_registrations TRTORCH_UNUSED = RegisterNodeConversionPatterns()
                 auto out_size = util::toVec(util::toDims(args[1].unwrapToIntList()));
 
                 TRTORCH_ASSERT(out_size.size() == 1, "aten::upsample_linear1d input Tensor and output size dimension mismatch");
-                
-                auto out_shape = in_shape; 
+
+                auto out_shape = in_shape;
                 std::copy(out_size.begin(), out_size.end(), out_shape.begin() + (in_shape.size() - out_size.size()));
 
-                if (!align_corners) {
-                    // align_corners not supported in TensorRT, create plugin and run layer through PyTorch
-                    create_plugin(ctx, n, in, "linear1d", in_shape, out_shape, out_size, std::string("linear"));
-                } else {
-                    resize_layer_size(ctx, n, in, out_shape, nvinfer1::ResizeMode::kLINEAR);
-                }
+                resize_layer_size(ctx, n, in, out_shape, nvinfer1::ResizeMode::kLINEAR, align_corners);
             } else {
                 TRTORCH_THROW_ERROR("Unable to convert node: " << util::node_info(n) << "\nScale factor parameter for upsample_linear1d not supported yet.");
             }
@@ -167,16 +142,11 @@ auto interpolate_registrations TRTORCH_UNUSED = RegisterNodeConversionPatterns()
                 auto out_size = util::toVec(util::toDims(args[1].unwrapToIntList()));
 
                 TRTORCH_ASSERT(out_size.size() == 2, "aten::upsample_bilinear2d input Tensor and output size dimension mismatch");
-                
+
                 auto out_shape = in_shape;
                 std::copy(out_size.begin(), out_size.end(), out_shape.begin() + (in_shape.size() - out_size.size()));
 
-                if (!align_corners) {
-                    // align_corners not supported in TensorRT, create plugin and run layer through PyTorch
-                    create_plugin(ctx, n, in, "bilinear2d", in_shape, out_shape, out_size, std::string("bilinear"));
-                } else {
-                    resize_layer_size(ctx, n, in, out_shape, nvinfer1::ResizeMode::kLINEAR);
-                }
+                resize_layer_size(ctx, n, in, out_shape, nvinfer1::ResizeMode::kLINEAR, align_corners);
             } else {
                 TRTORCH_THROW_ERROR("Unable to convert node: " << util::node_info(n) << "\nScale factor parameter for upsample_bilinear2d not supported yet.");
             }
@@ -195,16 +165,11 @@ auto interpolate_registrations TRTORCH_UNUSED = RegisterNodeConversionPatterns()
                 auto out_size = util::toVec(util::toDims(args[1].unwrapToIntList()));
 
                 TRTORCH_ASSERT(out_size.size() == 3, "aten::upsample_trilinear3d input Tensor and output size dimension mismatch");
-                
+
                 auto out_shape = in_shape;
                 std::copy(out_size.begin(), out_size.end(), out_shape.begin() + (in_shape.size() - out_size.size()));
 
-                if (!align_corners) {
-                    // align_corners not supported in TensorRT, create plugin and run layer through PyTorch
-                    create_plugin(ctx, n, in, "trilinear3d", in_shape, out_shape, out_size, std::string("trilinear"));
-                } else {
-                    resize_layer_size(ctx, n, in, out_shape, nvinfer1::ResizeMode::kLINEAR);
-                }
+                resize_layer_size(ctx, n, in, out_shape, nvinfer1::ResizeMode::kLINEAR, align_corners);
             } else {
                 TRTORCH_THROW_ERROR("Unable to convert node: " << util::node_info(n) << "\nScale factor parameter for upsample_trilinear3d not supported yet.");
             }

core/conversion/converters/impl/plugins/BUILD

@@ -24,6 +24,12 @@ cc_library(
         "//conditions:default":  ["@libtorch//:libtorch"],
     }),
     alwayslink = True,
+    copts = [
+        "-pthread"
+    ],
+    linkopts = [
+        "-lpthread",
+    ]
 )
 
 load("@rules_pkg//:pkg.bzl", "pkg_tar")

core/conversion/converters/impl/plugins/interpolate_plugin.cpp

@@ -9,57 +9,57 @@ namespace converters {
 namespace impl {
 namespace plugins {
 
-/* 
+/*
  * InterpolatePlugin class implementations
  */
 
-InterpolatePlugin::InterpolatePlugin(std::vector<int64_t> in_shape, std::vector<int64_t> out_shape, std::vector<int64_t> size, std::string mode, bool align_corners) : 
-    in_shape(in_shape), out_shape(out_shape), size(size), mode(mode), align_corners(align_corners) 
+InterpolatePlugin::InterpolatePlugin(std::vector<int64_t> in_shape, std::vector<int64_t> out_shape, std::vector<int64_t> size, std::string mode, bool align_corners) :
+    in_shape_(in_shape), out_shape_(out_shape), size_(size), mode_(mode), align_corners_(align_corners)
 {}
 
 InterpolatePlugin::InterpolatePlugin(const char *data, size_t length) {
     std::istringstream data_stream(std::string(data, length));
-    
+
     torch::serialize::InputArchive input_archive;
     input_archive.load_from(data_stream);
-    
+
     {
         torch::IValue value;
         input_archive.read("in_shape", value);
-        in_shape = value.toIntVector();
+        in_shape_ = value.toIntVector();
     }
     {
         torch::IValue value;
         input_archive.read("out_shape", value);
-        out_shape = value.toIntVector();
+        out_shape_ = value.toIntVector();
     }
     {
         torch::IValue value;
         input_archive.read("size", value);
-        size = value.toIntVector();
+        size_ = value.toIntVector();
     }
     {
         torch::IValue value;
         input_archive.read("mode", value);
-        mode = value.toStringRef();
+        mode_ = value.toStringRef();
     }
     {
         torch::IValue value;
         input_archive.read("align_corners", value);
-        align_corners = value.toBool();
+        align_corners_ = value.toBool();
     }
 }
 
 std::vector<int64_t> InterpolatePlugin::getInputShape() {
-    return in_shape;
+    return in_shape_;
 }
 
 std::vector<int64_t> InterpolatePlugin::getOutputShape() {
-    return out_shape;
+    return out_shape_;
 }
 
 std::vector<int64_t> InterpolatePlugin::getOutputSize() {
-    return size;
+    return size_;
 }
 
 int InterpolatePlugin::getNbOutputs() const {
@@ -80,14 +80,14 @@ const char* InterpolatePlugin::getPluginNamespace() const {
 
 
 nvinfer1::IPluginV2DynamicExt* InterpolatePlugin::clone() const {
-    return new InterpolatePlugin(in_shape, out_shape, size, mode, align_corners);
+    return new InterpolatePlugin(in_shape_, out_shape_, size_, mode_, align_corners_);
 }
 
 nvinfer1::DimsExprs InterpolatePlugin::getOutputDimensions(int outputIndex, const nvinfer1::DimsExprs *inputs, int nbInputs, nvinfer1::IExprBuilder &exprBuilder) {
    nvinfer1::DimsExprs output(inputs[0]);
 
-   for (unsigned int i = 0; i < out_shape.size(); i++) {
-       output.d[i] = exprBuilder.constant(out_shape[i]);
+   for (unsigned int i = 0; i < out_shape_.size(); i++) {
+       output.d[i] = exprBuilder.constant(out_shape_[i]);
    }
 
    return output;
@@ -98,10 +98,10 @@ nvinfer1::DataType InterpolatePlugin::getOutputDataType(int index, const nvinfer
 }
 
 int InterpolatePlugin::initialize() {
-    tensor_options = tensor_options.device(c10::kCUDA);
+    tensor_options_ = tensor_options_.device(c10::kCPU);
 
     // c10::kFloat = FLOAT32
-    tensor_options = tensor_options.dtype(c10::kFloat);
+    tensor_options_ = tensor_options_.dtype(c10::kFloat);
 
     return 0;
 }
@@ -117,11 +117,11 @@ void InterpolatePlugin::serialize(void* buffer) const {
 std::string InterpolatePlugin::serializeToString() const {
     torch::serialize::OutputArchive output_archive;
 
-    output_archive.write("in_shape", torch::IValue(in_shape));
-    output_archive.write("out_shape", torch::IValue(out_shape));
-    output_archive.write("size", torch::IValue(size));
-    output_archive.write("mode", torch::IValue(mode));
-    output_archive.write("align_corners", torch::IValue(align_corners));
+    output_archive.write("in_shape", torch::IValue(in_shape_));
+    output_archive.write("out_shape", torch::IValue(out_shape_));
+    output_archive.write("size", torch::IValue(size_));
+    output_archive.write("mode", torch::IValue(mode_));
+    output_archive.write("align_corners", torch::IValue(align_corners_));
 
     std::ostringstream data_str;
     output_archive.save_to(data_str);
@@ -146,56 +146,48 @@ bool InterpolatePlugin::supportsFormatCombination(int pos, const nvinfer1::Plugi
 
     // pos == 1, accessing information about output tensor
     const PluginTensorDesc& out = inOut[1];
-    
+
     return (in.type == out.type) && (in.format == out.format);
 }
 
 void InterpolatePlugin::configurePlugin(const nvinfer1::DynamicPluginTensorDesc* in, int nbInputs, const nvinfer1::DynamicPluginTensorDesc* out, int nbOutputs) {
-    dtype = DataType::kFLOAT;
+    dtype_ = DataType::kFLOAT;
 }
 
 size_t InterpolatePlugin::getWorkspaceSize(const nvinfer1::PluginTensorDesc* inputs, int nbInputs, const nvinfer1::PluginTensorDesc* outputs, int nbOutputs) const {
     return 0;
 }
 
-int InterpolatePlugin::enqueue(const nvinfer1::PluginTensorDesc* inputDesc, const nvinfer1::PluginTensorDesc* outputDesc, const void *const *inputs, 
-                                                                                                        void *const *outputs, void *workspace, 
+int InterpolatePlugin::enqueue(const nvinfer1::PluginTensorDesc* inputDesc, const nvinfer1::PluginTensorDesc* outputDesc, const void* const* inputs,
+                                                                                                        void* const* outputs, void* workspace,
                                                                                                         cudaStream_t stream) {
-    at::Tensor input = at::from_blob((void*) inputs[0], util::toVec(inputDesc->dims), [](void*){}, tensor_options);
-    at::Tensor output = at::from_blob(outputs[0], out_shape, [](void*){}, tensor_options);
-
-    at::cuda::CUDAStream torch_stream = at::cuda::getStreamFromPool();
-    at::cuda::CUDAStreamGuard torch_guard(torch_stream);
-
-    cudaEvent_t event;
-    cudaEventCreate(&event);
-    cudaEventRecord(event, stream);
-
-    cudaStreamWaitEvent(torch_stream.stream(), event, 0);
-
-    if (mode == "linear") {
-        at::upsample_linear1d_out(output, input, {size[0]}, align_corners);
-    } else if (mode == "bilinear") {
-        at::upsample_bilinear2d_out(output, input, {size[0], size[1]}, align_corners);
-    } else if (mode == "trilinear") {
-        at::upsample_trilinear3d_out(output, input, {size[0], size[1], size[2]}, align_corners);
-    } else if (mode == "adaptive_pool2d") {
-        at::adaptive_avg_pool2d_out(output, input, {size[0], size[1]});
+    // TODO: When PyTorch updates to cuDNN 8 try moving back to CUDA based ATen kernels
+    // HACK: WAR because there is a segfault if you try to create a CUDA Tensor in the context of TensorRT execution
+    float* input_blob = (float*) malloc(util::volume(inputDesc->dims) * sizeof(float));
+    cudaMemcpyAsync(input_blob, static_cast<const void*>(inputs[0]), util::volume(inputDesc->dims) * sizeof(float), cudaMemcpyDeviceToHost, stream);
+    cudaStreamSynchronize(stream);
+
+    at::Tensor input = at::from_blob((void*)input_blob, util::toVec(inputDesc->dims), tensor_options_);
+
+    at::Tensor output;
+    if (mode_ == "adaptive_pool2d") {
+        output = at::adaptive_avg_pool2d(input, {size_[0], size_[1]});
     }
 
-    cudaEvent_t torch_event;
-    cudaEventCreate(&torch_event);
-    cudaEventRecord(torch_event, torch_stream.stream());
+    output = output.contiguous();
+    for (int i = 0; i < util::volume(outputDesc->dims); i++) {
+        std::cout << ((float*)output.data_ptr())[i] << std::endl;
+    }
 
-    cudaStreamWaitEvent(stream, torch_event, 0);
+    cudaMemcpyAsync(outputs[0], output.data_ptr(), util::volume(outputDesc->dims) * sizeof(float), cudaMemcpyHostToDevice, stream);
+    cudaStreamSynchronize(stream);
 
-    cudaEventDestroy(event);
-    cudaEventDestroy(torch_event);
+    free(input_blob);
 
     return 0;
 }
 
-/* 
+/*
  * InterpolatePluginCreator class implementations
  */
 const char* InterpolatePluginCreator::getPluginNamespace() const {
@@ -214,15 +206,15 @@ nvinfer1::IPluginV2* InterpolatePluginCreator::createPlugin(const char* name, co
     return nullptr;
 }
 
-InterpolatePlugin* InterpolatePluginCreator::createPlugin(const char* name, std::vector<int64_t> in_shape, std::vector<int64_t> out_shape, 
-                                                                                                           std::vector<int64_t> size, 
+InterpolatePlugin* InterpolatePluginCreator::createPlugin(const char* name, std::vector<int64_t> in_shape, std::vector<int64_t> out_shape,
+                                                                                                           std::vector<int64_t> size,
                                                                                                            std::string mode, bool align_corners) {
-    name = name;
+    name_ = name;
     return new InterpolatePlugin(in_shape, out_shape, size, mode, align_corners);
 }
 
 nvinfer1::IPluginV2* InterpolatePluginCreator::deserializePlugin(const char* name, const void *serialData, size_t serialLength) {
-    name = name;
+    name_ = name;
     return new InterpolatePlugin((const char*) serialData, serialLength);
 }