mirror of
https://github.com/NVIDIA/TensorRT-LLM.git
synced 2026-01-14 06:27:45 +08:00
b2fd493c16
* Update TensorRT-LLM --------- Co-authored-by: Shixiaowei02 <39303645+Shixiaowei02@users.noreply.github.com>
221 lines
7.1 KiB
C++
221 lines
7.1 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 <gmock/gmock.h>
|
|
#include <gtest/gtest.h>
|
|
|
|
#include "tensorrt_llm/runtime/bufferManager.h"
|
|
#include "tensorrt_llm/runtime/iTensor.h"
|
|
|
|
using namespace tensorrt_llm::runtime;
|
|
namespace tc = tensorrt_llm::common;
|
|
|
|
TEST(ITensorTest, SqueezeTensor)
|
|
{
|
|
auto dims = ITensor::makeShape({16, 1, 4});
|
|
auto constexpr dataType = nvinfer1::DataType::kFLOAT;
|
|
ITensor::SharedPtr tensor{BufferManager::cpu(dims, dataType)};
|
|
|
|
auto squeezeDim = 0;
|
|
EXPECT_THROW(tensor->squeeze(squeezeDim), std::runtime_error);
|
|
squeezeDim = 1;
|
|
auto squeezed = ITensor::view(tensor, ITensor::squeeze(dims, squeezeDim));
|
|
|
|
EXPECT_EQ(squeezed->getSize(), tensor->getSize());
|
|
EXPECT_EQ(squeezed->getShape().nbDims, tensor->getShape().nbDims - 1);
|
|
EXPECT_EQ(squeezed->getShape().d[0], tensor->getShape().d[0]);
|
|
EXPECT_EQ(squeezed->getShape().d[1], tensor->getShape().d[2]);
|
|
|
|
EXPECT_NO_THROW(squeezed->release());
|
|
EXPECT_EQ(squeezed->data(), nullptr);
|
|
EXPECT_NE(tensor->data(), nullptr);
|
|
}
|
|
|
|
TEST(ITensorTest, UnsqueezeShape)
|
|
{
|
|
auto oldShape = ITensor::makeShape({2, 3, 4, 5});
|
|
{
|
|
auto shape = ITensor::unsqueeze(oldShape, 0);
|
|
|
|
EXPECT_EQ(shape.nbDims, 5);
|
|
EXPECT_EQ(shape.d[0], 1);
|
|
EXPECT_EQ(shape.d[1], 2);
|
|
EXPECT_EQ(shape.d[2], 3);
|
|
EXPECT_EQ(shape.d[3], 4);
|
|
EXPECT_EQ(shape.d[4], 5);
|
|
}
|
|
{
|
|
auto shape = ITensor::unsqueeze(oldShape, 1);
|
|
|
|
EXPECT_EQ(shape.nbDims, 5);
|
|
EXPECT_EQ(shape.d[0], 2);
|
|
EXPECT_EQ(shape.d[1], 1);
|
|
EXPECT_EQ(shape.d[2], 3);
|
|
EXPECT_EQ(shape.d[3], 4);
|
|
EXPECT_EQ(shape.d[4], 5);
|
|
}
|
|
|
|
{
|
|
auto shape = ITensor::unsqueeze(oldShape, 4);
|
|
|
|
EXPECT_EQ(shape.nbDims, 5);
|
|
EXPECT_EQ(shape.d[0], 2);
|
|
EXPECT_EQ(shape.d[1], 3);
|
|
EXPECT_EQ(shape.d[2], 4);
|
|
EXPECT_EQ(shape.d[3], 5);
|
|
EXPECT_EQ(shape.d[4], 1);
|
|
}
|
|
|
|
std::vector<int> invalidDims{-1, 5, 10};
|
|
for (auto invalidDim : invalidDims)
|
|
{
|
|
try
|
|
{
|
|
ITensor::unsqueeze(oldShape, invalidDim);
|
|
FAIL() << "Expected failure";
|
|
}
|
|
catch (tensorrt_llm::common::TllmException const& e)
|
|
{
|
|
EXPECT_THAT(e.what(), testing::HasSubstr("Invalid dim"));
|
|
}
|
|
catch (...)
|
|
{
|
|
FAIL() << "Expected TllmException";
|
|
}
|
|
}
|
|
}
|
|
|
|
TEST(ITensorTest, UnsqueezeTensor)
|
|
{
|
|
auto oldShape = ITensor::makeShape({2, 3, 4, 5});
|
|
|
|
{
|
|
auto tensor = BufferManager::cpu(oldShape, nvinfer1::DataType::kINT32);
|
|
tensor->unsqueeze(0);
|
|
auto shape = tensor->getShape();
|
|
|
|
EXPECT_EQ(shape.nbDims, 5);
|
|
EXPECT_EQ(shape.d[0], 1);
|
|
EXPECT_EQ(shape.d[1], 2);
|
|
EXPECT_EQ(shape.d[2], 3);
|
|
EXPECT_EQ(shape.d[3], 4);
|
|
EXPECT_EQ(shape.d[4], 5);
|
|
}
|
|
{
|
|
auto tensor = BufferManager::cpu(oldShape, nvinfer1::DataType::kINT32);
|
|
tensor->unsqueeze(1);
|
|
auto shape = tensor->getShape();
|
|
|
|
EXPECT_EQ(shape.nbDims, 5);
|
|
EXPECT_EQ(shape.d[0], 2);
|
|
EXPECT_EQ(shape.d[1], 1);
|
|
EXPECT_EQ(shape.d[2], 3);
|
|
EXPECT_EQ(shape.d[3], 4);
|
|
EXPECT_EQ(shape.d[4], 5);
|
|
}
|
|
|
|
{
|
|
auto tensor = BufferManager::cpu(oldShape, nvinfer1::DataType::kINT32);
|
|
tensor->unsqueeze(4);
|
|
auto shape = tensor->getShape();
|
|
|
|
EXPECT_EQ(shape.nbDims, 5);
|
|
EXPECT_EQ(shape.d[0], 2);
|
|
EXPECT_EQ(shape.d[1], 3);
|
|
EXPECT_EQ(shape.d[2], 4);
|
|
EXPECT_EQ(shape.d[3], 5);
|
|
EXPECT_EQ(shape.d[4], 1);
|
|
}
|
|
|
|
std::vector<int> invalidDims{-1, 5, 10};
|
|
for (auto invalidDim : invalidDims)
|
|
{
|
|
try
|
|
{
|
|
auto tensor = BufferManager::cpu(oldShape, nvinfer1::DataType::kINT32);
|
|
tensor->unsqueeze(invalidDim);
|
|
FAIL() << "Expected failure";
|
|
}
|
|
catch (tensorrt_llm::common::TllmException const& e)
|
|
{
|
|
EXPECT_THAT(e.what(), testing::HasSubstr("Invalid dim"));
|
|
}
|
|
catch (...)
|
|
{
|
|
FAIL() << "Expected TllmException";
|
|
}
|
|
}
|
|
}
|
|
|
|
TEST(ITensorTest, TensorView)
|
|
{
|
|
auto const dims = ITensor::makeShape({16, 1, 4});
|
|
auto constexpr dataType = nvinfer1::DataType::kFLOAT;
|
|
ITensor::SharedPtr tensor = BufferManager::cpu(dims, dataType);
|
|
|
|
auto const viewDims = ITensor::makeShape({16, 1, 2});
|
|
|
|
auto view = ITensor::view(tensor, viewDims);
|
|
EXPECT_EQ(view->getSize(), tensor->getSize() / 2);
|
|
EXPECT_EQ(view->getShape().nbDims, tensor->getShape().nbDims);
|
|
EXPECT_EQ(view->getShape().d[2], tensor->getShape().d[2] / 2);
|
|
|
|
EXPECT_NO_THROW(view->release());
|
|
EXPECT_EQ(view->data(), nullptr);
|
|
EXPECT_NE(tensor->data(), nullptr);
|
|
}
|
|
|
|
TEST(ITensorTest, TensorSlice)
|
|
{
|
|
auto dims = ITensor::makeShape({16, 8, 4});
|
|
auto constexpr dataType = nvinfer1::DataType::kFLOAT;
|
|
ITensor::SharedPtr tensor{BufferManager::cpu(dims, dataType)};
|
|
auto offset = dims.d[0] / 4;
|
|
auto slice = ITensor::slice(tensor, offset);
|
|
auto const sizeSlice = 3 * tensor->getSize() / 4;
|
|
EXPECT_EQ(slice->getShape().d[0], dims.d[0] - offset);
|
|
EXPECT_EQ(slice->getSize(), sizeSlice);
|
|
EXPECT_EQ(slice->getCapacity(), sizeSlice);
|
|
EXPECT_EQ(static_cast<std::uint8_t*>(slice->data()) - static_cast<std::uint8_t*>(tensor->data()),
|
|
offset * ITensor::volume(dims) / dims.d[0] * BufferDataType(dataType).getSize());
|
|
|
|
auto dimsNew = ITensor::makeShape({12, 32});
|
|
EXPECT_EQ(ITensor::volume(dimsNew), sizeSlice);
|
|
EXPECT_NO_THROW(slice->reshape(dimsNew));
|
|
EXPECT_EQ(slice->getShape().d[1], dimsNew.d[1]);
|
|
dimsNew.d[0] = 6;
|
|
EXPECT_LT(ITensor::volume(dimsNew), sizeSlice);
|
|
EXPECT_NO_THROW(slice->reshape(dimsNew));
|
|
EXPECT_EQ(slice->getShape().d[0], dimsNew.d[0]);
|
|
dimsNew.d[0] = 16;
|
|
EXPECT_GT(ITensor::volume(dimsNew), sizeSlice);
|
|
EXPECT_THROW(slice->reshape(dimsNew), std::runtime_error);
|
|
|
|
EXPECT_NO_THROW(slice->resize(sizeSlice));
|
|
EXPECT_NO_THROW(slice->resize(sizeSlice / 2));
|
|
EXPECT_EQ(slice->getShape().d[0], sizeSlice / 2);
|
|
EXPECT_THROW(slice->resize(sizeSlice * 2), std::runtime_error);
|
|
EXPECT_NO_THROW(slice->release());
|
|
EXPECT_EQ(slice->data(), nullptr);
|
|
EXPECT_NE(tensor->data(), nullptr);
|
|
|
|
std::shared_ptr<ITensor const> constTensor{tensor};
|
|
auto constSlice = ITensor::slice(constTensor, offset);
|
|
EXPECT_EQ(constSlice->getShape().d[0], dims.d[0] - offset);
|
|
auto uniqueSlice = ITensor::slice(std::move(constSlice), 1);
|
|
EXPECT_EQ(uniqueSlice->getShape().d[0], dims.d[0] - offset - 1);
|
|
}
|