perf: reduce TensorView metadata allocations

src/tensor_view.cc

@@ -16,8 +16,11 @@ static TensorView::Index GetEffectiveIndex(TensorView::Index index,
 TensorView::TensorView(void* data, std::initializer_list<Size> shape,
                        const DataType& dtype, const Device& device,
                        std::initializer_list<Stride> 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 {

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 <typename Shape>
   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 <typename Shape>
   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 <typename Shape>
   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 <typename Shape>
   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 <typename Shape, typename Strides>
   TensorView(void* data, const Shape& shape, const DataType& dtype,
              const Device& device, const Strides& strides)

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)
 

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];

tests/test_tensor_view_allocations.cc

@@ -0,0 +1,191 @@
+#include <infini/rt.h>
+
+#include <cstddef>
+#include <cstdlib>
+#include <new>
+#include <utility>
+
+#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 <typename Function>
+std::size_t CountAllocations(Function&& function) {
+  AllocationScope scope;
+  std::forward<Function>(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();
+}