From 759c35cb9290aebda62cfd85bc09f3c551196405 Mon Sep 17 00:00:00 2001
From: XiangyunYang <xiangyunyang@taichi.graphics>
Date: Wed, 26 Jan 2022 10:06:19 +0800
Subject: [PATCH 01/10] commit

---
 python/taichi/lang/__init__.py           | 16 ++++------------
 python/taichi/lang/impl.py               |  6 ++++++
 taichi/ir/expr.cpp                       |  2 +-
 tests/python/test_ast_refactor.py        |  7 ++++---
 tests/python/test_clear_all_gradients.py |  6 ++++--
 tests/python/test_div.py                 |  4 +++-
 tests/python/test_internal_func.py       | 14 ++++++++------
 tests/python/test_ndarray.py             | 17 +++++++++--------
 tests/python/test_types.py               |  3 ++-
 9 files changed, 41 insertions(+), 34 deletions(-)

diff --git a/python/taichi/lang/__init__.py b/python/taichi/lang/__init__.py
index 07bb2ed2b3783..e8744d59bdd73 100644
--- a/python/taichi/lang/__init__.py
+++ b/python/taichi/lang/__init__.py
@@ -24,15 +24,9 @@
                                    TaichiSyntaxError, TaichiTypeError)
 from taichi.lang.expr import Expr, make_expr_group
 from taichi.lang.field import Field, ScalarField
-from taichi.lang.impl import (axes, begin_frontend_if,
-                              begin_frontend_struct_for, call_internal,
-                              current_cfg, deactivate_all_snodes, expr_init,
-                              expr_init_func, expr_init_list, field,
-                              get_runtime, grouped,
-                              insert_expr_stmt_if_ti_func, ndarray, one, root,
-                              static, static_assert, static_print, stop_grad,
-                              subscript, ti_assert, ti_float, ti_format,
-                              ti_int, ti_print, zero)
+from taichi.lang.impl import (axes, deactivate_all_snodes, field, grouped,
+                              ndarray, one, root, static, static_assert,
+                              static_print, stop_grad, zero)
 from taichi.lang.kernel_arguments import SparseMatrixProxy
 from taichi.lang.kernel_impl import (KernelArgError, KernelDefError,
                                      data_oriented, func, kernel, pyfunc)
@@ -62,8 +56,6 @@
 
 from taichi import _logging
 
-runtime = impl.get_runtime()
-
 i = axes(0)
 j = axes(1)
 k = axes(2)
@@ -791,7 +783,7 @@ def Tape(loss, clear_gradients=True):
     from taichi._kernels import clear_loss  # pylint: disable=C0415
     clear_loss(loss)
 
-    return runtime.get_tape(loss)
+    return impl.get_runtime().get_tape(loss)
 
 
 def clear_all_gradients():
diff --git a/python/taichi/lang/impl.py b/python/taichi/lang/impl.py
index a7e52a9f96929..837b0879dd9bc 100644
--- a/python/taichi/lang/impl.py
+++ b/python/taichi/lang/impl.py
@@ -904,3 +904,9 @@ def mesh_relation_access(mesh, from_index, to_element_type):
     if isinstance(mesh, MeshInstance):
         return MeshRelationAccessProxy(mesh, from_index, to_element_type)
     raise RuntimeError("Relation access should be with a mesh instance!")
+
+
+__all__ = [
+    'axes', 'deactivate_all_snodes', 'field', 'grouped', 'ndarray', 'one',
+    'root', 'static', 'static_assert', 'static_print', 'stop_grad', 'zero'
+]
diff --git a/taichi/ir/expr.cpp b/taichi/ir/expr.cpp
index 1b80256eba88e..4f3dfa92883ee 100644
--- a/taichi/ir/expr.cpp
+++ b/taichi/ir/expr.cpp
@@ -76,7 +76,7 @@ void Expr::set_or_insert_assignment(const Expr &o) {
       TI_ERROR("Cannot assign to non-lvalue: {}", serialize());
     }
   } else {
-    set(o);  // Literally set this Expr to o
+    set(o);  // Literally set this Exp r to o
   }
 }
 
diff --git a/tests/python/test_ast_refactor.py b/tests/python/test_ast_refactor.py
index 29a0731a6da71..3a43484f4a546 100644
--- a/tests/python/test_ast_refactor.py
+++ b/tests/python/test_ast_refactor.py
@@ -1,5 +1,6 @@
 import numpy as np
 import pytest
+from taichi.lang.impl import subscript, ti_float
 
 import taichi as ti
 from taichi import approx
@@ -12,7 +13,7 @@ def foo(x: ti.i32, y: ti.i32, a: ti.template()):
         a[0] = x + y
         a[1] = x - y
         a[2] = x * y
-        a[3] = ti.ti_float(x) / y
+        a[3] = ti_float(x) / y
         a[4] = x // y
         a[5] = x % y
         a[6] = x**y
@@ -687,7 +688,7 @@ def bar(x: ti.template()):
             return ti.Matrix([[1, 0], [0, 1]])
 
     def fibonacci(x):
-        return ti.subscript(bar(x), 1, 0)
+        return subscript(bar(x), 1, 0)
 
     @ti.kernel
     def foo(x: ti.template()) -> ti.i32:
@@ -779,7 +780,7 @@ def bar(x: tc.template()):
             return tc.Matrix([[1, 0], [0, 1]])
 
     def fibonacci(x):
-        return tc.subscript(bar(x), 1, 0)
+        return subscript(bar(x), 1, 0)
 
     @tc.kernel
     def foo(x: tc.template()) -> tc.i32:
diff --git a/tests/python/test_clear_all_gradients.py b/tests/python/test_clear_all_gradients.py
index 8624d6ec52060..eb278dd801cd5 100644
--- a/tests/python/test_clear_all_gradients.py
+++ b/tests/python/test_clear_all_gradients.py
@@ -1,3 +1,5 @@
+from taichi.lang.impl import get_runtime
+
 import taichi as ti
 
 
