feat(//core/conversion/converters/impl/plugins): interpolate plugin compiles now. time to test it.

Signed-off-by: Abhiram Iyer <abhirami@nvidia.com>

Signed-off-by: Abhiram Iyer <abhi.iyer.ai@gmail.com>

Author: abhi-iyer

core/conversion/converters/BUILD

@@ -35,6 +35,7 @@ cc_library(
         "@tensorrt//:nvinfer",
         "//core/util:prelude",
         "//core/conversion/conversionctx",
+        "//core/conversion/converters/impl/plugins"
     ] + select({
         ":use_pre_cxx11_abi":  ["@libtorch_pre_cxx11_abi//:libtorch"],
         "//conditions:default":  ["@libtorch//:libtorch"],

core/conversion/converters/impl/plugins/BUILD

@@ -0,0 +1,33 @@
+package(default_visibility = ["//visibility:public"])
+
+config_setting(
+    name = "use_pre_cxx11_abi",
+    values = {
+        "define": "abi=pre_cxx11_abi",
+    }
+)
+
+cc_library(
+    name = "plugins",
+    hdrs = [],
+    srcs = [
+        "interpolate_plugin.cpp"
+    ],
+    deps = [
+        "@tensorrt//:nvinfer",
+        "//core/util:prelude",
+        "//core/conversion/conversionctx",
+    ] + select({
+        ":use_pre_cxx11_abi":  ["@libtorch_pre_cxx11_abi//:libtorch"],
+        "//conditions:default":  ["@libtorch//:libtorch"],
+    }),
+    alwayslink = True,
+)
+
+load("@rules_pkg//:pkg.bzl", "pkg_tar")
+
+pkg_tar(
+    name = "include",
+    package_dir = "core/conversion/converters/impl/plugins",
+    srcs = [],
+)

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

@@ -1,18 +1,17 @@
-#include <torch/extension.h>
-#include <torch/script.h>
 #include <string>
 #include <iostream>
 #include <sstream>
-#include <NvInfer.h>
 #include <ATen/ATen.h>
 #include <ATen/cuda/CUDAEvent.h>
-#include <torch/torch.h>
 #include <cuda_runtime_api.h>
-#include "NvInferVersion.h"
 #include <vector>
 #include <cudnn.h>
-#include <NVInferRuntime.h>
-#include <NVInferRuntimeCommon.h>
+
+#include "core/util/prelude.h"
+#include "torch/torch.h"
+#include "NvInfer.h"
+
+using namespace nvinfer1;
 
 namespace trtorch {
 namespace core {
@@ -25,24 +24,58 @@ namespace {
 class InterpolatePlugin : public nvinfer1::IPluginV2DynamicExt {
 private:
     at::TensorOptions tensor_options;
-    std::vector<int64_t> input_sizes;
-    std::vector<int64_t> output_sizes;
     DataType dtype;
 
+    std::vector<int64_t> in_shape;
+    std::vector<int64_t> out_shape;
     std::vector<int64_t> size;
     std::string mode;
     bool align_corners;
 
 public:
-    InterpolatePlugin(const char* name, std::vector<int64_t> in_shape, 
-                                        std::vector<int64_t> out_shape, 
-                                        std::string mode, 
-                                        bool align_corners) : name(name), in_shape(in_shape), out_shape(out_shape), mode(mode), align_corners(align_corners) {}
-
+    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(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();
+        }
+        {
+            torch::IValue value;
+            input_archive.read("out_shape", value);
+            out_shape = value.toIntVector();
+        }
+        {
+            torch::IValue value;
+            input_archive.read("size", value);
+            size = value.toIntVector();
+        }
+        {
+            torch::IValue value;
+            input_archive.read("mode", value);
+            mode = value.toStringRef();
+        }
+        {
+            torch::IValue value;
+            input_archive.read("align_corners", value);
+            align_corners = value.toBool();
+        }
+    }
 
+    int getNbOutputs() const override {
+        return 1;
+    }
 
     const char* getPluginType() const override {
-        return "Interpolate";
+        return "Interpolate_TRTorch";
     }
 
     const char* getPluginVersion() const override {
@@ -60,79 +93,125 @@ class InterpolatePlugin : public nvinfer1::IPluginV2DynamicExt {
     }
 
     nvinfer1::IPluginV2DynamicExt* clone() const override {
-        auto* plugin = new InterpolatePlugin(*this);
-        return plugin;
+        return new InterpolatePlugin(in_shape, out_shape, size, mode, align_corners);
     }
 
-    nvinfer::DimsExprs getOutputDimensions(int outputIndex, const nvinfer1::DimsExprs *inputs, int nbInputs, nvinfer1::IExprBuilder &exprBuilder) const override {
-       
-    }
+    nvinfer1::DimsExprs getOutputDimensions(int outputIndex, const nvinfer1::DimsExprs *inputs, int nbInputs, nvinfer1::IExprBuilder &exprBuilder) override {
+       //nvinfer1::DimsExprs output(inputs[0]);
 
-    nvinfer1::DataType getOutputDataType(int index, const nvinfer1::DataType* inputTypes, int nbInputs) const override {
+    //    output.nbDims = out_shape.size(); 
+
+    //    for (int i = 0; i < out_shape.size(); i++) {
+    //        output.d[i] = exprBuilder.getConstantValue(out_shape[i]);
+    //    }
 
+    //    return output; 
+        nvinfer1::DimsExprs empty;
+        return empty;
     }
 
-    int getNbOutputs() const override {
-        return 1;
+    nvinfer1::DataType getOutputDataType(int index, const nvinfer1::DataType* inputTypes, int nbInputs) const override {
+        return DataType::kFLOAT;
     }
 
     int initialize() override {
+        tensor_options = tensor_options.device(c10::kCUDA);
+        tensor_options = tensor_options.dtype(c10::kFloat);
 
+        return 0;
     }
 
-    void terminate() override {
-
-    }
+    void terminate() override {}
 
     void serialize(void* buffer) const override {
+        std::string data = serializeToString();
+        size_t size = getSerializationSize();
 
+        data.copy((char *) buffer, size);
     }
 
-    void size_t getSerializationSize() const override {
+    std::string 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));
 
-    void destroy() override {
+        std::ostringstream data_str;
+        output_archive.save_to(data_str);
 
+        return data_str.str();
     }
 
-    bool supportsFormatCombination(int pos, const nvinfer1::PluginTensorDesc* inOut, int nbInputs, int nbOutputs) override {
-
+    size_t getSerializationSize() const override {
+        return serializeToString().size();
     }
 
-    void configurePlugin(const nvinfer1::DynamicPluginTensorDesc* in, int nbInputs, const nvinfer1::DynamicPluginTensorDesc* out, int nbOutputs) override {
+    void destroy() override {}
 
-    }
+    bool supportsFormatCombination(int pos, const nvinfer1::PluginTensorDesc* inOut, int nbInputs, int nbOutputs) override {
+        if (inOut->format != nvinfer1::TensorFormat::kLINEAR) {
+            return false;
+        } 
 
-    size_t getWorkspaceSize(const nvinfer1::PluginTensorDesc* inputs, int nbInputs, const PluginTensorDesc* outputs, int nbOutputs) const override {
+        if (inOut->type == DataType::kINT32 || inOut->type == DataType::kINT8) {
+            return false;
+        }
 
+        return true;
     }
 
-    void attachToContext(nvinfer1::cudnnContext*, nvinfer1::cublasContext*, nvinfer1::IGpuAllocator*) override {}
+    void configurePlugin(const nvinfer1::DynamicPluginTensorDesc* in, int nbInputs, const nvinfer1::DynamicPluginTensorDesc* out, int nbOutputs) override {
+        dtype = DataType::kFLOAT;
+    }
 
-    void detachFromContext() override {}
+    size_t getWorkspaceSize(const nvinfer1::PluginTensorDesc* inputs, int nbInputs, const nvinfer1::PluginTensorDesc* outputs, int nbOutputs) const override {
+        return 0;
+    }
 
     int enqueue(const nvinfer1::PluginTensorDesc* inputDesc, const nvinfer1::PluginTensorDesc* outputDesc, const void *const *inputs, 
                                                                                                            void *const *outputs, void *workspace, 
                                                                                                            cudaStream_t stream) override {
-                                                                                                               
-    }
+        at::Tensor input = at::from_blob((void*) inputs[0], in_shape, [](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);
 
-private:
-    std::string name;
-    std::vector<int64_t> in_shape;
-    std::vector<int64_t> out_shape;
-    std::string mode;
-    bool align_corners;
+        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);
+        }
+
+        cudaEvent_t torch_event;
+        cudaEventCreate(&torch_event);
+        cudaEventRecord(torch_event, torch_stream.stream());
+
+        cudaStreamWaitEvent(stream, torch_event, 0);
+
+        cudaEventDestroy(event);
+        cudaEventDestroy(torch_event);
 
-    nvinfer1::DataType dtype;
-}
+        return 0;
+    }
+};
 
 
 class InterpolatePluginCreator : public nvinfer1::IPluginCreator {
+private:
+    std::string name;
+
 public:
     InterpolatePluginCreator() {}
 
@@ -158,18 +237,20 @@ class InterpolatePluginCreator : public nvinfer1::IPluginCreator {
         return nullptr;
     }
 
-    nvinfer1::IPluginV2* createPlugin(const char* name, std::vector<int64_t> in_shape, std::vector<int64_t> out_shape, std::string mode, bool align_corners) {
-        return new InterpolatePlugin(name, in_shape, out_shape, mode, align_corners);
+    nvinfer1::IPluginV2* 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;
+        return new InterpolatePlugin(in_shape, out_shape, size, mode, align_corners);
     }
 
     nvinfer1::IPluginV2* deserializePlugin(const char* name, const void *serialData, size_t serialLength) override {
-        return nullptr;
+        name = name;
+        return new InterpolatePlugin((const char*) serialData, serialLength);
     }
 
     const nvinfer1::PluginFieldCollection* getFieldNames() override {
         return nullptr;
     }
-}
+};
 
 REGISTER_TENSORRT_PLUGIN(InterpolatePluginCreator);