From 57e197475bff018bc0fbc57ddeaf567e86ef51f9 Mon Sep 17 00:00:00 2001
From: Megvii Engine Team <megengine@megvii.com>
Date: Wed, 17 Nov 2021 12:38:49 +0800
Subject: [PATCH] fix(imperative): check async error when getting value

GitOrigin-RevId: 3945a9bfa2b27e8f47f6ee00184f69f3d992bb22
---
 imperative/python/megengine/functional/vision.py     | 1 +
 imperative/python/src/tensor.cpp                     | 4 ++++
 imperative/src/impl/interpreter/interpreter_impl.cpp | 6 ++++++
 3 files changed, 11 insertions(+)

diff --git a/imperative/python/megengine/functional/vision.py b/imperative/python/megengine/functional/vision.py
index 7ab6bf24..0fd905ee 100644
--- a/imperative/python/megengine/functional/vision.py
+++ b/imperative/python/megengine/functional/vision.py
@@ -420,6 +420,7 @@ def warp_affine(
         Here all available options for params are listed,
         however it does not mean that you can use all the combinations.
         On different platforms, different combinations are supported.
+        ``warp_affine`` only support forward inference, Please refer to ``warp_perspective`` if backward is needed.
     """
     conv_format = _config._get_actual_op_param(format, _config.__conv_format)
 
diff --git a/imperative/python/src/tensor.cpp b/imperative/python/src/tensor.cpp
index 87f2459d..f415f572 100644
--- a/imperative/python/src/tensor.cpp
+++ b/imperative/python/src/tensor.cpp
@@ -1074,6 +1074,10 @@ void init_tensor(py::module m) {
             []() {
                 interpreter_for_py->sync();
                 CompNode::sync_all();
+                CompNode::foreach ([](CompNode cn) {
+                    auto err = cn.check_async_error();
+                    mgb_assert(!err, "%s", err->what());
+                });
                 sync_py_task_q();
             },
             py::call_guard<py::gil_scoped_release>());
diff --git a/imperative/src/impl/interpreter/interpreter_impl.cpp b/imperative/src/impl/interpreter/interpreter_impl.cpp
index bf122982..e8b6fb43 100644
--- a/imperative/src/impl/interpreter/interpreter_impl.cpp
+++ b/imperative/src/impl/interpreter/interpreter_impl.cpp
@@ -156,6 +156,8 @@ TensorInfo* ChannelImpl::put_impl(const HostTensorND& value, bool no_cache) {
     if (m_async_level == 0) {
         sync_impl();
         info->desc.comp_node.sync();
+        auto err = info->desc.comp_node.check_async_error();
+        mgb_assert(!err, "%s", err->what());
     }
     return info;
 }
@@ -336,6 +338,8 @@ void ChannelImpl::dispatch_kernel(
         for (auto&& oup : *outputs) {
             auto info = reinterpret_cast<TensorInfo*>(oup);
             info->ptr->comp_node().sync();
+            auto err = info->ptr->comp_node().check_async_error();
+            mgb_assert(!err, "%s", err->what());
         }
     }
 }
@@ -944,6 +948,8 @@ TensorPtr ChannelImpl::wait_tensor(TensorInfo* info, TensorProp prop) {
     });
     MGB_RECORD_EVENT(TensorWaitPropFinishEvent, info->id, m_waitee_id, prop);
     m_waitee = nullptr;
+    auto err = info->ptr->comp_node().check_async_error();
+    mgb_assert(!err, "%s", err->what());
     return info->ptr;
 }