Merge pull request #126 from NVIDIA/python_abi_option

Handle Dropout_ and ABI fixes for C++ and Python

Author: narendasan

core/conversion/conversion_blacklist.cpp

@@ -22,7 +22,8 @@ const std::unordered_set<std::string>& get_non_convertable_nodes() {
         "prim::GetAttr",
         "prim::CallMethod",
         "prim::Drop",
-        "aten:dropout",
+        "aten::dropout",
+        "aten::dropout_"
     };
     return nonconvertable_nodes;
 }

core/lowering/passes/remove_dropout.cpp

@@ -20,6 +20,20 @@ void RemoveDropout(std::shared_ptr<torch::jit::Graph>& graph) {
     remove_dropout.RegisterRewritePattern(
         dropout_pattern, no_dropout_pattern);
     remove_dropout.runOnGraph(graph);
+
+    std::string dropout_inplace_pattern = R"IR(
+        graph(%input, %4, %5):
+            %6 = aten::dropout_(%input, %4, %5)
+            return (%6))IR";
+    std::string no_dropout_inplace_pattern = R"IR(
+        graph(%input, %4, %5):
+            return (%input))IR";
+
+    torch::jit::SubgraphRewriter remove_dropout_inplace_pattern;
+    remove_dropout_inplace_pattern.RegisterRewritePattern(
+        dropout_inplace_pattern, no_dropout_inplace_pattern);
+    remove_dropout_inplace_pattern.runOnGraph(graph);
+
     LOG_GRAPH("Post remove dropout: " << *graph);
 }
 

cpp/trtorchc/BUILD

@@ -1,14 +1,28 @@
 package(default_visibility = ["//visibility:public"])
 
+config_setting(
+    name = "use_pre_cxx11_abi",
+    values = {
+        "define": "abi=pre_cxx11_abi",
+    }
+)
+
 cc_binary(
     name = "trtorchc",
     srcs = [
         "main.cpp"
     ],
     deps = [
-        "@libtorch//:libtorch",
-        "@libtorch//:caffe2",
         "//third_party/args",
         "//cpp/api:trtorch"
-    ],
-)
+    ] + select({
+        ":use_pre_cxx11_abi":  [
+            "@libtorch_pre_cxx11_abi//:libtorch",
+            "@libtorch_pre_cxx11_abi//:caffe2",
+        ],
+        "//conditions:default":  [
+            "@libtorch//:libtorch",
+            "@libtorch//:caffe2",
+        ],
+    }),
+)

py/setup.py

@@ -18,7 +18,13 @@
 
 __version__ = '0.0.2'
 
-def build_libtrtorch_pre_cxx11_abi(develop=True, use_dist_dir=True):
+CXX11_ABI = False
+
+if "--use-cxx11-abi" in sys.argv:
+    sys.argv.remove("--use-cxx11-abi")
+    CXX11_ABI = True
+
+def build_libtrtorch_pre_cxx11_abi(develop=True, use_dist_dir=True, cxx11_abi=False):
     cmd = ["/usr/bin/bazel", "build"]
     cmd.append("//cpp/api/lib:libtrtorch.so")
     if develop:
@@ -27,7 +33,10 @@ def build_libtrtorch_pre_cxx11_abi(develop=True, use_dist_dir=True):
         cmd.append("--compilation_mode=opt")
     if use_dist_dir:
         cmd.append("--distdir=third_party/dist_dir/x86_64-linux-gnu")
-    cmd.append("--config=python")
+    if not cxx11_abi:
+        cmd.append("--config=python")
+    else:
+        print("using CXX11 ABI build")
 
     print("building libtrtorch")
     status_code = subprocess.run(cmd).returncode
@@ -64,7 +73,8 @@ def finalize_options(self):
         develop.finalize_options(self)
 
     def run(self):
-        build_libtrtorch_pre_cxx11_abi(develop=True)
+        global CXX11_ABI
+        build_libtrtorch_pre_cxx11_abi(develop=True, cxx11_abi=CXX11_ABI)
         gen_version_file()
         copy_libtrtorch()
         develop.run(self)
@@ -80,7 +90,8 @@ def finalize_options(self):
         install.finalize_options(self)
 
     def run(self):
-        build_libtrtorch_pre_cxx11_abi(develop=False)
+        global CXX11_ABI
+        build_libtrtorch_pre_cxx11_abi(develop=False, cxx11_abi=CXX11_ABI)
         gen_version_file()
         copy_libtrtorch()
         install.run(self)
@@ -95,7 +106,8 @@ def finalize_options(self):
         bdist_wheel.finalize_options(self)
 
     def run(self):
-        build_libtrtorch_pre_cxx11_abi(develop=False)
+        global CXX11_ABI
+        build_libtrtorch_pre_cxx11_abi(develop=False, cxx11_abi=CXX11_ABI)
         gen_version_file()
         copy_libtrtorch()
         bdist_wheel.run(self)
@@ -138,15 +150,16 @@ def run(self):
                                     dir_path + "/../bazel-TRTorch/external/tensorrt/include",
                                 ],
                                 extra_compile_args=[
-                                    "-D_GLIBCXX_USE_CXX11_ABI=0",
-                                    "-Wno-deprecated-declaration",
-                                ],
+                                    "-Wno-deprecated",
+                                    "-Wno-deprecated-declarations",
+                                ] + ["-D_GLIBCXX_USE_CXX11_ABI=1"] if CXX11_ABI else ["-D_GLIBCXX_USE_CXX11_ABI=0"],
                                 extra_link_args=[
-                                    "-D_GLIBCXX_USE_CXX11_ABI=0"
+                                    "-Wno-deprecated",
+                                    "-Wno-deprecated-declarations",
                                     "-Wl,--no-as-needed",
                                     "-ltrtorch",
                                     "-Wl,-rpath,$ORIGIN/lib"
-                                ],
+                                ] + ["-D_GLIBCXX_USE_CXX11_ABI=1"] if CXX11_ABI else ["-D_GLIBCXX_USE_CXX11_ABI=0"],
                                 undef_macros=[ "NDEBUG" ]
                             )
 ]
@@ -178,7 +191,6 @@ def run(self):
     zip_safe=False,
     license="BSD",
     packages=find_packages(),
-    platform="Linux",
     classifiers=[
         "Development Status :: 3 - Alpha",
         "Environment :: GPU :: NVIDIA CUDA",