@@ -23,7 +25,7 @@ def test_clear_all_gradients():
             w.grad[i, j] = 6
 
     ti.clear_all_gradients()
-    assert ti.get_runtime().get_num_compiled_functions() == 3
+    assert get_runtime().get_num_compiled_functions() == 3
 
     assert x.grad[None] == 0
     for i in range(n):
@@ -34,4 +36,4 @@ def test_clear_all_gradients():
 
     ti.clear_all_gradients()
     # No more kernel compilation
-    assert ti.get_runtime().get_num_compiled_functions() == 3
+    assert get_runtime().get_num_compiled_functions() == 3
diff --git a/tests/python/test_div.py b/tests/python/test_div.py
index 65c8fb7c9750f..4bcaa8633730c 100644
--- a/tests/python/test_div.py
+++ b/tests/python/test_div.py
@@ -1,3 +1,5 @@
+from taichi.lang.impl import get_runtime
+
 import taichi as ti
 
 
@@ -58,7 +60,7 @@ def test_true_div():
 
 @ti.test()
 def test_div_default_ip():
-    ti.get_runtime().set_default_ip(ti.i64)
+    get_runtime().set_default_ip(ti.i64)
     z = ti.field(ti.f32, shape=())
 
     @ti.kernel
diff --git a/tests/python/test_internal_func.py b/tests/python/test_internal_func.py
index 784ca77e945d7..d97581daf249d 100644
--- a/tests/python/test_internal_func.py
+++ b/tests/python/test_internal_func.py
@@ -1,5 +1,7 @@
 import time
 
+from taichi.lang.impl import call_internal
+
 import taichi as ti
 
 
@@ -8,7 +10,7 @@ def test_basic():
     @ti.kernel
     def test():
         for _ in range(10):
-            ti.call_internal("do_nothing")
+            call_internal("do_nothing")
 
     test()
 
@@ -19,7 +21,7 @@ def test_host_polling():
 
     @ti.kernel
     def test():
-        ti.call_internal("refresh_counter")
+        call_internal("refresh_counter")
 
     for i in range(10):
         print('updating tail to', i)
@@ -31,7 +33,7 @@ def test():
 def test_list_manager():
     @ti.kernel
     def test():
-        ti.call_internal("test_list_manager")
+        call_internal("test_list_manager")
 
     test()
     test()
@@ -41,7 +43,7 @@ def test():
 def test_node_manager():
     @ti.kernel
     def test():
-        ti.call_internal("test_node_allocator")
+        call_internal("test_node_allocator")
 
     test()
     test()
@@ -51,7 +53,7 @@ def test():
 def test_node_manager_gc():
     @ti.kernel
     def test_cpu():
-        ti.call_internal("test_node_allocator_gc_cpu")
+        call_internal("test_node_allocator_gc_cpu")
 
     test_cpu()
 
@@ -60,7 +62,7 @@ def test_cpu():
 def test_return():
     @ti.kernel
     def test_cpu():
-        ret = ti.call_internal("test_internal_func_args", 1.0, 2.0, 3)
+        ret = call_internal("test_internal_func_args", 1.0, 2.0, 3)
         assert ret == 9
 
     test_cpu()
diff --git a/tests/python/test_ndarray.py b/tests/python/test_ndarray.py
index 6315a23f23c4f..aaea429a2aef3 100644
--- a/tests/python/test_ndarray.py
+++ b/tests/python/test_ndarray.py
@@ -2,6 +2,7 @@
 
 import numpy as np
 import pytest
+from taichi.lang.impl import get_runtime
 
 import taichi as ti
 
@@ -109,7 +110,7 @@ def test_default_fp_ndarray_torch(dtype):
 
     x = ti.Vector.ndarray(2, float, ())
 
-    assert x.dtype == ti.get_runtime().default_fp
+    assert x.dtype == get_runtime().default_fp
 
 
 @pytest.mark.parametrize('dtype', [ti.f32, ti.f64])
@@ -118,7 +119,7 @@ def test_default_fp_ndarray(dtype):
 
     x = ti.Vector.ndarray(2, float, ())
 
-    assert x.dtype == ti.get_runtime().default_fp
+    assert x.dtype == get_runtime().default_fp
 
 
 @pytest.mark.parametrize('dtype', [ti.i32, ti.i64])
@@ -128,7 +129,7 @@ def test_default_ip_ndarray_torch(dtype):
 
     x = ti.Vector.ndarray(2, int, ())
 
-    assert x.dtype == ti.get_runtime().default_ip
+    assert x.dtype == get_runtime().default_ip
 
 
 @pytest.mark.parametrize('dtype', [ti.i32, ti.i64])
@@ -137,7 +138,7 @@ def test_default_ip_ndarray(dtype):
 
     x = ti.Vector.ndarray(2, int, ())
 
-    assert x.dtype == ti.get_runtime().default_ip
+    assert x.dtype == get_runtime().default_ip
 
 
 # access
@@ -559,17 +560,17 @@ def func(a: ti.any_arr(element_dim=1)):
 
     v = ti.Vector.ndarray(10, ti.i32, 5)
     func(v)
-    assert ti.get_runtime().get_num_compiled_functions() == 1
+    assert get_runtime().get_num_compiled_functions() == 1
     v = np.zeros((6, 10), dtype=np.int32)
     func(v)
-    assert ti.get_runtime().get_num_compiled_functions() == 1
+    assert get_runtime().get_num_compiled_functions() == 1
     import torch
     v = torch.zeros((6, 11), dtype=torch.int32)
     func(v)
-    assert ti.get_runtime().get_num_compiled_functions() == 2
+    assert get_runtime().get_num_compiled_functions() == 2
     v = ti.Vector.ndarray(10, ti.i32, 5, layout=ti.Layout.SOA)
     func(v)
