Add at::Tensor -> TensorInfo

This commit adds a conversion from ATen tensor directly to TensorInfo.
Previously one needed to go through multiple conversions.
These conversions were expected to be fully memoized by keeping a reference
to the executor. However the PyTorch autograd function is stateless and we
cannot rely on this assumption.

Therefore we reduce the overhead by removing all the unnecessary conversions.
Notice that a we need to retrieve the dlpack type from the ATen type and
this function is not exposed on the ATen side so we currently copy the ATen
code. Once the function is exposed and shipped in the proper conda package
we can get rid of our copy.

Author: nicolasvasilache

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;

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

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

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;