Add pybindings for gflags

Author: prigoyal

docs/source/framework/pytorch_integration/debugging.rst

@@ -28,7 +28,7 @@ Example usage
     import tensor_comprehensions as tc
     import torch
 
-    tc.GlobalDebugInit(["--debug_tc_mapper=true", "--debug_lang=false"])
+    tc.GlobalDebugInit(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()
@@ -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.GlobalDebugInit(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/pybinds/pybind_engine.cc

@@ -39,21 +39,23 @@ namespace py = pybind11;
 using ATenCudaCompilationUnit = tc::ATenCompilationUnit<tc::CudaTcExecutor>;
 
 PYBIND11_MODULE(tc, m) {
+  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_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,25 @@
     "threads": 32, "generations": 5, "tuner_min_launch_total_threads": 1,
 }
 
+###############################################################################
+# Set global debugging flags
+###############################################################################
+class GlobalDebugInit(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_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

@@ -21,7 +21,7 @@
 
 # enable this to dump cuda code generated whenever tc layer runs: simple run or
 # autotuner run
-tc.GlobalDebugInit(["--dump_cuda=true"])
+tc.GlobalDebugInit(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.GlobalDebugInit(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.GlobalDebugInit(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.GlobalDebugInit(dump_cuda=False)
 
 
 MATMUL_LANG = """