mirror of
https://github.com/NVIDIA/TensorRT-LLM.git
synced 2026-01-14 06:27:45 +08:00
280 lines
8.7 KiB
C++
280 lines
8.7 KiB
C++
/*
|
|
* Copyright (c) 2022-2023, NVIDIA CORPORATION. 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 <unordered_map>
|
|
#include <vector>
|
|
|
|
#include <gtest/gtest.h>
|
|
|
|
#include "tensorrt_llm/common/tensor.h"
|
|
|
|
using namespace tensorrt_llm::common;
|
|
|
|
namespace
|
|
{
|
|
|
|
#define EXPECT_EQUAL_TENSORS(t1, t2) \
|
|
do \
|
|
{ \
|
|
EXPECT_TRUE(t1.where == t2.where); \
|
|
EXPECT_TRUE(t1.type == t2.type); \
|
|
EXPECT_TRUE(t1.shape == t2.shape); \
|
|
EXPECT_TRUE(t1.data == t2.data); \
|
|
} while (false)
|
|
|
|
TEST(TensorMapTest, HasKeyCorrectness)
|
|
{
|
|
bool* v1 = new bool(true);
|
|
float* v2 = new float[6]{1.0f, 1.1f, 1.2f, 1.3f, 1.4f, 1.5f};
|
|
Tensor t1 = Tensor{MEMORY_CPU, TYPE_BOOL, {1}, v1};
|
|
Tensor t2 = Tensor{MEMORY_CPU, TYPE_FP32, {3, 2}, v2};
|
|
|
|
TensorMap map({{"t1", t1}, {"t2", t2}});
|
|
EXPECT_TRUE(map.contains("t1"));
|
|
EXPECT_TRUE(map.contains("t2"));
|
|
EXPECT_FALSE(map.contains("t3"));
|
|
|
|
delete v1;
|
|
delete[] v2;
|
|
}
|
|
|
|
TEST(TensorMapTest, InsertCorrectness)
|
|
{
|
|
int* v1 = new int[4]{1, 10, 20, 30};
|
|
float* v2 = new float[2]{1.0f, 2.0f};
|
|
Tensor t1 = Tensor(MEMORY_CPU, TYPE_INT32, {4}, v1);
|
|
Tensor t2 = Tensor(MEMORY_CPU, TYPE_INT32, {2}, v2);
|
|
|
|
TensorMap map({{"t1", t1}});
|
|
EXPECT_TRUE(map.size() == 1);
|
|
EXPECT_TRUE(map.contains("t1"));
|
|
EXPECT_EQUAL_TENSORS(map.at("t1"), t1);
|
|
EXPECT_FALSE(map.contains("t2"));
|
|
|
|
delete[] v1;
|
|
delete[] v2;
|
|
}
|
|
|
|
TEST(TensorMapTest, InsertDoesNotAllowNoneTensor)
|
|
{
|
|
TensorMap map;
|
|
EXPECT_TRUE(map.size() == 0);
|
|
// forbid a none tensor.
|
|
EXPECT_THROW(map.insert("none", {}), std::runtime_error);
|
|
|
|
// forbid a tensor having null data pointer.
|
|
Tensor none_data_tensor = Tensor(MEMORY_CPU, TYPE_INT32, {}, nullptr);
|
|
EXPECT_THROW(map.insert("empty", none_data_tensor), std::runtime_error);
|
|
}
|
|
|
|
TEST(TensorMapTest, InsertDoesNotAllowDuplicatedKey)
|
|
{
|
|
int* v1 = new int[4]{1, 10, 20, 30};
|
|
Tensor t1 = Tensor(MEMORY_CPU, TYPE_INT32, {4}, v1);
|
|
Tensor t2 = Tensor(MEMORY_CPU, TYPE_INT32, {2}, v1);
|
|
TensorMap map({{"t1", t1}});
|
|
EXPECT_TRUE(map.size() == 1);
|
|
// forbid a duplicated key.
|
|
EXPECT_THROW(map.insert("t1", t2), std::runtime_error);
|
|
delete[] v1;
|
|
}
|
|
|
|
TEST(TensorMapTest, GetValCorrectness)
|
|
{
|
|
int* v1 = new int[4]{1, 10, 20, 30};
|
|
Tensor t1 = Tensor(MEMORY_CPU, TYPE_INT32, {4}, v1);
|
|
|
|
TensorMap map({{"t1", t1}});
|
|
EXPECT_TRUE(map.size() == 1);
|
|
// throw exception since the map doesn't have a key "t3".
|
|
EXPECT_THROW(map.getVal<int>("t3"), std::runtime_error);
|
|
EXPECT_TRUE(map.getVal<int>("t1") == 1);
|
|
EXPECT_TRUE(map.getVal<int>("t1", 3) == 1);
|
|
|
|
// map doesn't have t2 so return the default value 3.
|
|
EXPECT_TRUE(map.getVal<int>("t2", 3) == 3);
|
|
|
|
v1[0] += 1; // update value.
|
|
EXPECT_TRUE(map.getVal<int>("t1") == 2);
|
|
EXPECT_TRUE(map.getVal<int>("t1", 3) == 2);
|
|
|
|
size_t index = 2;
|
|
EXPECT_TRUE(map.getValWithOffset<int>("t1", index) == 20);
|
|
EXPECT_TRUE(map.getValWithOffset<int>("t1", index, 3) == 20);
|
|
EXPECT_TRUE(map.getValWithOffset<int>("t2", index, 3) == 3);
|
|
delete[] v1;
|
|
}
|
|
|
|
TEST(TensorMapTest, GetTensorCorrectness)
|
|
{
|
|
bool* t1_val = new bool(true);
|
|
float* t2_val = new float[6]{1.0f, 1.1f, 1.2f, 1.3f, 1.4f, 1.5f};
|
|
Tensor t1 = Tensor{MEMORY_CPU, TYPE_BOOL, {1}, t1_val};
|
|
Tensor t2 = Tensor{MEMORY_CPU, TYPE_FP32, {3, 2}, t2_val};
|
|
|
|
int* default_val = new int[4]{0, 1, 2, 3};
|
|
Tensor default_tensor = Tensor{MEMORY_CPU, TYPE_INT32, {4}, default_val};
|
|
|
|
TensorMap map({{"t1", t1}, {"t2", t2}});
|
|
EXPECT_THROW(map.at("t3"), std::runtime_error);
|
|
EXPECT_EQUAL_TENSORS(map.at("t1", default_tensor), t1);
|
|
EXPECT_EQUAL_TENSORS(map.at("t2", default_tensor), t2);
|
|
EXPECT_EQUAL_TENSORS(map.at("t3", default_tensor), default_tensor);
|
|
EXPECT_EQUAL_TENSORS(map.at("t3", {}), Tensor());
|
|
|
|
delete[] default_val;
|
|
delete[] t2_val;
|
|
delete[] t1_val;
|
|
}
|
|
|
|
TEST(TensorMapTest, GetTensorCorrectnessAtConstTensorMap)
|
|
{
|
|
bool* t1_val = new bool(true);
|
|
float* t2_val = new float[6]{1.0f, 1.1f, 1.2f, 1.3f, 1.4f, 1.5f};
|
|
Tensor t1 = Tensor{MEMORY_CPU, TYPE_BOOL, {1}, t1_val};
|
|
Tensor t2 = Tensor{MEMORY_CPU, TYPE_FP32, {3, 2}, t2_val};
|
|
|
|
int* default_val = new int[4]{0, 1, 2, 3};
|
|
Tensor default_tensor = Tensor{MEMORY_CPU, TYPE_INT32, {4}, default_val};
|
|
|
|
const TensorMap map({{"t1", t1}, {"t2", t2}});
|
|
EXPECT_THROW(map.at("t3"), std::runtime_error);
|
|
EXPECT_EQUAL_TENSORS(map.at("t1", default_tensor), t1);
|
|
EXPECT_EQUAL_TENSORS(map.at("t2", default_tensor), t2);
|
|
EXPECT_EQUAL_TENSORS(map.at("t3", default_tensor), default_tensor);
|
|
EXPECT_EQUAL_TENSORS(map.at("t3", {}), Tensor());
|
|
|
|
delete[] default_val;
|
|
delete[] t2_val;
|
|
delete[] t1_val;
|
|
}
|
|
|
|
TEST(TensorTest, EmptyTensorMinMaxRaiseError)
|
|
{
|
|
Tensor t1;
|
|
EXPECT_THROW(t1.min<int>(), std::runtime_error);
|
|
EXPECT_THROW(t1.max<int>(), std::runtime_error);
|
|
|
|
Tensor t2 = Tensor{MEMORY_CPU, TYPE_INT32, {}, nullptr};
|
|
EXPECT_THROW(t2.min<int>(), std::runtime_error);
|
|
EXPECT_THROW(t2.max<int>(), std::runtime_error);
|
|
}
|
|
|
|
using TensorTypes = testing::Types<int8_t, int, float>;
|
|
|
|
template <typename T>
|
|
class TensorFuncTest : public testing::Test
|
|
{
|
|
};
|
|
|
|
TYPED_TEST_SUITE(TensorFuncTest, TensorTypes);
|
|
|
|
TYPED_TEST(TensorFuncTest, MaxCorrectness)
|
|
{
|
|
using T = TypeParam;
|
|
|
|
size_t size = 4;
|
|
|
|
T* v1 = new T[size]{T(1), T(2), T(3), T(4)};
|
|
T* v2 = new T[size]{T(4), T(3), T(2), T(1)};
|
|
T* v3 = new T[size]{T(1), T(2), T(4), T(3)};
|
|
|
|
Tensor t1 = Tensor(MEMORY_CPU, getTensorType<T>(), {size}, v1);
|
|
Tensor t2 = Tensor(MEMORY_CPU, getTensorType<T>(), {size}, v2);
|
|
Tensor t3 = Tensor(MEMORY_CPU, getTensorType<T>(), {size}, v3);
|
|
|
|
EXPECT_EQ(t1.max<T>(), T(4));
|
|
EXPECT_EQ(t2.max<T>(), T(4));
|
|
EXPECT_EQ(t3.max<T>(), T(4));
|
|
|
|
delete[] v1;
|
|
delete[] v2;
|
|
delete[] v3;
|
|
}
|
|
|
|
TYPED_TEST(TensorFuncTest, MinCorrectness)
|
|
{
|
|
using T = TypeParam;
|
|
|
|
size_t size = 4;
|
|
|
|
T* v1 = new T[size]{T(1), T(2), T(3), T(4)};
|
|
T* v2 = new T[size]{T(4), T(3), T(2), T(1)};
|
|
T* v3 = new T[size]{T(1), T(2), T(4), T(3)};
|
|
|
|
Tensor t1 = Tensor(MEMORY_CPU, getTensorType<T>(), {size}, v1);
|
|
Tensor t2 = Tensor(MEMORY_CPU, getTensorType<T>(), {size}, v2);
|
|
Tensor t3 = Tensor(MEMORY_CPU, getTensorType<T>(), {size}, v3);
|
|
|
|
EXPECT_EQ(t1.min<T>(), T(1));
|
|
EXPECT_EQ(t2.min<T>(), T(1));
|
|
EXPECT_EQ(t3.min<T>(), T(1));
|
|
|
|
delete[] v1;
|
|
delete[] v2;
|
|
delete[] v3;
|
|
}
|
|
|
|
TYPED_TEST(TensorFuncTest, AnyCorrectness)
|
|
{
|
|
using T = TypeParam;
|
|
|
|
T* v = new T[4]{T(1), T(2), T(3), T(4)};
|
|
Tensor t = Tensor{MEMORY_CPU, getTensorType<T>(), {4}, v};
|
|
EXPECT_TRUE(t.any<T>(T(1)));
|
|
EXPECT_FALSE(t.any<T>(T(5)));
|
|
delete[] v;
|
|
}
|
|
|
|
TYPED_TEST(TensorFuncTest, AllCorrectness)
|
|
{
|
|
using T = TypeParam;
|
|
|
|
constexpr size_t size = 4;
|
|
T* v1 = new T[size]{T(1), T(1), T(1), T(1)};
|
|
T* v2 = new T[size]{T(1), T(1), T(1), T(2)};
|
|
Tensor t1 = Tensor{MEMORY_CPU, getTensorType<T>(), {size}, v1};
|
|
Tensor t2 = Tensor{MEMORY_CPU, getTensorType<T>(), {size}, v2};
|
|
EXPECT_TRUE(t1.all<T>(T(1)));
|
|
EXPECT_FALSE(t2.all<T>(T(2)));
|
|
delete[] v1;
|
|
delete[] v2;
|
|
}
|
|
|
|
TYPED_TEST(TensorFuncTest, SliceCorrectness)
|
|
{
|
|
using T = TypeParam;
|
|
|
|
constexpr int size = 12;
|
|
T* v = new T[size];
|
|
for (int i = 0; i < size; ++i)
|
|
{
|
|
v[i] = i;
|
|
}
|
|
|
|
DataType dtype = getTensorType<T>();
|
|
Tensor t1 = Tensor(MEMORY_CPU, dtype, {3, 4}, v);
|
|
Tensor t2 = t1.slice({2, 4}, 4);
|
|
|
|
EXPECT_EQUAL_TENSORS(t2, Tensor(MEMORY_CPU, dtype, {2, 4}, &v[4]));
|
|
// An overflowed tensor throws an exception.
|
|
EXPECT_THROW(t1.slice({2, 4}, 5), std::runtime_error);
|
|
|
|
delete[] v;
|
|
}
|
|
|
|
} // end of namespace
|