-    assert ti.get_runtime().get_num_compiled_functions() == 3
+    assert get_runtime().get_num_compiled_functions() == 3
 
 
 @pytest.mark.skipif(not ti.has_pytorch(), reason='Pytorch not installed.')
diff --git a/tests/python/test_types.py b/tests/python/test_types.py
index bd90fa43bed60..d5ece425012b3 100644
--- a/tests/python/test_types.py
+++ b/tests/python/test_types.py
@@ -1,4 +1,5 @@
 import pytest
+from taichi.lang.impl import get_runtime
 
 import taichi as ti
 
@@ -138,7 +139,7 @@ def test_overflow64(dt, n):
 @ti.test(require=ti.extension.data64)
 def test_uint_max(dt, val):
     # https://github.com/taichi-dev/taichi/issues/2060
-    ti.get_runtime().default_ip = dt
+    get_runtime().default_ip = dt
     N = 16
     f = ti.field(dt, shape=N)
 

From ab4292297105da0883795a23f202444128dc8cf0 Mon Sep 17 00:00:00 2001
From: Taichi Gardener <taichigardener@gmail.com>
Date: Wed, 26 Jan 2022 02:12:52 +0000
Subject: [PATCH 02/10] Auto Format

---
 tests/python/test_ast_refactor.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tests/python/test_ast_refactor.py b/tests/python/test_ast_refactor.py
index 4ea5cd74bd4ac..c24c1067155d7 100644
--- a/tests/python/test_ast_refactor.py
+++ b/tests/python/test_ast_refactor.py
@@ -1,7 +1,7 @@
 import numpy as np
 import pytest
-from taichi.lang.impl import subscript, ti_float
 from taichi._testing import approx
+from taichi.lang.impl import subscript, ti_float
 
 import taichi as ti
 

From 0783b183bb1f9d0fcd2d62e5e5b2a0ff434cc1be Mon Sep 17 00:00:00 2001
From: XiangyunYang <xiangyunyang@taichi.graphics>
Date: Wed, 26 Jan 2022 11:11:37 +0800
Subject: [PATCH 03/10] fix

---
 docs/lang/articles/misc/internal.md            |  2 +-
 misc/benchmark_rebuild_graph.py                |  3 ++-
 .../taichi/examples/algorithm/print_offset.py  |  3 ++-
 tests/python/test_ad_if.py                     |  4 ++--
 tests/python/test_ast_refactor.py              |  8 ++++----
 tests/python/test_clear_all_gradients.py       |  6 +++---
 tests/python/test_compare.py                   |  3 ++-
 tests/python/test_cuda_internals.py            |  7 ++++---
 tests/python/test_field.py                     |  5 +++--
 tests/python/test_internal_func.py             | 14 +++++++-------
 tests/python/test_matrix.py                    | 11 ++++++-----
 tests/python/test_ndarray.py                   | 18 +++++++++---------
 tests/python/test_tensor_reflection.py         |  4 ++--
 tests/python/test_torch_io.py                  |  7 ++++---
 tests/python/test_types.py                     |  4 ++--
 15 files changed, 53 insertions(+), 46 deletions(-)

diff --git a/docs/lang/articles/misc/internal.md b/docs/lang/articles/misc/internal.md
index 610f93292347a..692b50f26947f 100644
--- a/docs/lang/articles/misc/internal.md
+++ b/docs/lang/articles/misc/internal.md
@@ -330,7 +330,7 @@ a = ti.field(ti.f32, shape=(128, 32, 8))
 b = ti.field(ti.f32)
 ti.root.dense(ti.j, 32).dense(ti.i, 16).place(b)
 
-ti.get_runtime().materialize()
+ti.lang.impl.get_runtime().materialize() # This is an internal api for dev, we don't make sure it is stable for user.
 
 mapping_a = a.snode().physical_index_position()
 
diff --git a/misc/benchmark_rebuild_graph.py b/misc/benchmark_rebuild_graph.py
index 8d43870e2c83e..176d2ef250f15 100644
--- a/misc/benchmark_rebuild_graph.py
+++ b/misc/benchmark_rebuild_graph.py
@@ -1,4 +1,5 @@
 import taichi as ti
+from taichi.lang import impl
 
 ti.init(arch=ti.cuda, async_mode=True)
 
@@ -23,4 +24,4 @@ def foo():
 for i in range(1000):
     foo()
 
-ti.get_runtime().prog.benchmark_rebuild_graph()
+impl.get_runtime().prog.benchmark_rebuild_graph()
diff --git a/python/taichi/examples/algorithm/print_offset.py b/python/taichi/examples/algorithm/print_offset.py
index b905126bd4f9b..3c0a4230d6bba 100644
--- a/python/taichi/examples/algorithm/print_offset.py
+++ b/python/taichi/examples/algorithm/print_offset.py
@@ -1,4 +1,5 @@
 import taichi as ti
+from taichi.lang import impl
 
 ti.init(arch=ti.cpu, print_ir=True)
 
@@ -19,7 +20,7 @@ def fill():
 fill()
 print(a.to_numpy())
 
-ti.get_runtime().prog.visualize_layout('layout.pdf')
+impl.get_runtime().prog.visualize_layout('layout.pdf')
 
 gui = ti.GUI('layout', res=(256, 512), background_color=0xFFFFFF)
 
diff --git a/tests/python/test_ad_if.py b/tests/python/test_ad_if.py
index 0b4e0cae16196..e8b062fb7ac85 100644
--- a/tests/python/test_ad_if.py
+++ b/tests/python/test_ad_if.py
@@ -1,5 +1,5 @@
 import taichi as ti
-
+from taichi.lang import impl
 
 @ti.test(require=ti.extension.adstack)
 def test_ad_if_simple():
@@ -235,6 +235,6 @@ def func():
 def test_stack():
     @ti.kernel
     def func():
