fix(imperative): fix imperative exit segment fault

GitOrigin-RevId: dc498e3634a11aae4e3b3af69c0f398a55b3b0dc

CMakeLists.txt

@@ -1089,6 +1089,11 @@ if(MGE_WITH_JIT AND MGE_WITH_HALIDE)
   include(cmake/Halide.cmake)
 endif()
 
+if(MGE_ENABLE_EXCEPTIONS)
+  add_compile_definitions(CPP_REDIS_ENABLE_EXCEPTION=1)
+else()
+  add_compile_definitions(CPP_REDIS_ENABLE_EXCEPTION=0)
+endif()
 include(cmake/cpp_redis.cmake)
 
 # Thread
@@ -1133,6 +1138,7 @@ if(MSVC OR WIN32)
 endif()
 
 if(CMAKE_THREAD_LIBS_INIT
+   OR CMAKE_HAVE_LIBC_PTHREAD
    OR CMAKE_USE_WIN32_THREADS_INIT
    OR ANDROID)
   set(MGB_HAVE_THREAD 1)

dnn/test/aarch64/relayout.cpp

@@ -12,7 +12,7 @@ namespace test {
 namespace {
 template <typename tag>
 class AARCH64_RELAYOUT : public AARCH64 {};
-TYPED_TEST_CASE(AARCH64_RELAYOUT, relayout::test_types);
+TYPED_TEST_SUITE(AARCH64_RELAYOUT, relayout::test_types);
 TYPED_TEST(AARCH64_RELAYOUT, run) {
     relayout::run_test<TypeParam>(this->handle());
 }

dnn/test/arm_common/elemwise.cpp

@@ -12,14 +12,14 @@ using namespace test;
 
 template <typename tag>
 class ARM_ELEMWISE : public ARM_COMMON {};
-TYPED_TEST_CASE(ARM_ELEMWISE, elemwise::test_types);
+TYPED_TEST_SUITE(ARM_ELEMWISE, elemwise::test_types);
 TYPED_TEST(ARM_ELEMWISE, run) {
     elemwise::run_test<TypeParam>(this->handle());
 }
 
 template <typename tag>
 class ARM_ELEMWISE_MULTI_THREADS : public ARM_COMMON_MULTI_THREADS {};
-TYPED_TEST_CASE(ARM_ELEMWISE_MULTI_THREADS, elemwise::test_types);
+TYPED_TEST_SUITE(ARM_ELEMWISE_MULTI_THREADS, elemwise::test_types);
 TYPED_TEST(ARM_ELEMWISE_MULTI_THREADS, run) {
     elemwise::run_test<TypeParam>(this->handle());
 }

dnn/test/arm_common/elemwise_multi_type.cpp

@@ -14,7 +14,7 @@ using namespace test;
 namespace {
 template <typename tag>
 class ARM_COMMON_ELEMWISE_MULTI_TYPE : public ARM_COMMON {};
-TYPED_TEST_CASE(ARM_COMMON_ELEMWISE_MULTI_TYPE, elemwise_multi_type::test_types);
+TYPED_TEST_SUITE(ARM_COMMON_ELEMWISE_MULTI_TYPE, elemwise_multi_type::test_types);
 }  // anonymous namespace
 
 TYPED_TEST(ARM_COMMON_ELEMWISE_MULTI_TYPE, run) {

dnn/test/common/small_vector.cpp

@@ -231,7 +231,7 @@ typedef ::testing::Types<
         SmallVector<Constructable, 2>, SmallVector<Constructable, 4>,
         SmallVector<Constructable, 5>>
         SmallVectorTestTypes;
-TYPED_TEST_CASE(SmallVectorTest, SmallVectorTestTypes);
+TYPED_TEST_SUITE(SmallVectorTest, SmallVectorTestTypes);
 // Constructor test.
 TYPED_TEST(SmallVectorTest, ConstructorNonIterTest) {
     SCOPED_TRACE("ConstructorTest");
@@ -836,7 +836,7 @@ typedef ::testing::Types<
         std::pair<SmallVector<Constructable, 2>, SmallVector<Constructable, 2>>>
         DualSmallVectorTestTypes;
 
-TYPED_TEST_CASE(DualSmallVectorsTest, DualSmallVectorTestTypes);
+TYPED_TEST_SUITE(DualSmallVectorsTest, DualSmallVectorTestTypes);
 
 TYPED_TEST(DualSmallVectorsTest, MoveAssignment) {
     SCOPED_TRACE("MoveAssignTest-DualVectorTypes");

dnn/test/cpu/relayout.cpp

@@ -12,7 +12,7 @@ using namespace test;
 namespace {
 template <typename tag>
 class CPU_RELAYOUT : public CPU {};
-TYPED_TEST_CASE(CPU_RELAYOUT, relayout::test_types);
+TYPED_TEST_SUITE(CPU_RELAYOUT, relayout::test_types);
 TYPED_TEST(CPU_RELAYOUT, run) {
     relayout::run_test<TypeParam>(this->handle());
 }

dnn/test/cuda/elemwise.cpp

@@ -71,7 +71,7 @@ void run_tensor_add(
 
 template <typename tag>
 class CUDA_ELEMWISE : public CUDA {};
-TYPED_TEST_CASE(CUDA_ELEMWISE, elemwise::test_types);
+TYPED_TEST_SUITE(CUDA_ELEMWISE, elemwise::test_types);
 TYPED_TEST(CUDA_ELEMWISE, run) {
     elemwise::run_test<TypeParam>(this->handle_cuda());
 }

dnn/test/cuda/elemwise_multi_type.cpp

@@ -13,7 +13,7 @@ using namespace test;
 namespace {
 template <typename tag>
 class CUDA_ELEMWISE_MULTI_TYPE : public CUDA {};
-TYPED_TEST_CASE(CUDA_ELEMWISE_MULTI_TYPE, elemwise_multi_type::test_types);
+TYPED_TEST_SUITE(CUDA_ELEMWISE_MULTI_TYPE, elemwise_multi_type::test_types);
 }  // anonymous namespace
 
 TYPED_TEST(CUDA_ELEMWISE_MULTI_TYPE, run) {

dnn/test/cuda/relayout.cpp

@@ -11,7 +11,7 @@ using namespace test;
 namespace {
 template <typename tag>
 class CUDA_RELAYOUT : public CUDA {};
-TYPED_TEST_CASE(CUDA_RELAYOUT, relayout::test_types);
+TYPED_TEST_SUITE(CUDA_RELAYOUT, relayout::test_types);
 TYPED_TEST(CUDA_RELAYOUT, run) {
     relayout::run_test<TypeParam>(this->handle_cuda());
 }

dnn/test/fallback/elemwise.cpp

@@ -11,7 +11,7 @@ using namespace test;
 
 template <typename tag>
 class FALLBACK_ELEMWISE : public FALLBACK {};
-TYPED_TEST_CASE(FALLBACK_ELEMWISE, elemwise::test_types);
+TYPED_TEST_SUITE(FALLBACK_ELEMWISE, elemwise::test_types);
 TYPED_TEST(FALLBACK_ELEMWISE, run) {
     elemwise::run_test<TypeParam>(this->handle());
 }

dnn/test/fallback/elemwise_multi_type.cpp

@@ -8,7 +8,7 @@ using namespace test;
 namespace {
 template <typename tag>
 class FALLBACK_ELEMWISE_MULTI_TYPE : public FALLBACK {};
-TYPED_TEST_CASE(FALLBACK_ELEMWISE_MULTI_TYPE, elemwise_multi_type::test_types);
+TYPED_TEST_SUITE(FALLBACK_ELEMWISE_MULTI_TYPE, elemwise_multi_type::test_types);
 }  // anonymous namespace
 
 TYPED_TEST(FALLBACK_ELEMWISE_MULTI_TYPE, run) {

dnn/test/fallback/relayout.cpp

@@ -14,7 +14,7 @@ using namespace test;
 namespace {
 template <typename tag>
 class FALLBACK_RELAYOUT : public FALLBACK {};
-TYPED_TEST_CASE(FALLBACK_RELAYOUT, relayout::test_types);
+TYPED_TEST_SUITE(FALLBACK_RELAYOUT, relayout::test_types);
 TYPED_TEST(FALLBACK_RELAYOUT, run) {
     relayout::run_test<TypeParam>(this->handle());
 }

dnn/test/rocm/elemwise.cpp

@@ -52,7 +52,7 @@ void run_elemwise_benchmark(
 
 template <typename tag>
 class ROCM_ELEMWISE : public ROCM {};
-TYPED_TEST_CASE(ROCM_ELEMWISE, elemwise::test_types);
+TYPED_TEST_SUITE(ROCM_ELEMWISE, elemwise::test_types);
 TYPED_TEST(ROCM_ELEMWISE, run) {
     elemwise::run_test<TypeParam>(this->handle_rocm());
 }

dnn/test/rocm/relayout.cpp

@@ -13,7 +13,7 @@ using namespace test;
 namespace {
 template <typename tag>
 class ROCM_RELAYOUT : public ROCM {};
-TYPED_TEST_CASE(ROCM_RELAYOUT, relayout::test_types);
+TYPED_TEST_SUITE(ROCM_RELAYOUT, relayout::test_types);
 TYPED_TEST(ROCM_RELAYOUT, run) {
     relayout::run_test<TypeParam>(this->handle_rocm());
 }

dnn/test/x86/elemwise.cpp

@@ -212,7 +212,7 @@ TEST_F(X86, ELEMWISE_FORWARD_TERNARY) {
 
 template <typename tag>
 class X86_ELEMWISE : public X86 {};
-TYPED_TEST_CASE(X86_ELEMWISE, elemwise::test_types);
+TYPED_TEST_SUITE(X86_ELEMWISE, elemwise::test_types);
 TYPED_TEST(X86_ELEMWISE, run) {
     elemwise::run_test<TypeParam>(this->handle());
 }

dnn/test/x86/elemwise_multi_type.cpp

@@ -11,7 +11,7 @@ using namespace test;
 namespace {
 template <typename tag>
 class X86_ELEMWISE_MULTI_TYPE : public X86 {};
-TYPED_TEST_CASE(X86_ELEMWISE_MULTI_TYPE, elemwise_multi_type::test_types);
+TYPED_TEST_SUITE(X86_ELEMWISE_MULTI_TYPE, elemwise_multi_type::test_types);
 }  // anonymous namespace
 
 TYPED_TEST(X86_ELEMWISE_MULTI_TYPE, run) {

imperative/python/src/tensor.cpp

@@ -994,6 +994,15 @@ void init_tensor(py::module m) {
         sync_py_task_q();
     });
     m.def("close", [channel]() {
+        // sync channel and compnode before close to ensure all tasks have been completed
+        if (channel->check_available()) {
+            channel->sync();
+        }
+        CompNode::sync_all();
+        CompNode::foreach ([](CompNode cn) {
+            auto err = cn.check_async_error();
+            mgb_assert(!err, "%s", err->what());
+        });
         channel->close();
         sync_py_task_q();
     });

imperative/python/test/unit/core/test_imperative_rt.py

@@ -71,3 +71,30 @@ def test_opdef_path():
     assert Mode.__module__ == "megengine.core._imperative_rt.ops"
     assert Mode.__name__ == "Mode"
     assert Mode.__qualname__ == "Elemwise.Mode"
+
+
+def _exit_impl():
+    import numpy as np
+    import megengine
+    from megengine import functional as F
+
+    megengine.set_default_device("cpu0")
+    in_channel = 32
+    out_channel = 32
+    x = megengine.tensor(np.random.randn(32, in_channel, 224, 224).astype(np.float32))
+    w = megengine.tensor(
+        np.random.randn(out_channel, in_channel, 3, 3).astype(np.float32)
+    )
+    y = F.conv2d(x, w)
+
+
+def test_imperative_exit():
+    import multiprocessing as mp
+
+    recover = mp.get_start_method()
+    mp.set_start_method("spawn", force=True)
+    pro = mp.Process(target=_exit_impl)
+    pro.start()
+    pro.join()
+    assert pro.exitcode == 0, f"{pro.exitcode}"
+    mp.set_start_method(recover, force=True)

imperative/src/include/megbrain/imperative/utils/platform.h

 #pragma once
 
 #include <string>
-
+#include <typeinfo>
 namespace mgb::imperative {
 
 std::string demangle(std::string mangled);

src/core/test/comp_node.cpp

@@ -789,7 +789,7 @@ TEST(TestCompNode, DepedentObjectList) {
 namespace {
 template <typename tag>
 class TestCPUCompSeqRec : public ::testing::Test {};
-TYPED_TEST_CASE(TestCPUCompSeqRec, comp_node_test::seq_rec::test_types);
+TYPED_TEST_SUITE(TestCPUCompSeqRec, comp_node_test::seq_rec::test_types);
 TYPED_TEST(TestCPUCompSeqRec, run) {
     comp_node_test::seq_rec::run<TypeParam>(CompNode::load("cpux"));
 }

src/core/test/static_mem_alloc.cpp

@@ -124,7 +124,7 @@ TEST_P(BasicCorrectness, OverwriteSameEnd) {
     ASSERT_EQ(align(1 + padding()), allocator->tot_alloc_lower_bound());
 }
 
-INSTANTIATE_TEST_CASE_P(
+INSTANTIATE_TEST_SUITE_P(
         TestStaticMemAllocAlgo, BasicCorrectness,
         TestParam::make_values({1, 2}, {1, 2}, {1}));
 
@@ -207,7 +207,7 @@ TEST_P(RandomOpr, Main) {
             timer.get_secs(), sz_tot, sz_lower, double(sz_tot) / sz_lower - 1);
 }
 
-INSTANTIATE_TEST_CASE_P(
+INSTANTIATE_TEST_SUITE_P(
         TestStaticMemAllocAlgo, RandomOpr,
         TestParam::make_values(
                 {1, 256}, {1, 32}, {10, INTERVAL_MOVE_MAX_SIZE, 1000, 10000}));

src/gopt/test/basic_arith.cpp

@@ -107,11 +107,13 @@ TEST(TestGoptBasicArithInplace, Absorbing) {
     ASSERT_EQ(y.as_immutable_scalar()->get_cast<float>(), 0.f);
 }
 
+namespace {
 auto gen_postive = [](HostTensorND& dest) {
     HostTensorGenerator<dtype::Float32, RandomDistribution::UNIFORM> mask_generator{
             2.f, 4.f};
     dest = *mask_generator(dest.shape(), dest.comp_node());
 };
+};
 
 TEST(TestGoptBasicArithInplace, LogExpExpand) {
     // test log(exp(a) * (exp(b) / (exp(c) * d**2))) -> a + b - c - log(d**2)

src/jit/test/codegen.cpp

@@ -290,7 +290,7 @@ void run_mlir_mode(CompNode cn) {
 #if MGB_JIT_HALIDE
 template <typename tag>
 class TestJITHalideCodeGenCuda : public ::testing::Test {};
-TYPED_TEST_CASE(TestJITHalideCodeGenCuda, test_types);
+TYPED_TEST_SUITE(TestJITHalideCodeGenCuda, test_types);
 TYPED_TEST(TestJITHalideCodeGenCuda, run) {
     REQUIRE_GPU(1);
     run<TypeParam>(Backend::HALIDE, CompNode::load("gpu0"));
@@ -301,7 +301,7 @@ TYPED_TEST(TestJITHalideCodeGenCuda, run) {
 
 template <typename tag>
 class TestJITNvrtcCodeGen : public ::testing::Test {};
-TYPED_TEST_CASE(TestJITNvrtcCodeGen, test_types);
+TYPED_TEST_SUITE(TestJITNvrtcCodeGen, test_types);
 TYPED_TEST(TestJITNvrtcCodeGen, run) {
     REQUIRE_GPU(1);
     run<TypeParam>(Backend::NVRTC, CompNode::load("gpu0"));
@@ -344,7 +344,7 @@ FOREACH_UNARY_MODE(def_tag)
 #define t(n) n,
 using mlir_elemwise_unary_types = ::testing::Types<FOREACH_UNARY_MODE(t) ABS>;
 #undef t
-TYPED_TEST_CASE(TestJITMlirUnaryElemwise, mlir_elemwise_unary_types);
+TYPED_TEST_SUITE(TestJITMlirUnaryElemwise, mlir_elemwise_unary_types);
 
 #define SKIP_MODE(_mode)                                 \
     if (TypeParam::mode == opr::Elemwise::Mode::_mode) { \
@@ -390,7 +390,7 @@ FOREACH_BINARY_MODE(def_tag)
 #define t(n) n,
 using mlir_elemwise_binary_types = ::testing::Types<FOREACH_BINARY_MODE(t) ADD>;
 #undef t
-TYPED_TEST_CASE(TestJITMlirBinaryElemwise, mlir_elemwise_binary_types);
+TYPED_TEST_SUITE(TestJITMlirBinaryElemwise, mlir_elemwise_binary_types);
 TYPED_TEST(TestJITMlirBinaryElemwise, run) {
     auto cn = CompNode::load("cpu0");
     run_mlir_mode<TypeParam, 2>(cn);
@@ -422,7 +422,7 @@ FOREACH_TERNARY_MODE(def_tag)
 using mlir_elemwise_ternary_types =
         ::testing::Types<FOREACH_TERNARY_MODE(t) COND_LEQ_MOV>;
 #undef t
-TYPED_TEST_CASE(TestJITMlirTernaryElemwise, mlir_elemwise_ternary_types);
+TYPED_TEST_SUITE(TestJITMlirTernaryElemwise, mlir_elemwise_ternary_types);
 TYPED_TEST(TestJITMlirTernaryElemwise, run) {
     auto cn = CompNode::load("cpu0");
     run_mlir_mode<TypeParam, 3>(cn);

src/jit/test/fusion.cpp

@@ -1147,7 +1147,7 @@ TEST(TestJITFusionHalide, SameGradOpr) {
 
 template <typename tag>
 class TestJITHalideFusionCuda : public ::testing::Test {};
-TYPED_TEST_CASE(TestJITHalideFusionCuda, test_types);
+TYPED_TEST_SUITE(TestJITHalideFusionCuda, test_types);
 TYPED_TEST(TestJITHalideFusionCuda, run) {
     set_backend(Backend::NONE);
 
@@ -1160,7 +1160,7 @@ TYPED_TEST(TestJITHalideFusionCuda, run) {
 
 template <typename tag>
 class TestJITNvrtcFusion : public ::testing::Test {};
-TYPED_TEST_CASE(TestJITNvrtcFusion, test_types);
+TYPED_TEST_SUITE(TestJITNvrtcFusion, test_types);
 TYPED_TEST(TestJITNvrtcFusion, run) {
     set_backend(Backend::NONE);
 

src/opr/test/basic_arith/elemwise.cpp

@@ -812,7 +812,7 @@ typedef ::testing::Types<
         void  // extra void to consume last comma
         >
         UnaryTraitTypes;
-TYPED_TEST_CASE(TestOprBasicArithUnaryElemwise, UnaryTraitTypes);
+TYPED_TEST_SUITE(TestOprBasicArithUnaryElemwise, UnaryTraitTypes);
 
 typedef ::testing::Types<
 #define DEF_TRAIT(_mode, _expr) _mode,
@@ -821,7 +821,7 @@ typedef ::testing::Types<
         void  // extra void to consume last comma
         >
         BinaryTraitTypes;
-TYPED_TEST_CASE(TestOprBasicArithBinaryElemwise, BinaryTraitTypes);
+TYPED_TEST_SUITE(TestOprBasicArithBinaryElemwise, BinaryTraitTypes);
 
 typedef ::testing::Types<
 #define DEF_TRAIT(_mode, _expr) _mode,
@@ -830,7 +830,7 @@ typedef ::testing::Types<
         void  // extra void to consume last comma
         >
         TernaryTraitTypes;
-TYPED_TEST_CASE(TestOprBasicArithTernaryElemwise, TernaryTraitTypes);
+TYPED_TEST_SUITE(TestOprBasicArithTernaryElemwise, TernaryTraitTypes);
 
 }  // anonymous namespace