Function TEST

tensorflow/core/common_runtime/gpu/gpu_vmem_allocator_test.cc:51–81

}

TEST(GPUVMemAllocatorTest, VMemCreate) {
  PlatformGpuId platform_gpu_id(0);
  se::StreamExecutor* se =
      GpuIdUtil::ExecutorForPlatformGpuId(platform_gpu_id).ValueOrDie();
  GPUMemAllocator* sub_allocator = new GPUMemAllocator(
      se, platform_gpu_id, false /*use_unified_memory*/, {}, {});
  Allocator* device_allocator = new GPUBFCAllocator(sub_allocator, 1 << 30, "GPU_0_bfc");

  // VMEM disabled, expect OOM
  setenv("TF_GPU_VMEM", "false", 1);
  Allocator* allocator = maybe_create_gpu_vmem_allocator(device_allocator,
                                                         0, platform_gpu_id, 0, se);
  void* raw_ptr = allocator->AllocateRaw(1, (1 << 30) + 1024);
  CHECK_EQ(raw_ptr, nullptr);

  // VMEM enabled, expect success
  setenv("TF_GPU_VMEM", "true", 1);
  setenv("TF_CUDA_HOST_MEM_LIMIT_IN_MB", "1025", 1);
  allocator = maybe_create_gpu_vmem_allocator(device_allocator,
                                              0, platform_gpu_id, 0, se);
  raw_ptr = allocator->AllocateRaw(1, (1 << 30) + 1024);
  CHECK_NE(raw_ptr, nullptr);

  // VMEM enbaled, but exceeds host memory size, expect OOM
  setenv("TF_GPU_VMEM", "true", 1);
  setenv("TF_CUDA_HOST_MEM_LIMIT_IN_MB", "1024", 1);
  allocator = maybe_create_gpu_vmem_allocator(device_allocator,
                                              0, platform_gpu_id, 0, se);
  raw_ptr = allocator->AllocateRaw(1, (1 << 30) + 1024);
  CHECK_EQ(raw_ptr, nullptr);
}

TEST(GPUVMemAllocatorTest, VMemNoDups) {
  PlatformGpuId platform_gpu_id(0);