-        ti.call_internal("test_stack")
+        impl.call_internal("test_stack")
 
     func()
diff --git a/tests/python/test_ast_refactor.py b/tests/python/test_ast_refactor.py
index 3a43484f4a546..268c956d2bc94 100644
--- a/tests/python/test_ast_refactor.py
+++ b/tests/python/test_ast_refactor.py
@@ -1,6 +1,6 @@
 import numpy as np
 import pytest
-from taichi.lang.impl import subscript, ti_float
+from taichi.lang import impl
 
 import taichi as ti
 from taichi import approx
@@ -13,7 +13,7 @@ def foo(x: ti.i32, y: ti.i32, a: ti.template()):
         a[0] = x + y
         a[1] = x - y
         a[2] = x * y
-        a[3] = ti_float(x) / y
+        a[3] = impl.ti_float(x) / y
         a[4] = x // y
         a[5] = x % y
         a[6] = x**y
@@ -688,7 +688,7 @@ def bar(x: ti.template()):
             return ti.Matrix([[1, 0], [0, 1]])
 
     def fibonacci(x):
-        return subscript(bar(x), 1, 0)
+        return impl.subscript(bar(x), 1, 0)
 
     @ti.kernel
     def foo(x: ti.template()) -> ti.i32:
@@ -780,7 +780,7 @@ def bar(x: tc.template()):
             return tc.Matrix([[1, 0], [0, 1]])
 
     def fibonacci(x):
-        return subscript(bar(x), 1, 0)
+        return impl.subscript(bar(x), 1, 0)
 
     @tc.kernel
     def foo(x: tc.template()) -> tc.i32:
diff --git a/tests/python/test_clear_all_gradients.py b/tests/python/test_clear_all_gradients.py
index eb278dd801cd5..0e0b7cce22b7b 100644
--- a/tests/python/test_clear_all_gradients.py
+++ b/tests/python/test_clear_all_gradients.py
@@ -1,4 +1,4 @@
-from taichi.lang.impl import get_runtime
+from taichi.lang import impl
 
 import taichi as ti
 
@@ -25,7 +25,7 @@ def test_clear_all_gradients():
             w.grad[i, j] = 6
 
     ti.clear_all_gradients()
-    assert get_runtime().get_num_compiled_functions() == 3
+    assert impl.get_runtime().get_num_compiled_functions() == 3
 
     assert x.grad[None] == 0
     for i in range(n):
@@ -36,4 +36,4 @@ def test_clear_all_gradients():
 
     ti.clear_all_gradients()
     # No more kernel compilation
-    assert get_runtime().get_num_compiled_functions() == 3
+    assert impl.get_runtime().get_num_compiled_functions() == 3
diff --git a/tests/python/test_compare.py b/tests/python/test_compare.py
index 6b77b3a861207..441299a74cbeb 100644
--- a/tests/python/test_compare.py
+++ b/tests/python/test_compare.py
@@ -1,6 +1,7 @@
 import pytest
 
 import taichi as ti
+from taichi.lang import impl
 
 
 @ti.test(require=ti.extension.sparse)
@@ -106,7 +107,7 @@ def why_this_foo_fail(n):
         return ti.atomic_add(b[None], n)
 
     def foo(n):
-        return ti.atomic_add(ti.subscript(b, None), n)
+        return ti.atomic_add(impl.subscript(b, None), n)
 
     @ti.kernel
     def func():
diff --git a/tests/python/test_cuda_internals.py b/tests/python/test_cuda_internals.py
index dede77a7656ba..12012643ca7bd 100644
--- a/tests/python/test_cuda_internals.py
+++ b/tests/python/test_cuda_internals.py
@@ -1,4 +1,5 @@
 import taichi as ti
+from taichi.lang import impl
 
 # TODO: these are not really tests...
 
@@ -8,7 +9,7 @@ def test_do_nothing():
     @ti.kernel
     def test():
         for i in range(10):
-            ti.call_internal("do_nothing")
+            impl.call_internal("do_nothing")
 
     test()
 
@@ -19,7 +20,7 @@ def test_active_mask():
     def test():
         for i in range(48):
             if i % 2 == 0:
-                ti.call_internal("test_active_mask")
+                impl.call_internal("test_active_mask")
 
     test()
 
@@ -29,6 +30,6 @@ def test_shfl_down():
     @ti.kernel
     def test():
         for i in range(32):
-            ti.call_internal("test_shfl")
+            impl.call_internal("test_shfl")
 
     test()
diff --git a/tests/python/test_field.py b/tests/python/test_field.py
index 7013e47d65b13..d61091023519f 100644
--- a/tests/python/test_field.py
+++ b/tests/python/test_field.py
@@ -5,6 +5,7 @@
 import pytest
 
 import taichi as ti
+from taichi.lang import impl
 
 data_types = [ti.i32, ti.f32, ti.i64, ti.f64]
 field_shapes = [(), 8, (6, 12)]
@@ -121,7 +122,7 @@ def test_default_fp(dtype):
 
     x = ti.Vector.field(2, float, ())
 
-    assert x.dtype == ti.get_runtime().default_fp
+    assert x.dtype == impl.get_runtime().default_fp
 
 
 @pytest.mark.parametrize('dtype', [ti.i32, ti.i64])
@@ -130,7 +131,7 @@ def test_default_ip(dtype):
 
     x = ti.Vector.field(2, int, ())
 
-    assert x.dtype == ti.get_runtime().default_ip
+    assert x.dtype == impl.get_runtime().default_ip
 
 
 @ti.test()
diff --git a/tests/python/test_internal_func.py b/tests/python/test_internal_func.py
index d97581daf249d..1c2604c3b8a40 100644
--- a/tests/python/test_internal_func.py
+++ b/tests/python/test_internal_func.py
@@ -1,6 +1,6 @@
 import time
 
