polish error msg

python/jittor/compiler.py

@@ -872,7 +872,7 @@ def check_cache_compile():
     jit_utils_core_files = files
     recompile = compile(cc_path, cc_flags+f" {opt_flags} ", files, jit_utils.cache_path+'/jit_utils_core'+extension_suffix, True)
     if recompile and jit_utils.cc:
-        LOG.e("jit_utils updated, please restart jittor.")
+        LOG.e("jit_utils updated, please rerun your command.")
         sys.exit(0)
     if not jit_utils.cc:
         with jit_utils.import_scope(import_flags):

python/jittor/src/mem/mem_info.cc

@@ -128,14 +128,6 @@ void display_memory_info(const char* fileline, bool dump_var, bool red_color) {
     log << "cpu&gpu:" << FloatOutput{(double)all_total, " KMG", 1024, "B"}
         << "gpu:" << FloatOutput{(double)gpu_total, " KMG", 1024, "B"}
         << "cpu:" << FloatOutput{(double)cpu_total, " KMG", 1024, "B"} >> '\n';
-
-    if (red_color){
-        bool cuda_overflow = (double)gpu_total>(double)mem_info.total_cuda_ram;
-        bool cpu_overflow = (double)cpu_total>(double)mem_info.total_cpu_ram;
-        if(cuda_overflow || cpu_overflow){
-            LOGf<<"CUDA memory or CPU memory is overflow, please reduce your batch_size or data size!";
-        }
-    }
     
     size_t cpu_free = 0;
 #if defined(__linux__)
@@ -194,6 +186,20 @@ void display_memory_info(const char* fileline, bool dump_var, bool red_color) {
         }
     }
     log >> "===========================\n";
+
+    if (red_color) {
+        bool gpu_overflow = (double)gpu_total>(double)mem_info.total_cuda_ram*0.95;
+        bool cpu_overflow = (double)cpu_total>(double)mem_info.total_cpu_ram*0.95;
+        if(gpu_overflow || cpu_overflow) {
+            double used = gpu_overflow ? (double)gpu_total : (double)cpu_total;
+            double total = gpu_overflow ? (double)mem_info.total_cuda_ram : (double)mem_info.total_cpu_ram;
+            log.end();
+            LOGf << "\n*******************\n"
+                >> (gpu_overflow?"GPU":"CPU") << "memory is overflow, please reduce your batch_size or data size!\nTotal:" << FloatOutput{(double)total, " KMG", 1024, "B"} << "Used:" << FloatOutput{(double)used, " KMG", 1024, "B"};
+        } else
+            return;
+    }
+
     log.end();
 }
 

python/jittor/test/test_error_msg.py

@@ -0,0 +1,34 @@
+# ***************************************************************
+# Copyright (c) 2021 Jittor. All Rights Reserved. 
+# Maintainers: 
+#    Dun Liang <randonlang@gmail.com>. 
+# 
+# This file is subject to the terms and conditions defined in
+# file 'LICENSE.txt', which is part of this source code package.
+# ***************************************************************
+import unittest
+import jittor as jt
+import numpy as np
+
+class TestErrorMsg(unittest.TestCase):
+
+    def test_error_msg(self):
+        a = jt.array([3,2,1])
+        b = jt.code(a.shape, a.dtype, [a],
+            cpu_header="""
+                #include <algorithm>
+                @alias(a, in0)
+                @alias(b, out)
+            """,
+            cpu_src="""
+                for (int i=0; i<a_shape0; i++)
+                    @b(i) = @a(i);
+                std::sort(&@b(0), &@b(in0_shape0));
+                throw std::runtime_error("???");
+            """
+        )
+        print(b)
+
+
+if __name__ == "__main__":
+    unittest.main()