chore: import upstream snapshot with attribution

tensorflow/lite/experimental/shlo/legacy/bench/BUILD

@@ -0,0 +1,51 @@
+load("@rules_cc//cc:cc_binary.bzl", "cc_binary")
+load("@rules_cc//cc:cc_library.bzl", "cc_library")
+
+# copybara:uncomment package(default_applicable_licenses = ["//tensorflow:LICENSE"])
+
+cc_library(
+    name = "util",
+    srcs = [
+    ],
+    hdrs = [
+        "util.h",
+    ],
+    tags = ["no_oss"],
+    deps = [
+        "//tensorflow/lite/experimental/shlo/legacy:float",
+    ],
+)
+
+cc_binary(
+    name = "shlo_benchmark",
+    srcs = [
+        "shlo_benchmark.cc",
+    ],
+    tags = ["no_oss"],
+    deps = [
+        ":util",
+        "//tensorflow/lite/experimental/shlo/legacy:shlo",
+        "//tensorflow/lite/experimental/shlo/legacy/test:util",
+        "@com_google_absl//absl/log",
+        "@com_google_absl//absl/log:check",
+        "@com_google_absl//absl/status",
+        "@com_google_benchmark//:benchmark",
+    ],
+)
+
+cc_binary(
+    name = "xnn_benchmark",
+    srcs = [
+        "xnn_benchmark.cc",
+    ],
+    data = [
+    ],
+    tags = ["no_oss"],
+    deps = [
+        ":util",
+        "//tensorflow/lite/experimental/shlo/legacy:float",
+        "@XNNPACK",
+        "@com_google_absl//absl/log",
+        "@com_google_benchmark//:benchmark",
+    ],
+)

tensorflow/lite/experimental/shlo/legacy/bench/shlo_benchmark.cc