-from taichi.lang.impl import call_internal
+from taichi.lang import impl
 
 import taichi as ti
 
@@ -10,7 +10,7 @@ def test_basic():
     @ti.kernel
     def test():
         for _ in range(10):
-            call_internal("do_nothing")
+            impl.call_internal("do_nothing")
 
     test()
 
@@ -21,7 +21,7 @@ def test_host_polling():
 
     @ti.kernel
     def test():
-        call_internal("refresh_counter")
+        impl.call_internal("refresh_counter")
 
     for i in range(10):
         print('updating tail to', i)
@@ -33,7 +33,7 @@ def test():
 def test_list_manager():
     @ti.kernel
     def test():
-        call_internal("test_list_manager")
+        impl.call_internal("test_list_manager")
 
     test()
     test()
@@ -43,7 +43,7 @@ def test():
 def test_node_manager():
     @ti.kernel
     def test():
-        call_internal("test_node_allocator")
+        impl.call_internal("test_node_allocator")
 
     test()
     test()
@@ -53,7 +53,7 @@ def test():
 def test_node_manager_gc():
     @ti.kernel
     def test_cpu():
-        call_internal("test_node_allocator_gc_cpu")
+        impl.call_internal("test_node_allocator_gc_cpu")
 
     test_cpu()
 
@@ -62,7 +62,7 @@ def test_cpu():
 def test_return():
     @ti.kernel
     def test_cpu():
-        ret = call_internal("test_internal_func_args", 1.0, 2.0, 3)
+        ret = impl.call_internal("test_internal_func_args", 1.0, 2.0, 3)
         assert ret == 9
 
     test_cpu()
diff --git a/tests/python/test_matrix.py b/tests/python/test_matrix.py
index 6df0bcb4f4cac..a367ecdda4626 100644
--- a/tests/python/test_matrix.py
+++ b/tests/python/test_matrix.py
@@ -4,6 +4,7 @@
 import pytest
 
 import taichi as ti
+from taichi.lang import impl
 from taichi import approx
 
 operation_types = [operator.add, operator.sub, operator.matmul]
@@ -425,7 +426,7 @@ def test_matrix_field_dynamic_index_different_path_length():
     ti.root.dense(ti.i, 8).place(x)
     ti.root.dense(ti.i, 2).dense(ti.i, 4).place(y)
 
-    ti.get_runtime().materialize()
+    impl.get_runtime().materialize()
     assert v.dynamic_index_stride is None
 
 
@@ -438,7 +439,7 @@ def test_matrix_field_dynamic_index_not_pure_dense():
     ti.root.dense(ti.i, 2).pointer(ti.i, 4).place(x)
     ti.root.dense(ti.i, 2).dense(ti.i, 4).place(y)
 
-    ti.get_runtime().materialize()
+    impl.get_runtime().materialize()
     assert v.dynamic_index_stride is None
 
 
@@ -453,7 +454,7 @@ def test_matrix_field_dynamic_index_different_cell_size_bytes():
     ti.root.dense(ti.i, 8).place(x, temp)
     ti.root.dense(ti.i, 8).place(y)
 
-    ti.get_runtime().materialize()
+    impl.get_runtime().materialize()
     assert v.dynamic_index_stride is None
 
 
@@ -469,7 +470,7 @@ def test_matrix_field_dynamic_index_different_offset_bytes_in_parent_cell():
     ti.root.dense(ti.i, 8).place(temp_a, x)
     ti.root.dense(ti.i, 8).place(y, temp_b)
 
-    ti.get_runtime().materialize()
+    impl.get_runtime().materialize()
     assert v.dynamic_index_stride is None
 
 
@@ -484,7 +485,7 @@ def test_matrix_field_dynamic_index_different_stride():
 
     ti.root.dense(ti.i, 8).place(x, y, temp, z)
 
-    ti.get_runtime().materialize()
+    impl.get_runtime().materialize()
     assert v.dynamic_index_stride is None
 
 
diff --git a/tests/python/test_ndarray.py b/tests/python/test_ndarray.py
index aaea429a2aef3..77622df55988a 100644
--- a/tests/python/test_ndarray.py
+++ b/tests/python/test_ndarray.py
@@ -2,7 +2,7 @@
 
 import numpy as np
 import pytest
-from taichi.lang.impl import get_runtime
+from taichi.lang import impl
 
 import taichi as ti
 
@@ -110,7 +110,7 @@ def test_default_fp_ndarray_torch(dtype):
 
     x = ti.Vector.ndarray(2, float, ())
 
-    assert x.dtype == get_runtime().default_fp
+    assert x.dtype == impl.get_runtime().default_fp
 
 
 @pytest.mark.parametrize('dtype', [ti.f32, ti.f64])
@@ -119,7 +119,7 @@ def test_default_fp_ndarray(dtype):
 
     x = ti.Vector.ndarray(2, float, ())
 
-    assert x.dtype == get_runtime().default_fp
+    assert x.dtype == impl.get_runtime().default_fp
 
 
 @pytest.mark.parametrize('dtype', [ti.i32, ti.i64])
@@ -129,7 +129,7 @@ def test_default_ip_ndarray_torch(dtype):
 
     x = ti.Vector.ndarray(2, int, ())
 
-    assert x.dtype == get_runtime().default_ip
+    assert x.dtype == impl.get_runtime().default_ip
 
 
 @pytest.mark.parametrize('dtype', [ti.i32, ti.i64])
@@ -138,7 +138,7 @@ def test_default_ip_ndarray(dtype):
 
     x = ti.Vector.ndarray(2, int, ())
 
-    assert x.dtype == get_runtime().default_ip
+    assert x.dtype == impl.get_runtime().default_ip
 
 
 # access
@@ -560,17 +560,17 @@ def func(a: ti.any_arr(element_dim=1)):
 
     v = ti.Vector.ndarray(10, ti.i32, 5)
     func(v)
