Merge pull request #495 from nicolasvasilache/pr/redo-pytorch

nicolasvasilache · web-flow · commit f6265e0a1a6b · 2018-06-13T19:33:46.000-04:00
More generic and lightweight PyTorch integration
diff --git a/.jenkins/build.sh b/.jenkins/build.sh
@@ -70,7 +70,7 @@ python setup.py install
 ./test_python/run_test.sh
 
 for f in $(find ./python/examples -name "*.py"); do
-    python $f
+    python $f -v
 done
 
 FILTER_OUT="benchmark_MLP_model benchmark_kronecker" ./test.sh
diff --git a/python/examples/tc_pybind_example.py b/python/examples/tc_pybind_example.py
diff --git a/tc/aten/aten-inl.h b/tc/aten/aten-inl.h
@@ -25,6 +25,57 @@
 
 namespace tc {
 namespace aten {
+
+// Stolen from ATen, get rid of our copy when ATen exposes the functionality
+// Unfortunately we need to wait for updated conda packages so we just copy
+// for now.
+inline DLDataType getDLDataType(const at::Type& type) {
+  using at::ScalarType;
+
+  DLDataType dtype;
+  dtype.lanes = 1;
+  dtype.bits = type.elementSizeInBytes() * 8;
+  switch (type.scalarType()) {
+    case ScalarType::Byte:
+      dtype.code = DLDataTypeCode::kDLUInt;
+      break;
+    case ScalarType::Char:
+      dtype.code = DLDataTypeCode::kDLInt;
+      break;
+    case ScalarType::Double:
+      dtype.code = DLDataTypeCode::kDLFloat;
+      break;
+    case ScalarType::Float:
+      dtype.code = DLDataTypeCode::kDLFloat;
+      break;
+    case ScalarType::Int:
+      dtype.code = DLDataTypeCode::kDLInt;
+      break;
+    case ScalarType::Long:
+      dtype.code = DLDataTypeCode::kDLInt;
+      break;
+    case ScalarType::Short:
+      dtype.code = DLDataTypeCode::kDLInt;
+      break;
+    case ScalarType::Half:
+      dtype.code = DLDataTypeCode::kDLFloat;
+      break;
+    case ScalarType::Undefined:
+      throw std::logic_error("Undefined is not a valid ScalarType");
+    case ScalarType::NumOptions:
+      throw std::logic_error("NumOptions is not a valid ScalarType");
+  }
+  return dtype;
+}
+
+inline TensorInfo toTensorInfo(const at::Tensor& t) {
+  return TensorInfo(
+      getDLDataType(t.type()),
+      reinterpret_cast<std::uintptr_t>(t.data_ptr()) % TensorInfo::kAlignment,
+      t.sizes(),
+      t.strides());
+}
+
 inline std::vector<DLTensorUPtr> makeDLTensors(
     const std::vector<at::Tensor>& tensors) {
   std::vector<DLTensorUPtr> dlTensors;
diff --git a/tc/aten/aten.h b/tc/aten/aten.h
@@ -25,6 +25,8 @@
 namespace tc {
 namespace aten {
 
+inline TensorInfo toTensorInfo(const at::Tensor&);
+
 inline std::vector<DLTensorUPtr> makeDLTensors(
     const std::vector<at::Tensor>& tensors);
 
diff --git a/tc/autotuner/autotuner-inl.h b/tc/autotuner/autotuner-inl.h
@@ -458,6 +458,7 @@ Autotuner<Backend, SearchStrategy>::tune(
     std::this_thread::sleep_for(std::chrono::milliseconds(100));
     if (sigint_) {
       tuningHarness.stopAfterCurrentIteration();
+      break;
     }
     if (sigterm_) {
       std::cerr << "Autotuning aborted." << std::endl;
diff --git a/tc/autotuner/options_cache-inl.h b/tc/autotuner/options_cache-inl.h
@@ -336,10 +336,10 @@ std::vector<typename Backend::MappingOptionsType> loadTopKFromCacheFile(
 
 template <typename Backend>
 void appendTopKToCacheFile(
-    const std::shared_ptr<OptionsCache<Backend>>& cache,
+    const OptionsCache<Backend>& cache,
     const std::string& cacheFilename,
     uint32_t count) {
-  OptionsCache<Backend> copy(*cache);
+  OptionsCache<Backend> copy(cache);
   copy.pruneKeepTopK(count);
   auto proto = copy.toProtobuf();
   OptionsCache<Backend> optionsCache;
diff --git a/tc/autotuner/options_cache.h b/tc/autotuner/options_cache.h
@@ -166,9 +166,10 @@ struct OptionsCache {
       OptionsCacheKeyHash>
       store_;
 
+  // Make friend to access toProtobuf/fromProtobuf
   template <typename BackendType>
   friend void appendTopKToCacheFile(
-      const std::shared_ptr<OptionsCache<BackendType>>& cache,
+      const OptionsCache<BackendType>& cache,
       const std::string& cacheFilename,
       uint32_t count);
 };
@@ -199,7 +200,7 @@ std::vector<typename Backend::MappingOptionsType> loadTopKFromCacheFile(
 /// needed.
 template <typename Backend>
 void appendTopKToCacheFile(
-    const std::shared_ptr<OptionsCache<Backend>>& cache,
+    const OptionsCache<Backend>& cache,
     const std::string& cacheFilename,
     uint32_t count);
 
diff --git a/tc/core/tensor.cc b/tc/core/tensor.cc
@@ -28,7 +28,8 @@ namespace tc {
 namespace detail {
 template <typename DLTensorType>
 uint64_t getDLTensorAlignment(const DLTensorType* t) {
-  return (reinterpret_cast<std::uintptr_t>(t->data) + t->byte_offset) % 256;
+  return (reinterpret_cast<std::uintptr_t>(t->data) + t->byte_offset) %
+      TensorInfo::kAlignment;
 }
 
 std::vector<int64_t> toIntVector(const int64_t* ptr, size_t ndim) {
diff --git a/tc/core/tensor.h b/tc/core/tensor.h
@@ -77,6 +77,8 @@ std::vector<const DLTensor*> extractRawPtrs(
  * It is serializable to protobuf and stored directly in the cache.
  */
 struct TensorInfo {
+  static constexpr int kAlignment = 256;
+
   DLDataType dtype;
   uint64_t alignment;
   std::vector<int64_t> shape;
diff --git a/tc/examples/blockdiagperm.cc b/tc/examples/blockdiagperm.cc
@@ -148,7 +148,7 @@ def blockdiagperm2dfissioned_2(float(B, N) I, int32(N) Idx) -> (O) {
   // later if required
   if (not FLAGS_options_cache.empty()) {
     tc::autotune::appendTopKToCacheFile(
-        geneticAutotuneATen.optionsCache,
+        *geneticAutotuneATen.optionsCache,
         FLAGS_options_cache,
         tc::FLAGS_tuner_save_best_candidates_count);
   }
diff --git a/tensor_comprehensions/pybinds/CMakeLists.txt b/tensor_comprehensions/pybinds/CMakeLists.txt
@@ -26,3 +26,22 @@ set_target_properties(tc PROPERTIES DEBUG_POSTFIX "")
 # "Install" to source_dir/tensor_comprehensions for easy inclusion in python
 # package with setuptools. Without this, setuptools is hell.
 install(TARGETS tc DESTINATION ${PROJECT_SOURCE_DIR}/tensor_comprehensions)
+
+add_library(tclib MODULE tclib.cc)
+target_include_directories(tclib PUBLIC ${PROJECT_SOURCE_DIR}/tc)
+target_link_libraries(
+  tclib
+
+  tc_autotuner
+  tc_aten
+
+  ${PYTHON_LIBRARIES}
+)
+
+set_target_properties(tclib PROPERTIES PREFIX "")
+set_target_properties(tclib PROPERTIES SUFFIX ".so")
+set_target_properties(tclib PROPERTIES DEBUG_POSTFIX "")
+
+# "Install" to source_dir/tensor_comprehensions for easy inclusion in python
+# package with setuptools. Without this, setuptools is hell.
+install(TARGETS tclib DESTINATION ${PROJECT_SOURCE_DIR}/tensor_comprehensions)
diff --git a/tensor_comprehensions/pybinds/tc.cc b/tensor_comprehensions/pybinds/tc.cc
@@ -106,7 +106,7 @@ PYBIND11_MODULE(tc, m) {
   m.def(
       "set_dump_cuda", [](bool dump_cuda) { tc::FLAGS_dump_cuda = dump_cuda; });
 
-  py::class_<ATenCudaTuner>(m, "ATenCudaTuner")
+  py::class_<ATenCudaTuner>(m, "ATenCudaTuner", py::module_local())
       .def(py::init<const std::string>())
       .def(
           "pop_size",
@@ -204,7 +204,8 @@ PYBIND11_MODULE(tc, m) {
                 numCandidates);
           });
 
-  py::class_<ATenCudaCompilationUnit>(m, "ATenCompilationUnit")
+  py::class_<ATenCudaCompilationUnit>(
+      m, "ATenCompilationUnit", py::module_local())
       .def(py::init<>())
       .def("define", &ATenCudaCompilationUnit::define, "Define the TC language")
       .def(
@@ -255,7 +256,8 @@ PYBIND11_MODULE(tc, m) {
   py::class_<tc::CudaMappingOptions>(
       m,
       "CudaMappingOptions",
-      "MappingCudaMappingOptions for a Tensor Comprehensions (TC)")
+      "CudaMappingOptions for a Tensor Comprehensions (TC)",
+      py::module_local())
       .def(
           py::init([](std::string type) {
             if (type == "naive") {
diff --git a/tensor_comprehensions/pybinds/tclib.cc b/tensor_comprehensions/pybinds/tclib.cc

Original file line number	Diff line number	Diff line change
`@@ -458,6 +458,7 @@ Autotuner<Backend, SearchStrategy>::tune(`
`458`	`458`	`std::this_thread::sleep_for(std::chrono::milliseconds(100));`
`459`	`459`	`if (sigint_) {`
`460`	`460`	`tuningHarness.stopAfterCurrentIteration();`
	`461`	`+ break;`
`461`	`462`	`}`
`462`	`463`	`if (sigterm_) {`
`463`	`464`	`std::cerr << "Autotuning aborted." << std::endl;`
Original file line number	Diff line number	Diff line change
`@@ -28,7 +28,8 @@ namespace tc {`
`28`	`28`	`namespace detail {`
`29`	`29`	`template <typename DLTensorType>`
`30`	`30`	`uint64_t getDLTensorAlignment(const DLTensorType* t) {`
`31`		`- return (reinterpret_cast<std::uintptr_t>(t->data) + t->byte_offset) % 256;`
	`31`	`+ return (reinterpret_cast<std::uintptr_t>(t->data) + t->byte_offset) %`
	`32`	`+ TensorInfo::kAlignment;`
`32`	`33`	`}`
`33`	`34`
`34`	`35`	`std::vector<int64_t> toIntVector(const int64_t* ptr, size_t ndim) {`
Original file line number	Diff line number	Diff line change
`@@ -148,7 +148,7 @@ def blockdiagperm2dfissioned_2(float(B, N) I, int32(N) Idx) -> (O) {`
`148`	`148`	`// later if required`
`149`	`149`	`if (not FLAGS_options_cache.empty()) {`
`150`	`150`	`tc::autotune::appendTopKToCacheFile(`
`151`		`- geneticAutotuneATen.optionsCache,`
	`151`	`+ *geneticAutotuneATen.optionsCache,`
`152`	`152`	`FLAGS_options_cache,`
`153`	`153`	`tc::FLAGS_tuner_save_best_candidates_count);`
`154`	`154`	`}`