@@ -0,0 +1,180 @@
+/* Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include <utility>
+
+#include "absl/log/check.h"
+#include "absl/log/log.h"
+#include "absl/status/status.h"
+#include "benchmark/benchmark.h"  // from @com_google_benchmark
+#include "tensorflow/lite/experimental/shlo/legacy/bench/util.h"
+#include "tensorflow/lite/experimental/shlo/legacy/include/shlo.h"
+#include "tensorflow/lite/experimental/shlo/legacy/src/storage.h"
+#include "tensorflow/lite/experimental/shlo/legacy/test/util.h"
+
+namespace stablehlo {
+namespace benchmark {
+namespace {
+
+template <absl::Status (*op)(const Tensor&, Tensor&), ElementType element_type,
+          int size>
+void BM_SHLO(::benchmark::State& state) {
+  Shape shape = {size};
+
+  using ST = typename Storage<element_type>::Type;
+  auto operand_values = GenerateRandomVector<ST>(size);
+  decltype(operand_values) result_values(operand_values.size());
+
+  Tensor operand(TensorType(Shape(shape), element_type),
+                 std::data(operand_values));
+  Tensor result(TensorType(Shape(shape), element_type),
+                std::data(result_values));
+
+  for (auto _ : state) {
+    auto res = op(operand, result);
+    CHECK_OK(res);
+  }
+}
+
+template <absl::Status (*op)(const Tensor&, const Tensor&, Tensor&),
+          ElementType element_type, int size>
+void BM_SHLO(::benchmark::State& state) {
+  Shape shape = {size};
+
+  using ST = typename Storage<element_type>::Type;
+  auto lhs_values = GenerateRandomVector<ST>(size);
+  auto rhs_values = GenerateRandomVector<ST>(size);
+  decltype(rhs_values) result_values(rhs_values.size());
+
+  Tensor lhs(TensorType(Shape(shape), element_type), std::data(lhs_values));
+  Tensor rhs(TensorType(Shape(shape), element_type), std::data(rhs_values));
+  Tensor result(TensorType(Shape(shape), element_type),
+                std::data(result_values));
+
+  for (auto _ : state) {
+    auto res = op(lhs, rhs, result);
+    CHECK_OK(res);
+  }
+}
+
+template <absl::Status (*op)(const QuantizedTensor&, QuantizedTensor&),
+          ElementType storage_type, ElementType expressed_type, int size>
+void BM_SHLO(::benchmark::State& state) {
+  Shape shape = {size};
+  QuantizedParameter quantized_parameter = {.scale = 0.1, .zero_point = 0};
+
+  using ET = typename Storage<expressed_type>::Type;
+
+  auto operand_values = GenerateRandomVector<ET>(size);
+
+  auto operand_quant_values =
+      testing::QuantizeVector<storage_type, expressed_type>(
+          operand_values, quantized_parameter);
+  decltype(operand_quant_values) result_quant_values(
+      operand_quant_values.size());
+
+  QuantizedTensorElementType element_type(storage_type, expressed_type,
+                                          std::move(quantized_parameter));
+
+  QuantizedTensor operand(
+      QuantizedTensorType(Shape(shape),
+                          QuantizedTensorElementType(element_type)),
+      operand_quant_values.data());
+  QuantizedTensor result(
+      QuantizedTensorType(Shape(shape),
+                          QuantizedTensorElementType(element_type)),
+      result_quant_values.data());
+
+  for (auto _ : state) {
+    auto res = op(operand, result);
+    CHECK_OK(res);
+  }
+}
+
+template <absl::Status (*op)(const QuantizedTensor&, const QuantizedTensor&,
+                             QuantizedTensor&),
+          ElementType storage_type, ElementType expressed_type, int size>
+void BM_SHLO(::benchmark::State& state) {
+  Shape shape = {size};
+  QuantizedParameter quantized_parameter = {.scale = 0.1, .zero_point = 0};
+
+  using ET = typename Storage<expressed_type>::Type;
+
+  auto lhs_values = GenerateRandomVector<ET>(size);
+  auto rhs_values = GenerateRandomVector<ET>(size);
+
+  auto lhs_quant_values = testing::QuantizeVector<storage_type, expressed_type>(
+      lhs_values, quantized_parameter);
+  auto rhs_quant_values = testing::QuantizeVector<storage_type, expressed_type>(
+      rhs_values, quantized_parameter);
+  decltype(rhs_quant_values) result_quant_values(rhs_quant_values.size());
+
+  QuantizedTensorElementType element_type(storage_type, expressed_type,
+                                          std::move(quantized_parameter));
+
+  QuantizedTensor lhs(
+      QuantizedTensorType(Shape(shape),
+                          QuantizedTensorElementType(element_type)),
+      lhs_quant_values.data());
+  QuantizedTensor rhs(
+      QuantizedTensorType(Shape(shape),
+                          QuantizedTensorElementType(element_type)),
+      rhs_quant_values.data());
+  QuantizedTensor result(
+      QuantizedTensorType(Shape(shape),
+                          QuantizedTensorElementType(element_type)),
+      result_quant_values.data());
+
+  for (auto _ : state) {
+    auto res = op(lhs, rhs, result);
+    CHECK_OK(res);
+  }
+}
+
+#define BENCHMARK_OP_HELPER(Op, ElementType)    \
+  BENCHMARK(BM_SHLO<Op, ElementType, 8 * KB>);  \
+  BENCHMARK(BM_SHLO<Op, ElementType, 16 * KB>); \
+  BENCHMARK(BM_SHLO<Op, ElementType, 32 * KB>); \
+  BENCHMARK(BM_SHLO<Op, ElementType, 64 * KB>);
+
+#define BENCHMARK_QOP_HELPER(Op, StorageType, ExpressedType)   \
+  BENCHMARK(BM_SHLO<Op, StorageType, ExpressedType, 8 * KB>);  \
+  BENCHMARK(BM_SHLO<Op, StorageType, ExpressedType, 16 * KB>); \
+  BENCHMARK(BM_SHLO<Op, StorageType, ExpressedType, 32 * KB>); \
+  BENCHMARK(BM_SHLO<Op, StorageType, ExpressedType, 64 * KB>);
+
+#define BENCHMARK_OP(Op)                                            \
+  BENCHMARK_OP_HELPER(Op, ElementType::kF32);                       \
+  BENCHMARK_OP_HELPER(Op, ElementType::kF16);                       \
+  BENCHMARK_OP_HELPER(Op, ElementType::kBF16);                      \
+  BENCHMARK_QOP_HELPER(Op, ElementType::kSI8, ElementType::kF32);   \
+  BENCHMARK_QOP_HELPER(Op, ElementType::kSI16, ElementType::kF32);  \
+  BENCHMARK_QOP_HELPER(Op, ElementType::kSI32, ElementType::kF32);  \
+  BENCHMARK_QOP_HELPER(Op, ElementType::kSI8, ElementType::kF16);   \
+  BENCHMARK_QOP_HELPER(Op, ElementType::kSI16, ElementType::kF16);  \
+  BENCHMARK_QOP_HELPER(Op, ElementType::kSI32, ElementType::kF16);  \
+  BENCHMARK_QOP_HELPER(Op, ElementType::kSI8, ElementType::kBF16);  \
+  BENCHMARK_QOP_HELPER(Op, ElementType::kSI16, ElementType::kBF16); \
+  BENCHMARK_QOP_HELPER(Op, ElementType::kSI32, ElementType::kBF16);
+
+BENCHMARK_OP(Abs);
+BENCHMARK_OP(Add);
+
+}  // namespace
+}  // namespace benchmark
+}  // namespace stablehlo
+
+// Run the benchmark
+BENCHMARK_MAIN();

tensorflow/lite/experimental/shlo/legacy/bench/util.h

@@ -0,0 +1,64 @@
+/* Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_LITE_EXPERIMENTAL_SHLO_LEGACY_BENCH_UTIL_H_
+#define TENSORFLOW_LITE_EXPERIMENTAL_SHLO_LEGACY_BENCH_UTIL_H_
+
+#include <algorithm>
+#include <cstddef>
+#include <limits>
+#include <random>
+#include <type_traits>
+#include <vector>
+
+#include "tensorflow/lite/experimental/shlo/legacy/src/bf16.h"
+#include "tensorflow/lite/experimental/shlo/legacy/src/f16.h"
+
+namespace stablehlo {
+
+namespace benchmark {
+
+static constexpr auto KB = 1024;
+
+template <typename Number>
+std::vector<Number> GenerateRandomVector(
+    size_t size, Number min = std::numeric_limits<Number>::min(),
+    Number max = std::numeric_limits<Number>::max()) {
+  std::vector<Number> data(size);
+  if constexpr (std::is_integral_v<Number>) {
+    static std::uniform_int_distribution<Number> distribution(min, max);
+    static std::default_random_engine generator;
+    std::generate(data.begin(), data.end(),
+                  [&]() { return distribution(generator); });
+  } else if constexpr (std::is_floating_point_v<Number>) {
+    static std::uniform_real_distribution<Number> distribution(min, max);
+    static std::default_random_engine generator;
+    std::generate(data.begin(), data.end(),
+                  [&]() { return distribution(generator); });
+  } else {
+    static_assert(std::is_same_v<Number, BF16> or std::is_same_v<Number, F16>);
+    static std::uniform_real_distribution<float> distribution(min, max);
+    static std::default_random_engine generator;
+    std::generate(data.begin(), data.end(),
+                  [&]() { return distribution(generator); });
+  }
+  return data;
+}
+
+}  // namespace benchmark
+
+}  // namespace stablehlo
+
+#endif  // TENSORFLOW_LITE_EXPERIMENTAL_SHLO_LEGACY_BENCH_UTIL_H_

tensorflow/lite/experimental/shlo/legacy/bench/xnn_benchmark.cc

@@ -0,0 +1,241 @@
+/* Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include <algorithm>
+#include <cstddef>
+#include <cstdint>
+#include <limits>
+#include <vector>
+
+#include "xnnpack.h"  // from @XNNPACK
+#include "absl/log/log.h"
+#include "benchmark/benchmark.h"  // from @com_google_benchmark
+#include "tensorflow/lite/experimental/shlo/legacy/bench/util.h"
+#include "tensorflow/lite/experimental/shlo/legacy/src/f16.h"
+
+namespace stablehlo {
+namespace benchmark {
+namespace {
+
+template <xnn_datatype datatype>
+struct Storage;
+
+template <>
+struct Storage<xnn_datatype_fp32> {
+  using Type = float;
+};
+
+template <>
+struct Storage<xnn_datatype_fp16> {
+  using Type = F16;
+};
+
+size_t size_in_bytes(xnn_datatype datatype) {
+  switch (datatype) {
+    case xnn_datatype_qint32:
+    case xnn_datatype_qcint32:
+    case xnn_datatype_fp32:
+      return 4;
+    case xnn_datatype_fp16:
+      return 2;
+    case xnn_datatype_qint8:
+    case xnn_datatype_quint8:
+    case xnn_datatype_qcint8:
+      return 1;
+    default:
+      // Crash OK
+      LOG(FATAL) << "Unexpected datatype: " << static_cast<int>(datatype);
+  }
+}
+
+size_t tensor_size_in_bytes(xnn_datatype datatype,
+                            const std::vector<size_t>& shape) {
+  size_t size = size_in_bytes(datatype);
+  for (auto s : shape) {
+    size *= s;
+  }
+  return size;
+}
+
+using UnaryOp = xnn_status (*)(xnn_subgraph_t subgraph, uint32_t input_id,
+                               uint32_t output_id, uint32_t flags);
+
+using BinaryOp = xnn_status (*)(xnn_subgraph_t subgraph, uint32_t input1_id,
+                                uint32_t input2_id, uint32_t output_id,
+                                uint32_t flags);
+
+using ClampedBinaryOp = xnn_status (*)(xnn_subgraph_t subgraph,
+                                       float output_min, float output_max,
+                                       uint32_t input1_id, uint32_t input2_id,
+                                       uint32_t output_id, uint32_t flags);
+
+template <xnn_datatype datatype, typename Op>
+void BM_XNN_HELPER(::benchmark::State& state, Op op, size_t size) {
+  size_t num_operands;
+  if constexpr (std::is_same_v<Op, UnaryOp>) {
+    num_operands = 1;
+  } else {
+    static_assert(std::is_same_v<Op, BinaryOp> or
+                  std::is_same_v<Op, ClampedBinaryOp>);
+    num_operands = 2;
+  }
+
+  std::vector<size_t> shape = {size};
+
+  using ST = typename Storage<datatype>::Type;
+
+  std::vector<std::vector<ST>> operand_values;
+  operand_values.reserve(num_operands);
+  for (auto i = 0; i < num_operands; ++i) {
+    operand_values.emplace_back(GenerateRandomVector<ST>(size));
+  }
+
+  auto size_in_bytes = tensor_size_in_bytes(datatype, shape);
+  std::vector<uint8_t> result_values(size_in_bytes);
+
+  std::vector<xnn_external_value> external_values;
+  external_values.reserve(num_operands + 1);
+  for (auto i = 0; i < num_operands; ++i) {
+    external_values.push_back(
+        {.id = static_cast<uint32_t>(i), .data = operand_values[i].data()});
+  }
+  external_values.push_back({.id = static_cast<uint32_t>(num_operands),
+                             .data = result_values.data()});
+
+  xnn_status status = xnn_initialize(/*allocator=*/nullptr);
+  if (status != xnn_status_success) {
+    // Crash OK
+    LOG(FATAL) << "Failed to invoke XNNPACK runtime: " << status;
+  }
+
+  xnn_workspace_t workspace;
+  status = xnn_create_workspace(&workspace);
+  if (status != xnn_status_success) {
+    // Crash OK
+    LOG(FATAL) << "Failed to invoke XNNPACK runtime: " << status;
+  }
+
+  xnn_subgraph_t subgraph = nullptr;
+  auto max_external_value_id =
+      std::max_element(external_values.begin(), external_values.end(),
+                       [](auto x, auto y) { return x.id < y.id; })
+          ->id;
+  status =
+      xnn_create_subgraph(1 + max_external_value_id, /*flags=*/0, &subgraph);
+  if (status != xnn_status_success) {
+    // Crash OK
+    LOG(FATAL) << "Failed to invoke XNNPACK runtime: " << status;
+  }
+
+  for (auto i = 0; i < num_operands; ++i) {
+    uint32_t operand_id;
+    status = xnn_define_tensor_value(
+        subgraph, datatype, shape.size(), shape.data(),
+        /* data */ nullptr, /* xnn_external_id */ i,
+        /* flags */ XNN_VALUE_FLAG_EXTERNAL_INPUT, &operand_id);
+    if (status != xnn_status_success) {
+      // Crash OK
+      LOG(FATAL) << "Failed to invoke XNNPACK runtime: " << status;
+    }
+  }
+
+  uint32_t result_id;
+  status = xnn_define_tensor_value(
+      subgraph, datatype, shape.size(), shape.data(),
+      /* data */ nullptr, /* xnn_external_id */ num_operands,
+      /* flags */ XNN_VALUE_FLAG_EXTERNAL_OUTPUT, &result_id);
+  if (status != xnn_status_success) {
+    // Crash OK
+    LOG(FATAL) << "Failed to invoke XNNPACK runtime: " << status;
+  }
+
+  if constexpr (std::is_same_v<Op, UnaryOp>) {
+    status = op(subgraph, 0, 1, /* flags */ 0);
+  } else if constexpr (std::is_same_v<Op, BinaryOp>) {
+    status = op(subgraph, 0, 1, 2, /* flags */ 0);
+  } else {
+    static_assert(std::is_same_v<Op, ClampedBinaryOp>);
+    status = op(subgraph, /* output_min */ std::numeric_limits<float>::min(),
+                /* output_max */ std::numeric_limits<float>::max(), 0, 1, 2,
+                /* flags */ 0);
+  }
+  if (status != xnn_status_success) {
+    // Crash OK
+    LOG(FATAL) << "Failed to invoke XNNPACK runtime: " << status;
+  }
+
+  xnn_runtime_t runtime = nullptr;
+  status = xnn_create_runtime_v4(subgraph,
+                                 /*weights_buffer_cache*/ nullptr, workspace,
+                                 /*threadpool*/ nullptr,
+                                 /*flags*/ 0, &runtime);
+  if (status != xnn_status_success) {
+    // Crash OK
+    LOG(FATAL) << "Failed to invoke XNNPACK runtime: " << status;
+  }
+
+  for (auto _ : state) {
+    status = xnn_setup_runtime(runtime, external_values.size(),
+                               external_values.data());
+    if (status != xnn_status_success) {
+      // Crash OK
+      LOG(FATAL) << "Failed to invoke XNNPACK runtime: " << status;
+    }
+    status = xnn_invoke_runtime(runtime);
+    if (status != xnn_status_success) {
+      // Crash OK
+      LOG(FATAL) << "Failed to invoke XNNPACK runtime: " << status;
+    }
+  }
+
+  xnn_delete_runtime(runtime);
+  xnn_delete_subgraph(subgraph);
+  xnn_release_workspace(workspace);
+  xnn_deinitialize();
+}
+
+template <UnaryOp op, xnn_datatype datatype, size_t size>
+void BM_XNN(::benchmark::State& state) {
+  BM_XNN_HELPER<datatype>(state, op, size);
+}
+
+template <BinaryOp op, xnn_datatype datatype, size_t size>
+void BM_XNN(::benchmark::State& state) {
+  BM_XNN_HELPER<datatype>(state, op, size);
+}
+
+template <ClampedBinaryOp op, xnn_datatype datatype, size_t size>
+void BM_XNN(::benchmark::State& state) {
+  BM_XNN_HELPER<datatype>(state, op, size);
+}
+
+#define BENCHMARK_OP_HELPER(Op, DataType)   \
+  BENCHMARK(BM_XNN<Op, DataType, 8 * KB>);  \
+  BENCHMARK(BM_XNN<Op, DataType, 16 * KB>); \
+  BENCHMARK(BM_XNN<Op, DataType, 32 * KB>); \
+  BENCHMARK(BM_XNN<Op, DataType, 64 * KB>);
+
+#define BENCHMARK_OP(Op) BENCHMARK_OP_HELPER(Op, xnn_datatype_fp32);
+//  BENCHMARK_OP_HELPER(Op, xnn_datatype_fp16);
+
+BENCHMARK_OP(xnn_define_abs);
+BENCHMARK_OP(xnn_define_add2);
+
+}  // namespace
+}  // namespace benchmark
+}  // namespace stablehlo
+
+// Run the benchmark
+BENCHMARK_MAIN();