[CLI] Fix ti test --arch ... (#764)

taichi-dev · Apr 13, 2020 · fbc12ab · fbc12ab
1 parent 936cd97
commit fbc12ab
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 6 deletions.
diff --git a/python/taichi/lang/__init__.py b/python/taichi/lang/__init__.py
@@ -237,12 +237,8 @@ def benchmark(func, repeat=100, args=()):
     return elapsed / repeat
 
 
-wanted_archs = None
-
-
 def set_wanted_archs(archs):
-    global wanted_archs
-    wanted_archs = archs
+    os.environ['TI_WANTED_ARCHS'] = ','.join(archs)
 
 
 def supported_archs():
@@ -254,7 +250,8 @@ def supported_archs():
         archs.append(metal)
     if ti.core.with_opengl():
         archs.append(opengl)
-    if wanted_archs is not None:
+    wanted_archs = os.environ.get('TI_WANTED_ARCHS', '').split(',')
+    if len(wanted_archs):
         archs, old_archs = [], archs
         for arch in old_archs:
             if ti.core.arch_name(arch) in wanted_archs:

diff --git a/python/taichi/main.py b/python/taichi/main.py
@@ -109,6 +109,7 @@ def main(debug=False):
     print()
     import taichi as ti
     if args.arch is not None:
+        print(f'Running on Arch={args.arch}')
         ti.set_wanted_archs(args.arch.split(','))
 
     if mode == 'help':