From b1187cc3e2511d368dbf40ac1b37bd458bd2843a Mon Sep 17 00:00:00 2001 From: Jiacheng Huang Date: Thu, 23 Jul 2026 00:00:27 +0800 Subject: [PATCH] perf: reduce TensorView metadata allocations --- src/tensor_view.cc | 7 +- src/tensor_view.h | 37 +++++ tests/CMakeLists.txt | 4 + tests/performance/perf_tensor_view.cc | 29 ++++ tests/test_tensor_view_allocations.cc | 191 ++++++++++++++++++++++++++ 5 files changed, 266 insertions(+), 2 deletions(-) create mode 100644 tests/test_tensor_view_allocations.cc diff --git a/src/tensor_view.cc b/src/tensor_view.cc index 14a2919..a19fbaa 100644 --- a/src/tensor_view.cc +++ b/src/tensor_view.cc @@ -16,8 +16,11 @@ static TensorView::Index GetEffectiveIndex(TensorView::Index index, TensorView::TensorView(void* data, std::initializer_list shape, const DataType& dtype, const Device& device, std::initializer_list strides) - : TensorView{data, decltype(shape_){shape}, dtype, device, - decltype(strides_){strides}} {} + : data_{data}, + shape_{shape}, + dtype_{dtype}, + device_{device}, + strides_{strides} {} TensorView TensorView::operator[](const Index& index) const { return { diff --git a/src/tensor_view.h b/src/tensor_view.h index dcf7cc9..ea822af 100644 --- a/src/tensor_view.h +++ b/src/tensor_view.h @@ -51,6 +51,13 @@ class TensorView { device_{tensor.device()}, strides_{tensor.strides()} {} + TensorView(void* data, Shape shape) + : data_{data}, + shape_{std::move(shape)}, + dtype_{DefaultDataType()}, + device_{DefaultDevice()}, + strides_{DefaultStrides(shape_)} {} + template TensorView(void* data, const Shape& shape) : data_{data}, @@ -59,6 +66,13 @@ class TensorView { device_{DefaultDevice()}, strides_{DefaultStrides(shape)} {} + TensorView(void* data, Shape shape, const DataType& dtype) + : data_{data}, + shape_{std::move(shape)}, + dtype_{dtype}, + device_{DefaultDevice()}, + strides_{DefaultStrides(shape_)} {} + template TensorView(void* data, const Shape& shape, const DataType& dtype) : data_{data}, @@ -67,6 +81,13 @@ class TensorView { device_{DefaultDevice()}, strides_{DefaultStrides(shape)} {} + TensorView(void* data, Shape shape, const Device& device) + : data_{data}, + shape_{std::move(shape)}, + dtype_{DefaultDataType()}, + device_{device}, + strides_{DefaultStrides(shape_)} {} + template TensorView(void* data, const Shape& shape, const Device& device) : data_{data}, @@ -75,6 +96,14 @@ class TensorView { device_{device}, strides_{DefaultStrides(shape)} {} + TensorView(void* data, Shape shape, const DataType& dtype, + const Device& device) + : data_{data}, + shape_{std::move(shape)}, + dtype_{dtype}, + device_{device}, + strides_{DefaultStrides(shape_)} {} + template TensorView(void* data, const Shape& shape, const DataType& dtype, const Device& device) @@ -84,6 +113,14 @@ class TensorView { device_{device}, strides_{DefaultStrides(shape)} {} + TensorView(void* data, Shape shape, const DataType& dtype, + const Device& device, Strides strides) + : data_{data}, + shape_{std::move(shape)}, + dtype_{dtype}, + device_{device}, + strides_{std::move(strides)} {} + template TensorView(void* data, const Shape& shape, const DataType& dtype, const Device& device, const Strides& strides) diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index 8ffa4b3..4b71fc5 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -41,6 +41,10 @@ endfunction() add_infini_rt_test(test_smoke test_smoke.cc) add_infini_rt_test(test_core test_core.cc) +if(CMAKE_SYSTEM_NAME STREQUAL "Linux") + add_infini_rt_test(test_tensor_view_allocations + test_tensor_view_allocations.cc) +endif() set(INFINI_RT_TEST_HAS_RUNTIME_BACKEND OFF) diff --git a/tests/performance/perf_tensor_view.cc b/tests/performance/perf_tensor_view.cc index 7a78886..d9e9b87 100644 --- a/tests/performance/perf_tensor_view.cc +++ b/tests/performance/perf_tensor_view.cc @@ -32,11 +32,40 @@ int main() { perf::DoNotOptimize(tensor); }); + perf::RunBenchmark("perf_tensor_view.construct_rvalue_full_metadata", + {perf::NumberParam("ndim", 2)}, kIterations, "ns", [&] { + TensorView tensor{data.data(), TensorView::Shape{32, 64}, + DataType::kFloat32, cpu_device, + TensorView::Strides{64, 1}}; + perf::DoNotOptimize(tensor); + }); + + perf::RunBenchmark("perf_tensor_view.construct_rvalue_default_strides", + {perf::NumberParam("ndim", 2)}, kIterations, "ns", [&] { + TensorView tensor{data.data(), TensorView::Shape{32, 64}, + DataType::kFloat32, cpu_device}; + perf::DoNotOptimize(tensor); + }); + + perf::RunBenchmark( + "perf_tensor_view.construct_initializer_list", + {perf::NumberParam("ndim", 2)}, kIterations, "ns", [&] { + TensorView tensor{ + data.data(), {32, 64}, DataType::kFloat32, cpu_device, {64, 1}}; + perf::DoNotOptimize(tensor); + }); + TensorView contiguous{data.data(), shape, DataType::kFloat32, cpu_device, contiguous_strides}; TensorView transposed{data.data(), shape, DataType::kFloat32, cpu_device, transposed_strides}; + perf::RunBenchmark("perf_tensor_view.transpose", + {perf::NumberParam("ndim", 2)}, kIterations, "ns", [&] { + const auto result = contiguous.T(); + perf::DoNotOptimize(result); + }); + perf::RunBenchmark("perf_tensor_view.operator_index", {perf::NumberParam("ndim", 2)}, kIterations, "ns", [&] { const auto slice = contiguous[1]; diff --git a/tests/test_tensor_view_allocations.cc b/tests/test_tensor_view_allocations.cc new file mode 100644 index 0000000..b76c682 --- /dev/null +++ b/tests/test_tensor_view_allocations.cc @@ -0,0 +1,191 @@ +#include + +#include +#include +#include +#include + +#include "test_helper.h" + +namespace { + +thread_local bool count_allocations = false; +thread_local std::size_t allocation_count = 0; + +class AllocationScope { + public: + AllocationScope() { + allocation_count = 0; + count_allocations = true; + } + + AllocationScope(const AllocationScope&) = delete; + AllocationScope& operator=(const AllocationScope&) = delete; + + ~AllocationScope() { count_allocations = false; } + + std::size_t count() const { return allocation_count; } +}; + +template +std::size_t CountAllocations(Function&& function) { + AllocationScope scope; + std::forward(function)(); + return scope.count(); +} + +void ExpectAllocationCount(infini::rt::test::TestContext* context, + std::size_t actual, std::size_t expected, + const char* message) { + context->ExpectEqual(actual, expected, message); +} + +} // namespace + +void* operator new(std::size_t size) { + if (void* pointer = std::malloc(size == 0 ? 1 : size)) { + if (count_allocations) { + ++allocation_count; + } + return pointer; + } + throw std::bad_alloc{}; +} + +void* operator new[](std::size_t size) { return ::operator new(size); } + +void operator delete(void* pointer) noexcept { std::free(pointer); } + +void operator delete[](void* pointer) noexcept { ::operator delete(pointer); } + +void operator delete(void* pointer, std::size_t) noexcept { + ::operator delete(pointer); +} + +void operator delete[](void* pointer, std::size_t) noexcept { + ::operator delete[](pointer); +} + +namespace { + +using infini::rt::DataType; +using infini::rt::Device; +using infini::rt::TensorView; + +void TestTensorViewAllocations(infini::rt::test::TestContext* context) { + alignas(float) std::byte data[6 * sizeof(float)]{}; + const Device cpu{Device::Type::kCpu}; + const Device indexed_cpu{Device::Type::kCpu, 1}; + const TensorView::Shape shape{2, 3}; + const TensorView::Strides strides{3, 1}; + + bool shape_only_metadata_is_default = false; + ExpectAllocationCount( + context, CountAllocations([&] { + TensorView tensor{data, TensorView::Shape{2, 3}}; + shape_only_metadata_is_default = tensor.dtype() == DataType::kFloat32 && + tensor.device() == cpu && + tensor.strides() == strides; + }), + 2, "Rvalue shape construction should use default metadata directly."); + context->Expect(shape_only_metadata_is_default, + "Shape-only construction should keep default metadata."); + + bool dtype_only_metadata_is_default = false; + ExpectAllocationCount( + context, CountAllocations([&] { + TensorView tensor{data, TensorView::Shape{2, 3}, DataType::kFloat64}; + dtype_only_metadata_is_default = tensor.dtype() == DataType::kFloat64 && + tensor.device() == cpu && + tensor.strides() == strides; + }), + 2, "Rvalue shape and dtype should use default device and strides."); + context->Expect( + dtype_only_metadata_is_default, + "Shape and dtype construction should keep default device and strides."); + + bool device_only_metadata_is_default = false; + ExpectAllocationCount( + context, CountAllocations([&] { + TensorView tensor{data, TensorView::Shape{2, 3}, indexed_cpu}; + device_only_metadata_is_default = + tensor.dtype() == DataType::kFloat32 && + tensor.device() == indexed_cpu && tensor.strides() == strides; + }), + 2, "Rvalue shape and device should use default dtype and strides."); + context->Expect( + device_only_metadata_is_default, + "Shape and device construction should keep default dtype and strides."); + + ExpectAllocationCount( + context, CountAllocations([&] { + TensorView tensor{data, shape, DataType::kFloat32, cpu, strides}; + (void)tensor; + }), + 2, "Lvalue shape and strides should allocate only their owned copies."); + + ExpectAllocationCount( + context, CountAllocations([&] { + TensorView tensor{data, TensorView::Shape{2, 3}, DataType::kFloat32, + cpu, TensorView::Strides{3, 1}}; + (void)tensor; + }), + 2, "Rvalue shape and strides should transfer their allocations."); + + ExpectAllocationCount( + context, CountAllocations([&] { + TensorView tensor{data, TensorView::Shape{2, 3}, DataType::kFloat32, + cpu}; + (void)tensor; + }), + 2, + "Rvalue shape construction should allocate shape and default strides."); + + ExpectAllocationCount( + context, CountAllocations([&] { + TensorView tensor{data, {2, 3}, DataType::kFloat32, cpu, {3, 1}}; + (void)tensor; + }), + 2, "Initializer lists should construct owned metadata directly."); + + const TensorView source{data, shape, DataType::kFloat32, cpu, strides}; + + ExpectAllocationCount(context, CountAllocations([&] { + TensorView indexed = source[0]; + (void)indexed; + }), + 2, + "Indexing should allocate only the result metadata."); + + ExpectAllocationCount( + context, CountAllocations([&] { + TensorView transposed = source.T(); + (void)transposed; + }), + 2, "Transposing should allocate only the result metadata."); + + ExpectAllocationCount( + context, CountAllocations([&] { + TensorView copied{source}; + (void)copied; + }), + 2, "Copying should allocate one owned shape and one owned stride array."); + + TensorView move_source{data, shape, DataType::kFloat32, cpu, strides}; + ExpectAllocationCount( + context, CountAllocations([&] { + TensorView moved{std::move(move_source)}; + (void)moved; + }), + 0, "Moving should transfer owned metadata without allocating."); +} + +} // namespace + +int main() { + infini::rt::test::TestContext context; + + TestTensorViewAllocations(&context); + + return context.ExitCode(); +}