ATen/Dlpack functions -> aten/utils*h

Author: nicolasvasilache

include/tc/aten/aten_compiler.h

@@ -22,13 +22,13 @@
 #include <ATen/ATen.h>
 #include <ATen/DLConvertor.h>
 
+#include "tc/aten/utils.h"
 #include "tc/core/cuda/cuda.h"
 #include "tc/core/cuda/cuda_tc_executor.h"
 #include "tc/core/execution_engine.h"
 #include "tc/lang/parser.h"
 
 namespace tc {
-
 /// This provides the basic interface for writing ATen style tensor operations
 /// based on Tensor Comprehensions.
 
@@ -74,13 +74,4 @@ class ATenCompilationUnit {
  private:
   std::unique_ptr<ExecutionEngine<CudaTcExecutor>> executionEngine_;
 };
-
-std::pair<std::vector<DLTensor*>, std::vector<DLManagedTensor*>>
-toDlpackTensors(const std::vector<at::Tensor>& tensors);
-
-std::pair<std::vector<const DLTensor*>, std::vector<DLManagedTensor*>>
-toConstDlpackTensors(const std::vector<at::Tensor>& tensors);
-
-void deleteDlmTensors(std::vector<DLManagedTensor*>& tensors);
-
 } // namespace tc

include/tc/aten/utils-inl.h

@@ -0,0 +1,53 @@
+/**
+ * Copyright (c) 2017-present, Facebook, Inc.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+#pragma once
+
+#include <string>
+#include <vector>
+
+#include <ATen/ATen.h>
+#include <ATen/DLConvertor.h>
+namespace tc {
+std::pair<std::vector<DLTensor*>, std::vector<DLManagedTensor*>>
+toDlpackTensors(const std::vector<at::Tensor>& tensors) {
+  std::vector<DLTensor*> dlTensors;
+  std::vector<DLManagedTensor*> dlMTensors;
+  for (auto tensor : tensors) {
+    auto dlMTensor = at::toDLPack(tensor);
+    dlTensors.push_back(&(dlMTensor->dl_tensor));
+    dlMTensors.push_back(dlMTensor);
+  }
+  return make_pair(dlTensors, dlMTensors);
+}
+
+std::pair<std::vector<const DLTensor*>, std::vector<DLManagedTensor*>>
+toConstDlpackTensors(const std::vector<at::Tensor>& tensors) {
+  std::vector<const DLTensor*> dlTensors;
+  std::vector<DLManagedTensor*> dlMTensors;
+  for (auto tensor : tensors) {
+    auto dlMTensor = at::toDLPack(tensor);
+    dlTensors.push_back(&(dlMTensor->dl_tensor));
+    dlMTensors.push_back(dlMTensor);
+  }
+  return make_pair(dlTensors, dlMTensors);
+}
+
+void deleteDlmTensors(std::vector<DLManagedTensor*>& tensors) {
+  for (auto& tensor : tensors) {
+    tensor->deleter(tensor);
+  }
+}
+} // namespace tc

include/tc/aten/utils.h

@@ -0,0 +1,34 @@
+/**
+ * Copyright (c) 2017-present, Facebook, Inc.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+#pragma once
+
+#include <string>
+#include <vector>
+
+#include <ATen/ATen.h>
+#include <ATen/DLConvertor.h>
+
+namespace tc {
+std::pair<std::vector<DLTensor*>, std::vector<DLManagedTensor*>>
+toDlpackTensors(const std::vector<at::Tensor>& tensors);
+
+std::pair<std::vector<const DLTensor*>, std::vector<DLManagedTensor*>>
+toConstDlpackTensors(const std::vector<at::Tensor>& tensors);
+
+void deleteDlmTensors(std::vector<DLManagedTensor*>& tensors);
+} // namespace tc
+
+#include "tc/aten/utils-inl.h"

src/aten/aten_compiler.cc

@@ -32,12 +32,6 @@ void ATenCompilationUnit::define(const std::string& language) {
   executionEngine_->define(language);
 }
 
-void deleteDlmTensors(std::vector<DLManagedTensor*>& tensors) {
-  for (auto& tensor : tensors) {
-    tensor->deleter(tensor);
-  }
-}
-
 namespace {
 
 // given the tensor shape and DLType, allocate storage for the tensor output
@@ -73,30 +67,6 @@ void prepareOutputs(
 
 } // namespace
 
-std::pair<std::vector<DLTensor*>, std::vector<DLManagedTensor*>>
-toDlpackTensors(const std::vector<at::Tensor>& tensors) {
-  std::vector<DLTensor*> dlTensors;
-  std::vector<DLManagedTensor*> dlMTensors;
-  for (auto tensor : tensors) {
-    auto dlMTensor = at::toDLPack(tensor);
-    dlTensors.push_back(&(dlMTensor->dl_tensor));
-    dlMTensors.push_back(dlMTensor);
-  }
-  return make_pair(dlTensors, dlMTensors);
-}
-
-std::pair<std::vector<const DLTensor*>, std::vector<DLManagedTensor*>>
-toConstDlpackTensors(const std::vector<at::Tensor>& tensors) {
-  std::vector<const DLTensor*> dlTensors;
-  std::vector<DLManagedTensor*> dlMTensors;
-  for (auto tensor : tensors) {
-    auto dlMTensor = at::toDLPack(tensor);
-    dlTensors.push_back(&(dlMTensor->dl_tensor));
-    dlMTensors.push_back(dlMTensor);
-  }
-  return make_pair(dlTensors, dlMTensors);
-}
-
 size_t ATenCompilationUnit::compile(
     const std::string& name,
     const std::vector<at::Tensor>& inputs,