Implement runtime tensor desc API #2370

ctodTT · 2025-03-05T17:59:31Z

This change implements an API on runtime tensors to reflect tensor metadata and contents reflecting the underlying TTNN tensor. This functionality is pybound as well, allowing for easy casting to torch tensors. Please see runtime/test/python/ttnn/test_runtime_api.py for an example.

NOTE: This is not the most efficient implementation possible, as there is effectively a double copy to get the data in a pybindable row major format. There is probably some trickery that can be done later to avoid this, but I would like to get this functionality out ASAP and avoid premature optimization

Big thanks to @tapspatel for helping me figure out the namespace dispatch intricacies here!

Ticket

Closes #1957

Checklist

New/Existing tests provide coverage for changes

This change implements an API on runtime tensors to reflect tensor metadata and contents reflecting the underlying TTNN tensor. This functionality is pybound as well, allowing for easy casting to torch tensors. Please see `runtime/test/python/ttnn/test_runtime_api.py` for an example. NOTE: This is not the most efficient implementation possible, as there is effectively a double copy to get the data in a pybindable row major format. There is probably some trickery that can be done later to avoid this, but I would like to get this functionality out ASAP and avoid premature optimization Closes #1957

github-actions

⚠️ Clang-Tidy found issue(s) with the introduced code (1/1)

runtime/lib/ttnn/runtime.cpp

tapspatel

fantastic changes! Thanks for the continued push on this. Few comments inline.

Will this break the golden callback functionality? Or do we already handle the types dynamically.

runtime/include/tt/runtime/detail/ttmetal.h

runtime/include/tt/runtime/detail/ttnn.h

tapspatel · 2025-03-05T18:08:31Z

runtime/lib/ttnn/runtime.cpp

@@ -567,6 +567,83 @@ std::vector<float> getTensorData(Tensor tensor) {
                            static_cast<float *>(dataPtr) + nnTensor->volume());
 }

+std::vector<std::byte> getDataBuffer(::tt::runtime::Tensor tensor) {


this is extremely clean. I love it.

tapspatel · 2025-03-05T18:08:55Z

runtime/test/python/ttnn/test_runtime_api.py

@@ -10,6 +10,34 @@
 from utils import TT_MLIR_HOME, Helper, DeviceContext, assert_pcc


+@pytest.mark.parametrize("shape", [(64, 128)])
+@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16])
+def test_tensor_buffer_api(shape, dtype):


great way to test!

Thank you for adding tests!

jnie-TT · 2025-03-05T18:12:11Z

runtime/include/tt/runtime/detail/ttnn.h

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


We have a TensorDesc type under runtime/include/tt/runtime/types.h:

struct TensorDesc { std::vector<std::uint32_t> shape; std::vector<std::uint32_t> stride; std::uint32_t itemsize; ::tt::target::DataType dataType; };

Was wondering if we could merge these APIs into one and return the TensorDesc directly. It shouldn't be hard to bind this structure, and it'll have all the information in one place.

Also on a side note I think getDtype and getTensorDataType do the same thing here so is probably redundant.

Good point, yeah I can bind and have it return that object if you want!

Yeah that would be awesome!

There is now getTensorDesc() & get_tensor_desc() for this exact purpose!

runtime/lib/ttnn/runtime.cpp

jnie-TT · 2025-03-05T18:21:05Z

runtime/lib/ttnn/runtime.cpp

