Merge pull request #166 from facebookresearch/fix-flags

Add pybindings for gflags

Author: nicolasvasilache

docs/source/framework/pytorch_integration/debugging.rst

@@ -16,7 +16,7 @@ can use these flags to enable logging. Various types of flags exposed are:
 * :code:`debug_tuner`: print debug spew for the tuner multithreading behavior.
 
 
-In order to use enable these flags, you need to call :code:`tc.GlobalDebugInit`
+In order to use enable these flags, you need to call :code:`tc.SetDebugFlags`
 and set the proper flags to :code:`true`. All of these flags are :code:`boolean`
 flags that take values :code:`true` or :code:`false`.
 
@@ -28,14 +28,14 @@ Example usage
     import tensor_comprehensions as tc
     import torch
 
-    tc.GlobalDebugInit(["--debug_tc_mapper=true", "--debug_lang=false"])
+    tc.SetDebugFlags(debug_tc_mapper=True, debug_lang=False)
 
     matmul = tc.define(tc.database['matmul']['lang'], name='matmul')
     mat1, mat2 = torch.randn(3, 4).cuda(), torch.randn(4, 5).cuda()
     out = matmul(mat1, mat2)
 
 In above example, when the TC executes, we will see the TC mapper information.
-You can chose to set any number of flags but the :code:`tc.GlobalDebugInit` should
+You can chose to set any number of flags but the :code:`tc.SetDebugFlags` should
 only be called once.
 
 Printing TC generated CUDA code
@@ -50,7 +50,7 @@ and the generated CUDA code will be printed on command line.
     import tensor_comprehensions as tc
     import torch
 
-    tc.GlobalDebugInit(["--dump_cuda=true"])
+    tc.SetDebugFlags(dump_cuda=True)
 
     matmul = tc.define(tc.database['matmul']['lang'], name='matmul')
     mat1, mat2 = torch.randn(3, 4).cuda(), torch.randn(4, 5).cuda()

include/tc/core/flags.h

@@ -61,10 +61,4 @@ DECLARE_bool(schedule_tree_verbose_validation);
 // random seed setting for reproducibility and debugging purposes
 uint64_t initRandomSeed();
 const uint64_t& randomSeed();
-
-// python
-namespace python {
-bool globalDebugGflagsGlogInit(int* pargc, char*** pargv);
-} // namespace python
-
 } // namespace tc

src/core/flags.cc

@@ -111,34 +111,6 @@ DEFINE_int64(
     -1,
     "The number of best candidates to restore from the proto cache");
 
