mirror of
https://github.com/NVIDIA/TensorRT-LLM.git
synced 2026-01-14 06:27:45 +08:00
16d2467ea8
* Update TensorRT-LLM --------- Co-authored-by: Denis Kayshev <topenkoff@gmail.com> Co-authored-by: akhoroshev <arthoroshev@gmail.com> Co-authored-by: Patrick Reiter Horn <patrick.horn@gmail.com> Update
144 lines
8.7 KiB
Plaintext
144 lines
8.7 KiB
Plaintext
/*
|
|
* Copyright (c) 2022-2024, 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 "tensorrt_llm/common/memoryUtils.h"
|
|
#include <array>
|
|
|
|
#include <gtest/gtest.h>
|
|
|
|
namespace tc = tensorrt_llm::common;
|
|
|
|
TEST(memory_utils, flat_index)
|
|
{
|
|
// All testing happens at compile time.
|
|
std::array<int, 4> constexpr dims{4, 3, 2, 1};
|
|
static_assert(tc::flat_index(dims, 0) == 0);
|
|
static_assert(tc::flat_index(dims, 1) == 1);
|
|
static_assert(tc::flat_index(dims, 2) == 2);
|
|
static_assert(tc::flat_index(dims, 3) == 3);
|
|
static_assert(tc::flat_index(dims, 0, 0) == 0);
|
|
static_assert(tc::flat_index(dims, 0, 1) == 1);
|
|
static_assert(tc::flat_index(dims, 1, 0) == 3);
|
|
static_assert(tc::flat_index(dims, 1, 1) == 4);
|
|
static_assert(tc::flat_index(dims, 1, 2) == 5);
|
|
static_assert(tc::flat_index(dims, 2, 0) == 6);
|
|
static_assert(tc::flat_index(dims, 2, 1) == 7);
|
|
static_assert(tc::flat_index(dims, 2, 2) == 8);
|
|
static_assert(tc::flat_index(dims, 0, 0, 0) == 0);
|
|
static_assert(tc::flat_index(dims, 0, 0, 1) == 1);
|
|
static_assert(tc::flat_index(dims, 0, 1, 0) == 2);
|
|
static_assert(tc::flat_index(dims, 0, 1, 1) == 3);
|
|
static_assert(tc::flat_index(dims, 0, 2, 0) == 4);
|
|
static_assert(tc::flat_index(dims, 0, 2, 1) == 5);
|
|
static_assert(tc::flat_index(dims, 1, 0, 0) == 6);
|
|
static_assert(tc::flat_index(dims, 1, 0, 1) == 7);
|
|
static_assert(tc::flat_index(dims, 1, 1, 0) == 8);
|
|
static_assert(tc::flat_index(dims, 1, 1, 1) == 9);
|
|
static_assert(tc::flat_index(dims, 1, 2, 0) == 10);
|
|
static_assert(tc::flat_index(dims, 1, 2, 1) == 11);
|
|
static_assert(tc::flat_index(0, dims, 1, 2, 1) == 11);
|
|
static_assert(tc::flat_index(1, dims, 1, 2, 1) == 35);
|
|
static_assert(tc::flat_index(2, dims, 1, 2, 1) == 59);
|
|
static_assert(tc::flat_index(dims, 0, 0, 0) == tc::flat_index(0, &dims[1], 0, 0));
|
|
static_assert(tc::flat_index(dims, 0, 0, 1) == tc::flat_index(0, &dims[1], 0, 1));
|
|
static_assert(tc::flat_index(dims, 0, 1, 0) == tc::flat_index(0, &dims[1], 1, 0));
|
|
static_assert(tc::flat_index(dims, 0, 1, 1) == tc::flat_index(0, &dims[1], 1, 1));
|
|
static_assert(tc::flat_index(dims, 0, 2, 0) == tc::flat_index(0, &dims[1], 2, 0));
|
|
static_assert(tc::flat_index(dims, 0, 2, 1) == tc::flat_index(0, &dims[1], 2, 1));
|
|
static_assert(tc::flat_index(dims, 1, 0, 0) == tc::flat_index(1, &dims[1], 0, 0));
|
|
static_assert(tc::flat_index(dims, 1, 0, 1) == tc::flat_index(1, &dims[1], 0, 1));
|
|
static_assert(tc::flat_index(dims, 1, 1, 0) == tc::flat_index(1, &dims[1], 1, 0));
|
|
static_assert(tc::flat_index(dims, 1, 1, 1) == tc::flat_index(1, &dims[1], 1, 1));
|
|
static_assert(tc::flat_index(dims, 1, 2, 0) == tc::flat_index(1, &dims[1], 2, 0));
|
|
static_assert(tc::flat_index(dims, 1, 2, 1) == tc::flat_index(1, &dims[1], 2, 1));
|
|
|
|
static_assert(tc::flat_index({4, 3, 2, 1}, 1, 2, 1) == 11);
|
|
static_assert(tc::flat_index({4, 3, 2}, 1, 2, 1) == 11);
|
|
static_assert(tc::flat_index(0, {4, 3, 2, 1}, 1, 2, 1) == 11);
|
|
static_assert(tc::flat_index(1, {4, 3, 2, 1}, 1, 2, 1) == 35);
|
|
|
|
int constexpr dim_array[]{4, 3, 2, 1};
|
|
static_assert(tc::flat_index(dim_array, 1, 2, 1) == 11);
|
|
static_assert(tc::flat_index(dim_array, 1, 2, 1) == 11);
|
|
static_assert(tc::flat_index(0, dim_array, 1, 2, 1) == 11);
|
|
static_assert(tc::flat_index(1, dim_array, 1, 2, 1) == 35);
|
|
|
|
static_assert(tc::flat_index<1>(dims, 2, 1, 0) == 5);
|
|
static_assert(tc::flat_index<1>(0, dims, 2, 1, 0) == 5);
|
|
static_assert(tc::flat_index<1>(1, dims, 2, 1, 0) == 11);
|
|
static_assert(tc::flat_index<1>({4, 3, 2, 1}, 2, 1, 0) == 5);
|
|
static_assert(tc::flat_index<1>(0, {4, 3, 2, 1}, 2, 1, 0) == 5);
|
|
static_assert(tc::flat_index<1>(1, {4, 3, 2, 1}, 2, 1, 0) == 11);
|
|
}
|
|
|
|
TEST(memory_utils, flat_index_dim)
|
|
{
|
|
std::array<int, 4> constexpr dims{4, 3, 2, 1};
|
|
static_assert(tc::flat_index(dims, 0, 0) == tc::flat_index2(0, 0, dims[1]));
|
|
static_assert(tc::flat_index(dims, 0, 1) == tc::flat_index2(0, 1, dims[1]));
|
|
static_assert(tc::flat_index(dims, 1, 0) == tc::flat_index2(1, 0, dims[1]));
|
|
static_assert(tc::flat_index(dims, 1, 1) == tc::flat_index2(1, 1, dims[1]));
|
|
static_assert(tc::flat_index(dims, 1, 2) == tc::flat_index2(1, 2, dims[1]));
|
|
static_assert(tc::flat_index(dims, 2, 0) == tc::flat_index2(2, 0, dims[1]));
|
|
static_assert(tc::flat_index(dims, 2, 1) == tc::flat_index2(2, 1, dims[1]));
|
|
static_assert(tc::flat_index(dims, 2, 2) == tc::flat_index2(2, 2, dims[1]));
|
|
static_assert(tc::flat_index(dims, 0, 0, 0) == tc::flat_index3(0, 0, 0, dims[1], dims[2]));
|
|
static_assert(tc::flat_index(dims, 0, 0, 1) == tc::flat_index3(0, 0, 1, dims[1], dims[2]));
|
|
static_assert(tc::flat_index(dims, 0, 1, 0) == tc::flat_index3(0, 1, 0, dims[1], dims[2]));
|
|
static_assert(tc::flat_index(dims, 0, 1, 1) == tc::flat_index3(0, 1, 1, dims[1], dims[2]));
|
|
static_assert(tc::flat_index(dims, 0, 2, 0) == tc::flat_index3(0, 2, 0, dims[1], dims[2]));
|
|
static_assert(tc::flat_index(dims, 0, 2, 1) == tc::flat_index3(0, 2, 1, dims[1], dims[2]));
|
|
static_assert(tc::flat_index(dims, 1, 0, 0) == tc::flat_index3(1, 0, 0, dims[1], dims[2]));
|
|
static_assert(tc::flat_index(dims, 1, 0, 1) == tc::flat_index3(1, 0, 1, dims[1], dims[2]));
|
|
static_assert(tc::flat_index(dims, 1, 1, 0) == tc::flat_index3(1, 1, 0, dims[1], dims[2]));
|
|
static_assert(tc::flat_index(dims, 1, 1, 1) == tc::flat_index3(1, 1, 1, dims[1], dims[2]));
|
|
static_assert(tc::flat_index(dims, 1, 2, 0) == tc::flat_index3(1, 2, 0, dims[1], dims[2]));
|
|
static_assert(tc::flat_index(dims, 1, 2, 1) == tc::flat_index3(1, 2, 1, dims[1], dims[2]));
|
|
static_assert(tc::flat_index(dims, 1, 1, 0, 0) == tc::flat_index4(1, 1, 0, 0, dims[1], dims[2], dims[3]));
|
|
static_assert(tc::flat_index(dims, 1, 1, 1, 0) == tc::flat_index4(1, 1, 1, 0, dims[1], dims[2], dims[3]));
|
|
static_assert(tc::flat_index(dims, 1, 2, 0, 0) == tc::flat_index4(1, 2, 0, 0, dims[1], dims[2], dims[3]));
|
|
static_assert(tc::flat_index(dims, 1, 2, 1, 0) == tc::flat_index4(1, 2, 1, 0, dims[1], dims[2], dims[3]));
|
|
}
|
|
|
|
TEST(memory_utils, flat_index_strided)
|
|
{
|
|
std::array<int, 4> constexpr dims{4, 3, 2, 1};
|
|
auto constexpr dim_12 = dims[1] * dims[2];
|
|
static_assert(tc::flat_index_strided3(0, 0, 0, dim_12, dims[2]) == tc::flat_index3(0, 0, 0, dims[1], dims[2]));
|
|
static_assert(tc::flat_index_strided3(0, 0, 1, dim_12, dims[2]) == tc::flat_index3(0, 0, 1, dims[1], dims[2]));
|
|
static_assert(tc::flat_index_strided3(0, 1, 0, dim_12, dims[2]) == tc::flat_index3(0, 1, 0, dims[1], dims[2]));
|
|
static_assert(tc::flat_index_strided3(0, 1, 1, dim_12, dims[2]) == tc::flat_index3(0, 1, 1, dims[1], dims[2]));
|
|
static_assert(tc::flat_index_strided3(0, 2, 0, dim_12, dims[2]) == tc::flat_index3(0, 2, 0, dims[1], dims[2]));
|
|
static_assert(tc::flat_index_strided3(0, 2, 1, dim_12, dims[2]) == tc::flat_index3(0, 2, 1, dims[1], dims[2]));
|
|
static_assert(tc::flat_index_strided3(1, 0, 0, dim_12, dims[2]) == tc::flat_index3(1, 0, 0, dims[1], dims[2]));
|
|
static_assert(tc::flat_index_strided3(1, 0, 1, dim_12, dims[2]) == tc::flat_index3(1, 0, 1, dims[1], dims[2]));
|
|
static_assert(tc::flat_index_strided3(1, 1, 0, dim_12, dims[2]) == tc::flat_index3(1, 1, 0, dims[1], dims[2]));
|
|
static_assert(tc::flat_index_strided3(1, 1, 1, dim_12, dims[2]) == tc::flat_index3(1, 1, 1, dims[1], dims[2]));
|
|
static_assert(tc::flat_index_strided3(1, 2, 0, dim_12, dims[2]) == tc::flat_index3(1, 2, 0, dims[1], dims[2]));
|
|
static_assert(tc::flat_index_strided3(1, 2, 1, dim_12, dims[2]) == tc::flat_index3(1, 2, 1, dims[1], dims[2]));
|
|
|
|
auto constexpr dims_123 = dims[1] * dims[2] * dims[3];
|
|
auto constexpr dims_23 = dims[2] * dims[3];
|
|
static_assert(tc::flat_index_strided4(1, 1, 0, 0, dims_123, dims_23, dims[3])
|
|
== tc::flat_index4(1, 1, 0, 0, dims[1], dims[2], dims[3]));
|
|
static_assert(tc::flat_index_strided4(1, 1, 1, 0, dims_123, dims_23, dims[3])
|
|
== tc::flat_index4(1, 1, 1, 0, dims[1], dims[2], dims[3]));
|
|
static_assert(tc::flat_index_strided4(1, 2, 1, 0, dims_123, dims_23, dims[3])
|
|
== tc::flat_index4(1, 2, 1, 0, dims[1], dims[2], dims[3]));
|
|
static_assert(tc::flat_index_strided4(3, 2, 1, 0, dims_123, dims_23, dims[3])
|
|
== tc::flat_index4(3, 2, 1, 0, dims[1], dims[2], dims[3]));
|
|
}
|