[aot] Make arch an optional argument in ti.aot.Module (taichi-dev#6574)

Fixes taichi-dev#5859

### Brief Summary
We haven't supported compiling to a target different from the
initialized arch and it won't be available in the short term. So let's
make this an optional argument and default to the current arch to save
users from some typing work.

python/taichi/aot/module.py

@@ -82,16 +82,24 @@ class Module:
         # Now the module file '/path/to/module' contains the Metal kernels
         # for running ``foo`` and ``bar``.
     """
-    def __init__(self, arch, caps=None):
+    def __init__(self, arch=None, caps=None):
         """Creates a new AOT module instance
 
         Args:
-          arch: Target backend architecture. This must match the one specified
-            in :func:`~taichi.lang.init`.
+          arch: Target backend architecture. Default to the one initialized in :func:`~taichi.lang.init` if not specified.
           caps (List[str]): Enabled device capabilities.
         """
         if caps is None:
             caps = []
+        curr_arch = impl.current_cfg().arch
+        if arch is None:
+            arch = curr_arch
+        elif arch != curr_arch:
+            # TODO: we'll support this eventually but not yet...
+            warnings.warn(
+                f"AOT compilation to a different arch than the current one is not yet supported, switching to {curr_arch}"
+            )
+            arch = curr_arch
 
         self._arch = arch
         self._kernels = []

tests/cpp/aot/python_scripts/bitmasked_aot_test.py

@@ -60,7 +60,7 @@ def check_value_1():
     assert "TAICHI_AOT_FOLDER_PATH" in os.environ.keys()
     dir_name = str(os.environ["TAICHI_AOT_FOLDER_PATH"])
 
-    m = ti.aot.Module(arch)
+    m = ti.aot.Module()
 
     m.add_kernel(activate, template_args={})
     m.add_kernel(check_value_0, template_args={})

tests/cpp/aot/python_scripts/comet_aot.py

@@ -121,7 +121,7 @@ def initialize():
 
 
 def save_kernels(arch):
-    mod = ti.aot.Module(arch)
+    mod = ti.aot.Module()
 
     # Initialize
     g_init_builder = ti.graph.GraphBuilder()

tests/cpp/aot/python_scripts/dense_field_aot_test.py

@@ -33,7 +33,7 @@ def ret() -> ti.f32:
     assert "TAICHI_AOT_FOLDER_PATH" in os.environ.keys()
     dir_name = str(os.environ["TAICHI_AOT_FOLDER_PATH"])
 
-    m = ti.aot.Module(arch)
+    m = ti.aot.Module()
     m.add_kernel(simple_return)
     m.add_kernel(init)
     m.add_kernel(ret)

tests/cpp/aot/python_scripts/dynamic_aot_test.py

@@ -59,7 +59,7 @@ def check_value_1():
         assert x[2, 0] == 9
         assert x[3, 0] == 12
 
-    m = ti.aot.Module(arch)
+    m = ti.aot.Module()
 
     m.add_kernel(activate, template_args={})
     m.add_kernel(check_value_0, template_args={})

tests/cpp/aot/python_scripts/field_aot_test.py

@@ -104,15 +104,15 @@ def check_activate_pointer_fields():
 
         run_graph = g_builder.compile()
 
-        m = ti.aot.Module(arch)
+        m = ti.aot.Module()
 
         m.add_field("x", x)
         m.add_field("y", y)
 
         m.add_graph('run_graph', run_graph)
         m.save(dir_name, '')
     else:
-        m = ti.aot.Module(arch)
+        m = ti.aot.Module()
 
         m.add_kernel(init_fields, template_args={})
         m.add_kernel(check_init_x, template_args={})

tests/cpp/aot/python_scripts/graph_aot_test.py

@@ -48,7 +48,7 @@ def run2(base: int, arr: ti.types.ndarray(field_dim=1, dtype=ti.i32)):
     assert "TAICHI_AOT_FOLDER_PATH" in os.environ.keys()
     tmpdir = str(os.environ["TAICHI_AOT_FOLDER_PATH"])
 
-    mod = ti.aot.Module(arch)
+    mod = ti.aot.Module()
     mod.add_graph('run_graph', run_graph)
     mod.save(tmpdir, '')
 

tests/cpp/aot/python_scripts/kernel_aot_test1.py

@@ -20,7 +20,7 @@ def run(base: int, arr: ti.types.ndarray(), v: ti.types.vector(3, ti.i32)):
     assert "TAICHI_AOT_FOLDER_PATH" in os.environ.keys()
     dir_name = str(os.environ["TAICHI_AOT_FOLDER_PATH"])
 
-    m = ti.aot.Module(arch)
+    m = ti.aot.Module()
     m.add_kernel(run, template_args={'arr': arr})
     m.save(dir_name, 'whatever')
 

tests/cpp/aot/python_scripts/kernel_aot_test2.py

@@ -22,7 +22,7 @@ def ker2(arr: ti.types.ndarray(), n: ti.i32):
     assert "TAICHI_AOT_FOLDER_PATH" in os.environ.keys()
     dir_name = str(os.environ["TAICHI_AOT_FOLDER_PATH"])
 
-    m = ti.aot.Module(arch)
+    m = ti.aot.Module()
     if save_compute_graph:
         sym_arr = ti.graph.Arg(ti.graph.ArgKind.NDARRAY,
                                'arr',

tests/cpp/aot/python_scripts/mpm88_graph_aot.py

@@ -181,7 +181,7 @@ def init_particles(x: ti.any_arr(field_dim=1), v: ti.any_arr(field_dim=1),
         grid_v = ti.Vector.ndarray(2, ti.f32, shape=(n_grid, n_grid))
         grid_m = ti.ndarray(ti.f32, shape=(n_grid, n_grid))
 
-        mod = ti.aot.Module(arch)
+        mod = ti.aot.Module()
         mod.add_graph('init', g_init)
         mod.add_graph('update', g_update)
         mod.save(tmpdir, '')
@@ -195,7 +195,7 @@ def init_particles(x: ti.any_arr(field_dim=1), v: ti.any_arr(field_dim=1),
         grid_v = ti.Vector.ndarray(2, ti.f32, shape=(n_grid, n_grid))
         grid_m = ti.ndarray(ti.f32, shape=(n_grid, n_grid))
 
-        mod = ti.aot.Module(arch)
+        mod = ti.aot.Module()
         mod.add_kernel(init_particles, template_args={'x': x, 'v': v, 'J': J})
         mod.add_kernel(substep_reset_grid,
                        template_args={

tests/cpp/aot/python_scripts/sph_aot.py

@@ -202,7 +202,7 @@ def copy_data_from_ndarray_to_field(src: ti.template(), dst: ti.any_arr()):
     print('running in graph mode')
 
     # Serialize!
-    mod = ti.aot.Module(arch)
+    mod = ti.aot.Module()
 
     mod.add_kernel(initialize,
                    template_args={

tests/cpp/aot/python_scripts/taichi_sparse_test.py

@@ -72,7 +72,7 @@ def paint():
 
 
 def save_kernels(arch):
-    m = ti.aot.Module(arch)
+    m = ti.aot.Module()
 
     m.add_kernel(fill_img, template_args={})
     m.add_kernel(block1_deactivate_all, template_args={})

tests/cpp/aot/python_scripts/texture_aot_test.py

@@ -74,7 +74,7 @@ def run2(tex0: ti.types.texture(num_dimensions=2),
     assert "TAICHI_AOT_FOLDER_PATH" in os.environ.keys()
     tmpdir = str(os.environ["TAICHI_AOT_FOLDER_PATH"])
 
-    mod = ti.aot.Module(arch)
+    mod = ti.aot.Module()
     mod.add_graph('run_graph', run_graph)
     mod.save(tmpdir, '')
 

tests/python/test_aot.py

@@ -45,7 +45,7 @@ def init():
             density[i, j] = 1
 
     with tempfile.TemporaryDirectory() as tmpdir:
-        m = ti.aot.Module(ti.opengl)
+        m = ti.aot.Module()
         m.add_field('density', density)
         m.add_kernel(init)
         m.save(tmpdir)
@@ -68,7 +68,7 @@ def init(x: ti.f32, density1: ti.types.ndarray(field_dim=2)):
             density1[i, j] = x + 1
 
     with tempfile.TemporaryDirectory() as tmpdir:
-        m = ti.aot.Module(ti.lang.impl.current_cfg().arch)
+        m = ti.aot.Module()
         m.add_kernel(init, {'density1': density1})
         m.save(tmpdir)
         with open(os.path.join(tmpdir, 'metadata.json')) as json_file:
@@ -92,8 +92,7 @@ def init():
             density[i, j] = 1
 
     with tempfile.TemporaryDirectory() as tmpdir:
-        # note ti.aot.Module(ti.opengl) is no-op according to its docstring.
-        m = ti.aot.Module(ti.lang.impl.current_cfg().arch)
+        m = ti.aot.Module()
         m.add_field('density', density)
         m.add_kernel(init)
         m.save(tmpdir)
@@ -111,8 +110,7 @@ def foo(n: ti.template()):
             density[0, 0] += 1
 
     with tempfile.TemporaryDirectory() as tmpdir:
-        # note ti.aot.Module(ti.opengl) is no-op according to its docstring.
-        m = ti.aot.Module(ti.lang.impl.current_cfg().arch)
+        m = ti.aot.Module()
         m.add_field('density', density)
         with m.add_kernel_template(foo) as kt:
             kt.instantiate(n=6)
@@ -132,7 +130,7 @@ def test_non_dense_snode():
     blk.dense(ti.i, n).place(y)
 
     with pytest.raises(RuntimeError, match='AOT: only supports dense field'):
-        m = ti.aot.Module(ti.lang.impl.current_cfg().arch)
+        m = ti.aot.Module()
         m.add_field('x', x)
         m.add_field('y', y)
 
@@ -219,7 +217,7 @@ def init():
             J[i] = 1
 
     with tempfile.TemporaryDirectory() as tmpdir:
-        m = ti.aot.Module(ti.lang.impl.current_cfg().arch)
+        m = ti.aot.Module()
         m.add_field("x", x)
         m.add_field("v", v)
         m.add_field("C", C)
@@ -396,7 +394,7 @@ def initialize():
             Jp[i] = 1
 
     with tempfile.TemporaryDirectory() as tmpdir:
-        m = ti.aot.Module(ti.lang.impl.current_cfg().arch)
+        m = ti.aot.Module()
         m.add_field('x', x)
         m.add_field('v', v)
         m.add_field('C', C)
@@ -492,7 +490,7 @@ def substep(x: ti.types.ndarray(element_dim=1),
     grid_m = ti.ndarray(ti.f32, (n_grid, n_grid))
 
     with tempfile.TemporaryDirectory() as tmpdir:
-        m = ti.aot.Module(ti.opengl)
+        m = ti.aot.Module()
         template_args = {
             'x': x,
             'v': v,
@@ -517,7 +515,7 @@ def run(arr: ti.types.ndarray(), val1: ti.f32, val2: ti.template()):
 
     with tempfile.TemporaryDirectory() as tmpdir:
         x = ti.ndarray(dtype=ti.f32, shape=16)
-        m = ti.aot.Module(ti.opengl)
+        m = ti.aot.Module()
         m.add_kernel(run, template_args={'arr': x, 'val2': 42})
         m.save(tmpdir, '')
         with open(os.path.join(tmpdir, 'metadata.json')) as json_file:
@@ -536,8 +534,7 @@ def init():
             density[i, j] = 1
 
     with tempfile.TemporaryDirectory() as tmpdir:
-        # note ti.aot.Module(ti.opengl) is no-op according to its docstring.
-        m = ti.aot.Module(ti.lang.impl.current_cfg().arch)
+        m = ti.aot.Module()
         m.add_field('density', density)
         m.add_kernel(init)
         tcm_path = f"{tmpdir}/x.tcm"
@@ -567,8 +564,7 @@ def init_data(test_vec: ivec3):
     g_init = g_init_builder.compile()
 
     with tempfile.TemporaryDirectory() as tmpdir:
-        # note ti.aot.Module(ti.opengl) is no-op according to its docstring.
-        m = ti.aot.Module(ti.lang.impl.current_cfg().arch)
+        m = ti.aot.Module()
         m.add_graph("g_init", g_init)
         m.save(tmpdir)
         with open(os.path.join(tmpdir, 'metadata.json'), "r") as json_file:
@@ -593,7 +589,17 @@ def test(a: ti.types.ndarray(), c: ti.u8):
 
     g.run({'a': a, 'c': c})
 
-    m = ti.aot.Module(ti.lang.impl.current_cfg().arch, caps=['spirv_has_int8'])
+    m = ti.aot.Module(caps=['spirv_has_int8'])
     m.add_graph('g_init', g)
     with tempfile.TemporaryDirectory() as tmpdir:
         m.save(tmpdir)
+
+
+@test_utils.test(arch=[ti.vulkan])
+def test_module_arch_fallback():
+    with pytest.warns(
+            Warning,
+            match=
+            r'AOT compilation to a different arch than the current one is not yet supported, switching'
+    ):
+        m = ti.aot.Module(ti.cpu)

tests/python/test_deprecation.py

@@ -12,7 +12,7 @@ def test_deprecated_aot_save_filename():
     density = ti.field(float, shape=(4, 4))
 
     with tempfile.TemporaryDirectory() as tmpdir:
-        m = ti.aot.Module(ti.lang.impl.current_cfg().arch)
+        m = ti.aot.Module()
         m.add_field('density', density)
         with pytest.warns(
                 DeprecationWarning,