Testable Device-Side Assertion Failures on CPUs (#605)

* add dummy error code * embed into a struct, and change taichi_assert_format * change taichi_assert_runtime (not working now) * move ErrorMessage to LLVMRuntime * let error_code == 1 when assertion fails for further extension * add tests * fixed print_traceback overflow on OS X (#610) * add dummy error code * embed into a struct, and change taichi_assert_format * change taichi_assert_runtime (not working now) * move ErrorMessage to LLVMRuntime * let error_code == 1 when assertion fails for further extension * add tests Co-authored-by: Yuanming Hu <yuanming-hu@users.noreply.github.com>
taichi-dev · Mar 18, 2020 · 97e1b50 · 97e1b50
1 parent a4b557b
commit 97e1b50
Show file tree

Hide file tree

Showing 6 changed files with 108 additions and 18 deletions.
diff --git a/taichi/inc/constants.h b/taichi/inc/constants.h
@@ -9,6 +9,7 @@ constexpr int taichi_max_gpu_block_dim = 1024;
 constexpr std::size_t taichi_global_tmp_buffer_size = 1024 * 1024;
 constexpr int taichi_max_num_mem_requests = 1024 * 64;
 constexpr std::size_t taichi_page_size = 4096;
+constexpr std::size_t taichi_max_message_length = 2048;
 
 template <typename T, typename G>
 T taichi_union_cast_with_different_sizes(G g) {

diff --git a/taichi/program/kernel.cpp b/taichi/program/kernel.cpp
@@ -44,6 +44,9 @@ void Kernel::operator()() {
   }
   compiled(program.get_context());
   program.sync = (program.sync && arch_is_cpu(arch));
+  if (program.config.debug && arch_is_cpu(arch)) {
+    program.check_runtime_error();
+  }
 }
 
 void Kernel::set_arg_float(int i, float64 d) {

diff --git a/taichi/program/program.cpp b/taichi/program/program.cpp
@@ -256,6 +256,23 @@ void Program::materialize_layout() {
   }
 }
 
+void Program::check_runtime_error() {
+  TI_ASSERT(arch_is_cpu(config.arch));
+  auto tlctx = llvm_context_host.get();
+  auto runtime_jit_module = tlctx->runtime_jit_module;
+  runtime_jit_module->call<void *>("retrieve_error_code", llvm_runtime);
+  auto error_code = runtime_jit_module->fetch_result<int64>();
+  if (error_code) {
+    runtime_jit_module->call<void *>("retrieve_error_message", llvm_runtime);
+    auto error_message = runtime_jit_module->fetch_result<char *>();
+    if (error_code == 1) {
+      TI_ERROR("Assertion failure: {}", error_message);
+    } else {
+      TI_NOT_IMPLEMENTED
+    }
+  }
+}
+
 void Program::synchronize() {
   if (!sync) {
     if (config.arch == Arch::cuda) {

diff --git a/taichi/program/program.h b/taichi/program/program.h
@@ -136,6 +136,8 @@ class Program {
 
   void materialize_layout();
 
+  void check_runtime_error();
+
   inline Kernel &get_current_kernel() {
     TI_ASSERT(current_kernel);
     return *current_kernel;

diff --git a/taichi/runtime/llvm/runtime.cpp b/taichi/runtime/llvm/runtime.cpp
@@ -496,8 +496,12 @@ struct LLVMRuntime {
   void (*profiler_start)(Ptr, Ptr);
   void (*profiler_stop)(Ptr);
 
+  char error_message_buffer[taichi_max_message_length];
+  i32 error_message_lock = 0;
+  i64 error_code = 0;
+
   Ptr result_buffer;
-  i32 lock;
+  i32 allocator_lock;
 
   template <typename T>
   void set_result(T t) {
@@ -614,6 +618,14 @@ Ptr get_temporary_pointer(LLVMRuntime *runtime, u64 offset) {
   return runtime->temporaries + offset;
 }
 
+void retrieve_error_code(LLVMRuntime *runtime) {
+  runtime->set_result(runtime->error_code);
+}
+
+void retrieve_error_message(LLVMRuntime *runtime) {
+  runtime->set_result(runtime->error_message_buffer);
+}
+
 #if ARCH_cuda
 void __assertfail(const char *message,
                   const char *file,
@@ -630,31 +642,34 @@ void taichi_assert_runtime(LLVMRuntime *runtime, i32 test, const char *msg) {
 }
 #else
 void taichi_assert_runtime(LLVMRuntime *runtime, i32 test, const char *msg) {
-  if (enable_assert) {
-    if (test == 0) {
-      runtime->assert_failed(msg);
+  if (!enable_assert || test != 0 || runtime->error_code)
+    return;
+  locked_task(&runtime->error_message_lock, [&] {
+    if (!runtime->error_code) {
+      runtime->error_code = 1; // Assertion failure
+      memcpy(runtime->error_message_buffer, msg,
+             std::min(strlen(msg), taichi_max_message_length));
     }
-  }
+  });
 }
 #endif
 
 void taichi_assert(Context *context, i32 test, const char *msg) {
   taichi_assert_runtime(context->runtime, test, msg);
 }
 
-const std::size_t ASSERT_MSG_BUFFER_SIZE = 2048;
-char assert_msg_buffer[ASSERT_MSG_BUFFER_SIZE];
-i32 assert_msg_buffer_lock = 0;
 void taichi_assert_format(LLVMRuntime *runtime, i32 test, const char *format,
                           ...) {
-  if (!enable_assert || test != 0)
+  if (!enable_assert || test != 0 || runtime->error_code)
     return;
   std::va_list args;
   va_start(args, format);
-  locked_task(&assert_msg_buffer_lock, [&] {
-    runtime->host_vsnprintf(assert_msg_buffer, ASSERT_MSG_BUFFER_SIZE, format,
-                            args);
-    runtime->assert_failed(assert_msg_buffer);
+  locked_task(&runtime->error_message_lock, [&] {
+    if (!runtime->error_code) {
+      runtime->error_code = 1; // Assertion failure
+      runtime->host_vsnprintf(runtime->error_message_buffer,
+                              taichi_max_message_length, format, args);
+    }
   });
   va_end(args);
 }
@@ -668,7 +683,7 @@ Ptr LLVMRuntime::allocate_aligned(std::size_t size, std::size_t alignment) {
 
 Ptr LLVMRuntime::allocate_from_buffer(std::size_t size, std::size_t alignment) {
   Ptr ret;
-  locked_task(&lock, [&] {
+  locked_task(&allocator_lock, [&] {
     preallocated_head +=
         alignment - 1 -
         ((std::size_t)preallocated_head + alignment - 1) % alignment;

diff --git a/tests/python/test_assert.py b/tests/python/test_assert.py
@@ -1,12 +1,64 @@
 import taichi as ti
 
-@ti.all_archs
-def test_assert():
-  return
-
+@ti.must_throw(RuntimeError)
+def test_assert_minimal():
+  ti.init(debug=True)
+  ti.set_gdb_trigger(False)
+
+  @ti.kernel
+  def func():
+    assert 0
+
+  func()
+
+
+@ti.must_throw(RuntimeError)
+def test_assert_basic():
+  ti.init(debug=True)
+  ti.set_gdb_trigger(False)
+
   @ti.kernel
   def func():
     x = 20
     assert 10 <= x < 20
 
   func()
+
+
+@ti.all_archs
+def test_assert_ok():
+  ti.init(debug=True)
+  ti.set_gdb_trigger(False)
+
+  @ti.kernel
+  def func():
+    x = 20
+    assert 10 <= x <= 20
+
+  func()
+
+
+@ti.must_throw(RuntimeError)
+def test_out_of_bound():
+  ti.init(debug=True)
+  ti.set_gdb_trigger(False)
+  x = ti.var(ti.i32, shape=(8, 16))
+
+  @ti.kernel
+  def func():
+    x[3, 16] = 1
+
+  func()
+
+
+@ti.all_archs
+def test_not_out_of_bound():
+  ti.init(debug=True)
+  ti.set_gdb_trigger(False)
+  x = ti.var(ti.i32, shape=(8, 16))
+
+  @ti.kernel
+  def func():
+    x[7, 15] = 1
+
+  func()