-    assert get_runtime().get_num_compiled_functions() == 1
+    assert impl.get_runtime().get_num_compiled_functions() == 1
     v = np.zeros((6, 10), dtype=np.int32)
     func(v)
-    assert get_runtime().get_num_compiled_functions() == 1
+    assert impl.get_runtime().get_num_compiled_functions() == 1
     import torch
     v = torch.zeros((6, 11), dtype=torch.int32)
     func(v)
-    assert get_runtime().get_num_compiled_functions() == 2
+    assert impl.get_runtime().get_num_compiled_functions() == 2
     v = ti.Vector.ndarray(10, ti.i32, 5, layout=ti.Layout.SOA)
     func(v)
-    assert get_runtime().get_num_compiled_functions() == 3
+    assert impl.get_runtime().get_num_compiled_functions() == 3
 
 
 @pytest.mark.skipif(not ti.has_pytorch(), reason='Pytorch not installed.')
diff --git a/tests/python/test_tensor_reflection.py b/tests/python/test_tensor_reflection.py
index b9462fd877799..637b198f504e4 100644
--- a/tests/python/test_tensor_reflection.py
+++ b/tests/python/test_tensor_reflection.py
@@ -1,7 +1,7 @@
 import pytest
 
 import taichi as ti
-
+from taichi.lang import impl
 
 @ti.test()
 def test_POT():
@@ -59,7 +59,7 @@ def test_unordered():
     assert val.snode in blk3.get_children()
     assert blk3 in blk2.get_children()
     assert blk2 in blk1.get_children()
-    ti.get_runtime().materialize_root_fb(False)
+    impl.get_runtime().materialize_root_fb(False)
     assert blk1 in ti.FieldsBuilder.finalized_roots()[0].get_children()
 
     expected_str = f'ti.root => dense {[n]} => dense {[m, n]}' \
diff --git a/tests/python/test_torch_io.py b/tests/python/test_torch_io.py
index 75d2bdb6a812b..9c4dc37620a97 100644
--- a/tests/python/test_torch_io.py
+++ b/tests/python/test_torch_io.py
@@ -2,6 +2,7 @@
 import pytest
 
 import taichi as ti
+from taichi.lang import impl
 
 if ti.has_pytorch():
     import torch
@@ -200,11 +201,11 @@ def test_io_struct():
 def test_fused_kernels():
     n = 12
     X = ti.Matrix.field(3, 2, ti.f32, shape=(n, n, n))
-    s = ti.get_runtime().get_num_compiled_functions()
+    s = impl.get_runtime().get_num_compiled_functions()
     t = X.to_torch()
-    assert ti.get_runtime().get_num_compiled_functions() == s + 1
+    assert impl.get_runtime().get_num_compiled_functions() == s + 1
     X.from_torch(t)
-    assert ti.get_runtime().get_num_compiled_functions() == s + 2
+    assert impl.get_runtime().get_num_compiled_functions() == s + 2
 
 
 @pytest.mark.skipif(not ti.has_pytorch(), reason='Pytorch not installed.')
diff --git a/tests/python/test_types.py b/tests/python/test_types.py
index d5ece425012b3..713800b89b770 100644
--- a/tests/python/test_types.py
+++ b/tests/python/test_types.py
@@ -1,5 +1,5 @@
 import pytest
-from taichi.lang.impl import get_runtime
+from taichi.lang import impl
 
 import taichi as ti
 
@@ -139,7 +139,7 @@ def test_overflow64(dt, n):
 @ti.test(require=ti.extension.data64)
 def test_uint_max(dt, val):
     # https://github.com/taichi-dev/taichi/issues/2060
-    get_runtime().default_ip = dt
+    impl.get_runtime().default_ip = dt
     N = 16
     f = ti.field(dt, shape=N)
 

From 29d131d352d350e0fac5dbbde5d050c3f7a5b491 Mon Sep 17 00:00:00 2001
From: XiangyunYang <xiangyunyang@taichi.graphics>
Date: Wed, 26 Jan 2022 11:15:19 +0800
Subject: [PATCH 04/10] fix

---
 tests/python/test_ast_refactor.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tests/python/test_ast_refactor.py b/tests/python/test_ast_refactor.py
index af342550f99b6..0f71c9f92276d 100644
--- a/tests/python/test_ast_refactor.py
+++ b/tests/python/test_ast_refactor.py
@@ -2,7 +2,7 @@
 import pytest
 
 import taichi as ti
-from taichi import approx
+from taichi._testing import approx
 from taichi.lang import impl
 
 @ti.test()

From d57a4a4ebe790e0eb01ccf553cf3e582fe218194 Mon Sep 17 00:00:00 2001
From: XiangyunYang <xiangyunyang@taichi.graphics>
Date: Wed, 26 Jan 2022 11:16:00 +0800
Subject: [PATCH 05/10] format

---
 misc/benchmark_rebuild_graph.py                  | 3 ++-
 python/taichi/examples/algorithm/print_offset.py | 3 ++-
 tests/python/test_ad_if.py                       | 4 +++-
 tests/python/test_ast_refactor.py                | 5 +++--
 tests/python/test_compare.py                     | 2 +-
 tests/python/test_cuda_internals.py              | 3 ++-
 tests/python/test_field.py                       | 2 +-
 tests/python/test_matrix.py                      | 4 ++--
 tests/python/test_tensor_reflection.py           | 3 ++-
 tests/python/test_torch_io.py                    | 2 +-
 10 files changed, 19 insertions(+), 12 deletions(-)

diff --git a/misc/benchmark_rebuild_graph.py b/misc/benchmark_rebuild_graph.py
index 176d2ef250f15..416605fa0411e 100644
--- a/misc/benchmark_rebuild_graph.py
+++ b/misc/benchmark_rebuild_graph.py
@@ -1,6 +1,7 @@
-import taichi as ti
 from taichi.lang import impl
 