-namespace {
-bool parseCommandLineFlags(int* pargc, char*** pargv) {
-  if (*pargc == 0) {
-    return true;
-  }
-  // TODO: (prigoyal): we need to do some filtering on flags here,
-  // add option for displaying the help message
-  return ::gflags::ParseCommandLineFlags(pargc, pargv, true);
-}
-
-bool initGoogleLogging(int* pargc, char** argv) {
-  if (*pargc == 0) {
-    return true;
-  }
-  ::google::InitGoogleLogging(argv[0]);
-  return true;
-}
-} // namespace
-
-namespace python {
-bool globalDebugGflagsGlogInit(int* pargc, char*** pargv) {
-  bool success = true;
-  success &= parseCommandLineFlags(pargc, pargv);
-  success &= initGoogleLogging(pargc, *pargv);
-  return success;
-}
-} // namespace python
-
 uint64_t initRandomSeed() {
   static std::mutex mut;
   static bool inited = false;

tensor_comprehensions/__init__.py

@@ -18,7 +18,7 @@
 from tensor_comprehensions.tc_unit import TcUnit
 from tensor_comprehensions.tc_unit import TcAutotuner
 from tensor_comprehensions.tc_unit import TcCompilationUnit
-from tensor_comprehensions.tc_unit import GlobalDebugInit
+from tensor_comprehensions.tc_unit import SetDebugFlags
 from tensor_comprehensions.tc_unit import autotuner_settings
 from tensor_comprehensions.tc_unit import small_sizes_autotuner_settings
 from tensor_comprehensions.tc_unit import ATenCompilationUnit
@@ -27,6 +27,6 @@
 
 __all__ = [
     'define', 'TcUnit', 'TcAutotuner', 'TcCompilationUnit', 'autotuner_settings',
-    'small_sizes_autotuner_settings', 'GlobalDebugInit', 'ATenCompilationUnit',
+    'small_sizes_autotuner_settings', 'SetDebugFlags', 'ATenCompilationUnit',
     'Options', 'database', 'decode',
 ]

tensor_comprehensions/pybinds/pybind_engine.cc

@@ -39,21 +39,26 @@ namespace py = pybind11;
 using ATenCudaCompilationUnit = tc::ATenCompilationUnit<tc::CudaTcExecutor>;
 
 PYBIND11_MODULE(tc, m) {
+  m.def("set_logtostderr", [](bool logtostderr) {
+    FLAGS_logtostderr = logtostderr;
+  });
+  m.def("set_debug_lang", [](bool debug_lang) {
+    tc::FLAGS_debug_lang = debug_lang;
+  });
+  m.def("set_debug_halide", [](bool debug_halide) {
+    tc::FLAGS_debug_halide = debug_halide;
+  });
+  m.def("set_debug_tc_mapper", [](bool debug_tc_mapper) {
+    tc::FLAGS_debug_tc_mapper = debug_tc_mapper;
+  });
+  m.def("set_debug_cuda", [](bool debug_cuda) {
+    tc::FLAGS_debug_cuda = debug_cuda;
+  });
+  m.def("set_debug_tuner", [](bool debug_tuner) {
+    tc::FLAGS_debug_tuner = debug_tuner;
+  });
   m.def(
-      "global_debug_init", // exposing the debugging flags to people
-      [](std::vector<std::string> args) {
-        if (args.size() > 0) {
-          args.insert(args.begin(), "tc");
-        }
-        int numArgs = args.size();
-        // now we construct a char** argv type from args
-        std::vector<char*> vargs; // char* vector args
-        for (auto& arg : args) {
-          vargs.push_back(const_cast<char*>(arg.data()));
-        }
-        char** argv = vargs.data();
-        tc::python::globalDebugGflagsGlogInit(&numArgs, &argv);
-      });
+      "set_dump_cuda", [](bool dump_cuda) { tc::FLAGS_dump_cuda = dump_cuda; });
 
   py::object dlpack;
   try {

tensor_comprehensions/tc_unit.py

@@ -19,7 +19,7 @@
 from torch.autograd import Variable
 
 from tensor_comprehensions.tc import ATenCompilationUnit
-from tensor_comprehensions.tc import global_debug_init as GlobalDebugInit
+from tensor_comprehensions.tc import set_logtostderr, set_debug_lang, set_debug_halide, set_debug_tc_mapper, set_debug_cuda, set_debug_tuner, set_dump_cuda
 from tensor_comprehensions.torch_tc.tc_function import TCFunction, unpack_variables, get_tensors, make_contiguous
 from tensor_comprehensions.autotuner import ATenAutotuner
 from tensor_comprehensions.mapping_options import Options
@@ -38,6 +38,26 @@
     "threads": 32, "generations": 5, "tuner_min_launch_total_threads": 1,
 }
 
+###############################################################################
+# Set global debugging flags
+###############################################################################
+class SetDebugFlags(object):
+    def __init__(self, **kwargs):
+        self.set_gflags(**kwargs)
+
+    def set_gflags(
+        self, debug_lang=False, debug_halide=False, debug_tc_mapper=False,
+        debug_cuda=False, debug_tuner=False, dump_cuda=False, **kwargs
+    ):
+        set_logtostderr(True)
+        set_debug_lang(debug_lang)
+        set_debug_halide(debug_halide)
+        set_debug_tc_mapper(debug_tc_mapper)
+        set_debug_cuda(debug_cuda)
+        set_debug_tuner(debug_tuner)
+        set_dump_cuda(dump_cuda)
+
+
 ###############################################################################
 # Some helper functions
 ###############################################################################

test_python/layers/test_dump_cuda.py

@@ -16,12 +16,11 @@
 import tensor_comprehensions as tc
 
 import torch
-import torch.cuda
 import unittest
 
 # enable this to dump cuda code generated whenever tc layer runs: simple run or
 # autotuner run
-tc.GlobalDebugInit(["--dump_cuda=true"])
+tc.SetDebugFlags(dump_cuda=True)
 
 
 class TestDumpCuda(unittest.TestCase):

test_python/layers/test_layernorm.py

@@ -19,7 +19,7 @@
 import torch.cuda
 import unittest
 
-tc.GlobalDebugInit(["--debug_tuner=false", "--debug_tc_mapper=false"])
+tc.SetDebugFlags(debug_tuner=False, debug_tc_mapper=False)
 
 
 class TestLayerNorm(unittest.TestCase):

test_python/test_debug_init.py

@@ -19,7 +19,7 @@
 import torch.cuda
 import tensor_comprehensions as tc
 
-tc.GlobalDebugInit(["--dump_cuda=true", "--debug_tc_mapper=false"])
+tc.SetDebugFlags(dump_cuda=True, debug_tc_mapper=True)
 
 
 class TestDebugInit(unittest.TestCase):

test_python/test_tc_torch.py

@@ -24,7 +24,7 @@
 from tensor_comprehensions.mapping_options import Options
 from common import TestCase, run_tests
 
-tc.GlobalDebugInit(["--dump_cuda=false"])
+tc.SetDebugFlags(dump_cuda=False)
 
 
 MATMUL_LANG = """