From aa749dd0e0e41438ec2b70ee3f6428fb79c407e3 Mon Sep 17 00:00:00 2001
From: Dun Liang <randonlang@gmail.com>
Date: Tue, 31 Aug 2021 11:41:41 +0800
Subject: [PATCH] polish exit cleanup

---
 python/jittor/__init__.py                                | 2 +-
 python/jittor/extern/cuda/cublas/ops/cublas_matmul_op.cc | 4 ++--
 python/jittor/extern/cuda/cublas/src/cublas_warper.cc    | 3 ++-
 python/jittor/src/init.cc                                | 2 ++
 python/jittor/src/utils/log.cc                           | 1 +
 5 files changed, 8 insertions(+), 4 deletions(-)

diff --git a/python/jittor/__init__.py b/python/jittor/__init__.py
index 49abca40..d8840faf 100644
--- a/python/jittor/__init__.py
+++ b/python/jittor/__init__.py
@@ -9,7 +9,7 @@
 # file 'LICENSE.txt', which is part of this source code package.
 # ***************************************************************
 
-__version__ = '1.2.3.94'
+__version__ = '1.2.3.95'
 from jittor_utils import lock
 with lock.lock_scope():
     ori_int = int
diff --git a/python/jittor/extern/cuda/cublas/ops/cublas_matmul_op.cc b/python/jittor/extern/cuda/cublas/ops/cublas_matmul_op.cc
index 1bbc42e5..e2228ec5 100644
--- a/python/jittor/extern/cuda/cublas/ops/cublas_matmul_op.cc
+++ b/python/jittor/extern/cuda/cublas/ops/cublas_matmul_op.cc
@@ -20,6 +20,8 @@ namespace jittor {
 
 CublasMatmulOp::CublasMatmulOp(Var* a, Var* b, bool trans_a, bool trans_b)
     : a(a), b(b), trans_a(trans_a), trans_b(trans_b) {
+    flags.set(NodeFlags::_cuda, 1);
+    flags.set(NodeFlags::_cpu, 0);
     // TODO: support int8 * int8
     ASSERT(a->dtype().is_float() && b->dtype().is_float()) << "type of two inputs should be the same";
     // TODO: support diffrent input type
@@ -51,7 +53,6 @@ void CublasMatmulOp::jit_prepare(JK& jk) {
 }
 
 #else // JIT
-#ifdef JIT_cpu
 #pragma clang diagnostic ignored "-Wtautological-compare"
 void CublasMatmulOp::jit_run() {
     cublasHandle_t& handle_ = cublas_handle;
@@ -78,7 +79,6 @@ void CublasMatmulOp::jit_run() {
     a->ptr<T>(), '@Trans_a' == 'N' ? m : n, &beta, 
     c->ptr<T>(), k));
 }
-#endif
 #endif // JIT
 
 } // jittor
diff --git a/python/jittor/extern/cuda/cublas/src/cublas_warper.cc b/python/jittor/extern/cuda/cublas/src/cublas_warper.cc
index f23c9429..0f1b912e 100644
--- a/python/jittor/extern/cuda/cublas/src/cublas_warper.cc
+++ b/python/jittor/extern/cuda/cublas/src/cublas_warper.cc
@@ -19,11 +19,12 @@ struct cublas_initer {
 inline cublas_initer() {
     if (!get_device_count()) return;
     checkCudaErrors(cublasCreate(&cublas_handle));
-    LOGv << "cublasCreate finished";
+    LOGv << "cublasCreate finished" << (void*)cublas_handle;
 }
 
 inline ~cublas_initer() {
     if (!get_device_count()) return;
+    LOGv << "cublasDestroy:" <<  (void*)cublas_handle;
     checkCudaErrors(cublasDestroy(cublas_handle));
     LOGv << "cublasDestroy finished";
 }
diff --git a/python/jittor/src/init.cc b/python/jittor/src/init.cc
index e792269a..82ee53eb 100644
--- a/python/jittor/src/init.cc
+++ b/python/jittor/src/init.cc
@@ -29,8 +29,10 @@ int current_seed;
 // fron fetch_op.cc
 extern list<VarPtr> fetcher;
 extern list<VarPtr> fetcher_to_free;
+extern bool exited;
 
 void cleanup() {
+    exited = true;
     fetcher_to_free.clear();
     fetcher.clear();
 }
diff --git a/python/jittor/src/utils/log.cc b/python/jittor/src/utils/log.cc
index 3e15ab8e..dbac0d9a 100644
--- a/python/jittor/src/utils/log.cc
+++ b/python/jittor/src/utils/log.cc
@@ -235,6 +235,7 @@ void segfault_sigaction(int signal, siginfo_t *si, void *arg) {
         exited = true;
         do_exit();
     }
+    if (exited) do_exit();
     std::cerr << "Caught segfault at address " << si->si_addr << ", "
         << "thread_name: '" << thread_name << "', flush log..." << std::endl;
     std::cerr.flush();