From d2957d95ad737ec8b8d7b7ad570dddda00745993 Mon Sep 17 00:00:00 2001
From: damnkk <17611288725@163.com>
Date: Tue, 20 Dec 2022 17:37:51 +0800
Subject: [PATCH] [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>
---
 c_api/src/taichi_core_impl.cpp      |   2 +
 c_api/tests/c_api_behavior_test.cpp | 100 ++++++++++++++++++++--------
 taichi/rhi/vulkan/vulkan_api.cpp    |   4 +-
 3 files changed, 76 insertions(+), 30 deletions(-)

diff --git a/c_api/src/taichi_core_impl.cpp b/c_api/src/taichi_core_impl.cpp
index cebd882ff4c38..6589ebef706e3 100644
--- 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
index 1adb2698a5989..204c5b17629b1 100644
--- 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
index 4615590159362..fde2670d83e16 100644
--- 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;