tc::ATenCompilationUnit<tc::CudaTcExecutor> -> ATenCudaCompilationUnit

Author: nicolasvasilache

tensor_comprehensions/pybinds/pybind_engine.cc

@@ -36,6 +36,8 @@ namespace python {
 
 namespace py = pybind11;
 
+using ATenCudaCompilationUnit = tc::ATenCompilationUnit<tc::CudaTcExecutor>;
+
 PYBIND11_MODULE(tc, m) {
   m.def(
       "global_debug_init", // exposing the debugging flags to people
@@ -60,17 +62,13 @@ PYBIND11_MODULE(tc, m) {
     std::cerr << "\n PyTorch installation is missing, binary will be useless \n"
               << e.what() << std::endl;
   }
-  py::class_<tc::ATenCompilationUnit<tc::CudaTcExecutor>>(
-      m, "ATenCompilationUnit")
+  py::class_<ATenCudaCompilationUnit>(m, "ATenCompilationUnit")
       .def(py::init<>())
-      .def(
-          "define",
-          &tc::ATenCompilationUnit<tc::CudaTcExecutor>::define,
-          "Define the TC language")
+      .def("define", &ATenCudaCompilationUnit::define, "Define the TC language")
       .def(
           "compile",
           [dlpack](
-              tc::ATenCompilationUnit<tc::CudaTcExecutor>& instance,
+              ATenCudaCompilationUnit& instance,
               const std::string& name,
               py::list& inputs,
               const tc::MappingOptions& options) {
@@ -80,7 +78,7 @@ PYBIND11_MODULE(tc, m) {
       .def(
           "run",
           [dlpack](
-              tc::ATenCompilationUnit<tc::CudaTcExecutor>& instance,
+              ATenCudaCompilationUnit& instance,
               const std::string& name,
               py::list& inputs,
               py::list& outputs,
@@ -95,7 +93,7 @@ PYBIND11_MODULE(tc, m) {
       .def(
           "uncheckedRun",
           [dlpack](
-              tc::ATenCompilationUnit<tc::CudaTcExecutor>& instance,
+              ATenCudaCompilationUnit& instance,
               py::list& inputs,
               py::list& outputs,
               size_t handle) {
@@ -107,7 +105,7 @@ PYBIND11_MODULE(tc, m) {
       .def(
           "inject_cuda",
           [dlpack](
-              tc::ATenCompilationUnit<tc::CudaTcExecutor>& instance,
+              ATenCudaCompilationUnit& instance,
               const std::string& name,
               const std::string& injectedKernelName,
               const std::string& cudaSource,