Update benchmarks to newer API

PiperOrigin-RevId: 356584743
Change-Id: I3c50bc0be4482171533b948b2a44c8eade403578
This commit is contained in:
A. Unique TensorFlower 2021-02-09 13:57:57 -08:00 committed by TensorFlower Gardener
parent 57cdc3ae73
commit 65f038e14d

View File

@ -333,7 +333,7 @@ TEST_P(GPUBFCAllocatorTest, DISABLED_AllocatorReceivesZeroMemory) {
INSTANTIATE_TEST_SUITE_P(GPUBFCAllocatorTestSuite, GPUBFCAllocatorTest, INSTANTIATE_TEST_SUITE_P(GPUBFCAllocatorTestSuite, GPUBFCAllocatorTest,
TestSuiteValues()); TestSuiteValues());
static void BM_Allocation(int iters) { static void BM_Allocation(::testing::benchmark::State& state) {
GPUBFCAllocator a(CreateSubAllocator(1ul << 36), 1uLL << 33, "GPU_0_bfc"); GPUBFCAllocator a(CreateSubAllocator(1ul << 36), 1uLL << 33, "GPU_0_bfc");
// Exercise a few different allocation sizes // Exercise a few different allocation sizes
std::vector<size_t> sizes = {256, 4096, 16384, 524288, std::vector<size_t> sizes = {256, 4096, 16384, 524288,
@ -341,7 +341,7 @@ static void BM_Allocation(int iters) {
1048576000, 2048576000}; 1048576000, 2048576000};
int size_index = 0; int size_index = 0;
while (--iters > 0) { for (auto s : state) {
size_t bytes = sizes[size_index++ % sizes.size()]; size_t bytes = sizes[size_index++ % sizes.size()];
void* p = a.AllocateRaw(1, bytes); void* p = a.AllocateRaw(1, bytes);
a.DeallocateRaw(p); a.DeallocateRaw(p);
@ -349,21 +349,27 @@ static void BM_Allocation(int iters) {
} }
BENCHMARK(BM_Allocation); BENCHMARK(BM_Allocation);
static void BM_AllocationThreaded(int iters, int num_threads) { static void BM_AllocationThreaded(::testing::benchmark::State& state) {
int num_threads = state.range(0);
int sub_iters = 500; // Pick a reasonably large number.
for (auto s : state) {
state.PauseTiming();
GPUBFCAllocator a(CreateSubAllocator(1ul << 36), 1uLL << 33, "GPU_0_bfc"); GPUBFCAllocator a(CreateSubAllocator(1ul << 36), 1uLL << 33, "GPU_0_bfc");
thread::ThreadPool pool(Env::Default(), "test", num_threads); thread::ThreadPool pool(Env::Default(), "test", num_threads);
std::atomic_int_fast32_t count(iters);
std::atomic_int_fast32_t count(sub_iters);
mutex done_lock; mutex done_lock;
condition_variable done; condition_variable done;
bool done_flag = false; bool done_flag = false;
state.ResumeTiming();
for (int t = 0; t < num_threads; t++) { for (int t = 0; t < num_threads; t++) {
pool.Schedule([&a, &count, &done_lock, &done, &done_flag, iters]() { pool.Schedule([&a, &count, &done_lock, &done, &done_flag, sub_iters]() {
// Exercise a few different allocation sizes // Exercise a few different allocation sizes
std::vector<int> sizes = {256, 4096, 16384, 524288, std::vector<int> sizes = {256, 4096, 16384, 524288,
512, 1048576, 10485760, 104857600}; 512, 1048576, 10485760, 104857600};
int size_index = 0; int size_index = 0;
for (int i = 0; i < iters; i++) { for (int i = 0; i < sub_iters; i++) {
int bytes = sizes[size_index++ % sizes.size()]; int bytes = sizes[size_index++ % sizes.size()];
void* p = a.AllocateRaw(1, bytes); void* p = a.AllocateRaw(1, bytes);
a.DeallocateRaw(p); a.DeallocateRaw(p);
@ -376,16 +382,20 @@ static void BM_AllocationThreaded(int iters, int num_threads) {
} }
}); });
} }
mutex_lock l(done_lock); mutex_lock l(done_lock);
if (!done_flag) { if (!done_flag) {
done.wait(l); done.wait(l);
} }
} }
}
BENCHMARK(BM_AllocationThreaded)->Arg(1)->Arg(4)->Arg(16); BENCHMARK(BM_AllocationThreaded)->Arg(1)->Arg(4)->Arg(16);
// A more complex benchmark that defers deallocation of an object for // A more complex benchmark that defers deallocation of an object for
// "delay" allocations. // "delay" allocations.
static void BM_AllocationDelayed(int iters, int delay) { static void BM_AllocationDelayed(::testing::benchmark::State& state) {
int delay = state.range(0);
GPUBFCAllocator a(CreateSubAllocator(1ull << 32), 1 << 30, "GPU_0_bfc"); GPUBFCAllocator a(CreateSubAllocator(1ull << 32), 1 << 30, "GPU_0_bfc");
// Exercise a few different allocation sizes // Exercise a few different allocation sizes
std::vector<int> sizes = {256, 4096, 16384, 4096, 512, 1024, 1024}; std::vector<int> sizes = {256, 4096, 16384, 4096, 512, 1024, 1024};
@ -397,7 +407,7 @@ static void BM_AllocationDelayed(int iters, int delay) {
ptrs.push_back(nullptr); ptrs.push_back(nullptr);
} }
int pindex = 0; int pindex = 0;
while (--iters > 0) { for (auto s : state) {
if (ptrs[pindex] != nullptr) { if (ptrs[pindex] != nullptr) {
a.DeallocateRaw(ptrs[pindex]); a.DeallocateRaw(ptrs[pindex]);
ptrs[pindex] = nullptr; ptrs[pindex] = nullptr;