Skip to content

Commit

Permalink
Fix tensor post processing
Browse files Browse the repository at this point in the history
This change fixes an issue with processing output tensors after they are
pulled off the device when checking golden-ness. The resulting tensor
had the first 32 bytes overwritten with seemingly uninitialized memory,
resulting in nondeterministic failures during similarity checks. Now,
the unadulterated tensor is returned.

Fixes #1795
  • Loading branch information
ctodTT committed Jan 15, 2025
1 parent f30a198 commit 0b18496
Showing 1 changed file with 5 additions and 16 deletions.
21 changes: 5 additions & 16 deletions runtime/lib/ttnn/runtime.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -467,23 +467,12 @@ Tensor getOpOutputTensor(OpContext opContextHandle,
return createNullTensor();
}

::ttnn::Tensor hostTensor = ::ttnn::from_device(*outPtr);
::ttnn::Tensor outCopy =
::ttnn::to_layout(hostTensor, ::ttnn::ROW_MAJOR_LAYOUT, std::nullopt,
std::nullopt, static_cast<::ttnn::IDevice *>(nullptr));

void *src = ::tt::tt_metal::get_raw_host_data_ptr(outCopy);
std::uint32_t outCopySize = outCopy.volume() * outCopy.element_size();
std::shared_ptr<void> data = ::tt::runtime::utils::malloc_shared(outCopySize);
std::memcpy(data.get(), src, outCopySize);

auto tensor = std::make_shared<::ttnn::Tensor>(
ttnn::createStorage<BorrowedStorage>(data.get(), outCopy.volume(),
::tt::target::DataType::Float32),
outCopy.shape().value, ::ttnn::DataType::FLOAT32,
::ttnn::Layout::ROW_MAJOR);
std::shared_ptr<::ttnn::Tensor> hostTensor =
std::make_shared<::ttnn::Tensor>(::ttnn::to_layout(
::ttnn::from_device(*outPtr), ::ttnn::Layout::ROW_MAJOR, std::nullopt,
std::nullopt, static_cast<::ttnn::IDevice *>(nullptr)));

return Tensor(std::static_pointer_cast<void>(tensor), nullptr,
return Tensor(std::static_pointer_cast<void>(hostTensor), nullptr,
DeviceRuntime::TTNN);
}

Expand Down

0 comments on commit 0b18496

Please sign in to comment.