unify dtype api

tenstorrent · Mar 5, 2025 · 5ecf5cb · 5ecf5cb
1 parent a59fdfe
commit 5ecf5cb
Show file tree

Hide file tree

Showing 5 changed files with 3 additions and 14 deletions.
diff --git a/runtime/include/tt/runtime/detail/ttmetal.h b/runtime/include/tt/runtime/detail/ttmetal.h
@@ -35,7 +35,6 @@ std::vector<std::uint32_t> getShape(::tt::runtime::Tensor tensor);
 std::vector<std::uint32_t> getStride(::tt::runtime::Tensor tensor);
 std::uint32_t getElementSize(::tt::runtime::Tensor tensor);
 std::uint32_t getVolume(::tt::runtime::Tensor tensor);
-target::DataType getDtype(::tt::runtime::Tensor tensor);
 
 size_t getNumAvailableDevices();
 

diff --git a/runtime/include/tt/runtime/detail/ttnn.h b/runtime/include/tt/runtime/detail/ttnn.h
@@ -99,7 +99,6 @@ std::vector<std::uint32_t> getShape(::tt::runtime::Tensor tensor);
 std::vector<std::uint32_t> getStride(::tt::runtime::Tensor tensor);
 std::uint32_t getElementSize(::tt::runtime::Tensor tensor);
 std::uint32_t getVolume(::tt::runtime::Tensor tensor);
-target::DataType getDtype(::tt::runtime::Tensor tensor);
 
 size_t getNumAvailableDevices();
 

diff --git a/runtime/lib/runtime.cpp b/runtime/lib/runtime.cpp
@@ -579,13 +579,13 @@ std::vector<std::uint32_t> Tensor::getStride() {
 target::DataType Tensor::getDtype() {
 #if defined(TT_RUNTIME_ENABLE_TTNN)
   if (getCurrentRuntime() == DeviceRuntime::TTNN) {
-    return ::tt::runtime::ttnn::getDtype(*this);
+    return ::tt::runtime::ttnn::getTensorDataType(*this);
   }
 #endif
 
 #if defined(TT_RUNTIME_ENABLE_TTMETAL)
   if (getCurrentRuntime() == DeviceRuntime::TTMetal) {
-    return ::tt::runtime::ttmetal::getDtype(*this);
+    return ::tt::runtime::ttmetal::getTensorDataType(*this);
   }
 #endif
   LOG_FATAL("runtime is not enabled");

diff --git a/runtime/lib/ttmetal/runtime.cpp b/runtime/lib/ttmetal/runtime.cpp
@@ -387,9 +387,5 @@ std::uint32_t getVolume(::tt::runtime::Tensor tensor) {
   LOG_WARNING("getVolume not implemented for metal runtime");
   return 0;
 }
-target::DataType getDtype(::tt::runtime::Tensor tensor) {
-  LOG_WARNING("getDtype not implemented for metal runtime");
-  return {};
-}
 
 } // namespace tt::runtime::ttmetal
diff --git a/runtime/lib/ttnn/runtime.cpp b/runtime/lib/ttnn/runtime.cpp
@@ -575,7 +575,7 @@ std::vector<std::byte> getDataBuffer(::tt::runtime::Tensor tensor) {
 
   // Need to `memcpy` in each case because the vector will go out of scope if we
   // wait until after the switch case
-  switch (getDtype(tensor)) {
+  switch (getTensorDataType(tensor)) {
   case target::DataType::BFP_BFloat4:
   case target::DataType::BFP_BFloat8:
   case target::DataType::Float32:
@@ -639,11 +639,6 @@ std::uint32_t getVolume(::tt::runtime::Tensor tensor) {
   return ttnnTensor->volume();
 }
 
-target::DataType getDtype(::tt::runtime::Tensor tensor) {
-  auto ttnnTensor = static_cast<::ttnn::Tensor *>(tensor.handle.get());
-  return utils::fromTTNNDataType(ttnnTensor->dtype());
-}
-
 std::vector<Tensor> submit(Device deviceHandle, Binary executableHandle,
                            std::uint32_t programIndex,
                            std::vector<Tensor> const &inputHandles) {