+import taichi as ti
+
 ti.init(arch=ti.cuda, async_mode=True)
 
 a = ti.field(dtype=ti.i32, shape=())
diff --git a/python/taichi/examples/algorithm/print_offset.py b/python/taichi/examples/algorithm/print_offset.py
index 3c0a4230d6bba..d731a7175b8e7 100644
--- a/python/taichi/examples/algorithm/print_offset.py
+++ b/python/taichi/examples/algorithm/print_offset.py
@@ -1,6 +1,7 @@
-import taichi as ti
 from taichi.lang import impl
 
+import taichi as ti
+
 ti.init(arch=ti.cpu, print_ir=True)
 
 n = 4
diff --git a/tests/python/test_ad_if.py b/tests/python/test_ad_if.py
index e8b062fb7ac85..7602b7de8df86 100644
--- a/tests/python/test_ad_if.py
+++ b/tests/python/test_ad_if.py
@@ -1,6 +1,8 @@
-import taichi as ti
 from taichi.lang import impl
 
+import taichi as ti
+
+
 @ti.test(require=ti.extension.adstack)
 def test_ad_if_simple():
     x = ti.field(ti.f32, shape=())
diff --git a/tests/python/test_ast_refactor.py b/tests/python/test_ast_refactor.py
index 0f71c9f92276d..8b2db4911fea5 100644
--- a/tests/python/test_ast_refactor.py
+++ b/tests/python/test_ast_refactor.py
@@ -1,10 +1,11 @@
 import numpy as np
 import pytest
-
-import taichi as ti
 from taichi._testing import approx
 from taichi.lang import impl
 
+import taichi as ti
+
+
 @ti.test()
 def test_binop():
     @ti.kernel
diff --git a/tests/python/test_compare.py b/tests/python/test_compare.py
index 441299a74cbeb..16d24d77a8aec 100644
--- a/tests/python/test_compare.py
+++ b/tests/python/test_compare.py
@@ -1,7 +1,7 @@
 import pytest
+from taichi.lang import impl
 
 import taichi as ti
-from taichi.lang import impl
 
 
 @ti.test(require=ti.extension.sparse)
diff --git a/tests/python/test_cuda_internals.py b/tests/python/test_cuda_internals.py
index 12012643ca7bd..99c69002e82b6 100644
--- a/tests/python/test_cuda_internals.py
+++ b/tests/python/test_cuda_internals.py
@@ -1,6 +1,7 @@
-import taichi as ti
 from taichi.lang import impl
 
+import taichi as ti
+
 # TODO: these are not really tests...
 
 