+}
+
+std::uint32_t getVolume(::tt::runtime::Tensor tensor) {
+  auto ttnnTensor = static_cast<::ttnn::Tensor *>(tensor.handle.get());


We should be using tensor.as here since it'll also check that the device runtime matches

jnie-TT · 2025-03-05T18:26:29Z

runtime/lib/ttnn/runtime.cpp

+  case target::DataType::Float32:
+    dataPtr = ttnnTensor->to_vector<float>().data();
+    assert(dataPtr != nullptr);
+    std::memcpy(dataVec.data(), dataPtr, numBytes);


I think numBytes will contain the wrong value here and won't be compatible for BFloat4 and BFloat8

oh shoot, you're right! good catch. let me fix that

jameszianxuTT

Looks clean and the added tests provide a good reference for my implementation of golden checking with this API in tt-torch.

Having an API to return a TensorDesc would be great, so I second Jackson's request, but otherwise no complaints. Thanks!

Now all golden checks utilize the runtime tensor buffer API

jnie-TT · 2025-03-07T20:55:23Z

runtime/include/tt/runtime/types.h

+  std::vector<std::uint32_t> getShape();
+  std::vector<std::uint32_t> getStride();
+  target::DataType getDtype();
+  TensorDesc getTensorDesc();


Let's not add these as member functions for consistency reasons. I think for now we want to put all user-facing functions in runtime.h. This is open for discussion though, if we prefer member functions over standalone functions then this should be a larger scale change that covers all existing functions.

Also since we have getTensorDesc we can probably remove all other APIs since getTensorDesc will return all the information anyway.

jnie-TT

Thanks for the change Collin!

ctodTT requested review from LPanosTT and jameszianxuTT March 5, 2025 17:59

ctodTT requested review from tapspatel, jnie-TT, kmabeeTT, AleksKnezevic and pilkicTT as code owners March 5, 2025 17:59

github-actions bot reviewed Mar 5, 2025

View reviewed changes

tapspatel approved these changes Mar 5, 2025

View reviewed changes

jnie-TT reviewed Mar 5, 2025

View reviewed changes

runtime/lib/ttnn/runtime.cpp Outdated Show resolved Hide resolved

jnie-TT reviewed Mar 5, 2025

View reviewed changes

LPanosTT approved these changes Mar 5, 2025

View reviewed changes

jnie-TT reviewed Mar 5, 2025

View reviewed changes

jameszianxuTT approved these changes Mar 5, 2025

View reviewed changes

ctodTT added 4 commits March 5, 2025 20:08

Fix nits

f66be99

use tensordesc

197a1a3

Fix block-type sizes and dangling pointer

5b8177d

Remove now obsolete getTensorData

490e719

Now all golden checks utilize the runtime tensor buffer API

ctodTT force-pushed the ctod/issue-1957 branch from 5ecf5cb to 490e719 Compare March 6, 2025 19:10

use correct sizes for memcpy

a90028d

jnie-TT reviewed Mar 7, 2025

View reviewed changes

ctodTT added 2 commits March 7, 2025 21:52

Make API consistent by removing methods

ef8fd1d

Change API names in C++ for consistency

8045310

jnie-TT approved these changes Mar 7, 2025

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Implement runtime tensor desc API #2370

Implement runtime tensor desc API #2370

ctodTT commented Mar 5, 2025 •

edited

Loading

github-actions bot left a comment

tapspatel left a comment

tapspatel Mar 5, 2025

tapspatel Mar 5, 2025

jnie-TT Mar 5, 2025

jnie-TT Mar 5, 2025 •

edited

Loading

ctodTT Mar 5, 2025

jnie-TT Mar 5, 2025

ctodTT Mar 6, 2025 •

edited

Loading

jnie-TT Mar 5, 2025 •

edited

Loading

ctodTT Mar 6, 2025

jnie-TT Mar 5, 2025

ctodTT Mar 5, 2025

ctodTT Mar 6, 2025

jameszianxuTT left a comment

jnie-TT Mar 7, 2025 •

edited

Loading

jnie-TT left a comment

Implement runtime tensor desc API #2370

Are you sure you want to change the base?

Implement runtime tensor desc API #2370

Conversation

ctodTT commented Mar 5, 2025 • edited Loading

Ticket

Checklist

github-actions bot left a comment

Choose a reason for hiding this comment

tapspatel left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

jnie-TT Mar 5, 2025 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

ctodTT Mar 6, 2025 • edited Loading

Choose a reason for hiding this comment

jnie-TT Mar 5, 2025 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

jameszianxuTT left a comment

Choose a reason for hiding this comment

jnie-TT Mar 7, 2025 • edited Loading

Choose a reason for hiding this comment

jnie-TT left a comment

Choose a reason for hiding this comment

ctodTT commented Mar 5, 2025 •

edited

Loading

jnie-TT Mar 5, 2025 •

edited

Loading

ctodTT Mar 6, 2025 •

edited

Loading

jnie-TT Mar 5, 2025 •

edited

Loading

jnie-TT Mar 7, 2025 •

edited

Loading