[aot] C_API behavior test (#6904)

Issue: # ### Brief Summary Added run out of memory test case. Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
taichi-dev · Dec 20, 2022 · d2957d9 · d2957d9
1 parent d085075
commit d2957d9
Show file tree

Hide file tree

Showing 3 changed files with 76 additions and 30 deletions.
diff --git a/c_api/src/taichi_core_impl.cpp b/c_api/src/taichi_core_impl.cpp
@@ -81,6 +81,8 @@ const char *describe_error(TiError error) {
       return "invalid state";
     case TI_ERROR_INCOMPATIBLE_MODULE:
       return "incompatible module";
+    case TI_ERROR_OUT_OF_MEMORY:
+      return "out of memory";
     default:
       return "unknown error";
   }

diff --git a/c_api/tests/c_api_behavior_test.cpp b/c_api/tests/c_api_behavior_test.cpp
@@ -102,39 +102,81 @@ TEST_F(CapiTest, TestBehaviorGetRuntimeCapabilities) {
 
 TEST_F(CapiTest, TestBehaviorAllocateMemory) {
   auto inner = [&](TiArch arch) {
-    if (ti::is_arch_available(arch)) {
-      // Attempt to allocate memory with size of 1024
-      TiRuntime runtime = ti_create_runtime(arch);
-      for (int i = 0; i < 4; ++i) {
-        TiMemoryAllocateInfo allocate_info;
-        allocate_info.size = 1024;
-        allocate_info.usage = TI_MEMORY_USAGE_STORAGE_BIT << i;
-        TiMemory memory = ti_allocate_memory(runtime, &allocate_info);
-        TI_ASSERT(memory != TI_NULL_HANDLE);
-        ti_free_memory(runtime, memory);
-      }
+    if (!ti::is_arch_available(arch)) {
+      TI_WARN("arch {} is not supported, so this test is skipped", arch);
+      return;
+    }
+    // Attempt to allocate memory with size of 1024
+    TiRuntime runtime = ti_create_runtime(arch);
+    TiMemoryAllocateInfo allocateInfo;
+    allocateInfo.size = 1024;
+    allocateInfo.usage = TI_MEMORY_USAGE_STORAGE_BIT;
+    TiMemory memory = ti_allocate_memory(runtime, &allocateInfo);
+    TI_ASSERT(memory != TI_NULL_HANDLE);
+    ti_free_memory(runtime, memory);
+    ti_destroy_runtime(runtime);
+  };
+  inner(TI_ARCH_VULKAN);
+}
 
-      // runtime and allocate_info are both null
-      {
-        ti_allocate_memory(TI_NULL_HANDLE, nullptr);
-        EXPECT_TAICHI_ERROR(TI_ERROR_ARGUMENT_NULL);
-      }
+TEST_F(CapiTest, TestBehaviorAllocInvalidMemory) {
+  auto inner = [&](TiArch arch) {
+    if (!ti::is_arch_available(arch)) {
+      TI_WARN("arch {} is not supported, so this test is skipped", arch);
+      return;
+    }
+    // Attemp to run out of memory
+    TiRuntime runtime = ti_create_runtime(arch);
+    TiMemoryAllocateInfo allocateInfo;
+    allocateInfo.size = 1000000000000000000;
+    allocateInfo.usage = TI_MEMORY_USAGE_STORAGE_BIT;
+    TiMemory memory = ti_allocate_memory(runtime, &allocateInfo);
+    EXPECT_TAICHI_ERROR(TI_ERROR_OUT_OF_MEMORY);
+    TI_ASSERT(memory == TI_NULL_HANDLE);
+    ti_destroy_runtime(runtime);
+  };
+  inner(TI_ARCH_VULKAN);
+}
 
-      // runtime is not null, allocate_info is null
-      {
-        ti_allocate_memory(runtime, nullptr);
-        EXPECT_TAICHI_ERROR(TI_ERROR_ARGUMENT_NULL);
-      }
+TEST_F(CapiTest, TestBehaviorAllocMemoryNoArg) {
+  auto inner = [&](TiArch arch) {
+    if (!ti::is_arch_available(arch)) {
+      TI_WARN("arch {} is not supported, so this test is skipped", arch);
+      return;
+    }
+    // runtime and allocate_info are both null
+    ti_allocate_memory(TI_NULL_HANDLE, nullptr);
+    EXPECT_TAICHI_ERROR(TI_ERROR_ARGUMENT_NULL);
+  };
+  inner(TI_ARCH_VULKAN);
+}
 
-      // runtime is null, allocate is not null;
-      {
-        TiMemoryAllocateInfo allocate_info;
-        allocate_info.size = 1024;
-        ti_allocate_memory(TI_NULL_HANDLE, &allocate_info);
-        EXPECT_TAICHI_ERROR(TI_ERROR_ARGUMENT_NULL);
-      }
-      ti_destroy_runtime(runtime);
+TEST_F(CapiTest, TestBehaviorAllocMemoryNoAllocInfo) {
+  auto inner = [&](TiArch arch) {
+    if (!ti::is_arch_available(arch)) {
+      TI_WARN("arch {} is not supported, so this test is skipped", arch);
+      return;
     }
+    // runtime is not null, allocate_info is null
+    TiRuntime runtime = ti_create_runtime(arch);
+    ti_allocate_memory(runtime, nullptr);
+    EXPECT_TAICHI_ERROR(TI_ERROR_ARGUMENT_NULL);
+    ti_destroy_runtime(runtime);
+  };
+  inner(TI_ARCH_VULKAN);
+}
+
+TEST_F(CapiTest, TestBehaviorAllocMemoryNoRuntime) {
+  auto inner = [&](TiArch arch) {
+    if (!ti::is_arch_available(arch)) {
+      TI_WARN("arch {} is not supported, so this test is skipped", arch);
+      return;
+    }
+    // runtime is null, allocate is not null;
+    TiMemoryAllocateInfo allocateInfo;
+    allocateInfo.size = 1024;
+    ti_allocate_memory(TI_NULL_HANDLE, &allocateInfo);
+    EXPECT_TAICHI_ERROR(TI_ERROR_ARGUMENT_NULL);
   };
   inner(TI_ARCH_VULKAN);
 }

diff --git a/taichi/rhi/vulkan/vulkan_api.cpp b/taichi/rhi/vulkan/vulkan_api.cpp
@@ -547,7 +547,9 @@ IVkBuffer create_buffer(VkDevice device,
 
   VkResult res = vmaCreateBuffer(allocator, buffer_info, alloc_info,
                                  &buffer->buffer, &buffer->allocation, nullptr);
-  assert(res != VK_ERROR_OUT_OF_DEVICE_MEMORY);
+  if (res == VK_ERROR_OUT_OF_DEVICE_MEMORY) {
+    throw std::bad_alloc();
+  }  // FIXME: (damnkk) Should be removed when RHI error codes are ready
   BAIL_ON_VK_BAD_RESULT_NO_RETURN(res, "failed to create buffer");
 
   return buffer;