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

More generic and lightweight PyTorch integration

Author: nicolasvasilache

.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

python/examples/tc_pybind_example.py

@@ -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;

tc/aten/aten-inl.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);
 

tc/aten/aten.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;

tc/autotuner/autotuner-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;

tc/autotuner/options_cache-inl.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);
 

tc/autotuner/options_cache.h

@@ -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) {

tc/core/tensor.cc

@@ -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;

tc/core/tensor.h

@@ -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);
   }