added bitwise_not converter

Author: blchu

core/conversion/converters/BUILD

@@ -54,6 +54,7 @@ cc_library(
         "NodeConverterRegistry.cpp",
         "impl/activation.cpp",
         "impl/batch_norm.cpp",
+        "impl/bitwise.cpp",
         "impl/cast.cpp",
         "impl/concat.cpp",
         "impl/constant.cpp",

core/conversion/converters/impl/bitwise.cpp

@@ -0,0 +1,52 @@
+#include "core/conversion/converters/converters.h"
+#include "core/util/prelude.h"
+
+#include <torch/torch.h>
+
+namespace torch_tensorrt {
+namespace core {
+namespace conversion {
+namespace converters {
+namespace impl {
+
+
+auto bitwisenot TORCHTRT_UNUSED =
+    RegisterNodeConversionPatterns()
+        .pattern({"aten::bitwise_not(Tensor self) -> Tensor",
+                  [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
+                    auto in = args[0].ITensorOrFreeze(ctx);
+                    nvinfer1::ILayer* out;
+                    
+                    if(in->getType() == nvinfer1::DataType::kINT32) {
+                      // Integer case
+                      auto one = torch::tensor({1}, util::TRTDataTypeToScalarType(in->getType()));
+                      auto one_const = tensor_to_const(ctx, one);
+                      auto neg = ctx->net->addUnary(*in, nvinfer1::UnaryOperation::kNEG);
+                      TORCHTRT_CHECK(neg, "Unable to create neg unary layer from node: " << *n);
+                      out = add_elementwise(
+                          ctx, nvinfer1::ElementWiseOperation::kSUB, neg->getOutput(0),
+                          one_const, util::node_info(n));
+                      TORCHTRT_CHECK(out, "Unable to create sub layer from node: " << *n);
+                    } else if(in->getType() == nvinfer1::DataType::kBOOL) {
+                      // Boolean case
+                      out = ctx->net->addUnary(*in, nvinfer1::UnaryOperation::kNOT);
+                      TORCHTRT_CHECK(out, "Unable to create logical not layer from node: " << *n);
+                    } else {
+                      LOG_ERROR("Input tensor must be 32 bit integer or boolean");
+                      return false;
+                    }
+
+                    out->setName(util::node_info(n).c_str());
+                    auto out_tensor = ctx->AssociateValueAndTensor(n->outputs()[0],
+                                                                   out->getOutput(0));
+                    LOG_DEBUG("Output tensor shape: " << out_tensor->getDimensions());
+
+                    return true;
+                  }});
+
+
+} // namespace impl
+} // namespace converters
+} // namespace conversion
+} // namespace core
+} // namespace torch_tensorrt

tests/core/conversion/converters/BUILD

@@ -15,6 +15,10 @@ converter_test(
     name = "test_batch_norm",
 )
 
+converter_test(
+    name = "test_bitwise",
+)
+
 converter_test(
     name = "test_instance_norm",
 )
@@ -136,6 +140,7 @@ test_suite(
     tests = [
         ":test_activation",
         ":test_batch_norm",
+        ":test_bitwise",
         ":test_instance_norm",
         ":test_cast",
         ":test_clone",

tests/core/conversion/converters/test_bitwise.cpp

@@ -1,24 +1,39 @@
+#include <string>
 #include "core/compiler.h"
 #include "gtest/gtest.h"
 #include "tests/util/util.h"
 #include "torch/csrc/jit/ir/irparser.h"
 
-TEST(Converters, ATenBitwiseNotConvertsCorrectly) {
-  const auto graph = R"IR(
-      graph(%0 : Tensor):
-        %3 : Tensor = aten::bitwise_not(%0)
-        return (%3))IR";
+std::string gen_test_graph() {
+    return R"IR(
+        graph(%0: Tensor):
+          %3 : Tensor = aten::bitwise_not(%0)
+          return (%3))IR";
+}
 
-  auto g = std::make_shared<torch::jit::Graph>();
-  torch::jit::parseIR(graph, g.get());
+#define test_bitwise_not(dtype)                                                            \
+  TEST(Converters, ATenBitwiseNot##dtype##ConvertsCorrectly) {                             \
+    const auto graph = gen_test_graph();                                                   \
+                                                                                           \
+    auto g = std::make_shared<torch::jit::Graph>();                                        \
+    torch::jit::parseIR(graph, g.get());                                                   \
+                                                                                           \
+    at::Tensor in;                                                                         \
+    if (strcmp(#dtype, "Integer") == 0)                                                    \
+      in = at::randint(-128, 128, {10}, {at::kCUDA}).toType(at::kInt);                     \
+    if (strcmp(#dtype, "Boolean") == 0)                                                    \
+      in = at::randint(0, 1, {10}, {at::kCUDA}).toType(at::kBool);                         \
+    auto params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {});            \
+    auto jit_results = torch_tensorrt::tests::util::RunGraph(g, params, {in});             \
+                                                                                           \
+    in = at::clone(in);                                                                    \
+    params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {});                 \
+    auto trt_results = torch_tensorrt::tests::util::RunGraphEngine(g, params, {in});       \
+                                                                                           \
+    ASSERT_TRUE(torch_tensorrt::tests::util::exactlyEqual(jit_results[0], trt_results[0]));\
+  }
 
-  auto in = at::randint(-128, 128, {10}, {at::kCUDA});
-  auto params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {});
-  auto jit_results = torch_tensorrt::tests::util::RunGraph(g, params, {in});
+test_bitwise_not(Integer);
+test_bitwise_not(Boolean);
 
-  in = at::clone(in);
-  params = torch_tensorrt::core::ir::get_static_params(g->inputs, {});
-  auto trt_results = torch_tensorrt::tests::util::RunGraphEngine(g, params, {in});
-  
-  ASSERT_TRUE(torch_tensorrt::tests::util::exactlyEqual(jit_results[0], trt_results[0]);
-}
+#undef test_bitwise_not