[cuda] [test] Add TI_DEVICE_MEMORY_GB and TI_DEVICE_MEMORY_FRACTION environment variable

python/taichi/lang/__init__.py

@@ -121,22 +121,31 @@ def init(arch=None,
     for k, v in kwargs.items():
         setattr(ti.cfg, k, v)
 
-    def boolean_config(key, name=None):
-        if name is None:
-            name = 'TI_' + key.upper()
-        value = os.environ.get(name)
-        if value is not None:
-            setattr(ti.cfg, key, len(value) and bool(int(value)))
+    def bool_int(x):
+        return bool(int(x))
+
+    def environ_config(key, cast=bool_int):
+        name = 'TI_' + key.upper()
+        value = os.environ.get(name, '')
+        if len(value):
+            setattr(ti.cfg, key, cast(value))
+
+        # TI_ASYNC=   : not work
+        # TI_ASYNC=0  : False
+        # TI_ASYNC=1  : True
 
     # does override
-    boolean_config("print_ir")
-    boolean_config("verbose")
-    boolean_config("fast_math")
-    boolean_config("async")
-    boolean_config("print_benchmark_stat")
-    gdb_trigger = os.environ.get("TI_GDB_TRIGGER")
-    if gdb_trigger is not None:
-        ti.set_gdb_trigger(len(gdb_trigger) and bool(int(gdb_trigger)))
+    environ_config("print_ir")
+    environ_config("verbose")
+    environ_config("fast_math")
+    environ_config("async")
+    environ_config("print_benchmark_stat")
+    environ_config("device_memory_fraction", float)
+    environ_config("device_memory_GB", float)
+    environ_config("gdb_trigger")
+
+    # Q: Why not environ_config("arch", ti.core.arch_from_name)?
+    # A: We need adaptive_arch_select for all.
     env_arch = os.environ.get("TI_ARCH")
     if env_arch is not None:
         print(f'Following TI_ARCH setting up for arch={env_arch}')

python/taichi/main.py

@@ -41,6 +41,7 @@ def test_python(args):
             threads = min(8, cpu_count())  # To prevent running out of memory
         except:
             threads = 2
+        os.environ['TI_DEVICE_MEMORY_GB'] = '0.5'  # Discussion: #769
         arg_threads = None
         if args.threads is not None:
             arg_threads = int(args.threads)

taichi/program/program.cpp

@@ -171,8 +171,10 @@ void Program::initialize_runtime_system(StructCompiler *scomp) {
     if (config.device_memory_fraction == 0) {
       TI_ASSERT(config.device_memory_GB > 0);
       prealloc_size = std::size_t(config.device_memory_GB * (1UL << 30));
-    } else
+    } else {
       prealloc_size = std::size_t(config.device_memory_fraction * total_mem);
+    }
+    TI_ASSERT(prealloc_size <= total_mem);
 
     TI_TRACE("Allocating device memory {:.2f} GB",
              1.0 * prealloc_size / (1UL << 30));
@@ -485,8 +487,8 @@ void Program::finalize() {
       std::string file_name = current_test;
       auto py_pos = file_name.find(".py::");
       TI_ASSERT(py_pos != file_name.npos);
-      file_name = file_name.substr(0, py_pos) + "__" +
-                  file_name.substr(py_pos + 5);
+      file_name =
+          file_name.substr(0, py_pos) + "__" + file_name.substr(py_pos + 5);
       auto first_space_pos = file_name.find_first_of(' ');
       TI_ASSERT(first_space_pos != file_name.npos);
       file_name = file_name.substr(0, first_space_pos);