wehub-resource-sync
339 lines
11 KiB
C++
339 lines
11 KiB
C++
/* 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 <cstdint>
|
|
#include <cstdlib>
|
|
#include <initializer_list>
|
|
#include <limits>
|
|
#include <vector>
|
|
|
|
#include <gmock/gmock.h>
|
|
#include <gtest/gtest.h>
|
|
#include "absl/log/absl_check.h"
|
|
#include "tensorflow/lite/c/common.h"
|
|
#include "tensorflow/lite/core/c/builtin_op_data.h"
|
|
#include "tensorflow/lite/core/interpreter.h"
|
|
#include "tensorflow/lite/kernels/kernel_util.h"
|
|
#include "tensorflow/lite/kernels/subgraph_test_util.h"
|
|
#include "tensorflow/lite/kernels/test_util.h"
|
|
#include "tensorflow/lite/schema/schema_generated.h"
|
|
|
|
namespace tflite {
|
|
namespace {
|
|
|
|
using ::testing::ElementsAreArray;
|
|
|
|
template <typename T>
|
|
tflite::TensorType GetTTEnum();
|
|
|
|
// NOLINTBEGIN
|
|
|
|
template <>
|
|
tflite::TensorType GetTTEnum<float>() {
|
|
return tflite::TensorType_FLOAT32;
|
|
}
|
|
|
|
template <>
|
|
tflite::TensorType GetTTEnum<int8_t>() {
|
|
return tflite::TensorType_INT8;
|
|
}
|
|
|
|
template <>
|
|
tflite::TensorType GetTTEnum<int16_t>() {
|
|
return tflite::TensorType_INT16;
|
|
}
|
|
|
|
template <>
|
|
tflite::TensorType GetTTEnum<int32_t>() {
|
|
return tflite::TensorType_INT32;
|
|
}
|
|
|
|
// NOLINTEND
|
|
|
|
class StablehloCaseOpModel : public SingleOpModel {
|
|
public:
|
|
StablehloCaseOpModel(const TensorData& input, const TensorData& input1,
|
|
const TensorData& input2, const TensorData& output,
|
|
const TfLiteStablehloCaseParams& params) {
|
|
InitializeCommonInputs(input, input1, input2, output, params);
|
|
}
|
|
|
|
template <typename T>
|
|
void SetInput(int index, std::initializer_list<T> data) {
|
|
PopulateTensor<T>(index, data);
|
|
}
|
|
|
|
template <typename T>
|
|
std::vector<T> GetOutput() {
|
|
return ExtractVector<T>(output_);
|
|
}
|
|
|
|
template <typename integer_dtype>
|
|
std::vector<float> GetDequantizedOutput() {
|
|
return Dequantize<integer_dtype>(ExtractVector<integer_dtype>(output_),
|
|
GetScale(output_), GetZeroPoint(output_));
|
|
}
|
|
|
|
int input() { return input_; }
|
|
|
|
int subgraph_input1() { return subgraph_input1_; }
|
|
|
|
int subgraph_input2() { return subgraph_input2_; }
|
|
|
|
protected:
|
|
void InitializeCommonInputs(const TensorData& input, const TensorData& input1,
|
|
const TensorData& input2,
|
|
const TensorData& output,
|
|
const TfLiteStablehloCaseParams& params) {
|
|
input_ = AddInput(input);
|
|
subgraph_input1_ = AddInput(SymmetricInt16Scaling(input1));
|
|
subgraph_input2_ = AddInput(SymmetricInt16Scaling(input2));
|
|
output_ = AddOutput(SymmetricInt16Scaling(output));
|
|
SetBuiltinOp(BuiltinOperator_STABLEHLO_CASE,
|
|
BuiltinOptions2_StablehloCaseOptions,
|
|
CreateStablehloCaseOptions(
|
|
builder_,
|
|
builder_.CreateVector(std::vector<int>(
|
|
params.branch_subgraph_indices,
|
|
params.branch_subgraph_indices + params.num_branches)))
|
|
.Union());
|
|
BuildInterpreter({GetShape(input_)}, /*num_threads=*/-1,
|
|
/*allow_fp32_relax_to_fp16=*/false,
|
|
/*apply_delegate=*/false, /*allocate_and_delegate=*/false);
|
|
AddSubgraphs(params.num_branches);
|
|
}
|
|
|
|
TensorData SymmetricInt16Scaling(TensorData tensor) {
|
|
if (tensor.type == TensorType_INT16) {
|
|
ABSL_CHECK_EQ(std::abs(tensor.min), tensor.max);
|
|
tensor.scale = tensor.max / std::numeric_limits<int16_t>::max();
|
|
tensor.zero_point = 0;
|
|
tensor.min = 0;
|
|
tensor.max = 0;
|
|
}
|
|
return tensor;
|
|
}
|
|
|
|
int input_;
|
|
int subgraph_input1_;
|
|
int subgraph_input2_;
|
|
int output_;
|
|
subgraph_test_util::SubgraphBuilder subgraph_builder_;
|
|
};
|
|
|
|
class StablehloCaseStaticOpModel : public StablehloCaseOpModel {
|
|
public:
|
|
StablehloCaseStaticOpModel(const TensorData& input, const TensorData& input1,
|
|
const TensorData& input2, const TensorData& output,
|
|
const TfLiteStablehloCaseParams& params)
|
|
: StablehloCaseOpModel(input, input1, input2, output, params) {
|
|
TfLiteType type = interpreter_->tensor(subgraph_input1_)->type;
|
|
subgraph_builder_.BuildAddSubgraph(interpreter_->subgraph(1), type);
|
|
subgraph_builder_.BuildMulSubgraph(interpreter_->subgraph(2), type);
|
|
subgraph_builder_.BuildMaximumSubgraph(interpreter_->subgraph(3), type);
|
|
subgraph_builder_.BuildMinimumSubgraph(interpreter_->subgraph(4), type);
|
|
AllocateAndDelegate(true);
|
|
}
|
|
int output() { return output_; }
|
|
};
|
|
|
|
class StablehloCaseDynamicOpModel : public StablehloCaseOpModel {
|
|
public:
|
|
StablehloCaseDynamicOpModel(const TensorData& input, const TensorData& input1,
|
|
const TensorData& input2,
|
|
const TensorData& output,
|
|
const TfLiteStablehloCaseParams& params)
|
|
: StablehloCaseOpModel(input, input1, input2, output, params) {
|
|
TfLiteType type = interpreter_->tensor(subgraph_input1_)->type;
|
|
subgraph_builder_.BuildAddSubgraph(interpreter_->subgraph(1), type);
|
|
subgraph_builder_.BuildPadSubgraph(interpreter_->subgraph(2));
|
|
AllocateAndDelegate(true);
|
|
}
|
|
int output() { return output_; }
|
|
};
|
|
|
|
template <typename T>
|
|
float GetTolerance(float min, float max) {
|
|
float kQuantizedStep =
|
|
2.0 * (max - min) /
|
|
(std::numeric_limits<T>::max() - std::numeric_limits<T>::min());
|
|
return kQuantizedStep;
|
|
}
|
|
|
|
template <typename Float>
|
|
class StablehloCaseTestFloat : public ::testing::Test {
|
|
public:
|
|
using FloatType = Float;
|
|
};
|
|
|
|
using FloatTestTypes = ::testing::Types<float>;
|
|
|
|
TYPED_TEST_SUITE(StablehloCaseTestFloat, FloatTestTypes);
|
|
|
|
TYPED_TEST(StablehloCaseTestFloat, CaseFloatMul) {
|
|
using Float = typename TestFixture::FloatType;
|
|
TfLiteStablehloCaseParams params = {
|
|
{1, 2, 3, 4},
|
|
4,
|
|
};
|
|
|
|
StablehloCaseStaticOpModel model(
|
|
{TensorType_INT32, {}}, {GetTTEnum<Float>(), {1, 2}},
|
|
{GetTTEnum<Float>(), {1, 2}}, {GetTTEnum<Float>(), {1, 2}}, params);
|
|
model.SetInput<int>(model.input(), {1});
|
|
model.SetInput<Float>(model.subgraph_input1(),
|
|
{static_cast<Float>(5.5), static_cast<Float>(2.5)});
|
|
model.SetInput<Float>(model.subgraph_input2(),
|
|
{static_cast<Float>(5.5), static_cast<Float>(2.5)});
|
|
ASSERT_EQ(model.Invoke(), kTfLiteOk);
|
|
|
|
EXPECT_THAT(model.GetOutput<Float>(),
|
|
Pointwise(FloatingPointEq(), {Float(30.25), Float(6.25)}));
|
|
}
|
|
|
|
TYPED_TEST(StablehloCaseTestFloat, CaseFloatAdd) {
|
|
using Float = typename TestFixture::FloatType;
|
|
TfLiteStablehloCaseParams params = {
|
|
{1, 2, 3, 4},
|
|
4,
|
|
};
|
|
|
|
StablehloCaseStaticOpModel model(
|
|
{TensorType_INT32, {}}, {GetTTEnum<Float>(), {1, 2}},
|
|
{GetTTEnum<Float>(), {1, 2}}, {GetTTEnum<Float>(), {1, 2}}, params);
|
|
model.SetInput<int>(model.input(), {0});
|
|
model.SetInput<Float>(model.subgraph_input1(),
|
|
{static_cast<Float>(5.5), static_cast<Float>(2.4)});
|
|
model.SetInput<Float>(model.subgraph_input2(),
|
|
{static_cast<Float>(5), static_cast<Float>(2)});
|
|
ASSERT_EQ(model.Invoke(), kTfLiteOk);
|
|
|
|
EXPECT_THAT(model.GetOutput<Float>(),
|
|
Pointwise(FloatingPointEq(),
|
|
{static_cast<Float>(10.5), static_cast<Float>(4.4)}));
|
|
}
|
|
|
|
template <typename Int>
|
|
class StablehloCaseTestInt : public ::testing::Test {
|
|
public:
|
|
using IntType = Int;
|
|
};
|
|
|
|
using IntTestTypes = ::testing::Types<int, int16_t>;
|
|
|
|
TYPED_TEST_SUITE(StablehloCaseTestInt, IntTestTypes);
|
|
|
|
TYPED_TEST(StablehloCaseTestInt, CaseIntMaximum) {
|
|
using Int = typename TestFixture::IntType;
|
|
TfLiteStablehloCaseParams params = {
|
|
{1, 2, 3, 4},
|
|
4,
|
|
};
|
|
|
|
StablehloCaseStaticOpModel model(
|
|
{TensorType_INT32, {}}, {GetTTEnum<Int>(), {1, 2}},
|
|
{GetTTEnum<Int>(), {1, 2}}, {GetTTEnum<Int>(), {1, 2}}, params);
|
|
model.SetInput<int>(model.input(), {2});
|
|
model.SetInput<Int>(model.subgraph_input1(), {5, 20});
|
|
model.SetInput<Int>(model.subgraph_input2(), {15, 2});
|
|
ASSERT_EQ(model.Invoke(), kTfLiteOk);
|
|
|
|
EXPECT_THAT(model.GetOutput<Int>(), ElementsAreArray({15, 20}));
|
|
}
|
|
|
|
TYPED_TEST(StablehloCaseTestInt, CaseIntMinimum) {
|
|
using Int = typename TestFixture::IntType;
|
|
TfLiteStablehloCaseParams params = {
|
|
{1, 2, 3, 4},
|
|
4,
|
|
};
|
|
|
|
StablehloCaseStaticOpModel model(
|
|
{TensorType_INT32, {}}, {GetTTEnum<Int>(), {1, 2}},
|
|
{GetTTEnum<Int>(), {1, 2}}, {GetTTEnum<Int>(), {1, 2}}, params);
|
|
model.SetInput<int>(
|
|
model.input(),
|
|
{-1}); // when index is out of bounds, case op executes the last branch
|
|
model.SetInput<Int>(model.subgraph_input1(),
|
|
{static_cast<Int>(5), static_cast<Int>(20)});
|
|
model.SetInput<Int>(model.subgraph_input2(),
|
|
{static_cast<Int>(15), static_cast<Int>(2)});
|
|
ASSERT_EQ(model.Invoke(), kTfLiteOk);
|
|
|
|
EXPECT_THAT(model.GetOutput<Int>(),
|
|
ElementsAreArray({static_cast<Int>(5), static_cast<Int>(2)}));
|
|
}
|
|
|
|
TEST(StablehloCaseTest, CaseQuantizedMul) {
|
|
float kQuantizedTolerance = GetTolerance<int8_t>(-127.0, 127.0);
|
|
TfLiteStablehloCaseParams params = {
|
|
{1, 2, 3, 4},
|
|
4,
|
|
};
|
|
|
|
StablehloCaseStaticOpModel model(
|
|
{TensorType_INT32, {}}, {TensorType_INT8, {1, 2}, -127.0f, 127.0f},
|
|
{TensorType_INT8, {1, 2}, -127.0f, 127.0f},
|
|
{TensorType_INT8, {1, 2}, -127.0f, 127.0f}, params);
|
|
model.SetInput<int>(model.input(), {0});
|
|
model.QuantizeAndPopulate<int8_t>(model.subgraph_input1(), {5.0, 2.0});
|
|
model.QuantizeAndPopulate<int8_t>(model.subgraph_input2(), {5.0, 2.0});
|
|
ASSERT_EQ(model.Invoke(), kTfLiteOk);
|
|
|
|
EXPECT_THAT(model.GetDequantizedOutput<int8_t>(),
|
|
ElementsAreArray(ArrayFloatNear({10, 4}, kQuantizedTolerance)));
|
|
}
|
|
|
|
TEST(StablehloCaseTest, DynamicCaseTestAdd) {
|
|
TfLiteStablehloCaseParams params = {
|
|
{1, 2},
|
|
2,
|
|
};
|
|
|
|
StablehloCaseDynamicOpModel model(
|
|
{TensorType_INT32, {}}, {TensorType_INT32, {2}},
|
|
{TensorType_INT32, {1, 2}}, {TensorType_INT32, {}}, params);
|
|
model.SetInput<int>(model.input(), {0});
|
|
model.SetInput<int>(model.subgraph_input1(), {5, 7});
|
|
model.SetInput<int>(model.subgraph_input2(), {1, 2});
|
|
ASSERT_EQ(model.Invoke(), kTfLiteOk);
|
|
|
|
EXPECT_TRUE(IsDynamicTensor(model.GetOutputTensor(0)));
|
|
EXPECT_THAT(model.GetOutput<int>(), ElementsAreArray({6, 9}));
|
|
}
|
|
|
|
TEST(StablehloCaseTest, DynamicCaseTestPad) {
|
|
TfLiteStablehloCaseParams params = {
|
|
{1, 2},
|
|
2,
|
|
};
|
|
|
|
StablehloCaseDynamicOpModel model(
|
|
{TensorType_INT32, {}}, {TensorType_INT32, {2}},
|
|
{TensorType_INT32, {1, 2}}, {TensorType_INT32, {}}, params);
|
|
model.SetInput<int>(model.input(),
|
|
{-1}); // when index value is out of bounds, case op
|
|
// executes the last branch
|
|
model.SetInput<int>(model.subgraph_input1(), {5, 7});
|
|
model.SetInput<int>(model.subgraph_input2(), {1, 2});
|
|
ASSERT_EQ(model.Invoke(), kTfLiteOk);
|
|
|
|
EXPECT_TRUE(IsDynamicTensor(model.GetOutputTensor(0)));
|
|
EXPECT_THAT(model.GetOutput<int>(), ElementsAreArray({0, 5, 7, 0, 0}));
|
|
}
|
|
|
|
} // namespace
|
|
} // namespace tflite
|