Implement an XNNPack runner for Gemma 4.

litert/tensor/tensor.cc

@@ -15,6 +15,7 @@ limitations under the License.
 
 #include "litert/tensor/tensor.h"
 
+#include <cstddef>
 #include <cstdint>
 #include <memory>
 #include <string>
@@ -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<decltype(arg)>;
         if constexpr (std::is_same_v<T, std::shared_ptr<Buffer>>) {
           info.buffer = std::forward<decltype(arg)>(arg);
-        } else if constexpr (std::is_same_v<T, std::vector<float>>) {
-          info.buffer = OwningCpuBuffer::Copy<Type::kFP32>(arg);
-        } else if constexpr (std::is_same_v<T, std::vector<int32_t>>) {
-          info.buffer = OwningCpuBuffer::Copy<Type::kI32>(arg);
-        } else if constexpr (std::is_same_v<T, std::vector<int8_t>>) {
-          info.buffer = OwningCpuBuffer::Copy<Type::kI8>(arg);
-        } else if constexpr (std::is_same_v<T, std::vector<int4_t>>) {
-          info.buffer = OwningCpuBuffer::Copy<Type::kI4>(arg);
+        } else if constexpr (std::is_same_v<T, std::vector<float>> ||
+                             std::is_same_v<T, std::vector<int32_t>> ||
+                             std::is_same_v<T, std::vector<int8_t>> ||
+                             std::is_same_v<T, std::vector<int4_t>>) {
+          info.type = info.type == Type::kUnknown
+                          ? ApiType<typename T::value_type>::value
+                          : info.type;
+          info.buffer = OwningCpuBuffer::CopyAs(info.type, arg);
+        } else if constexpr (std::is_arithmetic_v<T>) {
+          if (info.type == Type::kUnknown) {
+            info.type = ApiType<T>::value;
+          }
+          if (const size_t size = info.GetSize(); size != 1) {
+            info.buffer =
+                OwningCpuBuffer::CopyAs(info.type, std::vector<T>(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.";

litert/tensor/tensor.h

@@ -44,7 +44,8 @@ struct TensorInit {
   Type type = Type::kUnknown;
   Shape shape;
   std::variant<std::shared_ptr<Buffer>, std::vector<float>,
-               std::vector<int32_t>, std::vector<int8_t>, std::vector<int4_t>>
+               std::vector<int32_t>, std::vector<int8_t>, std::vector<int4_t>,
+               float, double, int8_t, int32_t, int64_t>
       buffer;
   std::shared_ptr<Quantization> quantization;
 };

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<const int>(), 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<const float>(), 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<const double>(), 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<const float>(), AllOf(SizeIs(6), Each(3.14)));
+}
+
 TEST(TensorTest, DefaultConstructedTensorDontHaveAProducer) {
   Tensor a;
   // The input tensors don't have a producer.