Add API to move tensors to CPU

tenstorrent · Sep 8, 2024 · 4555dd8 · 4555dd8
1 parent 608ed89
commit 4555dd8
Show file tree

Hide file tree

Showing 4 changed files with 41 additions and 8 deletions.
diff --git a/runtime/include/tt/runtime/detail/ttnn.h b/runtime/include/tt/runtime/detail/ttnn.h
@@ -79,6 +79,10 @@ inline Tensor createTensor(std::shared_ptr<void> data, TensorDesc const &desc) {
 
 tt::target::DataType getTensorDataType(Tensor tensor);
 
+void deallocateTensor(Tensor tensor, bool force);
+
+Tensor toCpu(Tensor tensor);
+
 Device openDevice(std::vector<int> const &deviceIds = {0},
                   std::vector<std::uint8_t> const &numHWCQs = {});
 
@@ -100,7 +104,7 @@ void runProgram(::ttnn::Device &device,
                 ::tt::target::ttnn::Program const *program,
                 std::vector<::ttnn::Tensor *> const &inputs,
                 std::vector<::ttnn::Tensor *> const &outputs);
-                
+
 std::vector<Tensor> runProgram(::ttnn::Device &device,
                                ::tt::target::ttnn::Program const *program,
                                std::vector<::ttnn::Tensor *> const &inputs);

diff --git a/runtime/include/tt/runtime/types.h b/runtime/include/tt/runtime/types.h
@@ -132,15 +132,16 @@ struct Tensor : public detail::RuntimeCheckedObjectImpl {
       : detail::RuntimeCheckedObjectImpl(handle, runtime), data(data),
         event(Event(nullptr, runtime)) {}
 
-
   Tensor(std::shared_ptr<void> handle, std::shared_ptr<void> data,
          DeviceRuntime runtime, Event event)
       : detail::RuntimeCheckedObjectImpl(handle, runtime), data(data),
         event(event) {}
 
   // Users need to manually deallocate tensors returned from submit
   // As the storage is now owned instead of borrowed
-  void deallocate();
+  void deallocate(bool force = false);
+
+  Tensor cpu() const;
 };
 
 } // namespace tt::runtime

diff --git a/runtime/lib/ttnn/runtime.cpp b/runtime/lib/ttnn/runtime.cpp
@@ -56,6 +56,20 @@ tt::target::DataType getTensorDataType(Tensor tensor) {
   return utils::fromTTNNDataType(nnTensor.dtype());
 }
 
+void deallocateTensor(Tensor tensor, bool force) {
+  ::ttnn::Tensor &ttnnTensor = tensor.as<::ttnn::Tensor>(DeviceRuntime::TTNN);
+  ttnnTensor.deallocate(force);
+}
+
+Tensor toCpu(Tensor tensor) {
+  ::ttnn::Tensor &ttnnTensor = tensor.as<::ttnn::Tensor>(DeviceRuntime::TTNN);
+  std::shared_ptr<::ttnn::Tensor> cpuTensor =
+      std::make_shared<::ttnn::Tensor>(ttnnTensor.cpu());
+  void *dataPtr = ::tt::tt_metal::get_raw_host_data_ptr(*cpuTensor);
+  return Tensor(cpuTensor, ::tt::runtime::utils::unsafe_borrow_shared(dataPtr),
+                DeviceRuntime::TTNN);
+}
+
 Device openDevice(std::vector<int> const &deviceIds,
                   std::vector<std::uint8_t> const &numHWCQs) {
   assert(deviceIds.size() == 1 && "Only one device is supported for now");

diff --git a/runtime/lib/types.cpp b/runtime/lib/types.cpp
@@ -14,19 +14,33 @@
 
 namespace tt::runtime {
 
-void Tensor::deallocate() {
+void Tensor::deallocate(bool force) {
 #if defined(TT_RUNTIME_ENABLE_TTNN)
   if (this->matchesRuntime(DeviceRuntime::TTNN)) {
-    ::ttnn::Tensor &tensor = this->as<::ttnn::Tensor>(DeviceRuntime::TTNN);
-    tensor.deallocate();
-    return;
+    ::tt::runtime::ttnn::deallocateTensor(*this, force);
   }
-#elif defined(TT_RUNTIME_ENABLE_TTMETAL)
+#endif
+
+#if defined(TT_RUNTIME_ENABLE_TTMETAL)
   if (this->matchesRuntime(DeviceRuntime::TTMetal)) {
     throw std::runtime_error("Not implemented");
   }
 #endif
   throw std::runtime_error("Runtime not enabled");
 }
 
+Tensor Tensor::cpu() const {
+#if defined(TT_RUNTIME_ENABLE_TTNN)
+  if (this->matchesRuntime(DeviceRuntime::TTNN)) {
+    return ::tt::runtime::ttnn::toCpu(*this);
+  }
+#endif
+
+#if defined(TT_RUNTIME_ENABLE_TTMETAL)
+  if (this->matchesRuntime(DeviceRuntime::TTMetal)) {
+    throw std::runtime_error("Not implemented");
+  }
+#endif
+  throw std::runtime_error("Runtime not enabled");
+}
 } // namespace tt::runtime