diff --git a/litert/tensor/tensor.cc b/litert/tensor/tensor.cc index 5ef2dddb509..5e69d726a8f 100644 --- a/litert/tensor/tensor.cc +++ b/litert/tensor/tensor.cc @@ -15,6 +15,7 @@ limitations under the License. #include "litert/tensor/tensor.h" +#include #include #include #include @@ -31,6 +32,7 @@ limitations under the License. #include "litert/tensor/buffer.h" #include "litert/tensor/datatypes.h" #include "litert/tensor/internal/graph.h" +#include "litert/tensor/utils/macros.h" #include "litert/tensor/utils/source_location.h" namespace litert::tensor { @@ -52,14 +54,19 @@ TensorHandle TensorHandle::ShallowClone() const { } void TensorHandle::ShallowCloneTo(TensorHandle& other) const { - *GetInfo(other.impl_) = *GetInfo(impl_); + LRT_TENSOR_ASSIGN_OR_ABORT(graph::TensorInformation & other_info, + GetInfo(other.impl_)); + LRT_TENSOR_ASSIGN_OR_ABORT(other_info, GetInfo(impl_)); } TensorHandle& TensorHandle::Set(TensorInit init, source_location loc) & { if (!graph::GetStatus(impl_).ok()) { impl_ = graph::NewTensor(loc); } - graph::TensorInformation& info = *GetInfo(impl_); + // We just checked that the tensor exists so we access it directly to avoid + // triggering the absl::StatusOr linter that would force us to do a redundant + // check. + graph::TensorInformation& info = impl_.group->tensor_infos[impl_.index]; info.name = std::move(init.name); info.type = init.type; info.shape = std::move(init.shape); @@ -68,14 +75,27 @@ TensorHandle& TensorHandle::Set(TensorInit init, source_location loc) & { using T = std::decay_t; if constexpr (std::is_same_v>) { info.buffer = std::forward(arg); - } else if constexpr (std::is_same_v>) { - info.buffer = OwningCpuBuffer::Copy(arg); - } else if constexpr (std::is_same_v>) { - info.buffer = OwningCpuBuffer::Copy(arg); - } else if constexpr (std::is_same_v>) { - info.buffer = OwningCpuBuffer::Copy(arg); - } else if constexpr (std::is_same_v>) { - info.buffer = OwningCpuBuffer::Copy(arg); + } else if constexpr (std::is_same_v> || + std::is_same_v> || + std::is_same_v> || + std::is_same_v>) { + info.type = info.type == Type::kUnknown + ? ApiType::value + : info.type; + info.buffer = OwningCpuBuffer::CopyAs(info.type, arg); + } else if constexpr (std::is_arithmetic_v) { + if (info.type == Type::kUnknown) { + info.type = ApiType::value; + } + if (const size_t size = info.GetSize(); size != 1) { + info.buffer = + OwningCpuBuffer::CopyAs(info.type, std::vector(size, arg)); + } else { + info.buffer = OwningCpuBuffer::CopyAs(info.type, {arg}); + if (info.shape.empty()) { + info.shape = {1}; + } + } } else { ABSL_LOG(ERROR) << "Failed to create buffer from typed vector: " "unsupported datatype."; diff --git a/litert/tensor/tensor.h b/litert/tensor/tensor.h index 0f353cfe671..a526bfa345b 100644 --- a/litert/tensor/tensor.h +++ b/litert/tensor/tensor.h @@ -44,7 +44,8 @@ struct TensorInit { Type type = Type::kUnknown; Shape shape; std::variant, std::vector, - std::vector, std::vector, std::vector> + std::vector, std::vector, std::vector, + float, double, int8_t, int32_t, int64_t> buffer; std::shared_ptr quantization; }; diff --git a/litert/tensor/tensor_test.cc b/litert/tensor/tensor_test.cc index b18b0773c85..f88f49c69e1 100644 --- a/litert/tensor/tensor_test.cc +++ b/litert/tensor/tensor_test.cc @@ -33,10 +33,13 @@ namespace { using ::litert::tensor::IsOk; using ::litert::tensor::IsOkAndHolds; using ::testing::Address; +using ::testing::AllOf; +using ::testing::Each; using ::testing::ElementsAre; using ::testing::ElementsAreArray; using ::testing::Eq; using ::testing::Not; +using ::testing::SizeIs; using ::testing::StrEq; MATCHER(IsValidTensor, "") { @@ -102,6 +105,36 @@ TEST(TensorTest, SetBufferOnRValueWorks) { EXPECT_THAT(buffer.Lock(), ElementsAreArray(expected_buffer->Lock())); } +TEST(TensorTest, BuildWithIntegerScalarWorks) { + Tensor a = + TensorHandle(TensorInit{.type = Type::kI32, .shape = {1}, .buffer = 3}); + LRT_TENSOR_ASSERT_OK_AND_ASSIGN(Buffer & buffer, a.GetBuffer()); + EXPECT_THAT(buffer.Lock().As(), ElementsAre(3)); +} + +TEST(TensorTest, BuildWithFloatingScalarWorks) { + Tensor a = TensorHandle( + TensorInit{.type = Type::kFP32, .shape = {1}, .buffer = 3.14}); + LRT_TENSOR_ASSERT_OK_AND_ASSIGN(Buffer & buffer, a.GetBuffer()); + EXPECT_THAT(buffer.Lock().As(), ElementsAre(3.14)); +} + +TEST(TensorTest, BuildWithScalarInfersTypeAndShape) { + Tensor a = TensorHandle(TensorInit{.buffer = 3.14}); + EXPECT_THAT(a.GetType(), Type::kFP64); + EXPECT_THAT(a.GetShape(), ElementsAre(1)); + LRT_TENSOR_ASSERT_OK_AND_ASSIGN(Buffer & buffer, a.GetBuffer()); + EXPECT_THAT(buffer.Lock().As(), ElementsAre(3.14)); +} + +TEST(TensorTest, BuildWithScalarBroadcastsShape) { + Tensor a = TensorHandle(TensorInit{.shape = {2, 3}, .buffer = 3.14f}); + EXPECT_THAT(a.GetType(), Type::kFP32); + EXPECT_THAT(a.GetShape(), ElementsAre(2, 3)); + LRT_TENSOR_ASSERT_OK_AND_ASSIGN(Buffer & buffer, a.GetBuffer()); + EXPECT_THAT(buffer.Lock().As(), AllOf(SizeIs(6), Each(3.14))); +} + TEST(TensorTest, DefaultConstructedTensorDontHaveAProducer) { Tensor a; // The input tensors don't have a producer.