diff --git a/tests/python/test_field.py b/tests/python/test_field.py
index d61091023519f..d65d2720d9d97 100644
--- a/tests/python/test_field.py
+++ b/tests/python/test_field.py
@@ -3,9 +3,9 @@
 '''
 
 import pytest
+from taichi.lang import impl
 
 import taichi as ti
-from taichi.lang import impl
 
 data_types = [ti.i32, ti.f32, ti.i64, ti.f64]
 field_shapes = [(), 8, (6, 12)]
diff --git a/tests/python/test_matrix.py b/tests/python/test_matrix.py
index a477dcf44aeb1..fc0c7a2886dc6 100644
--- a/tests/python/test_matrix.py
+++ b/tests/python/test_matrix.py
@@ -3,10 +3,10 @@
 
 import numpy as np
 import pytest
+from taichi._testing import approx
+from taichi.lang import impl
 
 import taichi as ti
-from taichi.lang import impl
-from taichi._testing import approx
 
 operation_types = [operator.add, operator.sub, operator.matmul]
 test_matrix_arrays = [
diff --git a/tests/python/test_tensor_reflection.py b/tests/python/test_tensor_reflection.py
index 637b198f504e4..28bb31339ee6a 100644
--- a/tests/python/test_tensor_reflection.py
+++ b/tests/python/test_tensor_reflection.py
@@ -1,7 +1,8 @@
 import pytest
+from taichi.lang import impl
 
 import taichi as ti
-from taichi.lang import impl
+
 
 @ti.test()
 def test_POT():
diff --git a/tests/python/test_torch_io.py b/tests/python/test_torch_io.py
index 9c4dc37620a97..fd4763a8b4f45 100644
--- a/tests/python/test_torch_io.py
+++ b/tests/python/test_torch_io.py
@@ -1,8 +1,8 @@
 import numpy as np
 import pytest
+from taichi.lang import impl
 
 import taichi as ti
-from taichi.lang import impl
 
 if ti.has_pytorch():
     import torch

From 7ab79fe1f76afc0b73b5857cf1ad70e79e5d8d5a Mon Sep 17 00:00:00 2001
From: Xiangyun Yang <xiangyunyang@taichi.graphics>
Date: Wed, 26 Jan 2022 11:17:10 +0800
Subject: [PATCH 06/10] Update expr.cpp

---
 taichi/ir/expr.cpp | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/taichi/ir/expr.cpp b/taichi/ir/expr.cpp
index 4f3dfa92883ee..1b80256eba88e 100644
--- a/taichi/ir/expr.cpp
+++ b/taichi/ir/expr.cpp
@@ -76,7 +76,7 @@ void Expr::set_or_insert_assignment(const Expr &o) {
       TI_ERROR("Cannot assign to non-lvalue: {}", serialize());
     }
   } else {
-    set(o);  // Literally set this Exp r to o
+    set(o);  // Literally set this Expr to o
   }
 }
 

From 22e1f3e5713c2ff1ce3be6edc163a20456533364 Mon Sep 17 00:00:00 2001
From: XiangyunYang <xiangyunyang@taichi.graphics>
Date: Wed, 26 Jan 2022 11:18:21 +0800
Subject: [PATCH 07/10] format

---
 tests/python/test_ast_refactor.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/tests/python/test_ast_refactor.py b/tests/python/test_ast_refactor.py
index 8b2db4911fea5..5d1230bdcc2af 100644
--- a/tests/python/test_ast_refactor.py
+++ b/tests/python/test_ast_refactor.py
@@ -13,7 +13,7 @@ def foo(x: ti.i32, y: ti.i32, a: ti.template()):
         a[0] = x + y
         a[1] = x - y
         a[2] = x * y
-        a[3] = ti_float(x) / y
+        a[3] = impl.ti_float(x) / y
         a[4] = x // y
         a[5] = x % y
         a[6] = x**y
@@ -688,7 +688,7 @@ def bar(x: ti.template()):
             return ti.Matrix([[1, 0], [0, 1]])
 
     def fibonacci(x):
-        return subscript(bar(x), 1, 0)
+        return impl.subscript(bar(x), 1, 0)
 
     @ti.kernel
     def foo(x: ti.template()) -> ti.i32:
@@ -780,7 +780,7 @@ def bar(x: tc.template()):
             return tc.Matrix([[1, 0], [0, 1]])
 
     def fibonacci(x):
-        return subscript(bar(x), 1, 0)
+        return impl.subscript(bar(x), 1, 0)
 
     @tc.kernel
     def foo(x: tc.template()) -> tc.i32:

From cc0140969e766b0c8112ffab876e52442244f009 Mon Sep 17 00:00:00 2001
From: XiangyunYang <xiangyunyang@taichi.graphics>
Date: Wed, 26 Jan 2022 11:19:19 +0800
Subject: [PATCH 08/10] fix

---
 tests/python/test_div.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/tests/python/test_div.py b/tests/python/test_div.py
index 4bcaa8633730c..9f87cbd211210 100644
--- a/tests/python/test_div.py
+++ b/tests/python/test_div.py
@@ -1,4 +1,4 @@
-from taichi.lang.impl import get_runtime
+from taichi.lang import impl
 
 import taichi as ti
 
@@ -60,7 +60,7 @@ def test_true_div():
 
 @ti.test()
 def test_div_default_ip():
-    get_runtime().set_default_ip(ti.i64)
+    impl.get_runtime().set_default_ip(ti.i64)
     z = ti.field(ti.f32, shape=())
 
     @ti.kernel

From 4e14ca964b9fd87f60e874957a268f82caf16351 Mon Sep 17 00:00:00 2001
From: XiangyunYang <xiangyunyang@taichi.graphics>
Date: Wed, 26 Jan 2022 11:22:11 +0800
Subject: [PATCH 09/10] fix

---
 tests/python/test_matrix.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tests/python/test_matrix.py b/tests/python/test_matrix.py
index fc0c7a2886dc6..d54f649335c23 100644
--- a/tests/python/test_matrix.py
+++ b/tests/python/test_matrix.py
@@ -75,7 +75,7 @@ def test_python_scope_matrix_field(ops):
 
 @ti.test(arch=ti.get_host_arch_list())
 def test_constant_matrices():
-    assert ti.cos(ti.math.pi / 3) == approx(0.5)
+    assert ti.cos(math.pi / 3) == approx(0.5)
     assert np.allclose((-ti.Vector([2, 3])).to_numpy(), np.array([-2, -3]))
     assert ti.cos(ti.Vector([2,
                              3])).to_numpy() == approx(np.cos(np.array([2,

From de13e54ee0898f814f6d0813a03d44e9d82ab672 Mon Sep 17 00:00:00 2001
From: XiangyunYang <xiangyunyang@taichi.graphics>
Date: Wed, 26 Jan 2022 14:14:51 +0800
Subject: [PATCH 10/10] fix

---
 tests/python/test_matrix.py | 10 +++++-----
 1 file changed, 5 insertions(+), 5 deletions(-)

diff --git a/tests/python/test_matrix.py b/tests/python/test_matrix.py
index d54f649335c23..3bc99cfce5182 100644
--- a/tests/python/test_matrix.py
+++ b/tests/python/test_matrix.py
@@ -427,7 +427,7 @@ def test_matrix_field_dynamic_index_different_path_length():
     ti.root.dense(ti.i, 8).place(x)
     ti.root.dense(ti.i, 2).dense(ti.i, 4).place(y)
 
-    ti.get_runtime().materialize()
+    impl.get_runtime().materialize()
     assert v.dynamic_index_stride is None
 
 
@@ -440,7 +440,7 @@ def test_matrix_field_dynamic_index_not_pure_dense():
     ti.root.dense(ti.i, 2).pointer(ti.i, 4).place(x)
     ti.root.dense(ti.i, 2).dense(ti.i, 4).place(y)
 
-    ti.get_runtime().materialize()
+    impl.get_runtime().materialize()
     assert v.dynamic_index_stride is None
 
 
@@ -455,7 +455,7 @@ def test_matrix_field_dynamic_index_different_cell_size_bytes():
     ti.root.dense(ti.i, 8).place(x, temp)
     ti.root.dense(ti.i, 8).place(y)
 
-    ti.get_runtime().materialize()
+    impl.get_runtime().materialize()
     assert v.dynamic_index_stride is None
 
 
@@ -471,7 +471,7 @@ def test_matrix_field_dynamic_index_different_offset_bytes_in_parent_cell():
     ti.root.dense(ti.i, 8).place(temp_a, x)
     ti.root.dense(ti.i, 8).place(y, temp_b)
 
-    ti.get_runtime().materialize()
+    impl.get_runtime().materialize()
     assert v.dynamic_index_stride is None
 
 
@@ -486,7 +486,7 @@ def test_matrix_field_dynamic_index_different_stride():
 
     ti.root.dense(ti.i, 8).place(x, y, temp, z)
 
-    ti.get_runtime().materialize()
+    impl.get_runtime().materialize()
     assert v.dynamic_index_stride is None