Expose dump_ptx flag to Python

Author: nicolasvasilache

tensor_comprehensions/__init__.py

@@ -28,6 +28,7 @@
 from tensor_comprehensions.tclib import debug_tc_mapper
 from tensor_comprehensions.tclib import debug_tuner
 from tensor_comprehensions.tclib import dump_cuda
+from tensor_comprehensions.tclib import dump_ptx
 
 from tensor_comprehensions.tclib import CompilationCache
 from tensor_comprehensions.tclib import MappingOptions
@@ -604,6 +605,7 @@ def make_autograd(forward_fun: Callable[[Iterable[torch.Tensor]], Iterable[torch
     'debug_tc_mapper',
     'debug_tuner',
     'dump_cuda',
+    'dump_ptx',
     # Functions exposed by the tclib
     'compile',
     'autotune',

tensor_comprehensions/pybinds/tclib.cc

@@ -438,6 +438,7 @@ PYBIND11_MODULE(tclib, m) {
     tc::FLAGS_debug_tuner = debug_tuner;
   });
   m.def("dump_cuda", [](bool dump_cuda) { tc::FLAGS_dump_cuda = dump_cuda; });
+  m.def("dump_ptx", [](bool dump_ptx) { tc::FLAGS_dump_ptx = dump_ptx; });
 
   // Access the names of the defs in a TC string
   m.def("parse_defs", [](const std::string& tc) {