wehub-resource-sync
520 lines
18 KiB
C++
520 lines
18 KiB
C++
/* Copyright 2023 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 <cstddef>
|
|
#include <cstdint>
|
|
#include <functional>
|
|
#include <ostream>
|
|
#include <string>
|
|
#include <utility>
|
|
#include <vector>
|
|
|
|
#include <gmock/gmock.h>
|
|
#include <gtest/gtest.h>
|
|
#include "absl/algorithm/container.h"
|
|
#include "absl/random/bit_gen_ref.h"
|
|
#include "absl/random/random.h"
|
|
#include "absl/status/status.h"
|
|
#include "absl/status/status_matchers.h"
|
|
#include "absl/status/statusor.h"
|
|
#include "absl/strings/str_format.h"
|
|
#include "absl/types/span.h"
|
|
#include "tensorflow/lite/c/c_api_types.h"
|
|
#include "tensorflow/lite/c/common.h"
|
|
#include "tensorflow/lite/kernels/stablehlo_reduce_window_test_util.h"
|
|
#include "tensorflow/lite/kernels/test_util.h"
|
|
#include "tensorflow/lite/schema/schema_generated.h"
|
|
|
|
namespace tflite {
|
|
namespace ops {
|
|
namespace builtin {
|
|
namespace stablehlo_pad {
|
|
namespace {
|
|
|
|
using ::absl_testing::StatusIs;
|
|
using ::testing::ElementsAre;
|
|
using ::testing::ElementsAreArray;
|
|
using ::testing::HasSubstr;
|
|
|
|
template <class T>
|
|
class StablehloPadModel : public SingleOpModel {
|
|
public:
|
|
static constexpr TensorType kTensorType = GetTensorType<T>();
|
|
|
|
void SetEdgePadding(std::vector<int64_t> low, std::vector<int64_t> high) {
|
|
edge_padding_low_ = std::move(low);
|
|
edge_padding_high_ = std::move(high);
|
|
}
|
|
|
|
const std::vector<int64_t>& GetEdgePaddingLow() const {
|
|
return edge_padding_low_;
|
|
}
|
|
|
|
const std::vector<int64_t>& GetEdgePaddingHigh() const {
|
|
return edge_padding_high_;
|
|
}
|
|
|
|
void SetInteriorPadding(std::vector<int64_t> padding) {
|
|
interior_padding_ = std::move(padding);
|
|
}
|
|
|
|
const std::vector<int64_t>& GetInteriorPadding() const {
|
|
return interior_padding_;
|
|
}
|
|
|
|
void SetInput(std::vector<int64_t> shape) {
|
|
input_.shape = std::move(shape);
|
|
input_.data.resize(absl::c_accumulate(input_.shape, static_cast<int64_t>(1),
|
|
std::multiplies<>()));
|
|
absl::c_iota(input_.data, static_cast<T>(1));
|
|
}
|
|
|
|
void SetInput(std::vector<int64_t> shape, std::vector<T> data) {
|
|
input_.shape = std::move(shape);
|
|
input_.data = std::move(data);
|
|
}
|
|
|
|
void SetInput(absl::Span<const int64_t> shape, absl::BitGenRef bitgen, T min,
|
|
T max) {
|
|
input_.shape.assign(shape.begin(), shape.end());
|
|
input_.data.resize(absl::c_accumulate(shape, static_cast<int64_t>(1),
|
|
std::multiplies<>()));
|
|
absl::c_generate(input_.data, [&] {
|
|
return absl::Uniform(absl::IntervalClosed, bitgen, min, max);
|
|
});
|
|
}
|
|
|
|
const reduce_window::reference::Tensor<T>& GetInput() const { return input_; }
|
|
|
|
void SetPaddingValue(const T& v) { padding_value_ = v; }
|
|
|
|
T GetPaddingValue() const { return padding_value_; }
|
|
|
|
absl::Span<const T> GetOutputData() {
|
|
return absl::Span<const T>(interpreter_->typed_tensor<T>(output_tensor_id_),
|
|
GetTensorSize(output_tensor_id_));
|
|
}
|
|
|
|
absl::Span<const int> GetOutputShape() {
|
|
const TfLiteIntArray& shape =
|
|
*(interpreter_->tensor(output_tensor_id_)->dims);
|
|
return absl::Span<const int>(shape.data, shape.size);
|
|
}
|
|
|
|
absl::Status CheckPreconditions() {
|
|
const size_t rank = input_.shape.size();
|
|
if (rank == 0) {
|
|
return absl::FailedPreconditionError("Input rank is 0.");
|
|
}
|
|
if (edge_padding_low_.empty()) {
|
|
edge_padding_low_ = std::vector<int64_t>(rank, 0);
|
|
} else if (edge_padding_low_.size() != rank) {
|
|
return absl::FailedPreconditionError(
|
|
"Low edge padding does not have the right size.");
|
|
}
|
|
if (edge_padding_high_.empty()) {
|
|
edge_padding_high_ = std::vector<int64_t>(rank, 0);
|
|
} else if (edge_padding_high_.size() != rank) {
|
|
return absl::FailedPreconditionError(
|
|
"High edge padding does not have the right size.");
|
|
}
|
|
if (interior_padding_.empty()) {
|
|
interior_padding_ = std::vector<int64_t>(rank, 0);
|
|
} else if (interior_padding_.size() != rank) {
|
|
return absl::FailedPreconditionError(
|
|
"Interior padding does not have the right size.");
|
|
}
|
|
return absl::OkStatus();
|
|
}
|
|
|
|
absl::Status SetupInterpreter() {
|
|
input_tensor_id_ =
|
|
AddInput({kTensorType,
|
|
std::vector<int>(input_.shape.begin(), input_.shape.end())});
|
|
padding_value_tensor_id_ =
|
|
AddConstInput(kTensorType, /*data=*/{padding_value_}, /*shape=*/{1});
|
|
output_tensor_id_ = AddOutput(kTensorType);
|
|
|
|
SetBuiltinOp(BuiltinOperator_STABLEHLO_PAD,
|
|
BuiltinOptions2_StablehloPadOptions,
|
|
CreateStablehloPadOptions(
|
|
builder_, builder_.CreateVector(edge_padding_low_),
|
|
builder_.CreateVector(edge_padding_high_),
|
|
builder_.CreateVector(interior_padding_))
|
|
.Union());
|
|
BuildInterpreter(
|
|
/*input_shapes=*/{std::vector<int>(input_.shape.begin(),
|
|
input_.shape.end())},
|
|
/*num_threads=*/-1, /*allow_fp32_relax_to_fp16=*/false,
|
|
/*apply_delegate=*/true, /*allocate_and_delegate=*/false);
|
|
return absl::OkStatus();
|
|
}
|
|
|
|
absl::Status Build() {
|
|
if (absl::Status status = CheckPreconditions(); !status.ok()) {
|
|
return status;
|
|
}
|
|
if (absl::Status status = SetupInterpreter(); !status.ok()) {
|
|
return status;
|
|
}
|
|
AllocateAndDelegate(/*apply_delegate=*/true);
|
|
PopulateTensor(input_tensor_id_, input_.data);
|
|
return absl::OkStatus();
|
|
}
|
|
|
|
absl::Status BuildAndInvoke() {
|
|
if (absl::Status status = Build(); !status.ok()) {
|
|
return status;
|
|
}
|
|
if (TfLiteStatus status = Invoke(); status != kTfLiteOk) {
|
|
const std::string msg =
|
|
absl::StrFormat("Invoke failed with status %d.", status);
|
|
return absl::InternalError(msg);
|
|
}
|
|
return absl::OkStatus();
|
|
}
|
|
|
|
TfLiteStatus BuildAndPrepare() {
|
|
if (absl::Status status = CheckPreconditions(); !status.ok()) {
|
|
return kTfLiteError;
|
|
}
|
|
if (absl::Status status = SetupInterpreter(); !status.ok()) {
|
|
return kTfLiteError;
|
|
}
|
|
return interpreter_->AllocateTensors();
|
|
}
|
|
|
|
friend std::ostream& operator<<(std::ostream& os,
|
|
const StablehloPadModel& model) {
|
|
auto print_vec = [&os](const auto& vec) {
|
|
os << "[";
|
|
if (!vec.empty()) {
|
|
auto it = vec.begin();
|
|
os << +*(it++);
|
|
for (; it != vec.end(); ++it) {
|
|
os << ", " << +*it;
|
|
}
|
|
}
|
|
os << "]";
|
|
};
|
|
os << " edge_padding_low: ";
|
|
print_vec(model.GetEdgePaddingLow());
|
|
os << "\n edge_padding_high: ";
|
|
print_vec(model.GetEdgePaddingHigh());
|
|
os << "\n interior_padding: ";
|
|
print_vec(model.GetInteriorPadding());
|
|
os << "\n padding_value: " << +model.GetPaddingValue();
|
|
os << "\n input shape: ";
|
|
print_vec(model.GetInput().shape);
|
|
return os;
|
|
}
|
|
|
|
private:
|
|
std::vector<int64_t> edge_padding_low_;
|
|
std::vector<int64_t> edge_padding_high_;
|
|
std::vector<int64_t> interior_padding_;
|
|
reduce_window::reference::Tensor<T> input_;
|
|
T padding_value_ = 0;
|
|
|
|
int input_tensor_id_;
|
|
int padding_value_tensor_id_;
|
|
int output_tensor_id_;
|
|
};
|
|
|
|
template <class T>
|
|
absl::StatusOr<reduce_window::reference::Tensor<T>> ComputeReference(
|
|
StablehloPadModel<T>& model) {
|
|
if (absl::Status status = model.CheckPreconditions(); !status.ok()) {
|
|
return status;
|
|
}
|
|
std::vector<int64_t> dilations, padding;
|
|
for (size_t i = 0; i < model.GetInput().shape.size(); ++i) {
|
|
padding.push_back(model.GetEdgePaddingLow()[i]);
|
|
padding.push_back(model.GetEdgePaddingHigh()[i]);
|
|
dilations.push_back(model.GetInteriorPadding()[i] + 1);
|
|
}
|
|
|
|
auto dilated_tensor = reduce_window::reference::Dilate(
|
|
model.GetInput(), dilations, model.GetPaddingValue());
|
|
auto padded_tensor = reduce_window::reference::Pad(dilated_tensor, padding,
|
|
model.GetPaddingValue());
|
|
return reduce_window::reference::Crop(padded_tensor, padding);
|
|
}
|
|
|
|
TEST(StablehloPadModelTest, DefaultModelFails) {
|
|
StablehloPadModel<int> model;
|
|
const auto expected_status = ComputeReference(model);
|
|
EXPECT_THAT(expected_status.status(),
|
|
StatusIs(absl::StatusCode::kFailedPrecondition,
|
|
HasSubstr("Input rank is 0.")));
|
|
}
|
|
|
|
TEST(StablehloPadModelTest, DefaultModelReturnsIdentity) {
|
|
StablehloPadModel<int> model;
|
|
model.SetInput({3, 1});
|
|
EXPECT_THAT(model.GetInput().shape, ElementsAre(3, 1));
|
|
const auto expected_status = ComputeReference(model);
|
|
ASSERT_TRUE(expected_status.ok());
|
|
EXPECT_THAT(expected_status.value().data,
|
|
ElementsAreArray(model.GetInput().data));
|
|
}
|
|
|
|
TEST(StablehloPadModelTest, WrongEdgePaddingSizeIsAnError) {
|
|
StablehloPadModel<int> model;
|
|
model.SetInput({3, 1});
|
|
model.SetEdgePadding(/*low=*/{3, 4, 5}, /*high=*/{6, 7});
|
|
{
|
|
const auto expected_status = ComputeReference(model);
|
|
EXPECT_THAT(
|
|
expected_status.status(),
|
|
StatusIs(absl::StatusCode::kFailedPrecondition,
|
|
HasSubstr("Low edge padding does not have the right size.")));
|
|
}
|
|
model.SetEdgePadding(/*low=*/{3, 4}, /*high=*/{5, 6, 7});
|
|
{
|
|
const auto expected_status = ComputeReference(model);
|
|
EXPECT_THAT(
|
|
expected_status.status(),
|
|
StatusIs(absl::StatusCode::kFailedPrecondition,
|
|
HasSubstr("High edge padding does not have the right size.")));
|
|
}
|
|
}
|
|
|
|
TEST(StablehloPadModelTest, WrongInteriorPaddingSizeIsAnError) {
|
|
StablehloPadModel<int> model;
|
|
model.SetInput({3, 1});
|
|
model.SetInteriorPadding({3, 4, 5});
|
|
const auto expected_status = ComputeReference(model);
|
|
EXPECT_THAT(
|
|
expected_status.status(),
|
|
StatusIs(absl::StatusCode::kFailedPrecondition,
|
|
HasSubstr("Interior padding does not have the right size.")));
|
|
}
|
|
|
|
TEST(StablehloPadTest, IdentityParams) {
|
|
StablehloPadModel<int> model;
|
|
model.SetInput({3, 3});
|
|
ASSERT_TRUE(model.BuildAndInvoke().ok());
|
|
EXPECT_THAT(model.GetOutputShape(), ElementsAreArray(model.GetInput().shape));
|
|
EXPECT_THAT(model.GetOutputData(), ElementsAreArray(model.GetInput().data));
|
|
}
|
|
|
|
TEST(StablehloPadTest, InteriorPad) {
|
|
StablehloPadModel<int> model;
|
|
model.SetInput({3, 3});
|
|
model.SetInteriorPadding({1, 2});
|
|
const auto expected_status = ComputeReference(model);
|
|
ASSERT_TRUE(expected_status.ok());
|
|
const auto& expected = expected_status.value();
|
|
ASSERT_TRUE(model.BuildAndInvoke().ok());
|
|
EXPECT_THAT(model.GetOutputShape(), ElementsAreArray(expected.shape));
|
|
EXPECT_THAT(model.GetOutputData(), ElementsAreArray(expected.data));
|
|
}
|
|
|
|
TEST(StablehloPadTest, LowPad) {
|
|
StablehloPadModel<int> model;
|
|
model.SetInput({3, 3});
|
|
model.SetEdgePadding({1, 1}, {0, 0});
|
|
const auto expected_status = ComputeReference(model);
|
|
ASSERT_TRUE(expected_status.ok());
|
|
const auto& expected = expected_status.value();
|
|
ASSERT_TRUE(model.BuildAndInvoke().ok());
|
|
EXPECT_THAT(model.GetOutputShape(), ElementsAreArray(expected.shape));
|
|
EXPECT_THAT(model.GetOutputData(), ElementsAreArray(expected.data));
|
|
}
|
|
|
|
TEST(StablehloPadTest, HighPad) {
|
|
StablehloPadModel<int> model;
|
|
model.SetInput({3, 3});
|
|
model.SetEdgePadding({0, 0}, {1, 1});
|
|
const auto expected_status = ComputeReference(model);
|
|
ASSERT_TRUE(expected_status.ok());
|
|
const auto& expected = expected_status.value();
|
|
ASSERT_TRUE(model.BuildAndInvoke().ok());
|
|
EXPECT_THAT(model.GetOutputShape(), ElementsAreArray(expected.shape));
|
|
EXPECT_THAT(model.GetOutputData(), ElementsAreArray(expected.data));
|
|
}
|
|
|
|
TEST(StablehloPadTest, AllPad) {
|
|
StablehloPadModel<int> model;
|
|
model.SetInput({3, 3});
|
|
model.SetEdgePadding({1, 1}, {1, 1});
|
|
const auto expected_status = ComputeReference(model);
|
|
ASSERT_TRUE(expected_status.ok());
|
|
const auto& expected = expected_status.value();
|
|
ASSERT_TRUE(model.BuildAndInvoke().ok());
|
|
EXPECT_THAT(model.GetOutputShape(), ElementsAreArray(expected.shape));
|
|
EXPECT_THAT(model.GetOutputData(), ElementsAreArray(expected.data));
|
|
}
|
|
|
|
TEST(StablehloPadTest, LowCrop) {
|
|
StablehloPadModel<int> model;
|
|
model.SetInput({3, 3});
|
|
model.SetEdgePadding({-1, -1}, {0, 0});
|
|
const auto expected_status = ComputeReference(model);
|
|
ASSERT_TRUE(expected_status.ok());
|
|
const auto& expected = expected_status.value();
|
|
ASSERT_TRUE(model.BuildAndInvoke().ok());
|
|
EXPECT_THAT(model.GetOutputShape(), ElementsAreArray(expected.shape));
|
|
EXPECT_THAT(model.GetOutputData(), ElementsAreArray(expected.data));
|
|
}
|
|
|
|
TEST(StablehloPadTest, HighCrop) {
|
|
StablehloPadModel<int> model;
|
|
model.SetInput({3, 3});
|
|
model.SetEdgePadding({0, 0}, {-1, -1});
|
|
const auto expected_status = ComputeReference(model);
|
|
ASSERT_TRUE(expected_status.ok());
|
|
const auto& expected = expected_status.value();
|
|
ASSERT_TRUE(model.BuildAndInvoke().ok());
|
|
EXPECT_THAT(model.GetOutputShape(), ElementsAreArray(expected.shape));
|
|
EXPECT_THAT(model.GetOutputData(), ElementsAreArray(expected.data));
|
|
}
|
|
|
|
TEST(StablehloPadTest, AllCrop) {
|
|
StablehloPadModel<int> model;
|
|
model.SetInput({3, 3});
|
|
model.SetEdgePadding({-1, -1}, {-1, -1});
|
|
const auto expected_status = ComputeReference(model);
|
|
ASSERT_TRUE(expected_status.ok());
|
|
const auto& expected = expected_status.value();
|
|
ASSERT_TRUE(model.BuildAndInvoke().ok());
|
|
EXPECT_THAT(model.GetOutputShape(), ElementsAreArray(expected.shape));
|
|
EXPECT_THAT(model.GetOutputData(), ElementsAreArray(expected.data));
|
|
}
|
|
|
|
TEST(StablehloPadTest, PadCrop) {
|
|
StablehloPadModel<int> model;
|
|
model.SetInput({3, 3});
|
|
model.SetEdgePadding({1, -1}, {1, -1});
|
|
const auto expected_status = ComputeReference(model);
|
|
ASSERT_TRUE(expected_status.ok());
|
|
const auto& expected = expected_status.value();
|
|
ASSERT_TRUE(model.BuildAndInvoke().ok());
|
|
EXPECT_THAT(model.GetOutputShape(), ElementsAreArray(expected.shape));
|
|
EXPECT_THAT(model.GetOutputData(), ElementsAreArray(expected.data));
|
|
}
|
|
|
|
TEST(StablehloPadTest, InteriorEdgePadding) {
|
|
StablehloPadModel<int> model;
|
|
model.SetInput({3, 3});
|
|
model.SetEdgePadding({-1, -4}, {0, 0});
|
|
model.SetInteriorPadding({1, 2});
|
|
const auto expected_status = ComputeReference(model);
|
|
ASSERT_TRUE(expected_status.ok());
|
|
const auto& expected = expected_status.value();
|
|
ASSERT_TRUE(model.BuildAndInvoke().ok());
|
|
EXPECT_THAT(model.GetOutputShape(), ElementsAreArray(expected.shape));
|
|
EXPECT_THAT(model.GetOutputData(), ElementsAreArray(expected.data));
|
|
}
|
|
|
|
TEST(StablehloPadTest, CallPrepareTwiceDoesNotFail) {
|
|
StablehloPadModel<int> model;
|
|
model.SetInput({3, 3});
|
|
model.SetEdgePadding({-1, -4}, {0, 0});
|
|
model.SetInteriorPadding({1, 2});
|
|
const auto expected_status = ComputeReference(model);
|
|
ASSERT_TRUE(expected_status.ok());
|
|
const auto& expected = expected_status.value();
|
|
// Applying delegates forces Prepare to be called twice.
|
|
model.SetApplyDefaultDelegates();
|
|
ASSERT_TRUE(model.BuildAndInvoke().ok());
|
|
EXPECT_THAT(model.GetOutputShape(), ElementsAreArray(expected.shape));
|
|
EXPECT_THAT(model.GetOutputData(), ElementsAreArray(expected.data));
|
|
}
|
|
|
|
TEST(StablehloPadTest, ZeroRankFails) {
|
|
StablehloPadModel<int> model;
|
|
model.SetInput({});
|
|
EXPECT_EQ(model.BuildAndPrepare(), kTfLiteError);
|
|
}
|
|
|
|
TEST(StablehloPadTest, NegativeInteriorPaddingFails) {
|
|
StablehloPadModel<int> model;
|
|
model.SetInput({3, 3});
|
|
model.SetInteriorPadding({-1, 0});
|
|
EXPECT_EQ(model.BuildAndPrepare(), kTfLiteError);
|
|
}
|
|
|
|
TEST(StablehloPadTest, DimensionOverflowFails) {
|
|
StablehloPadModel<int> model;
|
|
model.SetInput({3, 3});
|
|
model.SetInteriorPadding({2000000000, 0});
|
|
EXPECT_EQ(model.BuildAndPrepare(), kTfLiteError);
|
|
}
|
|
|
|
TEST(StablehloPadTest, ExtremeNegativeEdgePaddingFails) {
|
|
StablehloPadModel<int> model;
|
|
model.SetInput({3, 3});
|
|
model.SetEdgePadding({-2147483648LL, 0}, {0, 0});
|
|
EXPECT_EQ(model.BuildAndPrepare(), kTfLiteError);
|
|
}
|
|
|
|
TEST(StablehloPadTest, EmptyTensorCropSucceeds) {
|
|
StablehloPadModel<int> model;
|
|
model.SetInput({3, 3});
|
|
model.SetEdgePadding({-5, 0}, {0, 0});
|
|
EXPECT_EQ(model.BuildAndPrepare(), kTfLiteOk);
|
|
}
|
|
|
|
// Returns a vector of given size with elements in the range [min, max].
|
|
template <class T>
|
|
std::vector<T> RandomVector(absl::BitGenRef bitgen, size_t size, T min, T max) {
|
|
std::vector<T> vec(size);
|
|
for (T& v : vec) {
|
|
v = absl::Uniform(absl::IntervalClosed, bitgen, min, max);
|
|
}
|
|
return vec;
|
|
}
|
|
|
|
template <class T>
|
|
class StablehloPadFuzzyTest : public testing::Test {};
|
|
|
|
using TestList =
|
|
testing::Types<int8_t, int16_t, int32_t, int64_t, uint8_t, float, double>;
|
|
TYPED_TEST_SUITE(StablehloPadFuzzyTest, TestList);
|
|
|
|
TYPED_TEST(StablehloPadFuzzyTest, FuzzyTest) {
|
|
absl::BitGen bitgen;
|
|
|
|
for (size_t iteration = 0; iteration < 200; ++iteration) {
|
|
const int rank = absl::Uniform(absl::IntervalClosed, bitgen, 1, 2);
|
|
|
|
StablehloPadModel<TypeParam> model;
|
|
model.SetInput(
|
|
/*shape=*/RandomVector<int64_t>(bitgen, rank, /*min=*/1, /*max=*/3),
|
|
bitgen, /*min=*/-5, /*max=*/5);
|
|
model.SetInteriorPadding(
|
|
RandomVector<int64_t>(bitgen, rank, /*min=*/0, /*max=*/2));
|
|
model.SetEdgePadding(
|
|
RandomVector<int64_t>(bitgen, rank, /*min=*/-5, /*max=*/5),
|
|
RandomVector<int64_t>(bitgen, rank, /*min=*/-5, /*max=*/5));
|
|
model.SetPaddingValue(
|
|
absl::Uniform(absl::IntervalClosed, bitgen, -127, 127));
|
|
|
|
const auto expected_status = ComputeReference(model);
|
|
ASSERT_TRUE(expected_status.ok());
|
|
const auto& expected = expected_status.value();
|
|
ASSERT_TRUE(model.BuildAndInvoke().ok());
|
|
EXPECT_THAT(model.GetOutputShape(), ElementsAreArray(expected.shape))
|
|
<< model;
|
|
EXPECT_THAT(model.GetOutputData(), ElementsAreArray(expected.data))
|
|
<< model;
|
|
}
|
|
}
|
|
|
|
} // namespace
|
|
} // namespace stablehlo_pad
|
|
} // namespace builtin
|
|
} // namespace ops
|
|
} // namespace tflite
|