diff --git a/tests/cpp/storage/storage_test.cc b/tests/cpp/storage/storage_test.cc index 269480b83c37..41cb5fc22f74 100644 --- a/tests/cpp/storage/storage_test.cc +++ b/tests/cpp/storage/storage_test.cc @@ -1,5 +1,4 @@ -/* - * Licensed to the Apache Software Foundation (ASF) under one +/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file @@ -22,9 +21,11 @@ * \file storage_test.cc * \brief cpu/gpu storage tests */ +#include #include #include #include +#include "storage/pooled_storage_manager.h" #include #include "test_util.h" @@ -60,5 +61,40 @@ TEST(Storage, Basic_GPU) { storage->Free(handle); } } + +TEST(GPUStorage, Round_GPU) { + if (mxnet::test::unitTestsWithCuda) { + putenv("MXNET_GPU_MEM_POOL_ROUND_LINEAR_CUTOFF=20"); + auto &&storage = mxnet::storage::GPUPooledRoundedStorageManager(); + mxnet::Context context_gpu = mxnet::Context::GPU(0); + auto &&handle = mxnet::Storage::Handle(); + auto &&handle2 = mxnet::Storage::Handle(); + handle.ctx = context_gpu; + handle2.ctx = context_gpu; + handle.size = 32; + handle2.size = 2097153; + storage.Alloc(&handle); + storage.Alloc(&handle2); + assert(handle.size == 32); + assert(handle2.size == 2097153); + auto ptr = handle.dptr; + auto ptr2 = handle2.dptr; + storage.Free(handle); + storage.Free(handle2); + + handle.size = 4095; + storage.Alloc(&handle); + EXPECT_EQ(handle.size, 4095); + EXPECT_EQ(handle.dptr, ptr); + storage.Free(handle); + + handle2.size = 3145728; + storage.Alloc(&handle2); + EXPECT_EQ(handle2.size, 3145728); + EXPECT_EQ(handle2.dptr, ptr2); + storage.Free(handle2); + unsetenv("MXNET_GPU_MEM_POOL_ROUND_LINEAR_CUTOFF"); + } +} #endif // MXNET_USE_CUDA