diff --git a/backends/aoti/slim/core/Storage.h b/backends/aoti/slim/core/Storage.h index da9bf0638c1..adf0bd0debb 100644 --- a/backends/aoti/slim/core/Storage.h +++ b/backends/aoti/slim/core/Storage.h @@ -205,6 +205,17 @@ class MaybeOwningStorage { } } + /// Constructs non-owning storage with external memory. + /// @param device The device where the data resides. + /// @param data Pointer to external memory (not owned by this storage). + /// @param nbytes Size of the external memory in bytes. + MaybeOwningStorage(const c10::Device& device, void* data, size_t nbytes) + : device_(device), + data_(data), + capacity_(nbytes), + deleter_(detail::noop), + is_owning_(false) {} + /// Default constructor is deleted - storage must have a device. MaybeOwningStorage() = delete; diff --git a/backends/aoti/slim/core/test/test_storage.cpp b/backends/aoti/slim/core/test/test_storage.cpp index 8a5e78ba058..a4e6859a836 100644 --- a/backends/aoti/slim/core/test/test_storage.cpp +++ b/backends/aoti/slim/core/test/test_storage.cpp @@ -72,6 +72,137 @@ TEST(DeviceTraitsCPUTest, MemcpyCPUToCPU) { DeviceTraits::free(dst); } +// ============================================================================= +// MaybeOwningStorage Tests - Non-Owning Mode +// ============================================================================= + +TEST(MaybeOwningStorageNonOwningTest, ConstructNonOwning) { + constexpr size_t kNumFloats = 64; + constexpr size_t kNbytes = kNumFloats * sizeof(float); + + // Allocate external memory + float* external_data = static_cast( + DeviceTraits::allocate(kNbytes)); + + // Initialize external data + for (size_t i = 0; i < kNumFloats; ++i) { + external_data[i] = static_cast(i) * 2.5f; + } + + { + // Create non-owning storage + MaybeOwningStorage storage(CPU_DEVICE, external_data, kNbytes); + + EXPECT_EQ(storage.data(), external_data); + EXPECT_EQ(storage.nbytes(), kNbytes); + EXPECT_TRUE(storage.device().is_cpu()); + EXPECT_FALSE(storage.is_owning()); + EXPECT_FALSE(storage.is_resizable()); + + // Verify data is accessible through storage + float* data = static_cast(storage.data()); + for (size_t i = 0; i < kNumFloats; ++i) { + EXPECT_FLOAT_EQ(data[i], static_cast(i) * 2.5f); + } + } + // After storage goes out of scope, external_data should still be valid + // because the storage did not own it + + // Verify external data is still accessible after storage is destroyed + for (size_t i = 0; i < kNumFloats; ++i) { + EXPECT_FLOAT_EQ(external_data[i], static_cast(i) * 2.5f); + } + + // Clean up external data manually + DeviceTraits::free(external_data); +} + +TEST(MaybeOwningStorageNonOwningTest, ModifyThroughStorage) { + constexpr size_t kNumFloats = 32; + constexpr size_t kNbytes = kNumFloats * sizeof(float); + + // Allocate and initialize external memory + float* external_data = static_cast( + DeviceTraits::allocate(kNbytes)); + for (size_t i = 0; i < kNumFloats; ++i) { + external_data[i] = 0.0f; + } + + { + MaybeOwningStorage storage(CPU_DEVICE, external_data, kNbytes); + + // Modify data through storage + float* data = static_cast(storage.data()); + for (size_t i = 0; i < kNumFloats; ++i) { + data[i] = static_cast(i) * 10.0f; + } + } + + // Verify external data was modified after storage is destroyed + for (size_t i = 0; i < kNumFloats; ++i) { + EXPECT_FLOAT_EQ(external_data[i], static_cast(i) * 10.0f); + } + + DeviceTraits::free(external_data); +} + +TEST(MaybeOwningStorageNonOwningTest, MoveConstruct) { + constexpr size_t kNbytes = 256; + float* external_data = static_cast( + DeviceTraits::allocate(kNbytes)); + + MaybeOwningStorage original(CPU_DEVICE, external_data, kNbytes); + + MaybeOwningStorage moved(std::move(original)); + + EXPECT_EQ(moved.data(), external_data); + EXPECT_EQ(moved.nbytes(), kNbytes); + EXPECT_FALSE(moved.is_owning()); + + EXPECT_EQ(original.data(), nullptr); + EXPECT_EQ(original.nbytes(), 0); + EXPECT_FALSE(original.is_owning()); + + DeviceTraits::free(external_data); +} + +TEST(MaybeOwningStorageNonOwningTest, MoveAssign) { + constexpr size_t kNbytes1 = 256; + constexpr size_t kNbytes2 = 512; + + // Create two external buffers + float* external_data1 = static_cast( + DeviceTraits::allocate(kNbytes1)); + float* external_data2 = static_cast( + DeviceTraits::allocate(kNbytes2)); + + MaybeOwningStorage storage1(CPU_DEVICE, external_data1, kNbytes1); + MaybeOwningStorage storage2(CPU_DEVICE, external_data2, kNbytes2); + + storage1 = std::move(storage2); + + EXPECT_EQ(storage1.data(), external_data2); + EXPECT_EQ(storage1.nbytes(), kNbytes2); + EXPECT_FALSE(storage1.is_owning()); + + EXPECT_EQ(storage2.data(), nullptr); + EXPECT_EQ(storage2.nbytes(), 0); + EXPECT_FALSE(storage2.is_owning()); + + // Clean up both external buffers + DeviceTraits::free(external_data1); + DeviceTraits::free(external_data2); +} + +TEST(MaybeOwningStorageNonOwningTest, ZeroBytes) { + // Non-owning with nullptr and zero bytes + MaybeOwningStorage storage(CPU_DEVICE, nullptr, 0); + + EXPECT_EQ(storage.data(), nullptr); + EXPECT_EQ(storage.nbytes(), 0); + EXPECT_FALSE(storage.is_owning()); +} + // ============================================================================= // MaybeOwningStorage Parameterized Tests (CPU and CUDA) // ============================================================================= @@ -462,6 +593,278 @@ TEST(DeviceTraitsCUDATest, MemcpyCUDAToCUDA) { DeviceTraits::free(cpu_verify); } +// ============================================================================= +// MaybeOwningStorage CUDA Tests +// ============================================================================= + +TEST(MaybeOwningStorageCUDATest, ConstructOwning) { + constexpr size_t kNbytes = 512; + MaybeOwningStorage storage(DEFAULT_CUDA_DEVICE, kNbytes); + + EXPECT_NE(storage.data(), nullptr); + EXPECT_EQ(storage.nbytes(), kNbytes); + EXPECT_TRUE(storage.device().is_cuda()); + EXPECT_FALSE(storage.device().is_cpu()); + EXPECT_TRUE(storage.is_owning()); + EXPECT_TRUE(storage.is_resizable()); + EXPECT_EQ(storage.device().index(), 0); +} + +TEST(MaybeOwningStorageCUDATest, ConstructOwningZeroBytes) { + MaybeOwningStorage storage(DEFAULT_CUDA_DEVICE, 0); + + EXPECT_EQ(storage.data(), nullptr); + EXPECT_EQ(storage.nbytes(), 0); + EXPECT_TRUE(storage.device().is_cuda()); + EXPECT_TRUE(storage.is_owning()); +} + +TEST(MaybeOwningStorageCUDATest, MoveConstruct) { + constexpr size_t kNbytes = 256; + MaybeOwningStorage original(DEFAULT_CUDA_DEVICE, kNbytes); + void* original_data = original.data(); + + MaybeOwningStorage moved(std::move(original)); + + EXPECT_EQ(moved.data(), original_data); + EXPECT_EQ(moved.nbytes(), kNbytes); + EXPECT_TRUE(moved.is_owning()); + EXPECT_TRUE(moved.device().is_cuda()); + + EXPECT_EQ(original.data(), nullptr); + EXPECT_EQ(original.nbytes(), 0); + EXPECT_FALSE(original.is_owning()); +} + +TEST(MaybeOwningStorageCUDATest, MoveAssign) { + constexpr size_t kNbytes1 = 256; + constexpr size_t kNbytes2 = 512; + MaybeOwningStorage storage1(DEFAULT_CUDA_DEVICE, kNbytes1); + MaybeOwningStorage storage2(DEFAULT_CUDA_DEVICE, kNbytes2); + void* storage2_data = storage2.data(); + + storage1 = std::move(storage2); + + EXPECT_EQ(storage1.data(), storage2_data); + EXPECT_EQ(storage1.nbytes(), kNbytes2); + EXPECT_TRUE(storage1.is_owning()); + + EXPECT_EQ(storage2.data(), nullptr); + EXPECT_EQ(storage2.nbytes(), 0); + EXPECT_FALSE(storage2.is_owning()); +} + +// ============================================================================= +// MaybeOwningStorage CUDA Tests - Non-Owning Mode +// ============================================================================= + +TEST(MaybeOwningStorageCUDANonOwningTest, ConstructNonOwning) { + constexpr size_t kNumFloats = 64; + constexpr size_t kNbytes = kNumFloats * sizeof(float); + + // Allocate external CUDA memory + float* external_data = + static_cast(DeviceTraits::allocate( + kNbytes, DEFAULT_CUDA_DEVICE)); + + // Initialize external data via CPU buffer + float* cpu_buffer = static_cast( + DeviceTraits::allocate(kNbytes)); + for (size_t i = 0; i < kNumFloats; ++i) { + cpu_buffer[i] = static_cast(i) * 2.5f; + } + DeviceTraits::memcpy( + external_data, cpu_buffer, kNbytes, DEFAULT_CUDA_DEVICE, CPU_DEVICE); + + { + // Create non-owning storage + MaybeOwningStorage storage(DEFAULT_CUDA_DEVICE, external_data, kNbytes); + + EXPECT_EQ(storage.data(), external_data); + EXPECT_EQ(storage.nbytes(), kNbytes); + EXPECT_TRUE(storage.device().is_cuda()); + EXPECT_FALSE(storage.is_owning()); + EXPECT_FALSE(storage.is_resizable()); + + // Verify data is accessible through storage by copying back to CPU + float* verify_buffer = static_cast( + DeviceTraits::allocate(kNbytes)); + DeviceTraits::memcpy( + verify_buffer, + storage.data(), + kNbytes, + CPU_DEVICE, + DEFAULT_CUDA_DEVICE); + for (size_t i = 0; i < kNumFloats; ++i) { + EXPECT_FLOAT_EQ(verify_buffer[i], static_cast(i) * 2.5f); + } + DeviceTraits::free(verify_buffer); + } + // After storage goes out of scope, external_data should still be valid + // because the storage did not own it + + // Verify external data is still accessible after storage is destroyed + float* verify_buffer2 = static_cast( + DeviceTraits::allocate(kNbytes)); + DeviceTraits::memcpy( + verify_buffer2, external_data, kNbytes, CPU_DEVICE, DEFAULT_CUDA_DEVICE); + for (size_t i = 0; i < kNumFloats; ++i) { + EXPECT_FLOAT_EQ(verify_buffer2[i], static_cast(i) * 2.5f); + } + + // Clean up + DeviceTraits::free(verify_buffer2); + DeviceTraits::free(cpu_buffer); + DeviceTraits::free(external_data); +} + +TEST(MaybeOwningStorageCUDANonOwningTest, ModifyThroughStorage) { + constexpr size_t kNumFloats = 32; + constexpr size_t kNbytes = kNumFloats * sizeof(float); + + // Allocate external CUDA memory + float* external_data = + static_cast(DeviceTraits::allocate( + kNbytes, DEFAULT_CUDA_DEVICE)); + + // Initialize to zeros + float* cpu_buffer = static_cast( + DeviceTraits::allocate(kNbytes)); + for (size_t i = 0; i < kNumFloats; ++i) { + cpu_buffer[i] = 0.0f; + } + DeviceTraits::memcpy( + external_data, cpu_buffer, kNbytes, DEFAULT_CUDA_DEVICE, CPU_DEVICE); + + { + MaybeOwningStorage storage(DEFAULT_CUDA_DEVICE, external_data, kNbytes); + + // Modify data through storage by copying new data + for (size_t i = 0; i < kNumFloats; ++i) { + cpu_buffer[i] = static_cast(i) * 10.0f; + } + DeviceTraits::memcpy( + storage.data(), cpu_buffer, kNbytes, DEFAULT_CUDA_DEVICE, CPU_DEVICE); + } + + // Verify external data was modified after storage is destroyed + float* verify_buffer = static_cast( + DeviceTraits::allocate(kNbytes)); + DeviceTraits::memcpy( + verify_buffer, external_data, kNbytes, CPU_DEVICE, DEFAULT_CUDA_DEVICE); + for (size_t i = 0; i < kNumFloats; ++i) { + EXPECT_FLOAT_EQ(verify_buffer[i], static_cast(i) * 10.0f); + } + + // Clean up + DeviceTraits::free(verify_buffer); + DeviceTraits::free(cpu_buffer); + DeviceTraits::free(external_data); +} + +TEST(MaybeOwningStorageCUDANonOwningTest, MoveConstruct) { + constexpr size_t kNbytes = 256; + float* external_data = + static_cast(DeviceTraits::allocate( + kNbytes, DEFAULT_CUDA_DEVICE)); + + MaybeOwningStorage original(DEFAULT_CUDA_DEVICE, external_data, kNbytes); + + MaybeOwningStorage moved(std::move(original)); + + EXPECT_EQ(moved.data(), external_data); + EXPECT_EQ(moved.nbytes(), kNbytes); + EXPECT_FALSE(moved.is_owning()); + EXPECT_TRUE(moved.device().is_cuda()); + + EXPECT_EQ(original.data(), nullptr); + EXPECT_EQ(original.nbytes(), 0); + EXPECT_FALSE(original.is_owning()); + + DeviceTraits::free(external_data); +} + +TEST(MaybeOwningStorageCUDANonOwningTest, MoveAssign) { + constexpr size_t kNbytes1 = 256; + constexpr size_t kNbytes2 = 512; + + // Create two external CUDA buffers + float* external_data1 = + static_cast(DeviceTraits::allocate( + kNbytes1, DEFAULT_CUDA_DEVICE)); + float* external_data2 = + static_cast(DeviceTraits::allocate( + kNbytes2, DEFAULT_CUDA_DEVICE)); + + MaybeOwningStorage storage1(DEFAULT_CUDA_DEVICE, external_data1, kNbytes1); + MaybeOwningStorage storage2(DEFAULT_CUDA_DEVICE, external_data2, kNbytes2); + + storage1 = std::move(storage2); + + EXPECT_EQ(storage1.data(), external_data2); + EXPECT_EQ(storage1.nbytes(), kNbytes2); + EXPECT_FALSE(storage1.is_owning()); + EXPECT_TRUE(storage1.device().is_cuda()); + + EXPECT_EQ(storage2.data(), nullptr); + EXPECT_EQ(storage2.nbytes(), 0); + EXPECT_FALSE(storage2.is_owning()); + + // Clean up both external buffers + DeviceTraits::free(external_data1); + DeviceTraits::free(external_data2); +} + +TEST(MaybeOwningStorageCUDANonOwningTest, ZeroBytes) { + // Non-owning with nullptr and zero bytes + MaybeOwningStorage storage(DEFAULT_CUDA_DEVICE, nullptr, 0); + + EXPECT_EQ(storage.data(), nullptr); + EXPECT_EQ(storage.nbytes(), 0); + EXPECT_FALSE(storage.is_owning()); + EXPECT_TRUE(storage.device().is_cuda()); +} + +// ============================================================================= +// Storage (SharedPtr) CUDA Tests +// ============================================================================= + +TEST(StorageSharedPtrCUDATest, BasicUsage) { + constexpr size_t kNbytes = 128; + Storage storage(new MaybeOwningStorage(DEFAULT_CUDA_DEVICE, kNbytes)); + + EXPECT_NE(storage.get(), nullptr); + EXPECT_NE(storage->data(), nullptr); + EXPECT_EQ(storage->nbytes(), kNbytes); + EXPECT_TRUE(storage->device().is_cuda()); + EXPECT_EQ(storage.use_count(), 1); +} + +TEST(StorageSharedPtrCUDATest, SharedOwnership) { + constexpr size_t kNbytes = 128; + Storage storage1(new MaybeOwningStorage(DEFAULT_CUDA_DEVICE, kNbytes)); + void* data_ptr = storage1->data(); + + Storage storage2 = storage1; + + EXPECT_EQ(storage1.use_count(), 2); + EXPECT_EQ(storage2.use_count(), 2); + EXPECT_EQ(storage1->data(), storage2->data()); + EXPECT_EQ(storage2->data(), data_ptr); +} + +TEST(StorageSharedPtrCUDATest, MoveSemantics) { + constexpr size_t kNbytes = 64; + Storage storage1(new MaybeOwningStorage(DEFAULT_CUDA_DEVICE, kNbytes)); + void* data_ptr = storage1->data(); + + Storage storage2 = std::move(storage1); + + EXPECT_EQ(storage1.get(), nullptr); + EXPECT_EQ(storage2->data(), data_ptr); + EXPECT_EQ(storage2.use_count(), 1); +} + #endif // CUDA_AVAILABLE } // namespace executorch::backends::aoti::slim