Finalize code changes, Fix all tests, clean up code

Signed-off-by: Dheeraj Peri <[email protected]>

Author: peri044

core/conversion/conversionctx/BUILD

@@ -19,12 +19,10 @@ cc_library(
         "@tensorrt//:nvinfer",
         "@tensorrt//:nvinferplugin",
         "//core/util:prelude",
-        #"//core/plugins:trtorch_plugins",
     ] + select({
         ":use_pre_cxx11_abi":  ["@libtorch_pre_cxx11_abi//:libtorch"],
         "//conditions:default":  ["@libtorch//:libtorch"],
     }),
-    copts = ["-fpic"],
 )
 
 load("@rules_pkg//:pkg.bzl", "pkg_tar")

core/conversion/conversionctx/ConversionCtx.cpp

@@ -5,10 +5,7 @@
 #include "NvInferPluginUtils.h"
 
 #include <utility>
-// #include "core/plugins/plugin_prelude.h"
 #include "core/conversion/conversionctx/ConversionCtx.h"
-// #include "NvInferPlugin.h"
-// #include "NvInferPluginUtils.h"
 
 namespace trtorch {
 namespace core {
@@ -52,7 +49,6 @@ ConversionCtx::ConversionCtx(BuilderSettings build_settings)
           "[TRTorch Conversion Context] - ",
           util::logging::get_logger().get_reportable_severity(),
           util::logging::get_logger().get_is_colored_output_on()) {
-
   // TODO: Support FP16 and FP32 from JIT information
   if (settings.device.gpu_id) {
     TRTORCH_CHECK(

core/conversion/converters/BUILD

@@ -29,7 +29,7 @@ cc_library(
 cc_library(
     name = "converters",
     hdrs = [
-        "converters.h",
+        "converters.h"
     ],
     srcs = [
         "NodeConverterRegistry.cpp",
@@ -67,7 +67,6 @@ cc_library(
         ":use_pre_cxx11_abi":  ["@libtorch_pre_cxx11_abi//:libtorch"],
         "//conditions:default":  ["@libtorch//:libtorch"],
     }),
-    #copts = ["-fPIC"],
     alwayslink = True,
 )
 

core/conversion/converters/impl/batch_norm.cpp

@@ -10,13 +10,6 @@ namespace converters {
 namespace impl {
 namespace {
 
-static inline at::Tensor repeat_if_defined(const at::Tensor& t, int64_t repeat) {
-  if (t.defined()) {
-    return t.repeat(repeat);
-  }
-  return t;
-}
-
 auto batch_norm_registrations TRTORCH_UNUSED = RegisterNodeConversionPatterns().pattern({
     R"SIG(aten::batch_norm(Tensor input, Tensor? gamma, Tensor? beta,
                             Tensor? mean, Tensor? var,

core/conversion/converters/impl/interpolate.cpp

@@ -53,7 +53,7 @@ void create_plugin(
 
   fc.nbFields = f.size();
   fc.fields = f.data();
-  auto creator = getPluginRegistry()->getPluginCreator("Interpolate", "1", "");
+  auto creator = getPluginRegistry()->getPluginCreator("Interpolate", "1", "trtorch");
   auto interpolate_plugin = creator->createPlugin(name, &fc);
 
   auto resize_layer = ctx->net->addPluginV2(reinterpret_cast<nvinfer1::ITensor* const*>(&in), 1, *interpolate_plugin);

core/conversion/converters/impl/normalize.cpp

@@ -2,8 +2,6 @@
 #include "NvInferRuntimeCommon.h"
 #include "core/conversion/converters/converters.h"
 #include "core/util/prelude.h"
-// #include "core/plugins/impl/normalize_plugin.h"
-// #include "core/plugins/plugin_prelude.h"
 #include "torch/torch.h"
 
 namespace trtorch {
@@ -42,7 +40,8 @@ void create_plugin(
       TRTORCH_THROW_ERROR("Axis of normalization layer cannot exceed input rank");
     }
   }
-  auto creator = getPluginRegistry()->getPluginCreator("NormalizePlugintrtorch", "1", "");
+
+  auto creator = getPluginRegistry()->getPluginCreator("NormalizePlugin", "1", "trtorch");
   auto plugin = creator->createPlugin(name, &fc);
   auto normalize_layer = ctx->net->addPluginV2(reinterpret_cast<nvinfer1::ITensor* const*>(&in), 1, *plugin);
   TRTORCH_CHECK(normalize_layer, "Unable to create normalization plugin from node" << *n);

core/conversion/converters/impl/pooling.cpp

@@ -313,7 +313,7 @@ auto pooling_registrations TRTORCH_UNUSED =
                  auto out_shape = in_shape;
                  std::copy(out_size.begin(), out_size.end(), out_shape.begin() + (in_shape.size() - out_size.size()));
 
-                 auto creator = getPluginRegistry()->getPluginCreator("Interpolate", "1", "");
+                 auto creator = getPluginRegistry()->getPluginCreator("Interpolate", "1", "trtorch");
 
                  // Configure the plugin fields
                  nvinfer1::PluginFieldCollection fc;

core/plugins/BUILD

@@ -10,14 +10,15 @@ config_setting(
 cc_library(
     name = "trtorch_plugins",
     hdrs = [
-        "plugins.h",
         "impl/interpolate_plugin.h",
         "impl/normalize_plugin.h",
+        "plugins.h",
+
     ],
     srcs = [
-        "register_plugins.cpp",
         "impl/interpolate_plugin.cpp",
         "impl/normalize_plugin.cpp",
+        "register_plugins.cpp",
     ],
     deps = [
         "@tensorrt//:nvinfer",
@@ -28,6 +29,12 @@ cc_library(
         "//conditions:default":  ["@libtorch//:libtorch"],
     }),
     alwayslink = True,
+    copts = [
+        "-pthread"
+    ],
+    linkopts = [
+        "-lpthread",
+    ]
 )
 
 load("@rules_pkg//:pkg.bzl", "pkg_tar")

core/plugins/impl/interpolate_plugin.cpp

@@ -2,11 +2,10 @@
 #include "core/plugins/plugins.h"
 #include "core/util/prelude.h"
 
-using namespace nvinfer1;
-using namespace trtorch::core;
-// namespace trtorch {
-// namespace core {
-// namespace plugins {
+namespace trtorch {
+namespace core {
+namespace plugins {
+namespace impl {
 
 /*
  * InterpolatePlugin class implementations
@@ -119,7 +118,7 @@ const char* InterpolatePlugin::getPluginVersion() const {
 }
 
 const char* InterpolatePlugin::getPluginNamespace() const {
-  return "";
+  return "trtorch";
 }
 
 nvinfer1::IPluginV2DynamicExt* InterpolatePlugin::clone() const {
@@ -163,7 +162,7 @@ nvinfer1::DimsExprs InterpolatePlugin::getOutputDimensions(
 
 nvinfer1::DataType InterpolatePlugin::getOutputDataType(int index, const nvinfer1::DataType* inputTypes, int nbInputs)
     const {
-  return DataType::kFLOAT;
+  return nvinfer1::DataType::kFLOAT;
 }
 
 int InterpolatePlugin::initialize() {
@@ -216,14 +215,14 @@ bool InterpolatePlugin::supportsFormatCombination(
   TRTORCH_ASSERT(nbInputs == 1, "Expected a single tensor as input to interpolate plugin");
   TRTORCH_ASSERT(nbOutputs == 1, "Expected a single tensor as output to interpolate plugin");
 
-  const PluginTensorDesc& in = inOut[0];
+  const nvinfer1::PluginTensorDesc& in = inOut[0];
 
   if (pos == 0) {
     return (in.type == nvinfer1::DataType::kFLOAT) && (in.format == nvinfer1::TensorFormat::kLINEAR);
   }
 
   // pos == 1, accessing information about output tensor
-  const PluginTensorDesc& out = inOut[1];
+  const nvinfer1::PluginTensorDesc& out = inOut[1];
 
   return (in.type == out.type) && (in.format == out.format);
 }
@@ -233,7 +232,7 @@ void InterpolatePlugin::configurePlugin(
     int nbInputs,
     const nvinfer1::DynamicPluginTensorDesc* out,
     int nbOutputs) {
-  dtype_ = DataType::kFLOAT;
+  dtype_ = nvinfer1::DataType::kFLOAT;
 }
 
 size_t InterpolatePlugin::getWorkspaceSize(
@@ -344,20 +343,20 @@ int InterpolatePlugin::enqueue(
  */
 
 InterpolatePluginCreator::InterpolatePluginCreator() {
-  mPluginAttributes.emplace_back(PluginField("in_shape", nullptr, PluginFieldType::kINT32, 1));
-  mPluginAttributes.emplace_back(PluginField("out_shape", nullptr, PluginFieldType::kINT32, 1));
-  mPluginAttributes.emplace_back(PluginField("out_size", nullptr, PluginFieldType::kINT32, 1));
-  mPluginAttributes.emplace_back(PluginField("scales", nullptr, PluginFieldType::kFLOAT32, 1));
-  mPluginAttributes.emplace_back(PluginField("mode", nullptr, PluginFieldType::kCHAR, 1));
-  mPluginAttributes.emplace_back(PluginField("align_corners", nullptr, PluginFieldType::kINT32, 1));
-  mPluginAttributes.emplace_back(PluginField("use_scales", nullptr, PluginFieldType::kINT32, 1));
+  mPluginAttributes.emplace_back(nvinfer1::PluginField("in_shape", nullptr, nvinfer1::PluginFieldType::kINT32, 1));
+  mPluginAttributes.emplace_back(nvinfer1::PluginField("out_shape", nullptr, nvinfer1::PluginFieldType::kINT32, 1));
+  mPluginAttributes.emplace_back(nvinfer1::PluginField("out_size", nullptr, nvinfer1::PluginFieldType::kINT32, 1));
+  mPluginAttributes.emplace_back(nvinfer1::PluginField("scales", nullptr, nvinfer1::PluginFieldType::kFLOAT32, 1));
+  mPluginAttributes.emplace_back(nvinfer1::PluginField("mode", nullptr, nvinfer1::PluginFieldType::kCHAR, 1));
+  mPluginAttributes.emplace_back(nvinfer1::PluginField("align_corners", nullptr, nvinfer1::PluginFieldType::kINT32, 1));
+  mPluginAttributes.emplace_back(nvinfer1::PluginField("use_scales", nullptr, nvinfer1::PluginFieldType::kINT32, 1));
 
   mFC.nbFields = mPluginAttributes.size();
   mFC.fields = mPluginAttributes.data();
 }
 
 const char* InterpolatePluginCreator::getPluginNamespace() const {
-  return "";
+  return "trtorch";
 }
 
 const char* InterpolatePluginCreator::getPluginName() const {
@@ -418,8 +417,7 @@ const nvinfer1::PluginFieldCollection* InterpolatePluginCreator::getFieldNames()
   return nullptr;
 }
 
-REGISTER_TRTORCH_PLUGIN(InterpolatePluginCreator);
-
-// } // namespace plugins
-// } // namespace core
-// } // namespace trtorch
+} // namespace impl
+} // namespace plugins
+} // namespace core
+} // namespace trtorch

core/plugins/impl/interpolate_plugin.h

@@ -13,16 +13,15 @@
 #include "core/util/prelude.h"
 #include "torch/torch.h"
 
-using namespace nvinfer1;
-//
-// namespace trtorch {
-// namespace core {
-// namespace plugins {
+namespace trtorch {
+namespace core {
+namespace plugins {
+namespace impl {
 
 class InterpolatePlugin : public nvinfer1::IPluginV2DynamicExt {
  private:
   at::TensorOptions tensor_options_;
-  DataType dtype_;
+  nvinfer1::DataType dtype_;
 
   std::vector<int64_t> in_shape_;
   std::vector<int64_t> out_shape_;
@@ -121,8 +120,8 @@ class InterpolatePlugin : public nvinfer1::IPluginV2DynamicExt {
 class InterpolatePluginCreator : public nvinfer1::IPluginCreator {
  private:
   std::string name_;
-  std::vector<PluginField> mPluginAttributes;
-  PluginFieldCollection mFC;
+  std::vector<nvinfer1::PluginField> mPluginAttributes;
+  nvinfer1::PluginFieldCollection mFC;
 
  public:
   InterpolatePluginCreator();
@@ -137,21 +136,12 @@ class InterpolatePluginCreator : public nvinfer1::IPluginCreator {
 
   nvinfer1::IPluginV2* createPlugin(const char* name, const nvinfer1::PluginFieldCollection* fc) override;
 
-  // InterpolatePlugin* createPlugin(
-  //     const char* name,
-  //     std::vector<int64_t> in_shape,
-  //     std::vector<int64_t> out_shape,
-  //     std::vector<int64_t> size,
-  //     std::vector<double> scales,
-  //     std::string mode,
-  //     bool align_corners,
-  //     bool use_scales);
-
   nvinfer1::IPluginV2* deserializePlugin(const char* name, const void* serialData, size_t serialLength) override;
 
   const nvinfer1::PluginFieldCollection* getFieldNames() override;
 };
 
-// } // namespace plugins
-// } // namespace core
-// } // namespace trtorch
+} // namespace impl
+} // namespace plugins
+} // namespace core
+} // namespace trtorch