Add shared_memory_size function to PyBinds

This function might be useful in the future for compilation options search.
One of the available options allows us to set the maximum shared memory used.

Author: Jules Pondard

tensor_comprehensions/pybinds/tclib.cc

@@ -30,6 +30,7 @@
 #include "tc/aten/aten_autotuner.h"
 #include "tc/autotuner/genetic_search.h"
 #include "tc/autotuner/options_cache.h"
+#include "tc/core/cuda/cuda.h"
 #include "tc/core/cuda/cuda_backend.h"
 #include "tc/core/cuda/cuda_tc_executor.h"
 #include "tc/core/flags.h"
@@ -447,6 +448,11 @@ PYBIND11_MODULE(tclib, m) {
     return res;
   });
 
+  // Get GPU shared memory size
+  m.def("shared_memory_size", []() {
+    return CudaGPUInfo::GPUInfo().SharedMemorySize();
+  });
+
   // Low-level stateful API compile returns an executor on which run and
   // unchecked_run can be called.
   py::class_<TcExecutor>(m, "TcExecutor")