Merge branch 'master' of github.com:taichi-dev/taichi into matrix_typ…

…e_bug_fix_15

.github/workflows/release.yml

@@ -272,55 +272,6 @@ jobs:
           PY: ${{ matrix.python }}
           GPU_TEST: ON
 
-  build_and_test_macos_1014:
-    name: Build and Upload (macos 1014)
-    needs: matrix_prep
-    strategy:
-      fail-fast: false
-      matrix: ${{ fromJson(needs.matrix_prep.outputs.matrix) }}
-    runs-on: [self-hosted, macos10.14]
-    steps:
-      - uses: actions/checkout@v3
-        with:
-          submodules: 'recursive'
-          fetch-depth: '0'
-
-      - name: Prepare Environment
-        run: |
-          . .github/workflows/scripts/common-utils.sh
-          prepare-build-cache
-
-      - name: Build
-        run: |
-          # TODO: should embed LLVM as submodule
-          export LLVM_DIR=/Users/buildbot6/taichi-llvm-10.0.0-macos
-          export PATH=$LLVM_DIR/bin:$PATH
-          .github/workflows/scripts/unix-build.sh
-        env:
-          PROJECT_NAME: ${{ matrix.name }}
-          PY: ${{ matrix.python }}
-          CXX: clang++
-          TAICHI_CMAKE_ARGS: >-
-            -DTI_WITH_OPENGL:BOOL=OFF
-            -DTI_WITH_CUDA:BOOL=OFF
-            -DTI_WITH_CC:BOOL=OFF
-            -DTI_WITH_VULKAN:BOOL=OFF
-            -DTI_BUILD_TESTS:BOOL=ON
-
-      - name: Archive Wheel Artifacts
-        uses: actions/upload-artifact@v3
-        with:
-          name: dist
-          path: dist/*.whl
-          retention-days: 20
-
-      - name: Test
-        run: |
-          .github/workflows/scripts/unix_test.sh
-        env:
-          TI_WANTED_ARCHS: 'cpu'
-          PY: ${{ matrix.python }}
-
   build_and_test_windows:
     name: Build and Upload (Windows only)
     needs: matrix_prep
@@ -365,7 +316,7 @@ jobs:
 
   upload_to_pypi:
     name: Upload release to PyPI
-    needs: [build_and_test_linux, build_and_test_manylinux2014, build_and_test_mac, build_and_test_m1, build_and_test_macos_1014, build_and_test_windows]
+    needs: [build_and_test_linux, build_and_test_manylinux2014, build_and_test_mac, build_and_test_m1, build_and_test_windows]
     runs-on: ubuntu-latest
     steps:
       - uses: actions/checkout@v3

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

@@ -104,7 +104,7 @@ function build-and-test-headless-demo {
     popd
 
     rm -rf taichi-aot-demo
-    git clone --recursive --depth=1 -b add_test_script https://github.com/taichi-dev/taichi-aot-demo
+    git clone --recursive --depth=1 https://github.com/taichi-dev/taichi-aot-demo
     cd taichi-aot-demo
 
     . $(pwd)/ci/test_utils.sh
@@ -139,8 +139,6 @@ function build-and-test-headless-demo {
     rm -rf output
     mkdir output
 
-    adb shell "ls /data/local/tmp/output"
-
     adb pull /data/local/tmp/output .
     compare_to_groundtruth android
 }

.github/workflows/scripts/unix_test.sh

@@ -78,13 +78,23 @@ if [ -z "$TI_SKIP_CPP_TESTS" ]; then
     python3 tests/run_tests.py --cpp
 fi
 
+function run-it {
+    ARCH=$1
+    PARALLELISM=$2
+    KEYS=${3:-"not torch and not paddle"}
+
+    if [[ $TI_WANTED_ARCHS == *"$1"* ]]; then
+        python3 tests/run_tests.py -vr2 -t$PARALLELISM -k "$KEYS" -m "not run_in_serial" -a $ARCH
+        python3 tests/run_tests.py -vr2 -t1 -k "$KEYS" -m "run_in_serial" -a $ARCH
+    fi
+}
+
 if [ -z "$GPU_TEST" ]; then
     if [[ $PLATFORM == *"m1"* ]]; then
-	# Split per arch to avoid flaky test
-        python3 tests/run_tests.py -vr2 -t4 -k "not torch and not paddle" -a cpu
-        # Run metal and vulkan separately so that they don't use M1 chip simultaneously.
-        python3 tests/run_tests.py -vr2 -t4 -k "not torch and not paddle" -a vulkan
-        python3 tests/run_tests.py -vr2 -t2 -k "not torch and not paddle" -a metal
+        run-it cpu 4
+        run-it vulkan 4
+        run-it metal 2
+
         python3 tests/run_tests.py -vr2 -t1 -k "torch" -a "$TI_WANTED_ARCHS"
     else
         # Fail fast, give priority to the error-prone tests
@@ -94,21 +104,11 @@ if [ -z "$GPU_TEST" ]; then
         python3 tests/run_tests.py -vr2 -t4 -k "not paddle" -a "$TI_WANTED_ARCHS"
     fi
 else
-    # Split per arch to increase parallelism for linux GPU tests
-    if [[ $TI_WANTED_ARCHS == *"cuda"* ]]; then
-        # FIXME: suddenly tests exibit OOM on nvidia driver 470 + RTX2060 cards, lower parallelism by 1 (4->3)
-        python3 tests/run_tests.py -vr2 -t3 -k "not torch and not paddle" -m "not run_in_serial" -a cuda
-        python3 tests/run_tests.py -vr2 -t1 -k "not torch and not paddle" -m "run_in_serial" -a cuda
-    fi
-    if [[ $TI_WANTED_ARCHS == *"cpu"* ]]; then
-        python3 tests/run_tests.py -vr2 -t8 -k "not torch and not paddle" -a cpu
-    fi
-    if [[ $TI_WANTED_ARCHS == *"vulkan"* ]]; then
-        python3 tests/run_tests.py -vr2 -t8 -k "not torch and not paddle" -a vulkan
-    fi
-    if [[ $TI_WANTED_ARCHS == *"opengl"* ]]; then
-        python3 tests/run_tests.py -vr2 -t4 -k "not torch and not paddle" -a opengl
-    fi
+    run-it cuda   6
+    run-it cpu    $(nproc)
+    run-it vulkan 8
+    run-it opengl 4
+
     python3 tests/run_tests.py -vr2 -t1 -k "torch" -a "$TI_WANTED_ARCHS"
     # Paddle's paddle.fluid.core.Tensor._ptr() is only available on develop branch, and CUDA version on linux will get error `Illegal Instruction`
 fi

.github/workflows/scripts/win_test.ps1

@@ -36,6 +36,13 @@ Invoke python tests/run_tests.py -vr2 -t1 -k "paddle" -a cpu
 # Disable paddle for the remaining test
 $env:TI_ENABLE_PADDLE = "0"
 
+function RunIt($arch, $parallelism) {
+    if ("$env:TI_WANTED_ARCHS".Contains("cuda")) {
+        Invoke python tests/run_tests.py -vr2 -t"$parallelism" -k "not torch and not paddle" -m "not run_in_serial" -a $arch
+        Invoke python tests/run_tests.py -vr2 -t1 -k "not torch and not paddle" -m "run_in_serial" -a $arch
+    }
+}
+
 if ("$env:TI_WANTED_ARCHS".Contains("cpu")) {
   # NOTE: Always test CPU with non-CUDA version of PyTorch,
   #       since CUDA version of PyTorch will load a lot of CUDA libraries,
@@ -45,20 +52,18 @@ if ("$env:TI_WANTED_ARCHS".Contains("cpu")) {
   #       This is a non-issue on Linux, since Linux overcommits.
   # TODO relax this when torch supports 3.10
   Invoke pip install "torch==1.12.1; python_version < '3.10'"
-  Invoke python tests/run_tests.py -vr2 "-t$(EstimateNumProcs)" -k "not torch and not paddle" -a cpu
+  RunIt cpu (EstimateNumProcs)
 }
+
 if ("$env:TI_WANTED_ARCHS".Contains("cuda")) {
   # TODO relax this when torch supports 3.10
   Invoke pip install "torch==1.10.1+cu113; python_version < '3.10'" -f https://download.pytorch.org/whl/cu113/torch_stable.html
-  Invoke python tests/run_tests.py -vr2 -t4 -k "not torch and not paddle" -m "not run_in_serial" -a cuda
-  Invoke python tests/run_tests.py -vr2 -t1 -k "not torch and not paddle" -m "run_in_serial" -a cuda
-}
-if ("$env:TI_WANTED_ARCHS".Contains("opengl")) {
-  Invoke python tests/run_tests.py -vr2 -t4 -k "not torch and not paddle" -a opengl
-}
-if ("$env:TI_WANTED_ARCHS".Contains("vulkan")) {
-  Invoke python tests/run_tests.py -vr2 -t4 -k "not torch and not paddle" -a vulkan
+  RunIt cuda 6
 }
+
+RunIt opengl 4
+RunIt vulkan 4
+
 Invoke python tests/run_tests.py -vr2 -t1 -k "torch" -a "$env:TI_WANTED_ARCHS"
 
 if ("$env:TI_RUN_RELEASE_TESTS" -eq "1" -and -not "$env:TI_LITE_TEST") {

.github/workflows/testing.yml

@@ -176,7 +176,7 @@ jobs:
   build_and_test_gpu_linux:
     name: Build and Test (GPU)
     needs: check_files
-    timeout-minutes: ${{ github.event.schedule != '0 18 * * *' && 60 || 120 }}
+    timeout-minutes: ${{ github.event.schedule != '0 18 * * *' && 90 || 120 }}
     strategy:
       matrix:
         tags:

.gitignore

@@ -88,3 +88,4 @@ imgui.ini
 /_skbuild/
 .cache
 .do-not-clean
+*.dylib

cmake/TaichiCAPI.cmake

@@ -18,6 +18,9 @@ endif()
 
 if (TI_WITH_VULKAN)
   list(APPEND C_API_SOURCE "c_api/src/taichi_vulkan_impl.cpp")
+  if (APPLE)
+    install(FILES ${MoltenVK_LIBRARY} DESTINATION c_api/lib)
+  endif()
 endif()
 
 if(TI_BUILD_TESTS)

docs/lang/articles/basic/sparse.md

@@ -276,7 +276,7 @@ You can use `ti.rescale_index(descendant_snode/field, ancestor_snode, index)` to
 print(ti.rescale_index(x, block1, ti.Vector([7, 3]))) # output: [1, 0]
 print(ti.rescale_index(x, block2, [7, 3]))            # output: [3, 1]
 print(ti.rescale_index(x, pixel,  [7, 3]))            # output: [7, 3]
-print(ti.rescale_index(block1, block2, [3, 1]))       # output: [1, 0]
+print(ti.rescale_index(block2, block1, [3, 1]))       # output: [1, 0]
 ```
 
 Regarding line 1, you can also compute the `block1` index given `pixel` index `[7, 3]` as `[7//2//2, 3//2//2]`. However, doing so couples computation code with the internal configuration of data structures (in this case, the size of `block1` containers). By using `ti.rescale_index()`, you can avoid hard-coding internal information of data structures.

misc/test_build_cusm_from_coo.py

@@ -26,7 +26,7 @@
 A = ti.linalg.SparseMatrix(n=4, m=4, dtype=ti.float32)
 A.build_coo(d_coo_row, d_coo_col, d_coo_val)
 
-A.spmv(x, y)
+y = A @ x
 
 # Check if the results are correct
 equal = True

misc/test_coo_cusolver.py

@@ -48,7 +48,7 @@ def print_x(x: ti.types.ndarray(), ncols: ti.i32):
 A_ti.build_coo(d_row_coo, d_col_coo, d_val_coo)
 x_ti = ti.ndarray(shape=ncols, dtype=ti.float32)
 solver = ti.linalg.SparseSolver()
-solver.solve_cu(A_ti, b, x_ti)
+x_ti = solver.solve_cu(A_ti, b)
 ti.sync()
 print_x(x_ti, ncols)
 ti.sync()
@@ -65,7 +65,7 @@ def print_x(x: ti.types.ndarray(), ncols: ti.i32):
 solver = ti.linalg.SparseSolver()
 solver.analyze_pattern(A_ti)
 solver.factorize(A_ti)
-solver.solve_rf(A_ti, b, x_ti)
+x_ti = solver.solve(b)
 
 ti.sync()
 print_x(x_ti, ncols)

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 = []

python/taichi/lang/ast/ast_transformer.py

@@ -513,6 +513,9 @@ def build_Call(ctx, node):
         keywords = dict(ChainMap(*[keyword.ptr for keyword in node.keywords]))
         func = node.func.ptr
 
+        if id(func) in [id(print), id(impl.ti_print)]:
+            ctx.func.has_print = True
+
         if isinstance(node.func, ast.Attribute) and isinstance(
                 node.func.value.ptr, str) and node.func.attr == 'format':
             args.insert(0, node.func.value.ptr)
@@ -537,6 +540,9 @@ def build_Call(ctx, node):
         node.ptr = func(*args, **keywords)
         ASTTransformer.warn_if_is_external_func(ctx, node)
 
+        if getattr(func, "_is_taichi_function", False):
+            ctx.func.has_print |= func.func.has_print
+
         return node.ptr
 
     @staticmethod

python/taichi/lang/expr.py

@@ -55,6 +55,22 @@ def get_shape(self):
                 f"Getting shape of non-tensor type: {self.ptr.get_ret_type()}")
         return tuple(self.ptr.get_shape())
 
+    @property
+    def n(self):
+        shape = self.get_shape()
+        if len(shape) < 1:
+            raise TaichiCompilationError(
+                f"Getting n of tensor type < 1D: {self.ptr.get_ret_type()}")
+        return shape[0]
+
+    @property
+    def m(self):
+        shape = self.get_shape()
+        if len(shape) < 2:
+            raise TaichiCompilationError(
+                f"Getting m of tensor type < 2D: {self.ptr.get_ret_type()}")
+        return shape[1]
+
     def __hash__(self):
         return self.ptr.get_raw_address()
 

python/taichi/lang/kernel_impl.py

@@ -69,6 +69,7 @@ def decorated(*args, **kwargs):
 
     decorated._is_taichi_function = True
     decorated._is_real_function = is_real_function
+    decorated.func = fun
     return decorated
 
 
@@ -99,6 +100,7 @@ def decorated(*args, **kwargs):
         return fun.__call__(*args, **kwargs)
 
     decorated._is_taichi_function = True
+    decorated.func = fun
     return decorated
 
 
@@ -202,6 +204,7 @@ def __init__(self,
         self.mapper = TaichiCallableTemplateMapper(
             self.arguments, self.template_slot_locations)
         self.taichi_functions = {}  # The |Function| class in C++
+        self.has_print = False
 
     def __call__(self, *args, **kwargs):
         args = _process_args(self, args, kwargs)
@@ -467,6 +470,7 @@ def __init__(self, _func, autodiff_mode, _classkernel=False):
         # TODO[#5114]: get rid of compiled_functions and use compiled_kernels instead.
         # Main motivation is that compiled_kernels can be potentially serialized in the AOT scenario.
         self.compiled_kernels = {}
+        self.has_print = False
 
     def reset(self):
         self.runtime = impl.get_runtime()
@@ -676,7 +680,7 @@ def func__(*args):
                 elif isinstance(needed, sparse_matrix_builder):
                     # Pass only the base pointer of the ti.types.sparse_matrix_builder() argument
                     launch_ctx.set_arg_uint(actual_argument_slot,
-                                            v._get_addr())
+                                            v._get_ndarray_addr())
                 elif isinstance(needed,
                                 ndarray_type.NdarrayType) and isinstance(
                                     v, taichi.lang._ndarray.Ndarray):
@@ -798,7 +802,7 @@ def callback(original, updated):
             ret_dt = self.return_type
             has_ret = ret_dt is not None
 
-            if has_ret:
+            if has_ret or self.has_print:
                 runtime_ops.sync()
 
             if has_ret: