/* * Copyright (C) 2017 The Android Open Source Project * * 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 "HashtableLookup.h" #include "NeuralNetworksWrapper.h" #include "gmock/gmock-matchers.h" #include "gtest/gtest.h" using ::testing::FloatNear; using ::testing::Matcher; namespace android { namespace nn { namespace wrapper { namespace { std::vector<Matcher<float>> ArrayFloatNear(const std::vector<float>& values, float max_abs_error=1.e-6) { std::vector<Matcher<float>> matchers; matchers.reserve(values.size()); for (const float& v : values) { matchers.emplace_back(FloatNear(v, max_abs_error)); } return matchers; } } // namespace using ::testing::ElementsAreArray; #define FOR_ALL_INPUT_AND_WEIGHT_TENSORS(ACTION) \ ACTION(Lookup, int) \ ACTION(Key, int) \ ACTION(Value, float) // For all output and intermediate states #define FOR_ALL_OUTPUT_TENSORS(ACTION) \ ACTION(Output, float) \ ACTION(Hits, uint8_t) class HashtableLookupOpModel { public: HashtableLookupOpModel(std::initializer_list<uint32_t> lookup_shape, std::initializer_list<uint32_t> key_shape, std::initializer_list<uint32_t> value_shape) { auto it_vs = value_shape.begin(); rows_ = *it_vs++; features_ = *it_vs; std::vector<uint32_t> inputs; // Input and weights OperandType LookupTy(Type::TENSOR_INT32, lookup_shape); inputs.push_back(model_.addOperand(&LookupTy)); OperandType KeyTy(Type::TENSOR_INT32, key_shape); inputs.push_back(model_.addOperand(&KeyTy)); OperandType ValueTy(Type::TENSOR_FLOAT32, value_shape); inputs.push_back(model_.addOperand(&ValueTy)); // Output and other intermediate state std::vector<uint32_t> outputs; std::vector<uint32_t> out_dim(lookup_shape.begin(), lookup_shape.end()); out_dim.push_back(features_); OperandType OutputOpndTy(Type::TENSOR_FLOAT32, out_dim); outputs.push_back(model_.addOperand(&OutputOpndTy)); OperandType HitsOpndTy(Type::TENSOR_QUANT8_ASYMM, lookup_shape, 1.f, 0); outputs.push_back(model_.addOperand(&HitsOpndTy)); auto multiAll = [](const std::vector<uint32_t> &dims) -> uint32_t { uint32_t sz = 1; for (uint32_t d : dims) { sz *= d; } return sz; }; Value_.insert(Value_.end(), multiAll(value_shape), 0.f); Output_.insert(Output_.end(), multiAll(out_dim), 0.f); Hits_.insert(Hits_.end(), multiAll(lookup_shape), 0); model_.addOperation(ANEURALNETWORKS_HASHTABLE_LOOKUP, inputs, outputs); model_.identifyInputsAndOutputs(inputs, outputs); model_.finish(); } void Invoke() { ASSERT_TRUE(model_.isValid()); Compilation compilation(&model_); compilation.finish(); Execution execution(&compilation); #define SetInputOrWeight(X, T) \ ASSERT_EQ(execution.setInput(HashtableLookup::k##X##Tensor, X##_.data(), \ sizeof(T) * X##_.size()), \ Result::NO_ERROR); FOR_ALL_INPUT_AND_WEIGHT_TENSORS(SetInputOrWeight); #undef SetInputOrWeight #define SetOutput(X, T) \ ASSERT_EQ(execution.setOutput(HashtableLookup::k##X##Tensor, X##_.data(), \ sizeof(T) * X##_.size()), \ Result::NO_ERROR); FOR_ALL_OUTPUT_TENSORS(SetOutput); #undef SetOutput ASSERT_EQ(execution.compute(), Result::NO_ERROR); } #define DefineSetter(X, T) \ void Set##X(const std::vector<T>& f) { \ X##_.insert(X##_.end(), f.begin(), f.end()); \ } FOR_ALL_INPUT_AND_WEIGHT_TENSORS(DefineSetter); #undef DefineSetter void SetHashtableValue(const std::function<float(uint32_t, uint32_t)>& function) { for (uint32_t i = 0; i < rows_; i++) { for (uint32_t j = 0; j < features_; j++) { Value_[i * features_ + j] = function(i, j); } } } const std::vector<float>& GetOutput() const { return Output_; } const std::vector<uint8_t>& GetHits() const { return Hits_; } private: Model model_; uint32_t rows_; uint32_t features_; #define DefineTensor(X, T) std::vector<T> X##_; FOR_ALL_INPUT_AND_WEIGHT_TENSORS(DefineTensor); FOR_ALL_OUTPUT_TENSORS(DefineTensor); #undef DefineTensor }; TEST(HashtableLookupOpTest, BlackBoxTest) { HashtableLookupOpModel m({4}, {3}, {3, 2}); m.SetLookup({1234, -292, -11, 0}); m.SetKey({-11, 0, 1234}); m.SetHashtableValue([](int i, int j) { return i + j / 10.0f; }); m.Invoke(); EXPECT_THAT(m.GetOutput(), ElementsAreArray(ArrayFloatNear({ 2.0, 2.1, // 2-rd item 0, 0, // Not found 0.0, 0.1, // 0-th item 1.0, 1.1, // 1-st item }))); EXPECT_EQ(m.GetHits(), std::vector<uint8_t>({ 1, 0, 1, 1, })); } } // namespace wrapper } // namespace nn } // namespace android