Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

#0: Improve error messages in Tensor storage classes #18321

Open
wants to merge 8 commits into
base: main
Choose a base branch
from
Open

#0: Improve error messages in Tensor storage classes #18321

Changes from all commits
Commits
Show all changes
8 commits
Select commit Hold shift + click to select a range
ada5273
Clarify error message
jjiangTT Feb 25, 2025
39b8ce1
improve error clarification
jjiangTT Feb 25, 2025
1699e7d
change assert to a fatal
jjiangTT Feb 26, 2025
f485fd1
Change asserts to fatals where reasonable
jjiangTT Feb 27, 2025
530e024
Search buffer for device id once only
jjiangTT Feb 27, 2025
ee92213
fix pointer error
jjiangTT Mar 3, 2025
d4227ab
fix pair access
jjiangTT Mar 3, 2025
4eeb969
fix ->second error, replace other redundant searches
jjiangTT Mar 3, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
25 changes: 14 additions & 11 deletions ttnn/cpp/ttnn/tensor/storage.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -188,19 +188,19 @@ struct MultiDeviceHostStorage {

OwnedBuffer get_buffer(int buffer_index) const {
std::lock_guard<std::mutex> lock(mtx);
TT_ASSERT(buffer_index < buffers.size(), "Buffer not found for buffer_index {}", buffer_index);
TT_FATAL(buffer_index < buffers.size(), "Buffer not found for buffer_index {}", buffer_index);
return buffers[buffer_index];
}

OwnedBuffer& get_buffer(int buffer_index) {
std::lock_guard<std::mutex> lock(mtx);
TT_ASSERT(buffer_index < buffers.size(), "Buffer not found for buffer_index {}", buffer_index);
TT_FATAL(buffer_index < buffers.size(), "Buffer not found for buffer_index {}", buffer_index);
return buffers[buffer_index];
}

TensorSpec get_tensor_spec(int spec_index) const {
std::lock_guard<std::mutex> lock(mtx);
TT_ASSERT(spec_index < specs.size(), "Buffer not found for device {}", spec_index);
TT_FATAL(spec_index < specs.size(), "Spec for device {} not found in spec list", spec_index);
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think you can make the same change in other places in the file, for consistency

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This looks good!

return specs[spec_index];
}

Expand Down Expand Up @@ -325,20 +325,22 @@ struct MultiDeviceStorage {

inline std::shared_ptr<Buffer> get_buffer_for_device(IDevice* device) const {
std::lock_guard<std::mutex> lock(buffer_mtx);
TT_ASSERT(buffers.find(device->id()) != buffers.end(), "Buffer not found for device {}", device->id());
auto buffer_it = buffers.find(device->id());
TT_FATAL(buffer_it != buffers.end(), "Buffer not found for device {}", device->id());
TT_ASSERT(
buffers.at(device->id())->device() == device,
buffer_it->second->device() == device,
"Mismatch between device derived from buffer and device derived from MultiDeviceStorage.");
return buffers.at(device->id());
return buffer_it->second;
}

inline std::shared_ptr<Buffer>& get_buffer_for_device(IDevice* device) {
std::lock_guard<std::mutex> lock(buffer_mtx);
TT_ASSERT(buffers.find(device->id()) != buffers.end(), "Buffer not found for device {}", device->id());
auto buffer_it = buffers.find(device->id());
TT_FATAL(buffer_it != buffers.end(), "Buffer not found for device {}", device->id());
TT_ASSERT(
buffers.at(device->id())->device() == device,
buffer_it->second->device() == device,
"Mismatch between device derived from buffer and device derived from MultiDeviceStorage.");
return buffers.at(device->id());
return buffer_it->second;
}

inline std::shared_ptr<Buffer> get_buffer_for_device_id(uint32_t device_id) const {
Expand All @@ -348,8 +350,9 @@ struct MultiDeviceStorage {

inline TensorSpec get_tensor_spec_for_device(IDevice* device) const {
std::lock_guard<std::mutex> lock(shape_mtx);
TT_ASSERT(specs.find(device->id()) != specs.end(), "Shape not found for device {}", device->id());
return specs.at(device->id());
auto spec_it = specs.find(device->id());
TT_FATAL(spec_it != specs.end(), "Shape not found for device {}", device->id());
return spec_it->second;
}

inline uint32_t num_buffers() const {
Expand Down
Loading