[misc] Cherry-pick 6415, 6391, 6416, 6422 (#6421)

Issue: #

### Brief Summary

Co-authored-by: Yi Xu <[email protected]>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Proton <[email protected]>
Co-authored-by: Lin Jiang <[email protected]>

.github/workflows/scripts/aot-demo.sh

@@ -103,8 +103,11 @@ function build-and-test-headless-demo {
     popd
 
     rm -rf taichi-aot-demo
-    git clone --recursive --depth=1 https://github.com/taichi-dev/taichi-aot-demo
+    git clone --recursive --depth=2 https://github.com/taichi-dev/taichi-aot-demo
     cd taichi-aot-demo
+
+    git checkout 28f1df9a5b33bd21b4b46b6ca081275dfd037a16
+
     mkdir build
     pushd build
     export TAICHI_C_API_INSTALL_DIR=$(find $TAICHI_REPO_DIR -name cmake-install -type d | head -n 1)/c_api

python/taichi/_kernels.py

@@ -1,5 +1,4 @@
-from typing import Iterable
-
+from taichi._funcs import field_fill_taichi_scope
 from taichi._lib.utils import get_os_name
 from taichi.lang import ops
 from taichi.lang._ndrange import ndrange
@@ -237,20 +236,8 @@ def clear_loss(l: template()):
 
 
 @kernel
-def fill_matrix(mat: template(), vals: template()):
-    for I in grouped(mat):
-        for p in static(range(mat.n)):
-            for q in static(range(mat.m)):
-                if static(mat[I].ndim == 2):
-                    if static(isinstance(vals[p], Iterable)):
-                        mat[I][p, q] = vals[p][q]
-                    else:
-                        mat[I][p, q] = vals[p]
-                else:
-                    if static(isinstance(vals[p], Iterable)):
-                        mat[I][p] = vals[p][q]
-                    else:
-                        mat[I][p] = vals[p]
+def field_fill_python_scope(F: template(), val: template()):
+    field_fill_taichi_scope(F, val)
 
 
 @kernel

python/taichi/_lib/utils.py

@@ -1,5 +1,6 @@
 import os
 import platform
+import re
 import sys
 
 from colorama import Fore, Style
@@ -92,6 +93,12 @@ def print_red_bold(*args, **kwargs):
     print(Style.RESET_ALL, end='')
 
 
+def print_yellow_bold(*args, **kwargs):
+    print(Fore.YELLOW + Style.BRIGHT, end='')
+    print(*args, **kwargs)
+    print(Style.RESET_ALL, end='')
+
+
 def check_exists(src):
     if not os.path.exists(src):
         raise FileNotFoundError(
@@ -164,3 +171,77 @@ def _print_taichi_header():
 
 
 _print_taichi_header()
+
+
+def try_get_wheel_tag(module):
+    try:
+        import wheel.metadata  # pylint: disable=import-outside-toplevel
+        wheel_path = f'{module.__path__[0]}-{".".join(map(str, module.__version__))}.dist-info/WHEEL'
+        meta = wheel.metadata.read_pkg_info(wheel_path)
+        return meta.get('Tag')
+    except Exception:
+        return None
+
+
+def try_get_loaded_libc_version():
+    assert platform.system() == "Linux"
+    with open('/proc/self/maps') as f:
+        content = f.read()
+
+    try:
+        libc_path = next(v for v in content.split() if 'libc-' in v)
+        ver = re.findall(r'\d+\.\d+', libc_path)
+        if not ver:
+            return None
+        return tuple([int(v) for v in ver[0].split('.')])
+    except StopIteration:
+        return None
+
+
+def try_get_pip_version():
+    try:
+        import pip  # pylint: disable=import-outside-toplevel
+        return tuple([int(v) for v in pip.__version__.split('.')])
+    except ImportError:
+        return None
+
+
+def warn_restricted_version():
+    if os.environ.get('TI_MANYLINUX2014_OK', ''):
+        return
+
+    if get_os_name() == 'linux':
+        try:
+            import taichi as ti  # pylint: disable=import-outside-toplevel
+            wheel_tag = try_get_wheel_tag(ti)
+            if wheel_tag and 'manylinux2014' in wheel_tag:
+                print_yellow_bold(
+                    "You have installed a restricted version of taichi, certain features (e.g. Vulkan & GGUI) will not work."
+                )
+                libc_ver = try_get_loaded_libc_version()
+                if libc_ver and libc_ver < (2, 27):
+                    print_yellow_bold(
+                        "!! Taichi requires glibc >= 2.27 to run, please try upgrading your OS to a recent one (e.g. Ubuntu 18.04 or later) if possible."
+                    )
+
+                pip_ver = try_get_pip_version()
+                if pip_ver and pip_ver < (20, 3, 0):
+                    print_yellow_bold(
+                        f"!! Your pip (version {'.'.join(map(str, pip_ver))}) is outdated (20.3.0 or later required), "
+                        "try upgrading pip and install taichi again.")
+                    print()
+                    print_yellow_bold(
+                        "    $ python3 -m pip install --upgrade pip")
+                    print_yellow_bold(
+                        "    $ python3 -m pip install --force-reinstall taichi"
+                    )
+                    print()
+
+                print_yellow_bold(
+                    "You can suppress this warning by setting the environment variable TI_MANYLINUX2014_OK=1."
+                )
+        except Exception:
+            pass
+
+
+warn_restricted_version()

python/taichi/lang/ast/ast_transformer.py

@@ -1024,6 +1024,16 @@ def build_ndrange_for(ctx, node):
                                                      ndrange_end.ptr)
             I = impl.expr_init(ndrange_loop_var)
             targets = ASTTransformer.get_for_loop_targets(node)
+            if len(targets) != len(ndrange_var.dimensions):
+                warnings.warn_explicit(
+                    'Ndrange for loop with number of the loop variables not equal to '
+                    'the dimension of the ndrange is deprecated. '
+                    'Please check if the number of arguments of ti.ndrange() is equal to '
+                    'the number of the loop variables.',
+                    DeprecationWarning,
+                    ctx.file,
+                    node.lineno + ctx.lineno_offset,
+                    module="taichi")
             for i, target in enumerate(targets):
                 if i + 1 < len(targets):
                     target_tmp = impl.expr_init(

python/taichi/lang/matrix.py

@@ -1630,34 +1630,24 @@ def fill(self, val):
         """Fills this matrix field with specified values.
 
         Args:
-            val (Union[Number, List, Tuple, Matrix]): Values to fill,
+            val (Union[Number, Expr, List, Tuple, Matrix]): Values to fill,
                 should have consistent dimension consistent with `self`.
         """
-        if isinstance(val, numbers.Number):
-            val = tuple(
-                [tuple([val for _ in range(self.m)]) for _ in range(self.n)])
-        elif isinstance(val,
-                        (list, tuple)) and isinstance(val[0], numbers.Number):
-            assert self.m == 1
-            val = tuple(val)
-        elif is_vector(val):
-            val = tuple([(val(i), ) for i in range(self.n * self.m)])
+        if isinstance(val, numbers.Number) or (isinstance(val, expr.Expr)
+                                               and not val.is_tensor()):
+            val = list(list(val for _ in range(self.m)) for _ in range(self.n))
         elif isinstance(val, Matrix):
-            val_tuple = []
-            for i in range(val.n):
-                row = []
-                for j in range(val.m):
-                    row.append(val(i, j))
-                row = tuple(row)
-                val_tuple.append(row)
-            val = tuple(val_tuple)
+            val = val.to_list()
+        else:
+            assert isinstance(val, (list, tuple))
+        val = tuple(tuple(x) if isinstance(x, list) else x for x in val)
         assert len(val) == self.n
         if self.ndim != 1:
             assert len(val[0]) == self.m
-
         if in_python_scope():
-            from taichi._kernels import fill_matrix  # pylint: disable=C0415
-            fill_matrix(self, val)
+            from taichi._kernels import \
+                field_fill_python_scope  # pylint: disable=C0415
+            field_fill_python_scope(self, val)
         else:
             from taichi._funcs import \
                 field_fill_taichi_scope  # pylint: disable=C0415

python/taichi/ui/utils.py

@@ -1,8 +1,8 @@
 import platform
-import re
 from math import acos, asin, cos, sin
 
 from taichi._lib import core as _ti_core
+from taichi._lib.utils import try_get_wheel_tag
 from taichi.lang.impl import default_cfg
 from taichi.lang.matrix import Vector
 
@@ -66,31 +66,6 @@ def vec_to_euler(v):
     return yaw, pitch
 
 
-def try_get_wheel_tag(module):
-    try:
-        import wheel.metadata  # pylint: disable=import-outside-toplevel
-        wheel_path = f'{module.__path__[0]}-{".".join(map(str, module.__version__))}.dist-info/WHEEL'
-        meta = wheel.metadata.read_pkg_info(wheel_path)
-        return meta.get('Tag')
-    except Exception:
-        return None
-
-
-def try_get_loaded_libc_version():
-    assert platform.system() == "Linux"
-    with open('/proc/self/maps') as f:
-        content = f.read()
-
-    try:
-        libc_path = next(v for v in content.split() if 'libc-' in v)
-        ver = re.findall(r'\d+\.\d+', libc_path)
-        if not ver:
-            return None
-        return tuple([int(v) for v in ver[0].split('.')])
-    except StopIteration:
-        return None
-
-
 class GGUINotAvailableException(Exception):
     pass
 
@@ -107,26 +82,9 @@ def check_ggui_availability():
         wheel_tag = try_get_wheel_tag(taichi)
         if platform.system(
         ) == "Linux" and wheel_tag and 'manylinux2014' in wheel_tag:
-            libc_ver = try_get_loaded_libc_version()
-            if libc_ver and libc_ver < (2, 27):
-                raise GGUINotAvailableException(
-                    "GGUI is not available since you have installed a restricted version of taichi. "
-                    "Your OS is outdated, try upgrading to a recent one (e.g. Ubuntu 18.04 or later) if possible."
-                )
-
-            try:
-                import pip  # pylint: disable=import-outside-toplevel
-                ver = tuple([int(v) for v in pip.__version__.split('.')])
-                if ver < (20, 3, 0):
-                    raise GGUINotAvailableException(
-                        "GGUI is not available since you have installed a restricted version of taichi. "
-                        f"Your pip (version {pip.__version__}) is outdated (20.3.0 or later required), "
-                        "try upgrading pip and install taichi again.")
-            except ImportError:
-                pass
-
             raise GGUINotAvailableException(
-                "GGUI is not available since you have installed a restricted version of taichi."
+                "GGUI is not available since you have installed a restricted version of taichi. "
+                "Please see yellow warning messages printed during startup for details."
             )
 
     except GGUINotAvailableException:

taichi/rhi/vulkan/vulkan_device_creator.cpp

@@ -738,10 +738,15 @@ void VulkanDeviceCreator::create_logical_device(bool manual_create) {
       if (CHECK_VERSION(1, 3) ||
           buffer_device_address_feature.bufferDeviceAddress) {
         if (device_supported_features.shaderInt64) {
+// Temporarily disable it on macOS:
+// https://github.com/taichi-dev/taichi/issues/6295
+#if !defined(__APPLE__)
           ti_device_->set_cap(
               DeviceCapability::spirv_has_physical_storage_buffer, true);
+#endif
         }
       }
+
       *pNextEnd = &buffer_device_address_feature;
       pNextEnd = &buffer_device_address_feature.pNext;
     }

tests/python/test_field.py

@@ -304,26 +304,6 @@ def test_indexing_mat_field_with_np_int():
     val[idx][idx, idx]
 
 
-@test_utils.test(exclude=[ti.cc], debug=True)
-def test_field_fill():
-    x = ti.field(int, shape=(3, 3))
-    x.fill(2)
-
-    y = ti.field(float, shape=(3, 3))
-    y.fill(2.0)
-
-    z = ti.Vector.field(3, float, shape=(3, 3))
-    z.fill([1, 2, 3])
-
-    @ti.kernel
-    def test():
-        x.fill(3)
-        y.fill(3.0)
-        z.fill([4, 5, 6])
-
-    test()
-
-
 @test_utils.test()
 def test_python_for_in():
     x = ti.field(int, shape=3)