diff --git a/.github/ISSUE_TEMPLATE/bug_report.md b/.github/ISSUE_TEMPLATE/bug_report.md
index 16db75a75fdca..0c144f2eb4082 100644
--- a/.github/ISSUE_TEMPLATE/bug_report.md
+++ b/.github/ISSUE_TEMPLATE/bug_report.md
@@ -1,5 +1,5 @@
---
-name: Bug report
+name: Bug Report
about: Create a report to help us improve
title: ''
labels: potential bug
diff --git a/.github/ISSUE_TEMPLATE/document.md b/.github/ISSUE_TEMPLATE/document.md
new file mode 100644
index 0000000000000..ddb5001154da3
--- /dev/null
+++ b/.github/ISSUE_TEMPLATE/document.md
@@ -0,0 +1,12 @@
+---
+name: Doc
+about: Have doubts or suggestions about Taichi Docs
+title: ''
+labels: doc
+assignees: ''
+
+---
+
+
diff --git a/.github/ISSUE_TEMPLATE/feature_request.md b/.github/ISSUE_TEMPLATE/feature_request.md
index 9a7b1c1797446..1bf3415e010e7 100644
--- a/.github/ISSUE_TEMPLATE/feature_request.md
+++ b/.github/ISSUE_TEMPLATE/feature_request.md
@@ -1,5 +1,5 @@
---
-name: Feature request
+name: Feature Request
about: Suggest an idea for this project
title: ''
labels: feature request
diff --git a/.github/ISSUE_TEMPLATE/question.md b/.github/ISSUE_TEMPLATE/question.md
new file mode 100644
index 0000000000000..accca01a54df9
--- /dev/null
+++ b/.github/ISSUE_TEMPLATE/question.md
@@ -0,0 +1,17 @@
+---
+name: Ask a Question
+about: Ask anything about Taichi
+title: ''
+labels: question
+assignees: ''
+
+---
+
+
diff --git a/.github/workflows/cancel.yml b/.github/workflows/cancel.yml
deleted file mode 100644
index da36ae34e9405..0000000000000
--- a/.github/workflows/cancel.yml
+++ /dev/null
@@ -1,14 +0,0 @@
-name: Cancel
-on:
- workflow_run:
- workflows: ["Presubmit Checks"]
- types:
- - requested
-
-jobs:
- cancel:
- runs-on: ubuntu-latest
- steps:
- - uses: styfle/cancel-workflow-action@0.9.0
- with:
- workflow_id: ${{ github.event.workflow.id }}
diff --git a/.github/workflows/issue_comment.yml b/.github/workflows/issue_comment.yml
index 7398d128a61fe..5f8f46081f383 100644
--- a/.github/workflows/issue_comment.yml
+++ b/.github/workflows/issue_comment.yml
@@ -18,3 +18,5 @@ jobs:
commands: |
format
benchmark
+ rebase
+ rerun
diff --git a/.github/workflows/performance_monitoring.yml b/.github/workflows/performance_monitoring.yml
new file mode 100644
index 0000000000000..e2bbded40c117
--- /dev/null
+++ b/.github/workflows/performance_monitoring.yml
@@ -0,0 +1,33 @@
+name: Performance Monitoring
+on:
+ push:
+ branches:
+ - master
+
+jobs:
+ gpu_backends:
+ name: Performance monitoring (NVIDIA GPU)
+ timeout-minutes: 60
+ # Disable this workflow on forks
+ if: github.repository_owner == 'taichi-dev'
+ runs-on: [self-hosted, x64, cuda, linux, benchmark]
+ steps:
+ - uses: actions/checkout@v2
+ with:
+ submodules: "recursive"
+
+ - name: Build & Install
+ run: |
+ .github/workflows/scripts/unix_build.sh
+ python3 -m pip install dist/*.whl
+
+ - name: Run performance-monitoring
+ run: |
+ cd ..
+ rm -rf performance-monitoring
+ git clone git@github.com:taichi-dev/performance-monitoring.git
+ cd performance-monitoring
+ export WORKFLOW_MODE=postsubmit
+ ./run.sh
+ env:
+ GITHUB_CONTEXT: ${{ toJson(github) }}
diff --git a/.github/workflows/postsubmit.yml b/.github/workflows/postsubmit.yml
deleted file mode 100644
index e7ad3b5bebc28..0000000000000
--- a/.github/workflows/postsubmit.yml
+++ /dev/null
@@ -1,233 +0,0 @@
-name: Postsubmit Checks
-on:
- push:
- branches:
- - master
-
-jobs:
- build_and_test_cpu_required:
- # This job will be required to pass before merging to master branch.
- name: Required Build and Test (CPU)
- timeout-minutes: 60
- strategy:
- matrix:
- include:
- - os: ubuntu-latest
- python: 3.6
- runs-on: ${{ matrix.os }}
- steps:
- - uses: actions/checkout@v2
- with:
- submodules: 'recursive'
-
- - uses: actions/setup-python@v2
- with:
- python-version: ${{ matrix.python }}
-
- - name: Download Pre-Built LLVM 10.0.0
- run: |
- python misc/ci_download.py
- mkdir taichi-llvm
- cd taichi-llvm
- unzip ../taichi-llvm.zip
- env:
- CI_PLATFORM: ${{ matrix.os }}
-
- - name: Build & Install
- run: .github/workflows/scripts/unix_build.sh
- env:
- CI_SETUP_CMAKE_ARGS: -DTI_WITH_OPENGL:BOOL=OFF -DTI_WITH_CC:BOOL=ON -DTI_WITH_VULKAN:BOOL=OFF -DTI_BUILD_TESTS:BOOL=ON
- CXX: clang++
-
- - name: Test
- run: .github/workflows/scripts/unix_test.sh
-
- build_and_test_cpu:
- name: Build and Test (CPU)
- needs: build_and_test_cpu_required
- timeout-minutes: 60
- strategy:
- matrix:
- include:
- - os: macos-latest
- python: 3.7
- with_cc: OFF
- with_cpp_tests: ON
- - os: ubuntu-latest
- python: 3.9
- with_cc: OFF
- with_cpp_tests: OFF
- - os: ubuntu-latest
- python: 3.8
- with_cc: ON
- with_cpp_tests: OFF
- runs-on: ${{ matrix.os }}
- steps:
- - uses: actions/checkout@v2
- with:
- submodules: 'recursive'
-
- - uses: actions/setup-python@v2
- with:
- python-version: ${{ matrix.python }}
-
- - name: Download Pre-Built LLVM 10.0.0
- run: |
- python misc/ci_download.py
- mkdir taichi-llvm
- cd taichi-llvm
- unzip ../taichi-llvm.zip
- env:
- CI_PLATFORM: ${{ matrix.os }}
-
- - name: Build & Install
- run: .github/workflows/scripts/unix_build.sh
- env:
- CI_SETUP_CMAKE_ARGS: -DTI_WITH_OPENGL:BOOL=OFF -DTI_WITH_CC:BOOL=${{ matrix.with_cc }} -DTI_WITH_VULKAN:BOOL=OFF -DTI_BUILD_TESTS:BOOL=${{ matrix.with_cpp_tests }}
- CXX: clang++
- # [DEBUG] Copy this step around to enable debugging inside Github Action instances.
- #- name: Setup tmate session
- # uses: mxschmitt/action-tmate@v3
- # with:
- # limit-access-to-actor: true
-
- - name: Test
- run: .github/workflows/scripts/unix_test.sh
- env:
- RUN_CPP_TESTS: ${{ matrix.with_cpp_tests }}
-
- build_and_test_gpu_linux:
- name: Build and Test (GPU)
- runs-on: [self-hosted, cuda, vulkan, cn]
- timeout-minutes: 60
- steps:
- - uses: actions/checkout@v2
- with:
- submodules: 'recursive'
-
- - name: Build
- run: |
- export PATH=$PATH:/usr/local/cuda/bin
- .github/workflows/scripts/unix_build.sh
- env:
- LLVM_LIB_ROOT_DIR: /opt/taichi-llvm-10.0.0
- CI_SETUP_CMAKE_ARGS: -DTI_WITH_OPENGL:BOOL=ON -DTI_WITH_CC:BOOL=OFF -DTI_WITH_VULKAN:BOOL=ON
- BUILD_NUM_THREADS: 8
- LLVM_PATH: /opt/taichi-llvm-10.0.0/bin
- LLVM_DIR: /opt/taichi-llvm-10.0.0/lib/cmake/llvm
- CXX: clang++-8
-
- - name: Test
- run: .github/workflows/scripts/unix_test.sh
- env:
- DISPLAY: :1
- GPU_TEST: ON
-
- build_and_test_windows:
- name: Build and Test (Windows)
- runs-on: windows-latest
- timeout-minutes: 60
- steps:
-
- - name: Install 7Zip PowerShell
- shell: powershell
- run: Install-Module 7Zip4PowerShell -Force -Verbose
-
- - uses: actions/checkout@v2
- with:
- submodules: 'recursive'
-
- - uses: actions/setup-python@v2
- with:
- python-version: 3.7
-
- - name: Add msbuild to PATH
- uses: microsoft/setup-msbuild@v1.0.2
-
- - name: Download And Install Vulkan
- shell: powershell
- run: |
- Invoke-WebRequest -Uri "https://sdk.lunarg.com/sdk/download/1.2.189.0/windows/VulkanSDK-1.2.189.0-Installer.exe" -OutFile VulkanSDK.exe
- $installer = Start-Process -FilePath VulkanSDK.exe -Wait -PassThru -ArgumentList @("/S");
- $installer.WaitForExit();
-
- - name: Build
- shell: powershell
- run: |
- $env:Path += ";C:/VulkanSDK/1.2.189.0/Bin"
- cd C:\
- Remove-item alias:curl
- curl --retry 10 --retry-delay 5 https://github.com/taichi-dev/taichi_assets/releases/download/llvm10/taichi-llvm-10.0.0-msvc2019.zip -LO
- 7z x taichi-llvm-10.0.0-msvc2019.zip -otaichi_llvm
- curl --retry 10 --retry-delay 5 https://github.com/taichi-dev/taichi_assets/releases/download/llvm10/clang-10.0.0-win.zip -LO
- 7z x clang-10.0.0-win.zip -otaichi_clang
- $env:PATH = ";C:\taichi_llvm\bin;C:\taichi_clang\bin;" + $env:PATH
- clang --version
- cd D:\a\taichi\taichi
- python -m pip install -r requirements_dev.txt
- cd python
- git fetch origin master
- $env:TAICHI_CMAKE_ARGS = $env:CI_SETUP_CMAKE_ARGS
- python build.py build
- cd ..\dist
- $env:WHL = $(dir *.whl)
- python -m pip install $env:WHL
- env:
- PYTHON: C:\hostedtoolcache\windows\Python\3.7.9\x64\python.exe
- CI_SETUP_CMAKE_ARGS: -G "Visual Studio 16 2019" -A x64 -DLLVM_DIR=C:\taichi_llvm\lib\cmake\llvm -DTI_WITH_VULKAN:BOOL=ON
- VULKAN_SDK: C:/VulkanSDK/1.2.189.0
-
- - name: Test
- shell: powershell
- run: |
- $env:PATH = ";C:\taichi_llvm\bin;C:\taichi_clang\bin;" + $env:PATH
- python -c "import taichi"
- python examples/algorithm/laplace.py
- python bin/taichi diagnose
- python bin/taichi changelog
- python bin/taichi test -vr2 -t2
- env:
- PYTHON: C:\hostedtoolcache\windows\Python\3.7.9\x64\python.exe
-
- build_and_test_m1:
- name: Build and Test (Apple M1)
- timeout-minutes: 60
- strategy:
- matrix:
- include:
- - os: macos-latest
- python: 3.8
- defaults:
- run:
- # https://github.com/actions/runner/issues/805#issuecomment-844426478
- shell: "/usr/bin/arch -arch arm64e /bin/bash --noprofile --norc -eo pipefail {0}"
- runs-on: [self-hosted, m1]
- steps:
- - uses: actions/checkout@v2
- with:
- submodules: 'recursive'
-
- - name: Build
- run: |
- python3 -m pip uninstall taichi -y
- rm -rf $HOME/Library/Python/3.8/lib/python/site-packages/taichi
- git --version
- export CXX=clang++
- python3 -m pip install -r requirements_dev.txt
- cd python
- git fetch origin master
- TAICHI_CMAKE_ARGS=$CI_SETUP_CMAKE_ARGS python3 build.py build
- cd ..
- export NUM_WHL=`ls dist/*.whl | wc -l`
- if [ $NUM_WHL -ne 1 ]; then echo 'ERROR: created more than 1 whl.' && exit 1; fi
- python3 -m pip install dist/*.whl
- env:
- CI_SETUP_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: Test
- run: |
- export PATH=$PATH:$HOME/Library/Python/3.8/bin
- python3 examples/algorithm/laplace.py
- TI_LIB_DIR=`python3 -c "import taichi;print(taichi.__path__[0])" | tail -1`
- TI_LIB_DIR="$TI_LIB_DIR/lib" ./build/taichi_cpp_tests
- ti test -vr2 -t4 -x
diff --git a/.github/workflows/presubmit.yml b/.github/workflows/presubmit.yml
deleted file mode 100644
index 398bfed6631aa..0000000000000
--- a/.github/workflows/presubmit.yml
+++ /dev/null
@@ -1,298 +0,0 @@
-name: Presubmit Checks
-on:
- pull_request:
- types: [opened, synchronize, reopened]
-
-jobs:
- title_format:
- name: Check PR Title
- if: ${{ github.event.pull_request }}
- runs-on: ubuntu-latest
- steps:
- - uses: actions/checkout@v2
- - uses: actions/setup-python@v2
- with:
- python-version: 3.8
-
- - name: Run PR Title Checker
- run: |
- pip install semver GitPython
- python misc/ci_check_pr_title.py "$PR_TITLE"
- env:
- PR_TITLE: ${{ github.event.pull_request.title }}
-
- check_code_format:
- name: Check Code Format
- runs-on: ubuntu-latest
- # This job will be required to pass before merging to master branch.
- steps:
- - uses: actions/checkout@v2
- - name: Set up Python ${{ matrix.python-version }}
- uses: actions/setup-python@v2
- with:
- python-version: 3.8
-
- - name: Setup git & clang-format
- run: |
- git config user.email "taichigardener@gmail.com"
- git config user.name "Taichi Gardener"
- git checkout -b _fake_squash
- git remote add upstream https://github.com/taichi-dev/taichi.git
- git fetch upstream master
- sudo apt install clang-format-10
-
- - name: Cache PIP
- uses: actions/cache@v2
- with:
- path: ~/.cache/pip
- key: ${{ hashFiles('setup.py') }}-${{ hashFiles('requirements_dev.txt') }}
-
- - name: Install requirements
- run: |
- python3 -m pip install --user -r requirements_dev.txt
-
- - name: Check code format
- run: |
- python3 misc/code_format.py
- git checkout -b _enforced_format
- git commit -am "enforce code format" || true
- # exit with 1 if there were differences:
- git diff _fake_squash _enforced_format --exit-code
-
- - name: Pylint
- run: |
- # Make sure pylint doesn't regress
- pylint python/taichi/ --disable=all --enable=C0121,C0415
- if [ $? -eq 0 ]
- then
- echo "PASSED: pylint is happy"
- exit 0
- else
- echo "FAILED: please run the pylint command above and make sure it passes"
- exit 1
- fi
-
- build_and_test_cpu_required:
- # This job will be required to pass before merging to master branch.
- name: Required Build and Test (CPU)
- needs: check_code_format
- timeout-minutes: 60
- strategy:
- matrix:
- include:
- - os: ubuntu-latest
- python: 3.6
- runs-on: ${{ matrix.os }}
- steps:
- - uses: actions/checkout@v2
- with:
- submodules: 'recursive'
-
- - uses: actions/setup-python@v2
- with:
- python-version: ${{ matrix.python }}
-
- - name: Download Pre-Built LLVM 10.0.0
- run: |
- python misc/ci_download.py
- mkdir taichi-llvm
- cd taichi-llvm
- unzip ../taichi-llvm.zip
- env:
- CI_PLATFORM: ${{ matrix.os }}
-
- - name: Build & Install
- run: .github/workflows/scripts/unix_build.sh
- env:
- CI_SETUP_CMAKE_ARGS: -DTI_WITH_OPENGL:BOOL=OFF -DTI_WITH_CC:BOOL=ON -DTI_WITH_VULKAN:BOOL=OFF -DTI_BUILD_TESTS:BOOL=ON
- CXX: clang++
-
- - name: Test
- run: .github/workflows/scripts/unix_test.sh
- env:
- RUN_CPP_TESTS: ON
-
- build_and_test_cpu:
- name: Build and Test (CPU)
- needs: build_and_test_cpu_required
- timeout-minutes: 60
- strategy:
- matrix:
- include:
- - os: macos-latest
- python: 3.7
- with_cc: OFF
- with_cpp_tests: ON
- - os: ubuntu-latest
- python: 3.9
- with_cc: OFF
- with_cpp_tests: OFF
- - os: ubuntu-latest
- python: 3.8
- with_cc: ON
- with_cpp_tests: OFF
- runs-on: ${{ matrix.os }}
- steps:
- - uses: actions/checkout@v2
- with:
- submodules: 'recursive'
-
- - uses: actions/setup-python@v2
- with:
- python-version: ${{ matrix.python }}
-
- - name: Download Pre-Built LLVM 10.0.0
- run: |
- python misc/ci_download.py
- mkdir taichi-llvm
- cd taichi-llvm
- unzip ../taichi-llvm.zip
- env:
- CI_PLATFORM: ${{ matrix.os }}
-
- - name: Build & Install
- run: .github/workflows/scripts/unix_build.sh
- env:
- CI_SETUP_CMAKE_ARGS: -DTI_WITH_OPENGL:BOOL=OFF -DTI_WITH_CC:BOOL=${{ matrix.with_cc }} -DTI_WITH_VULKAN:BOOL=OFF -DTI_BUILD_TESTS:BOOL=${{ matrix.with_cpp_tests }}
- CXX: clang++
- # [DEBUG] Copy this step around to enable debugging inside Github Action instances.
- #- name: Setup tmate session
- # uses: mxschmitt/action-tmate@v3
- # with:
- # limit-access-to-actor: true
-
- - name: Test
- run: .github/workflows/scripts/unix_test.sh
- env:
- RUN_CPP_TESTS: ${{ matrix.with_cpp_tests }}
-
- build_and_test_gpu_linux:
- name: Build and Test (GPU)
- needs: check_code_format
- runs-on: [self-hosted, cuda, vulkan, cn]
- timeout-minutes: 60
- steps:
- - uses: actions/checkout@v2
- with:
- submodules: 'recursive'
-
- - name: Build
- run: |
- export PATH=$PATH:/usr/local/cuda/bin
- .github/workflows/scripts/unix_build.sh
- env:
- LLVM_LIB_ROOT_DIR: /opt/taichi-llvm-10.0.0
- CI_SETUP_CMAKE_ARGS: -DTI_WITH_OPENGL:BOOL=ON -DTI_WITH_CC:BOOL=OFF -DTI_WITH_VULKAN:BOOL=ON -DTI_BUILD_TESTS:BOOL=ON
- BUILD_NUM_THREADS: 8
- LLVM_PATH: /opt/taichi-llvm-10.0.0/bin
- LLVM_DIR: /opt/taichi-llvm-10.0.0/lib/cmake/llvm
- CXX: clang++-8
-
- - name: Test
- run: .github/workflows/scripts/unix_test.sh
- env:
- DISPLAY: :1
- GPU_TEST: ON
- RUN_CPP_TESTS: ON
-
- build_and_test_windows:
- name: Build and Test (Windows)
- needs: check_code_format
- runs-on: windows-latest
- timeout-minutes: 90
- steps:
-
- - name: Install 7Zip PowerShell
- shell: powershell
- run: Install-Module 7Zip4PowerShell -Force -Verbose
-
- - uses: actions/checkout@v2
- with:
- submodules: 'recursive'
-
- - uses: actions/setup-python@v2
- with:
- python-version: 3.7
-
- - name: Add msbuild to PATH
- uses: microsoft/setup-msbuild@v1.0.2
-
- - name: Download And Install Vulkan
- shell: powershell
- run: |
- Invoke-WebRequest -Uri "https://sdk.lunarg.com/sdk/download/1.2.189.0/windows/VulkanSDK-1.2.189.0-Installer.exe" -OutFile VulkanSDK.exe
- $installer = Start-Process -FilePath VulkanSDK.exe -Wait -PassThru -ArgumentList @("/S");
- $installer.WaitForExit();
-
- - name: Build
- shell: powershell
- run: |
- $env:Path += ";C:/VulkanSDK/1.2.189.0/Bin"
- cd C:\
- Remove-item alias:curl
- curl --retry 10 --retry-delay 5 https://github.com/taichi-dev/taichi_assets/releases/download/llvm10/taichi-llvm-10.0.0-msvc2019.zip -LO
- 7z x taichi-llvm-10.0.0-msvc2019.zip -otaichi_llvm
- curl --retry 10 --retry-delay 5 https://github.com/taichi-dev/taichi_assets/releases/download/llvm10/clang-10.0.0-win.zip -LO
- 7z x clang-10.0.0-win.zip -otaichi_clang
- $env:PATH = ";C:\taichi_llvm\bin;C:\taichi_clang\bin;" + $env:PATH
- clang --version
- cd D:\a\taichi\taichi
- python -m pip install -r requirements_dev.txt
- cd python
- git fetch origin master
- $env:TAICHI_CMAKE_ARGS = $env:CI_SETUP_CMAKE_ARGS
- python build.py build
- cd ..\dist
- $env:WHL = $(dir *.whl)
- python -m pip install $env:WHL
- env:
- PYTHON: C:\hostedtoolcache\windows\Python\3.7.9\x64\python.exe
- CI_SETUP_CMAKE_ARGS: -G "Visual Studio 16 2019" -A x64 -DLLVM_DIR=C:\taichi_llvm\lib\cmake\llvm -DTI_WITH_VULKAN:BOOL=ON
- VULKAN_SDK: C:/VulkanSDK/1.2.189.0
-
- - name: Test
- shell: powershell
- run: |
- $env:PATH = ";C:\taichi_llvm\bin;C:\taichi_clang\bin;" + $env:PATH
- python -c "import taichi"
- python examples/algorithm/laplace.py
- python bin/taichi diagnose
- python bin/taichi changelog
- python bin/taichi test -vr2 -t2
- env:
- PYTHON: C:\hostedtoolcache\windows\Python\3.7.9\x64\python.exe
-
- build_and_test_m1:
- name: Build and Test (Apple M1)
- needs: check_code_format
- timeout-minutes: 60
- strategy:
- matrix:
- include:
- - os: macos-latest
- python: 3.8
- defaults:
- run:
- # https://github.com/actions/runner/issues/805#issuecomment-844426478
- shell: "/usr/bin/arch -arch arm64e /bin/bash --noprofile --norc -eo pipefail {0}"
- runs-on: [self-hosted, m1]
- steps:
- - uses: actions/checkout@v2
- with:
- submodules: 'recursive'
-
- - name: Build
- run: |
- rm -rf $HOME/Library/Python/3.8/lib/python/site-packages/taichi
- .github/workflows/scripts/unix_build.sh
- env:
- CI_SETUP_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
- CXX: clang++
-
- - name: Test
- run: |
- export PATH=$PATH:$HOME/Library/Python/3.8/bin
- python3 examples/algorithm/laplace.py
- TI_LIB_DIR=`python3 -c "import taichi;print(taichi.__path__[0])" | tail -1`
- TI_LIB_DIR="$TI_LIB_DIR/lib" ./build/taichi_cpp_tests
- ti test -vr2 -t4 -x
diff --git a/.github/workflows/pull_request.yml b/.github/workflows/pull_request.yml
new file mode 100644
index 0000000000000..290a2961bf0a7
--- /dev/null
+++ b/.github/workflows/pull_request.yml
@@ -0,0 +1,30 @@
+name: Presubmit Title Checks
+on:
+ pull_request_target:
+ types: [opened, synchronize, reopened, edited]
+
+jobs:
+ pre_submit:
+ name: Presubmit Title Checks
+ if: ${{ github.event.pull_request }}
+ runs-on: ubuntu-latest
+ steps:
+ - uses: actions/checkout@v2
+ - uses: actions/setup-python@v2
+ with:
+ python-version: 3.8
+ - name: Install Dependencies
+ run: pip install semver GitPython PyGithub
+
+ - name: Run PR Title Checker
+ run: |
+ python misc/ci_check_pr_title.py "$PR_TITLE"
+ env:
+ PR_TITLE: ${{ github.event.pull_request.title }}
+
+ - name: PR Project Card Creation
+ if: github.event.action == 'opened' || github.event.action == 'edited'
+ run: python misc/ci_create_pr_card.py
+ env:
+ GITHUB_TOKEN: ${{ secrets.GARDENER_PAT }}
+ GH_EVENT: ${{ toJson(github.event) }}
diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml
index ee651749ca5a8..a448e4ab1d939 100644
--- a/.github/workflows/release.yml
+++ b/.github/workflows/release.yml
@@ -1,35 +1,65 @@
name: Publishing Release
on:
- release:
- # https://docs.github.com/en/free-pro-team@latest/actions/reference/events-that-trigger-workflows#release
- types: [published]
- # When triggered by schedule and workflow_dispatch, github.event.action is an empty string.
- # We use this to distinguish which taichi to release.
schedule:
- cron: "0 0 * * *"
workflow_dispatch:
+ # Manually trigger the release workflow, a version must be provided
+ inputs:
+ version:
+ description: "The version to release (e.g. v0.8.0)"
+ type: string
+ required: true
+
+env:
+ PROD_PWD: ${{ secrets.PYPI_PWD_PROD }}
+ NIGHT_PWD: ${{ secrets.PYPI_PWD_NIGHTLY }}
+ METADATA_USERNAME: ${{ secrets.METADATA_USERNAME }}
+ METADATA_PASSWORD: ${{ secrets.METADATA_PASSWORD }}
+ METADATA_URL: ${{ secrets.METADATA_URL }}
+ RELEASE_VERSION: ${{ github.event.inputs.version }}
jobs:
+ add_version_to_database:
+ name: Add version to database
+ # Skip running release workflow on forks
+ if: github.repository_owner == 'taichi-dev'
+ runs-on: ubuntu-latest
+ steps:
+ - uses: actions/checkout@v2
+
+ - name: Save new version
+ run: |
+ if [ -z "$RELEASE_VERSION" ]; then
+ echo "Not a production release"
+ exit 0
+ fi
+ python3 -m pip install requests==2.26
+ python3 misc/save_new_version.py
+
# This job set environment matrix with respect to production release and nightly release.
matrix_prep:
runs-on: ubuntu-latest
+ needs: add_version_to_database
outputs:
matrix: ${{ steps.set-matrix.outputs.matrix }}
matrix_osx: ${{ steps.set-matrix.outputs.matrix_osx }}
steps:
- id: set-matrix
run: |
- # For nightly release, we only run on python 3.8
- [ -z "${{ github.event.action }}" ] && matrix="[{\"name\":\"taichi-nightly\",\"python\":\"3.8\"}]"
- # For production release, we run on four python versions.
- [ -z "${{ github.event.action }}" ] || matrix="[{\"name\":\"taichi\",\"python\":\"3.6\"},{\"name\":\"taichi\",\"python\":\"3.7\"},{\"name\":\"taichi\",\"python\":\"3.8\"},{\"name\":\"taichi\",\"python\":\"3.9\"}]"
- echo ::set-output name=matrix::{\"include\":$(echo $matrix)}\"
- # M1 only supports py38 and py39(conda), so change matrix.
- [ -z "${{ github.event.action }}" ] && matrix_osx="[{\"name\":\"taichi-nightly\",\"python\":\"3.8\"}]"
- [ -z "${{ github.event.action }}" ] || matrix_osx="[{\"name\":\"taichi\",\"python\":\"3.8\"},{\"name\":\"taichi\",\"python\":\"3.9\"}]"
- echo ::set-output name=matrix_osx::{\"include\":$(echo $matrix_osx)}\"
-
- build_and_upload_linux:
+ if [[ "$GITHUB_EVENT_NAME" == "workflow_dispatch" ]]; then
+ # For production release, we run on four python versions.
+ echo '::set-output name=matrix::{"include":[{"name":"taichi","python":"3.6","conda_python":"py36"},{"name":"taichi","python":"3.7","conda_python":"py37"},{"name":"taichi","python":"3.8","conda_python":"py38"},{"name":"taichi","python":"3.9","conda_python":"py39"}]}"'
+
+ echo '::set-output name=matrix_osx::{"include":[{"name":"taichi","python":"3.8"},{"name":"taichi","python":"3.9"}]}"'
+ else
+ # For nightly release, we only run on python 3.8
+ echo '::set-output name=matrix::{"include":[{"name":"taichi-nightly","python":"3.8","conda_python":"py38"},{"name":"taichi-nightly","python":"3.10","conda_python":"py310"}]}"'
+
+ # M1 only supports py38 and py39(conda), so change matrix.
+ echo '::set-output name=matrix_osx::{"include":[{"name":"taichi-nightly","python":"3.8"},{"name":"taichi-nightly","python":"3.10"}]}"'
+ fi
+
+ build_and_test_linux:
name: Build and Upload (linux only)
needs: matrix_prep
strategy:
@@ -39,138 +69,115 @@ jobs:
steps:
- uses: actions/checkout@v2
with:
- submodules: 'recursive'
+ submodules: "recursive"
- - name: Create Python Wheel
+ - name: Get sccache cache
+ uses: actions/cache@v2
+ with:
+ path: sccache_cache
+ key: sccache-linux-gpu-${{ github.sha }}
+ restore-keys: |
+ sccache-linux-gpu-
+
+ - name: Build
run: |
- # We hacked here because conda activate in CI won't update python PATH
- # automatically. So we don't activate and use desired python version
- # directly.
- export PATH=/home/buildbot/miniconda3/envs/$PYTHON/bin:$PATH
- TAICHI_REPO_DIR=`pwd`
- export PATH=$LLVM_LIB_ROOT_DIR/bin:/usr/local/cuda/bin:$PATH
- export LLVM_DIR=$LLVM_LIB_ROOT_DIR/lib/cmake/llvm
- export CXX=clang++-8
- python3 -m pip uninstall taichi taichi-nightly -y
- python3 -m pip install -r requirements_dev.txt
- python3 -m pip install twine
- cd python
- git fetch origin master
- export TAICHI_CMAKE_ARGS=$CI_SETUP_CMAKE_ARGS
- python3 build.py build --project_name $PROJECT_NAME
- cd ..
- NUM_WHL=`ls dist/*.whl | wc -l`
- if [ $NUM_WHL -ne 1 ]; then echo 'ERROR: created more than 1 whl.' && exit 1; fi
- python3 -m pip install dist/*.whl
+ mkdir -m777 shared
+ docker create --user dev --name taichi_build --gpus all -v /tmp/.X11-unix:/tmp/.X11-unix \
+ -e DISPLAY -e PY -e GPU_BUILD -e TAICHI_CMAKE_ARGS -e PROJECT_NAME \
+ registry.taichigraphics.com/taichidev-ubuntu18.04:v0.2.1 \
+ /home/dev/${{ github.event.repository.name }}/.github/workflows/scripts/unix_build.sh
+ tar -cf - ../${{ github.event.repository.name }} --mode u=+rwx,g=+rwx,o=+rwx --owner 1000 --group 1000 | docker cp - taichi_build:/home/dev/
+ docker start -a taichi_build
+ docker cp taichi_build:/home/dev/${{ github.event.repository.name }}/dist shared/dist
+ docker cp taichi_build:/home/dev/${{ github.event.repository.name }}/build shared/build
env:
- LLVM_LIB_ROOT_DIR: /opt/taichi-llvm-10.0.0
- BUILD_NUM_THREADS: 8
- CI_SETUP_CMAKE_ARGS: -DTI_WITH_VULKAN:BOOL=ON -DTI_WITH_OPENGL:BOOL=ON -DTI_WITH_CC:BOOL=OFF -DTI_BUILD_TESTS:BOOL=${{ matrix.with_cpp_tests }}
+ PY: ${{ matrix.conda_python }}
+ GPU_BUILD: ON
+ TAICHI_CMAKE_ARGS: -DTI_WITH_OPENGL:BOOL=ON -DTI_WITH_CC:BOOL=OFF -DTI_WITH_VULKAN:BOOL=ON -DTI_BUILD_TESTS:BOOL=ON -DCMAKE_C_COMPILER_LAUNCHER=sccache -DCMAKE_CXX_COMPILER_LAUNCHER=sccache
PROJECT_NAME: ${{ matrix.name }}
- PYTHON: ${{ matrix.python }}
+ DISPLAY: ":1"
- name: Archive Wheel Artifacts
- uses: actions/upload-artifact@v2
+ uses: actions/upload-artifact@v3
with:
- name: ${{ matrix.name }}-py${{ matrix.python }}-linux.whl
- path: dist/*.whl
+ name: dist
+ path: shared/dist/*.whl
+ retention-days: 20
- name: Test
run: |
- export PATH=/home/buildbot/miniconda3/envs/$PYTHON/bin:$PATH
- python3 examples/algorithm/laplace.py
- export DISPLAY=:1
- hash -r
- glewinfo
- ti diagnose
- ti changelog
- ti test -vr2 -t2 -k "not ndarray and not torch"
- # ndarray test might OOM if run with -t2.
- # FIXME: unify this with presubmit.yml to avoid further divergence
- ti test -vr2 -t1 -k "ndarray or torch"
+ docker create --user dev --name taichi_test --gpus all -v /tmp/.X11-unix:/tmp/.X11-unix \
+ -e DISPLAY -e PY -e GPU_TEST registry.taichigraphics.com/taichidev-ubuntu18.04:v0.2.1 \
+ /home/dev/unix_test.sh
+ docker cp .github/workflows/scripts/unix_test.sh taichi_test:/home/dev/unix_test.sh
+ docker cp ./requirements_test.txt taichi_test:/home/dev/requirements_test.txt
+ docker cp shared/dist/ taichi_test:/home/dev/
+ docker cp shared/build/ taichi_test:/home/dev/
+ docker cp tests/ taichi_test:/home/dev/
+ docker start -a taichi_test
env:
- PYTHON: ${{ matrix.python }}
+ PY: ${{ matrix.conda_python }}
+ GPU_TEST: ON
+ DISPLAY: ":1"
- - name: Upload PyPI
- env:
- # https://docs.github.com/en/free-pro-team@latest/actions/reference/encrypted-secrets#using-encrypted-secrets-in-a-workflow
- PROD_PWD: ${{ secrets.PYPI_PWD_PROD }}
- NIGHT_PWD: ${{ secrets.PYPI_PWD_NIGHTLY }}
- PROJECT_NAME: ${{ matrix.name }}
- PYTHON: ${{ matrix.python }}
+ - name: clean docker container
+ if: always()
run: |
- export PATH=/home/buildbot/miniconda3/envs/$PYTHON/bin:$PATH
- cd python
- if [ $PROJECT_NAME == "taichi-nightly" ]; then export PYPI_PWD="$NIGHT_PWD" && python3 build.py upload --skip_build --testpypi --project_name $PROJECT_NAME
- elif [ $PROJECT_NAME == "taichi" ]; then export PYPI_PWD="$PROD_PWD" && python3 build.py upload --skip_build; fi
+ docker rm taichi_build taichi_test -f
- build_and_upload_mac:
+ build_and_test_mac:
name: Build and Upload (macOS only)
needs: matrix_prep
strategy:
fail-fast: false
matrix: ${{ fromJson(needs.matrix_prep.outputs.matrix) }}
- runs-on: macos-latest
+ runs-on: macos-10.15
steps:
- uses: actions/checkout@v2
with:
- submodules: 'recursive'
+ submodules: "recursive"
+
+ - name: Get sccache cache
+ uses: actions/cache@v2
+ with:
+ path: sccache_cache
+ key: sccache-mac-${{ github.sha }}
+ restore-keys: |
+ sccache-mac-
- uses: actions/setup-python@v2
with:
python-version: ${{ matrix.python }}
- name: Download Pre-Built LLVM 10.0.0
- run: |
- python misc/ci_download.py
- mkdir taichi-llvm
- cd taichi-llvm
- unzip ../taichi-llvm.zip
+ run: python misc/ci_download.py
env:
- CI_PLATFORM: macos-latest
+ CI_PLATFORM: macos-10.15
- name: Create Python Wheel
run: |
- TAICHI_REPO_DIR=`pwd`
- export PATH=$TAICHI_REPO_DIR/taichi-llvm/bin/:$PATH
- export CXX=clang++
- python -m pip install -r requirements_dev.txt
- cd python
- git fetch origin master
- export TAICHI_CMAKE_ARGS=$CI_SETUP_CMAKE_ARGS
- python build.py build --project_name $PROJECT_NAME
- cd ..
- NUM_WHL=`ls dist/*.whl | wc -l`
- if [ $NUM_WHL -ne 1 ]; then echo 'ERROR: created more than 1 whl.' && exit 1; fi
- pip install dist/*.whl
+ brew install molten-vk
+ export PATH=$(pwd)/taichi-llvm/bin/:$PATH
+ bash .github/workflows/scripts/unix_build.sh
+ brew uninstall molten-vk
env:
- CI_SETUP_CMAKE_ARGS: -DTI_WITH_VULKAN:BOOL=OFF -DTI_WITH_OPENGL:BOOL=OFF -DTI_WITH_CC:BOOL=OFF -DTI_BUILD_TESTS:BOOL=${{ matrix.with_cpp_tests }}
+ TAICHI_CMAKE_ARGS: -DTI_WITH_VULKAN:BOOL=ON -DTI_WITH_OPENGL:BOOL=OFF -DTI_WITH_CC:BOOL=OFF -DTI_BUILD_TESTS:BOOL=ON -DCMAKE_C_COMPILER_LAUNCHER=sccache -DCMAKE_CXX_COMPILER_LAUNCHER=sccache
PROJECT_NAME: ${{ matrix.name }}
+ CXX: clang++
- name: Archive Wheel Artifacts
- uses: actions/upload-artifact@v2
+ uses: actions/upload-artifact@v3
with:
- name: ${{ matrix.name }}-py${{ matrix.python }}-macos.whl
+ name: dist
path: dist/*.whl
+ retention-days: 20
- name: Test
- run: |
- python examples/algorithm/laplace.py
- ti diagnose
- ti test -vr2 -t2
-
- - name: Upload PyPI
+ run: .github/workflows/scripts/unix_test.sh
env:
- # https://docs.github.com/en/free-pro-team@latest/actions/reference/encrypted-secrets#using-encrypted-secrets-in-a-workflow
- PROD_PWD: ${{ secrets.PYPI_PWD_PROD }}
- NIGHT_PWD: ${{ secrets.PYPI_PWD_NIGHTLY }}
- PROJECT_NAME: ${{ matrix.name }}
- run: |
- cd python
- if [ $PROJECT_NAME == "taichi-nightly" ]; then export PYPI_PWD="$NIGHT_PWD" && python build.py upload --skip_build --testpypi --project_name $PROJECT_NAME
- elif [ $PROJECT_NAME == "taichi" ]; then export PYPI_PWD="$PROD_PWD" && python build.py upload --skip_build; fi
+ TI_WANTED_ARCHS: "cpu"
- build_and_upload_m1:
+ build_and_test_m1:
name: Build and Upload (Apple M1)
needs: matrix_prep
strategy:
@@ -183,51 +190,48 @@ jobs:
steps:
- uses: actions/checkout@v2
with:
- submodules: 'recursive'
+ submodules: "recursive"
+
+ - name: Get sccache cache
+ uses: actions/cache@v2
+ with:
+ path: sccache_cache
+ key: sccache-m1-${{ github.sha }}
+ restore-keys: |
+ sccache-m1-
- name: Build
run: |
+ brew install molten-vk
# We hacked here because conda activate in CI won't update python PATH
# automatically. So we don't activate and use desired python version
# directly.
export PATH=/Users/github/miniforge3/envs/$PYTHON/bin:$PATH
- python3 -m pip uninstall taichi taichi-nightly -y
- git --version
- export CXX=clang++
- python3 -m pip install -r requirements_dev.txt
- cd python
- git fetch origin master
- export TAICHI_CMAKE_ARGS=$CI_SETUP_CMAKE_ARGS
- python3 build.py build --project_name $PROJECT_NAME
- cd ..
- export NUM_WHL=`ls dist/*.whl | wc -l`
- if [ $NUM_WHL -ne 1 ]; then echo 'ERROR: created more than 1 whl.' && exit 1; fi
- python3 -m pip install dist/*.whl
+ bash .github/workflows/scripts/unix_build.sh
+ brew uninstall molten-vk
env:
- CI_SETUP_CMAKE_ARGS: -DTI_WITH_OPENGL:BOOL=OFF -DTI_WITH_CUDA:BOOL=OFF -DTI_WITH_CC:BOOL=OFF -DTI_WITH_VULKAN:BOOL=OFF -DTI_WITH_TESTS:BOOL=ON
+ TAICHI_CMAKE_ARGS: -DTI_WITH_OPENGL:BOOL=OFF -DTI_WITH_CUDA:BOOL=OFF -DTI_WITH_CC:BOOL=OFF -DTI_WITH_VULKAN:BOOL=ON -DTI_BUILD_TESTS:BOOL=ON -DCMAKE_C_COMPILER_LAUNCHER=sccache -DCMAKE_CXX_COMPILER_LAUNCHER=sccache
PROJECT_NAME: ${{ matrix.name }}
PYTHON: ${{ matrix.python }}
+ CXX: clang++
- name: Archive Wheel Artifacts
- uses: actions/upload-artifact@v2
+ uses: actions/upload-artifact@v3
with:
- name: ${{ matrix.name }}-py${{ matrix.python }}-macos-m1.whl
+ name: dist
path: dist/*.whl
+ retention-days: 20
- - name: Upload PyPI
- env:
- # https://docs.github.com/en/free-pro-team@latest/actions/reference/encrypted-secrets#using-encrypted-secrets-in-a-workflow
- PROD_PWD: ${{ secrets.PYPI_PWD_PROD }}
- NIGHT_PWD: ${{ secrets.PYPI_PWD_NIGHTLY }}
- PROJECT_NAME: ${{ matrix.name }}
- PYTHON: ${{ matrix.python }}
+ - name: Test
run: |
export PATH=/Users/github/miniforge3/envs/$PYTHON/bin:$PATH
- cd python
- if [ $PROJECT_NAME == "taichi-nightly" ]; then export PYPI_PWD="$NIGHT_PWD" && python3 build.py upload --skip_build --testpypi --project_name $PROJECT_NAME
- elif [ $PROJECT_NAME == "taichi" ]; then export PYPI_PWD="$PROD_PWD" && python3 build.py upload --skip_build; fi
+ .github/workflows/scripts/unix_test.sh
+ env:
+ TI_WANTED_ARCHS: "metal,vulkan,cpu"
+ PYTHON: ${{ matrix.python }}
+ GPU_TEST: ON
- build_and_upload_macos_1014:
+ build_and_test_macos_1014:
name: Build and Upload (macos 1014)
needs: matrix_prep
strategy:
@@ -237,7 +241,15 @@ jobs:
steps:
- uses: actions/checkout@v2
with:
- submodules: 'recursive'
+ submodules: "recursive"
+
+ - name: Get sccache cache
+ uses: actions/cache@v2
+ with:
+ path: sccache_cache
+ key: sccache-1014-${{ github.sha }}
+ restore-keys: |
+ sccache-1014-
- name: Build
run: |
@@ -247,53 +259,29 @@ jobs:
export PATH=/Users/buildbot6/miniconda3/envs/$PYTHON/bin:$PATH
export LLVM_DIR=/Users/buildbot6/taichi-llvm-10.0.0-macos
export PATH=$LLVM_DIR/bin:$PATH
- python3 -m pip uninstall taichi taichi-nightly -y
- git --version
- export CXX=clang++
- python3 -m pip install -r requirements_dev.txt
- cd python
- git fetch origin master
- export TAICHI_CMAKE_ARGS=$CI_SETUP_CMAKE_ARGS
- python3 build.py build --project_name $PROJECT_NAME
- cd ..
- export NUM_WHL=`ls dist/*.whl | wc -l`
- if [ $NUM_WHL -ne 1 ]; then echo 'ERROR: created more than 1 whl.' && exit 1; fi
- python3 -m pip install dist/*.whl
+ bash .github/workflows/scripts/unix_build.sh
env:
- CI_SETUP_CMAKE_ARGS: -DTI_WITH_OPENGL:BOOL=OFF -DTI_WITH_CUDA:BOOL=OFF -DTI_WITH_CC:BOOL=OFF -DTI_WITH_VULKAN:BOOL=OFF -DTI_WITH_TESTS:BOOL=ON
+ 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 -DCMAKE_C_COMPILER_LAUNCHER=sccache -DCMAKE_CXX_COMPILER_LAUNCHER=sccache
PROJECT_NAME: ${{ matrix.name }}
PYTHON: ${{ matrix.python }}
+ CXX: clang++
- name: Archive Wheel Artifacts
- uses: actions/upload-artifact@v2
+ uses: actions/upload-artifact@v3
with:
- name: ${{ matrix.name }}-py${{ matrix.python }}-macos-1014.whl
+ name: dist
path: dist/*.whl
+ retention-days: 20
- name: Test
run: |
export PATH=/Users/buildbot6/miniconda3/envs/$PYTHON/bin:$PATH
- python examples/algorithm/laplace.py
- ti diagnose
- ti test -vr2 -t2 -a cpu
+ .github/workflows/scripts/unix_test.sh
env:
+ TI_WANTED_ARCHS: "cpu"
PYTHON: ${{ matrix.python }}
- - name: Upload PyPI
- env:
- # https://docs.github.com/en/free-pro-team@latest/actions/reference/encrypted-secrets#using-encrypted-secrets-in-a-workflow
- PROD_PWD: ${{ secrets.PYPI_PWD_PROD }}
- NIGHT_PWD: ${{ secrets.PYPI_PWD_NIGHTLY }}
- PROJECT_NAME: ${{ matrix.name }}
- PYTHON: ${{ matrix.python }}
- run: |
- export PATH=/Users/buildbot6/miniconda3/envs/$PYTHON/bin:$PATH
- cd python
- if [ $PROJECT_NAME == "taichi-nightly" ]; then export PYPI_PWD="$NIGHT_PWD" && python3 build.py upload --skip_build --testpypi --project_name $PROJECT_NAME
- elif [ $PROJECT_NAME == "taichi" ]; then export PYPI_PWD="$PROD_PWD" && python3 build.py upload --skip_build; fi
-
-
- build_and_upload_windows:
+ build_and_test_windows:
name: Build and Upload (Windows only)
needs: matrix_prep
strategy:
@@ -301,79 +289,138 @@ jobs:
matrix: ${{ fromJson(needs.matrix_prep.outputs.matrix) }}
runs-on: windows-latest
steps:
- - name: Install 7Zip PowerShell
- shell: powershell
- run: Install-Module 7Zip4PowerShell -Force -Verbose
-
- uses: actions/checkout@v2
with:
- submodules: 'recursive'
+ submodules: "recursive"
- uses: actions/setup-python@v2
with:
python-version: ${{ matrix.python }}
- - name: Add msbuild to PATH
- uses: microsoft/setup-msbuild@v1.0.2
+ - name: Add Visual Studio Shell to ENV
+ uses: egor-tensin/vs-shell@v2
+ with:
+ arch: x64
- - name: Download And Install Vulkan
- shell: powershell
- run: |
- Invoke-WebRequest -Uri "https://sdk.lunarg.com/sdk/download/1.2.189.0/windows/VulkanSDK-1.2.189.0-Installer.exe" -OutFile VulkanSDK.exe
- $installer = Start-Process -FilePath VulkanSDK.exe -Wait -PassThru -ArgumentList @("/S");
- $installer.WaitForExit();
+ - name: Get sccache cache
+ uses: actions/cache@v2
+ with:
+ path: ccache_cache
+ key: ccache-win64-clang-${{ github.sha }}
+ restore-keys: |
+ ccache-win64-clang-
- name: Build Python Wheel
shell: powershell
run: |
- $env:Path += ";C:/VulkanSDK/1.2.189.0/Bin"
- cd C:\
- Remove-item alias:curl
- curl --retry 10 --retry-delay 5 https://github.com/taichi-dev/taichi_assets/releases/download/llvm10/taichi-llvm-10.0.0-msvc2019.zip -LO
- 7z x taichi-llvm-10.0.0-msvc2019.zip -otaichi_llvm
- curl --retry 10 --retry-delay 5 https://github.com/taichi-dev/taichi_assets/releases/download/llvm10/clang-10.0.0-win.zip -LO
- 7z x clang-10.0.0-win.zip -otaichi_clang
- $env:PATH = ";C:\taichi_llvm\bin;C:\taichi_clang\bin;" + $env:PATH
- clang --version
- cd D:\a\taichi\taichi
- python -m pip install -r requirements_dev.txt
- cd python
- git fetch origin master
- $env:TAICHI_CMAKE_ARGS = $env:CI_SETUP_CMAKE_ARGS
- python build.py build --project_name $env:PROJECT_NAME
- cd ..\dist
- $env:WHL = $(dir *.whl)
- python -m pip install $env:WHL
+ .\.github\workflows\scripts\win_build.ps1 -installVulkan -libsDir C:\
+ venv\Scripts\python -m pip install $(dir dist\*.whl)
env:
- CI_SETUP_CMAKE_ARGS: -G "Visual Studio 16 2019" -A x64 -DLLVM_DIR=C:\taichi_llvm\lib\cmake\llvm -DTI_WITH_VULKAN:BOOL=ON
PROJECT_NAME: ${{ matrix.name }}
- VULKAN_SDK: C:/VulkanSDK/1.2.189.0
- name: Archive Wheel Artifacts
uses: actions/upload-artifact@v2
with:
- name: ${{ matrix.name }}-py${{ matrix.python }}-windows.whl
+ name: dist
path: dist/*.whl
+ retention-days: 20
- name: Test
shell: powershell
run: |
$env:PATH = ";C:\taichi_llvm\bin;C:\taichi_clang\bin;" + $env:PATH
+ . venv\Scripts\activate.ps1
python -c "import taichi"
- python examples/algorithm/laplace.py
- python bin/taichi diagnose
- python bin/taichi test -vr2 -t2
-
- - name: Upload PyPI
- shell: powershell
+ pip install torch
+ ti diagnose
+ python tests/run_tests.py -vr2 -t2
env:
- # https://docs.github.com/en/free-pro-team@latest/actions/reference/encrypted-secrets#using-encrypted-secrets-in-a-workflow
- PROD_PWD: ${{ secrets.PYPI_PWD_PROD }}
- NIGHT_PWD: ${{ secrets.PYPI_PWD_NIGHTLY }}
- PROJECT_NAME: ${{ matrix.name }}
+ TI_SKIP_VERSION_CHECK: ON
+
+ upload_to_pypi:
+ name: Upload release to PyPI
+ needs:
+ [
+ build_and_test_linux,
+ build_and_test_mac,
+ build_and_test_m1,
+ build_and_test_macos_1014,
+ build_and_test_windows,
+ ]
+ runs-on: ubuntu-latest
+ steps:
+ - uses: actions/checkout@v2
+ - uses: actions/setup-python@v2
+ with:
+ python-version: 3.9
+
+ - name: Get dist files
+ uses: actions/download-artifact@v3
+ with:
+ name: dist
+ path: dist
+
+ - name: Upload to PyPI
run: |
- cd python
- if ( $env:PROJECT_NAME -eq "taichi-nightly" ) {$env:PYPI_PWD = "$env:NIGHT_PWD"}
- if ( $env:PROJECT_NAME -eq "taichi-nightly" ) {python build.py upload --skip_build --testpypi --project_name $env:PROJECT_NAME}
- if ( $env:PROJECT_NAME -eq "taichi" ) {$env:PYPI_PWD = "$env:PROD_PWD"}
- if ( $env:PROJECT_NAME -eq "taichi" ) {python build.py upload --skip_build}
+ ls -l dist/
+ if [ -z "$RELEASE_VERSION" ]; then
+ export PROJECT_NAME="taichi-nightly"
+ else
+ export PROJECT_NAME="taichi"
+ fi
+ python -m pip install requests twine
+ python misc/upload_release.py
+
+ create_release:
+ name: Create tag and publish release
+ needs: upload_to_pypi
+ runs-on: ubuntu-latest
+ if: github.event_name == 'workflow_dispatch'
+ steps:
+ - uses: actions/checkout@v2
+ - uses: actions/setup-python@v2
+ with:
+ python-version: 3.9
+
+ - name: Generate Changelog
+ id: changelog
+ run: |
+ pip3 install gitpython
+ content=$(python3 misc/make_changelog.py)
+ echo $content
+ # Escape multiline strings:
+ # https://renehernandez.io/snippets/multiline-strings-as-a-job-output-in-github-actions/
+ content="${content//'%'/'%25'}"
+ content="${content//$'\n'/'%0A'}"
+ content="${content//$'\r'/'%0D'}"
+ echo "::set-output name=content::$content"
+
+ - name: Create tag
+ run: |
+ git config user.email "taichigardener@gmail.com"
+ git config user.name "Taichi Gardener"
+ git tag -a ${RELEASE_VERSION} -m "Release ${RELEASE_VERSION}"
+ git push origin --tags
+
+ - name: Publish release
+ uses: softprops/action-gh-release@v1
+ with:
+ body: ${{ steps.changelog.outputs.content }}
+ tag_name: ${{ github.event.inputs.version }}
+
+ - name: Bump version
+ run: |
+ version_parts=(${RELEASE_VERSION//./ })
+ version_parts[2]=$(expr ${version_parts[2]} + 1)
+ next_version=$(IFS=.; echo "${version_parts[*]}")
+ # Update version.txt
+ git checkout -b "bump/$next_version"
+ echo "$next_version" > version.txt
+ git add version.txt
+ # Commit and push changes
+ git commit -m "Bump version to $next_version"
+ git push origin "bump/$next_version"
+ # Create pull request
+ gh pr create -B master -t "[misc] Bump version to $next_version"
+ env:
+ GITHUB_TOKEN: ${{ secrets.GARDENER_PAT }}
diff --git a/.github/workflows/scripts/check_clang_tidy.sh b/.github/workflows/scripts/check_clang_tidy.sh
new file mode 100755
index 0000000000000..d9db1c9a3433f
--- /dev/null
+++ b/.github/workflows/scripts/check_clang_tidy.sh
@@ -0,0 +1,12 @@
+#!/bin/bash
+
+CI_SETUP_CMAKE_ARGS=$1
+
+cd taichi
+python3 -m pip install -r requirements_dev.txt
+
+rm -rf build && mkdir build && cd build
+cmake $CI_SETUP_CMAKE_ARGS ..
+
+cd ..
+python3 ./scripts/run_clang_tidy.py $PWD/taichi -clang-tidy-binary clang-tidy-10 -checks=-*,performance-inefficient-string-concatenation,readability-identifier-naming -header-filter=$PWD/taichi -p $PWD/build -j2
diff --git a/.github/workflows/scripts/unix_build.sh b/.github/workflows/scripts/unix_build.sh
index f715a8146bf13..a4960f1d2e1ff 100755
--- a/.github/workflows/scripts/unix_build.sh
+++ b/.github/workflows/scripts/unix_build.sh
@@ -1,11 +1,76 @@
+#!/bin/bash
set -ex
-export PATH=`pwd`/taichi-llvm/bin/:$LLVM_PATH:$PATH
-python3 -m pip uninstall taichi taichi-nightly -y
-python3 -m pip install -r requirements_dev.txt
-cd python
-git fetch origin master
-TAICHI_CMAKE_ARGS=$CI_SETUP_CMAKE_ARGS python3 build.py build
-cd ..
-export NUM_WHL=`ls dist/*.whl | wc -l`
-if [ $NUM_WHL -ne 1 ]; then echo `ERROR: created more than 1 whl.` && exit 1; fi
-python3 -m pip install dist/*.whl
+
+check_in_docker() {
+ # This is a temporary solution to detect in a docker, but it should work
+ if [[ $(whoami) == "dev" ]]; then
+ echo "true"
+ else
+ echo "false"
+ fi
+}
+
+IN_DOCKER=$(check_in_docker)
+[[ "$IN_DOCKER" == "true" ]] && cd taichi
+
+setup_sccache() {
+ export SCCACHE_DIR=$(pwd)/sccache_cache
+ export SCCACHE_CACHE_SIZE="128M"
+ export SCCACHE_LOG=error
+ export SCCACHE_ERROR_LOG=$(pwd)/sccache_error.log
+ mkdir -p "$SCCACHE_DIR"
+ echo "sccache dir: $SCCACHE_DIR"
+ ls -la "$SCCACHE_DIR"
+
+ if [[ $OSTYPE == "linux-"* ]]; then
+ wget https://github.com/mozilla/sccache/releases/download/v0.2.15/sccache-v0.2.15-x86_64-unknown-linux-musl.tar.gz
+ tar -xzf sccache-v0.2.15-x86_64-unknown-linux-musl.tar.gz
+ chmod +x sccache-v0.2.15-x86_64-unknown-linux-musl/sccache
+ export PATH=$(pwd)/sccache-v0.2.15-x86_64-unknown-linux-musl:$PATH
+ elif [[ $(uname -m) == "arm64" ]]; then
+ wget https://github.com/mozilla/sccache/releases/download/v0.2.15/sccache-v0.2.15-aarch64-apple-darwin.tar.gz
+ tar -xzf sccache-v0.2.15-aarch64-apple-darwin.tar.gz
+ chmod +x sccache-v0.2.15-aarch64-apple-darwin/sccache
+ export PATH=$(pwd)/sccache-v0.2.15-aarch64-apple-darwin:$PATH
+ else
+ wget https://github.com/mozilla/sccache/releases/download/v0.2.15/sccache-v0.2.15-x86_64-apple-darwin.tar.gz
+ tar -xzf sccache-v0.2.15-x86_64-apple-darwin.tar.gz
+ chmod +x sccache-v0.2.15-x86_64-apple-darwin/sccache
+ export PATH=$(pwd)/sccache-v0.2.15-x86_64-apple-darwin:$PATH
+ fi
+}
+
+setup_python() {
+ if [[ "$IN_DOCKER" == "true" ]]; then
+ source $HOME/miniconda/etc/profile.d/conda.sh
+ conda activate "$PY"
+ fi
+ python3 -m pip uninstall taichi taichi-nightly -y
+ python3 -m pip install -r requirements_dev.txt
+}
+
+build() {
+ git fetch origin master
+ PROJECT_TAGS=""
+ EXTRA_ARGS=""
+ if [ "$PROJECT_NAME" = "taichi-nightly" ]; then
+ PROJECT_TAGS="egg_info --tag-date"
+ fi
+
+ if [[ $OSTYPE == "linux-"* ]]; then
+ EXTRA_ARGS="-p manylinux1_x86_64"
+ fi
+ python3 misc/make_changelog.py origin/master ./ True
+ python3 setup.py $PROJECT_TAGS bdist_wheel $EXTRA_ARGS
+ sccache -s
+}
+
+setup_sccache
+setup_python
+build
+cat "$SCCACHE_ERROR_LOG"
+NUM_WHL=$(ls dist/*.whl | wc -l)
+if [ $NUM_WHL -ne 1 ]; then echo "ERROR: created more than 1 whl." && exit 1; fi
+
+chmod -R 777 "$SCCACHE_DIR"
+rm -f python/CHANGELOG.md
diff --git a/.github/workflows/scripts/unix_test.sh b/.github/workflows/scripts/unix_test.sh
index 907c52160b8be..fca398ae663af 100755
--- a/.github/workflows/scripts/unix_test.sh
+++ b/.github/workflows/scripts/unix_test.sh
@@ -1,15 +1,63 @@
+#!/bin/bash
set -ex
-TAICHI_REPO_DIR=`pwd`
-TI_LIB_DIR=`python3 -c "import taichi;print(taichi.__path__[0])" | tail -1`
-[[ $RUN_CPP_TESTS == "ON" ]] && TI_LIB_DIR="$TI_LIB_DIR/lib" ./build/taichi_cpp_tests
-export PATH=$TAICHI_REPO_DIR/taichi-llvm/bin/:$PATH
-## Only GPU machine uses system python.
-[ -z $GPU_TEST ] || export PATH=$PATH:$HOME/.local/bin
-hash -r
-python3 examples/algorithm/laplace.py
+
+check_in_docker() {
+ # This is a temporary solution to detect in a docker, but it should work
+ if [[ $(whoami) == "dev" ]]; then
+ echo "true"
+ else
+ echo "false"
+ fi
+}
+
+export TI_SKIP_VERSION_CHECK=ON
+export TI_IN_DOCKER=$(check_in_docker)
+
+if [[ "$TI_IN_DOCKER" == "true" ]]; then
+ source $HOME/miniconda/etc/profile.d/conda.sh
+ conda activate "$PY"
+fi
+python3 -m pip install dist/*.whl
+if [ -z "$GPU_TEST" ]; then
+ python3 -m pip install -r requirements_test.txt
+ python3 -m pip install "torch; python_version < '3.10'"
+else
+ ## Only GPU machine uses system python.
+ export PATH=$PATH:$HOME/.local/bin
+ # pip will skip packages if already installed
+ python3 -m pip install -r requirements_test.txt
+fi
ti diagnose
ti changelog
-[ -z $GPU_TEST ] && ti test -vr2 -t2
+echo "wanted archs: $TI_WANTED_ARCHS"
+
+TI_PATH=$(python3 -c "import taichi;print(taichi.__path__[0])" | tail -1)
+TI_LIB_DIR="$TI_PATH/_lib/runtime" ./build/taichi_cpp_tests
-[ -z $GPU_TEST ] || ti test -vr2 -t2 -k "not ndarray and not torch"
-[ -z $GPU_TEST ] || ti test -vr2 -t1 -k "ndarray or torch"
+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" -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" -a vulkan
+ python3 tests/run_tests.py -vr2 -t2 -k "not torch" -a metal
+ python3 tests/run_tests.py -vr2 -t1 -k "torch" -a "$TI_WANTED_ARCHS"
+ else
+ python3 tests/run_tests.py -vr2 -t4 -a "$TI_WANTED_ARCHS"
+ fi
+else
+ # Split per arch to increase parallelism for linux GPU tests
+ if [[ $TI_WANTED_ARCHS == *"cuda"* ]]; then
+ python3 tests/run_tests.py -vr2 -t4 -k "not torch" -a cuda
+ fi
+ if [[ $TI_WANTED_ARCHS == *"cpu"* ]]; then
+ python3 tests/run_tests.py -vr2 -t8 -k "not torch" -a cpu
+ fi
+ if [[ $TI_WANTED_ARCHS == *"vulkan"* ]]; then
+ python3 tests/run_tests.py -vr2 -t8 -k "not torch" -a vulkan
+ fi
+ if [[ $TI_WANTED_ARCHS == *"opengl"* ]]; then
+ python3 tests/run_tests.py -vr2 -t4 -k "not torch" -a opengl
+ fi
+ python3 tests/run_tests.py -vr2 -t1 -k "torch" -a "$TI_WANTED_ARCHS"
+fi
diff --git a/.github/workflows/scripts/win_build.ps1 b/.github/workflows/scripts/win_build.ps1
new file mode 100644
index 0000000000000..e58c179ee6952
--- /dev/null
+++ b/.github/workflows/scripts/win_build.ps1
@@ -0,0 +1,101 @@
+# Build script for windows
+
+param (
+ [switch]$clone = $false,
+ [switch]$installVulkan = $false,
+ [switch]$develop = $false,
+ [switch]$install = $false,
+ [string]$libsDir = "."
+)
+
+$ErrorActionPreference = "Stop"
+
+$RepoURL = 'https://github.com/taichi-dev/taichi'
+
+function WriteInfo($text) {
+ Write-Host -ForegroundColor Green "[BUILD] $text"
+}
+
+# Get sccache
+$env:CCACHE_DIR="${pwd}/ccache_cache"
+$env:CCACHE_MAXSIZE="128M"
+$env:CCACHE_LOGFILE="${pwd}/ccache_error.log"
+WriteInfo("ccache dir: $Env:CCACHE_DIR")
+md "$Env:CCACHE_DIR" -ea 0
+if (-not (Test-Path "ccache-4.5.1-windows-64")) {
+ curl.exe --retry 10 --retry-delay 5 https://github.com/ccache/ccache/releases/download/v4.5.1/ccache-4.5.1-windows-64.zip -LO
+ 7z x ccache-4.5.1-windows-64.zip
+ $env:PATH += ";${pwd}/ccache-4.5.1-windows-64"
+}
+ccache -v -s
+
+# WriteInfo("Install 7Zip")
+# Install-Module 7Zip4PowerShell -Force -Verbose -Scope CurrentUser
+
+if ($clone) {
+ WriteInfo("Clone the repository")
+ git clone --recurse-submodules $RepoURL
+ Set-Location .\taichi
+}
+
+$libsDir = (Resolve-Path $libsDir).Path
+
+if (-not (Test-Path $libsDir)) {
+ New-Item -ItemType Directory -Path $libsDir
+}
+Push-Location $libsDir
+if (-not (Test-Path "taichi_llvm")) {
+ WriteInfo("Download and extract LLVM")
+ curl.exe --retry 10 --retry-delay 5 https://github.com/taichi-dev/taichi_assets/releases/download/llvm10/taichi-llvm-10.0.0-msvc2019.zip -LO
+ 7z x taichi-llvm-10.0.0-msvc2019.zip -otaichi_llvm
+}
+if (-not (Test-Path "taichi_clang")) {
+ WriteInfo("Download and extract Clang")
+ curl.exe --retry 10 --retry-delay 5 https://github.com/taichi-dev/taichi_assets/releases/download/llvm10/clang-10.0.0-win.zip -LO
+ 7z x clang-10.0.0-win.zip -otaichi_clang
+}
+$env:LLVM_DIR = "$libsDir\taichi_llvm"
+$env:TAICHI_CMAKE_ARGS += " -DCMAKE_CXX_COMPILER=clang++ -DCMAKE_C_COMPILER=clang"
+if ($installVulkan) {
+ WriteInfo("Download and install Vulkan")
+ if (-not (Test-Path "VulkanSDK")) {
+ curl.exe --retry 10 --retry-delay 5 https://sdk.lunarg.com/sdk/download/1.2.189.0/windows/VulkanSDK-1.2.189.0-Installer.exe -Lo VulkanSDK.exe
+ $installer = Start-Process -FilePath VulkanSDK.exe -Wait -PassThru -ArgumentList @("/S");
+ $installer.WaitForExit();
+ }
+ $env:VULKAN_SDK = "$libsDir\VulkanSDK\1.2.189.0"
+ $env:PATH += ";$env:VULKAN_SDK\Bin"
+ $env:TAICHI_CMAKE_ARGS += " -DTI_WITH_VULKAN:BOOL=ON"
+}
+
+Pop-Location
+clang --version
+
+WriteInfo("Setting up Python environment")
+python -m venv venv
+. venv\Scripts\activate.ps1
+python -m pip install wheel
+python -m pip install -r requirements_dev.txt
+python -m pip install -r requirements_test.txt
+if (-not $?) { exit 1 }
+WriteInfo("Building Taichi")
+$env:TAICHI_CMAKE_ARGS += " -DCLANG_EXECUTABLE=$libsDir\\taichi_clang\\bin\\clang++.exe"
+$env:TAICHI_CMAKE_ARGS += " -DLLVM_AS_EXECUTABLE=$libsDir\\taichi_llvm\\bin\\llvm-as.exe"
+if ($install) {
+ if ($develop) {
+ python setup.py develop
+ } else {
+ python setup.py install
+ }
+ if (-not $?) { exit 1 }
+ WriteInfo("Build and install finished")
+} else {
+ if ($env:PROJECT_NAME -eq "taichi-nightly") {
+ python setup.py egg_info --tag-date bdist_wheel
+ } else {
+ python setup.py bdist_wheel
+ }
+ if (-not $?) { exit 1 }
+ WriteInfo("Build finished")
+}
+ccache -s -v
diff --git a/.github/workflows/scripts/win_test.ps1 b/.github/workflows/scripts/win_test.ps1
new file mode 100644
index 0000000000000..40ab79826257d
--- /dev/null
+++ b/.github/workflows/scripts/win_test.ps1
@@ -0,0 +1,28 @@
+$ErrorActionPreference = "Stop"
+
+. venv\Scripts\activate.ps1
+python -c "import taichi"
+ti diagnose
+ti changelog
+echo wanted arch: $env:TI_WANTED_ARCHS
+pip install -r requirements_test.txt
+# TODO relax this when torch supports 3.10
+if ("$env:TI_WANTED_ARCHS".Contains("cuda")) {
+ pip install "torch==1.10.1+cu113; python_version < '3.10'" -f https://download.pytorch.org/whl/cu113/torch_stable.html
+} else {
+ pip install "torch; python_version < '3.10'"
+}
+if ("$env:TI_WANTED_ARCHS".Contains("cuda")) {
+ python tests/run_tests.py -vr2 -t4 -k "not torch" -a cuda
+ if (-not $?) { exit 1 }
+}
+if ("$env:TI_WANTED_ARCHS".Contains("cpu")) {
+ python tests/run_tests.py -vr2 -t6 -k "not torch" -a cpu
+ if (-not $?) { exit 1 }
+}
+if ("$env:TI_WANTED_ARCHS".Contains("opengl")) {
+ python tests/run_tests.py -vr2 -t4 -k "not torch" -a opengl
+ if (-not $?) { exit 1 }
+}
+python tests/run_tests.py -vr2 -t2 -k "torch" -a "$env:TI_WANTED_ARCHS"
+if (-not $?) { exit 1 }
diff --git a/.github/workflows/testing.yml b/.github/workflows/testing.yml
new file mode 100644
index 0000000000000..dc8b63a3f86ec
--- /dev/null
+++ b/.github/workflows/testing.yml
@@ -0,0 +1,471 @@
+name: Build and Test
+on:
+ pull_request:
+ types: [opened, synchronize, reopened]
+ push:
+ branches: [master]
+
+concurrency:
+ group: ${{ github.event.number || github.run_id }}
+ cancel-in-progress: true
+
+jobs:
+ check_files:
+ name: Check files
+ # Disable this workflow on forks
+ if: github.repository_owner == 'taichi-dev'
+ outputs:
+ run_job: ${{ steps.check_files.outputs.run_job }}
+ runs-on: ubuntu-latest
+ steps:
+ - name: Checkout code
+ uses: actions/checkout@v2
+ with:
+ fetch-depth: 2
+
+ - name: check modified files
+ id: check_files
+ run: |
+ echo "Concurrency group: ${{ github.event.number || github.run_id }}"
+ echo "=============== list modified files ==============="
+ git diff --name-only @^
+
+ chore_files=( LICENSE CONTRIBUTING.md README.md netlify.toml )
+ chore_dirs=( docs )
+ run_job=false
+
+ for file in $(git diff --name-only @^); do
+ is_chore=false
+
+ for chore_file in ${chore_files[*]}; do
+ [[ ${file} == ${chore_file} ]] && is_chore=true && break
+ done
+
+ for chore_dir in ${chore_dirs[*]}; do
+ [[ ${file} == ${chore_dir}/* ]] && is_chore=true && break
+ done
+
+ if ! ${is_chore}; then
+ run_job=true
+ break
+ fi
+ done
+
+ if ${run_job}; then
+ echo "::set-output name=run_job::true"
+ else
+ echo "::set-output name=run_job::false"
+ fi
+
+ check_code_format:
+ name: Check Code Format
+ runs-on: ubuntu-latest
+ needs: check_files
+ # This job will be required to pass before merging to master branch.
+ steps:
+ - uses: actions/checkout@v2
+ - name: Set up Python ${{ matrix.python-version }}
+ uses: actions/setup-python@v2
+ with:
+ python-version: 3.9
+
+ - name: Setup git & clang-format
+ run: |
+ git config user.email "taichigardener@gmail.com"
+ git config user.name "Taichi Gardener"
+ git checkout -b _fake_squash
+ git remote add upstream https://github.com/taichi-dev/taichi.git
+ git fetch upstream master
+ sudo apt install clang-format-10
+
+ - name: Cache PIP
+ uses: actions/cache@v2
+ with:
+ path: ~/.cache/pip
+ key: ${{ hashFiles('setup.py') }}-${{ hashFiles('requirements_dev.txt') }}
+
+ - name: Install requirements
+ run: |
+ python3 -m pip install --user -r requirements_dev.txt
+
+ - name: Check code format
+ run: |
+ python3 misc/code_format.py
+ git checkout -b _enforced_format
+ git commit -am "enforce code format" || true
+ # exit with 1 if there were differences:
+ git diff _fake_squash _enforced_format --exit-code
+
+ check_static_analyzer:
+ name: Check Static Analyzer
+ runs-on: ubuntu-latest
+ needs: check_files
+ steps:
+ - uses: actions/checkout@v2
+ with:
+ submodules: "recursive"
+
+ - name: Set up Python ${{ matrix.python-version }}
+ uses: actions/setup-python@v2
+ with:
+ python-version: 3.9
+
+ - name: Pylint
+ run: |
+ if [[ ${{needs.check_files.outputs.run_job}} == false ]]; then
+ exit 0
+ fi
+ python3 -m pip install --user pylint
+ # Make sure pylint doesn't regress
+ pylint python/taichi/ --disable=all --enable=$(python scripts/generate_pylint_tags.py)
+ if [ $? -eq 0 ]
+ then
+ echo "PASSED: pylint is happy"
+ exit 0
+ else
+ echo "FAILED: please run the pylint command above and make sure it passes"
+ exit 1
+ fi
+
+ - name: clang-tidy
+ run: |
+ if [[ ${{needs.check_files.outputs.run_job}} == false ]]; then
+ exit 0
+ fi
+ # https://docs.github.com/en/packages/managing-github-packages-using-github-actions-workflows/publishing-and-installing-a-package-with-github-actions#upgrading-a-workflow-that-accesses-ghcrio
+ echo $CR_PAT | docker login ghcr.io -u ${{ github.actor }} --password-stdin
+ docker pull ghcr.io/taichi-dev/taichidev-cpu-ubuntu18.04:v0.2.2
+ docker run -id --user dev --name check_clang_tidy ghcr.io/taichi-dev/taichidev-cpu-ubuntu18.04:v0.2.2 /bin/bash
+ tar -cf - ../${{ github.event.repository.name }} --mode u=+rwx,g=+rwx,o=+rwx --owner 1000 --group 1000 | docker cp - check_clang_tidy:/home/dev/
+ docker exec --user root check_clang_tidy apt install -y clang-tidy-10
+ docker exec --user dev check_clang_tidy /home/dev/taichi/.github/workflows/scripts/check_clang_tidy.sh "$CI_SETUP_CMAKE_ARGS"
+ env:
+ CR_PAT: ${{ secrets.GITHUB_TOKEN }}
+ CI_SETUP_CMAKE_ARGS: -DCMAKE_EXPORT_COMPILE_COMMANDS=ON -DTI_WITH_OPENGL:BOOL=OFF -DTI_WITH_CC:BOOL=ON -DTI_WITH_VULKAN:BOOL=OFF -DTI_BUILD_TESTS:BOOL=OFF
+
+ build_and_test_cpu_linux:
+ name: Build and Test linux (CPU)
+ needs: [check_code_format, check_files]
+ timeout-minutes: 60
+ strategy:
+ matrix:
+ include:
+ - os: ubuntu-latest
+ python: py39
+ with_cc: ON
+ wanted_archs: "cpu,cc"
+ - os: ubuntu-latest
+ python: py310
+ with_cc: ON
+ wanted_archs: "cpu,cc"
+ runs-on: ${{ matrix.os }}
+ permissions:
+ packages: read
+ contents: read
+ steps:
+ - uses: actions/checkout@v2
+ with:
+ submodules: "recursive"
+
+ - name: Get sccache cache
+ uses: actions/cache@v2
+ with:
+ path: sccache_cache
+ key: sccache-linux-${{matrix.with_cc}}-${{ github.sha }}
+ restore-keys: |
+ sccache-linux-${{matrix.with_cc}}-
+
+ - name: Get docker images
+ run: |
+ if [[ ${{needs.check_files.outputs.run_job}} == false ]]; then
+ exit 0
+ fi
+ # https://docs.github.com/en/packages/managing-github-packages-using-github-actions-workflows/publishing-and-installing-a-package-with-github-actions#upgrading-a-workflow-that-accesses-ghcrio
+ echo $CR_PAT | docker login ghcr.io -u ${{ github.actor }} --password-stdin
+ docker pull ghcr.io/taichi-dev/taichidev-cpu-ubuntu18.04:v0.2.2
+ env:
+ CR_PAT: ${{ secrets.GITHUB_TOKEN }}
+
+ - name: Build
+ run: |
+ if [[ ${{needs.check_files.outputs.run_job}} == false ]]; then
+ exit 0
+ fi
+ mkdir -m777 shared
+ docker create --user dev --name taichi_build \
+ -e PY -e PROJECT_NAME -e TAICHI_CMAKE_ARGS \
+ ghcr.io/taichi-dev/taichidev-cpu-ubuntu18.04:v0.2.2 \
+ /home/dev/taichi/.github/workflows/scripts/unix_build.sh
+ # A tarball is needed because sccache needs some permissions that only the file owner has.
+ # 1000 is the uid and gid of user "dev" in the container.
+ # If the uid or gid of the user inside the docker changes, please change the uid and gid in the following line.
+ tar -cf - ../${{ github.event.repository.name }} --mode u=+rwx,g=+rwx,o=+rwx --owner 1000 --group 1000 | docker cp - taichi_build:/home/dev/
+ docker start -a taichi_build
+ rm -rf sccache_cache
+ docker cp taichi_build:/home/dev/taichi/sccache_cache sccache_cache
+ docker cp taichi_build:/home/dev/taichi/dist shared/dist
+ docker cp taichi_build:/home/dev/taichi/build shared/build
+ env:
+ PY: ${{ matrix.python }}
+ PROJECT_NAME: taichi
+ TAICHI_CMAKE_ARGS: -DTI_WITH_OPENGL:BOOL=OFF -DTI_WITH_CC:BOOL=${{ matrix.with_cc }} -DTI_WITH_VULKAN:BOOL=OFF -DTI_BUILD_TESTS:BOOL=ON -DCMAKE_C_COMPILER_LAUNCHER=sccache -DCMAKE_CXX_COMPILER_LAUNCHER=sccache
+
+ - name: Test
+ run: |
+ if [[ ${{needs.check_files.outputs.run_job}} == false ]]; then
+ exit 0
+ fi
+ docker create --user dev --name taichi_test -e PY -e TI_WANTED_ARCHS ghcr.io/taichi-dev/taichidev-cpu-ubuntu18.04:v0.2.2 /home/dev/unix_test.sh
+ docker cp .github/workflows/scripts/unix_test.sh taichi_test:/home/dev/unix_test.sh
+ docker cp shared/dist/ taichi_test:/home/dev/
+ docker cp shared/build/ taichi_test:/home/dev/
+ docker cp ./requirements_test.txt taichi_test:/home/dev/requirements_test.txt
+ docker cp tests/ taichi_test:/home/dev/
+ docker start -a taichi_test
+ env:
+ PY: ${{ matrix.python }}
+ TI_WANTED_ARCHS: ${{ matrix.wanted_archs }}
+
+ - name: clean docker container
+ if: always()
+ run: |
+ docker rm taichi_build taichi_test -f
+
+ build_and_test_cpu_mac:
+ name: Build and Test macos (CPU)
+ needs: [check_code_format, check_files]
+ timeout-minutes: 60
+ strategy:
+ matrix:
+ include:
+ - os: macos-10.15
+ python: 3.7
+ with_cc: OFF
+ with_cpp_tests: ON
+ wanted_archs: "cpu"
+ runs-on: ${{ matrix.os }}
+ steps:
+ - uses: actions/checkout@v2
+ with:
+ submodules: "recursive"
+
+ - uses: actions/setup-python@v2
+ with:
+ python-version: ${{ matrix.python }}
+
+ - name: Get sccache cache
+ uses: actions/cache@v2
+ with:
+ path: sccache_cache
+ key: sccache-mac-${{ github.sha }}
+ restore-keys: |
+ sccache-mac-
+
+ - name: Download Pre-Built LLVM 10.0.0
+ run: |
+ if [[ ${{needs.check_files.outputs.run_job}} == false ]]; then
+ exit 0
+ fi
+ python misc/ci_download.py
+ env:
+ CI_PLATFORM: ${{ matrix.os }}
+
+ - name: Build & Install
+ run: |
+ brew install molten-vk
+ if [[ ${{needs.check_files.outputs.run_job}} == false ]]; then
+ exit 0
+ fi
+ mkdir -p sccache_cache
+ export PATH=`pwd`/taichi-llvm/bin/:$PATH
+ .github/workflows/scripts/unix_build.sh
+ brew uninstall molten-vk
+ env:
+ TAICHI_CMAKE_ARGS: -DTI_WITH_OPENGL:BOOL=OFF -DTI_WITH_CC:BOOL=${{ matrix.with_cc }} -DTI_WITH_VULKAN:BOOL=ON -DTI_BUILD_TESTS:BOOL=${{ matrix.with_cpp_tests }} -DCMAKE_C_COMPILER_LAUNCHER=sccache -DCMAKE_CXX_COMPILER_LAUNCHER=sccache
+ CXX: clang++
+ # [DEBUG] Copy this step around to enable debugging inside Github Action instances.
+ #- name: Setup tmate session
+ # uses: mxschmitt/action-tmate@v3
+ # with:
+ # limit-access-to-actor: true
+
+ - name: Test
+ run: |
+ if [[ ${{needs.check_files.outputs.run_job}} == false ]]; then
+ exit 0
+ fi
+ .github/workflows/scripts/unix_test.sh
+ env:
+ TI_WANTED_ARCHS: ${{ matrix.wanted_archs }}
+
+ build_and_test_gpu_linux:
+ name: Build and Test (GPU)
+ needs: [check_code_format, check_files]
+ runs-on: [self-hosted, cuda, vulkan, cn]
+ timeout-minutes: 60
+ steps:
+ - uses: actions/checkout@v2
+ with:
+ submodules: "recursive"
+
+ - name: Get sccache cache
+ uses: actions/cache@v2
+ with:
+ path: sccache_cache
+ key: sccache-linux-gpu-${{ github.sha }}
+ restore-keys: |
+ sccache-linux-gpu-
+
+ - name: Build & Install
+ run: |
+ if [[ ${{needs.check_files.outputs.run_job}} == false ]]; then
+ exit 0
+ fi
+ mkdir -m777 shared
+ docker create --user dev --name taichi_build --gpus all -v /tmp/.X11-unix:/tmp/.X11-unix \
+ -e PY -e GPU_BUILD -e PROJECT_NAME -e TAICHI_CMAKE_ARGS -e DISPLAY \
+ registry.taichigraphics.com/taichidev-ubuntu18.04:v0.2.1 \
+ /home/dev/taichi/.github/workflows/scripts/unix_build.sh
+ # A tarball is needed because sccache needs some permissions that only the file owner has.
+ # 1000 is the uid and gid of user "dev" in the container.
+ # If the uid or gid of the user inside the docker changes, please change the uid and gid in the following line.
+ tar -cf - ../${{ github.event.repository.name }} --mode u=+rwx,g=+rwx,o=+rwx --owner 1000 --group 1000 | docker cp - taichi_build:/home/dev/
+ docker start -a taichi_build
+ rm -rf sccache_cache
+ docker cp taichi_build:/home/dev/taichi/sccache_cache sccache_cache
+ docker cp taichi_build:/home/dev/taichi/dist shared/dist
+ docker cp taichi_build:/home/dev/taichi/build shared/build
+ env:
+ PY: py38
+ GPU_BUILD: ON
+ PROJECT_NAME: taichi
+ TAICHI_CMAKE_ARGS: -DTI_WITH_OPENGL:BOOL=ON -DTI_WITH_CC:BOOL=OFF -DTI_WITH_VULKAN:BOOL=ON -DTI_BUILD_TESTS:BOOL=ON -DCMAKE_C_COMPILER_LAUNCHER=sccache -DCMAKE_CXX_COMPILER_LAUNCHER=sccache
+ DISPLAY: :1
+
+ - name: Test
+ run: |
+ if [[ ${{needs.check_files.outputs.run_job}} == false ]]; then
+ exit 0
+ fi
+ docker create --user dev --name taichi_test --gpus all -v /tmp/.X11-unix:/tmp/.X11-unix \
+ -e DISPLAY -e PY -e GPU_TEST -e TI_WANTED_ARCHS \
+ registry.taichigraphics.com/taichidev-ubuntu18.04:v0.2.1 \
+ /home/dev/unix_test.sh
+ docker cp .github/workflows/scripts/unix_test.sh taichi_test:/home/dev/unix_test.sh
+ docker cp shared/dist/ taichi_test:/home/dev/
+ docker cp shared/build/ taichi_test:/home/dev/
+ docker cp tests/ taichi_test:/home/dev/
+ docker cp requirements_test.txt taichi_test:/home/dev/requirements_test.txt
+ docker start -a taichi_test
+ env:
+ PY: py38
+ GPU_TEST: ON
+ DISPLAY: :1
+ TI_WANTED_ARCHS: "cpu,cuda,vulkan,opengl"
+
+ - name: clean docker container
+ if: always()
+ run: |
+ docker rm taichi_build taichi_test -f
+
+ build_and_test_windows:
+ name: Build and Test Windows
+ needs: [check_code_format, check_files]
+ runs-on: [self-hosted, windows, gpu]
+ timeout-minutes: 90
+ steps:
+ # See also https://github.com/taichi-dev/taichi/issues/4161
+ - name: Cleanup
+ shell: powershell
+ run: |
+ remove-item '${{ github.workspace }}\*' -recurse -force
+
+ - uses: actions/checkout@v2
+ with:
+ submodules: "recursive"
+
+ - uses: actions/setup-python@v2
+ with:
+ python-version: 3.7
+
+ - name: Add Visual Studio Shell to ENV
+ uses: egor-tensin/vs-shell@v2
+ with:
+ arch: x64
+
+ - name: Get sccache cache
+ uses: actions/cache@v2
+ with:
+ path: ccache_cache
+ key: ccache-win64-${{ github.sha }}
+ restore-keys: |
+ ccache-win64-
+
+ - name: Build
+ shell: powershell
+ if: ${{ needs.check_files.outputs.run_job != 'false' }}
+ run: |
+ .\.github\workflows\scripts\win_build.ps1 -installVulkan -install -libsDir C:\
+
+ - name: Test
+ shell: powershell
+ if: ${{ needs.check_files.outputs.run_job != 'false' }}
+ run: |
+ .\.github\workflows\scripts\win_test.ps1
+ env:
+ TI_WANTED_ARCHS: cpu,cuda,opengl
+ TAICHI_CMAKE_ARGS: -DTI_WITH_OPENGL:BOOL=ON -DTI_WITH_CC:BOOL=OFF
+ TI_SKIP_VERSION_CHECK: ON
+ PYTHON: "3.7"
+
+ build_and_test_m1:
+ name: Build and Test (Apple M1)
+ needs: [check_code_format, check_files]
+ timeout-minutes: 60
+ strategy:
+ matrix:
+ include:
+ - os: macos-latest
+ python: 3.8
+ defaults:
+ run:
+ # https://github.com/actions/runner/issues/805#issuecomment-844426478
+ shell: "/usr/bin/arch -arch arm64e /bin/bash --noprofile --norc -eo pipefail {0}"
+ runs-on: [self-hosted, m1]
+ steps:
+ - uses: actions/checkout@v2
+ with:
+ submodules: "recursive"
+
+ - name: Get sccache cache
+ uses: actions/cache@v2
+ with:
+ path: sccache_cache
+ key: sccache-m1-${{ github.sha }}
+ restore-keys: |
+ sccache-m1-
+
+ - name: Build
+ run: |
+ if [[ ${{needs.check_files.outputs.run_job}} == false ]]; then
+ exit 0
+ fi
+ export PATH=/Users/github/miniforge3/envs/$PY/bin:$PATH
+ brew install molten-vk
+ .github/workflows/scripts/unix_build.sh
+ env:
+ TAICHI_CMAKE_ARGS: -DTI_WITH_OPENGL:BOOL=OFF -DTI_WITH_CUDA:BOOL=OFF -DTI_WITH_CC:BOOL=OFF -DTI_WITH_VULKAN:BOOL=ON -DTI_BUILD_TESTS:BOOL=ON -DCMAKE_C_COMPILER_LAUNCHER=sccache -DCMAKE_CXX_COMPILER_LAUNCHER=sccache
+ PY: ${{ matrix.python }}
+ CXX: clang++
+
+ - name: Test
+ run: |
+ if [[ ${{needs.check_files.outputs.run_job}} == false ]]; then
+ exit 0
+ fi
+ export PATH=/Users/github/miniforge3/envs/$PY/bin:$PATH
+ .github/workflows/scripts/unix_test.sh
+ env:
+ TI_WANTED_ARCHS: "metal,vulkan,cpu"
+ PY: ${{ matrix.python }}
+ PLATFORM: "m1"
diff --git a/.gitignore b/.gitignore
index 902a589b6ca7c..fd39d08f9acea 100644
--- a/.gitignore
+++ b/.gitignore
@@ -57,6 +57,7 @@ __pycache__
*.jpg
!docs/**/*.jpg
!docs/**/*.png
+!tests/python/expected/*.png
*.egg-info
.tlang_cache
/taichi/common/version.h
@@ -77,6 +78,10 @@ _build
*.ll
*.bc
*.yml
+!.github/**/*.yml
*.dot
*.json
+!tests/**/*.json
!docs/**/*.json
+imgui.ini
+/venv/
diff --git a/.gitmodules b/.gitmodules
index e9d6d00209671..be8c04fa168b0 100644
--- a/.gitmodules
+++ b/.gitmodules
@@ -50,3 +50,9 @@
[submodule "external/SPIRV-Cross"]
path = external/SPIRV-Cross
url = https://github.com/KhronosGroup/SPIRV-Cross
+[submodule "external/Vulkan-Headers"]
+ path = external/Vulkan-Headers
+ url = https://github.com/KhronosGroup/Vulkan-Headers
+[submodule "external/FP16"]
+ path = external/FP16
+ url = https://github.com/Maratyszcza/FP16
diff --git a/CMakeLists.txt b/CMakeLists.txt
index 5b5beab468380..ecfb0117f3e0f 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -6,9 +6,12 @@ cmake_minimum_required(VERSION 3.12)
project(taichi)
-SET(TI_VERSION_MAJOR ${TI_VERSION_MAJOR})
-SET(TI_VERSION_MINOR ${TI_VERSION_MINOR})
-SET(TI_VERSION_PATCH ${TI_VERSION_PATCH})
+if (NOT DEFINED TI_VERSION_MAJOR)
+ message(WARNING "It seems that you are running cmake manually, which may cause issues. Please use setup.py to build taichi from source, see https://docs.taichi.graphics/lang/articles/contribution/dev_install for more details.")
+ set(TI_VERSION_MAJOR 0)
+ set(TI_VERSION_MINOR 0)
+ set(TI_VERSION_PATCH 0)
+endif()
set(CMAKE_CXX_STANDARD 17)
@@ -48,10 +51,17 @@ endif ()
set(CMAKE_RUNTIME_OUTPUT_DIRECTORY "${CMAKE_CURRENT_SOURCE_DIR}/build")
set(CMAKE_LIBRARY_OUTPUT_DIRECTORY "${CMAKE_CURRENT_SOURCE_DIR}/build")
-include(cmake/PythonNumpyPybind11.cmake)
+find_program(CCACHE_PROGRAM ccache)
+if(CCACHE_PROGRAM)
+ set(CMAKE_CXX_COMPILER_LAUNCHER "${CCACHE_PROGRAM}")
+endif()
+
+# No support of Python for Android build
+if (NOT ANDROID)
+ include(cmake/PythonNumpyPybind11.cmake)
+endif()
include(cmake/TaichiCXXFlags.cmake)
include(cmake/TaichiCore.cmake)
-include(cmake/TaichiMain.cmake)
option(TI_BUILD_TESTS "Build the CPP tests" OFF)
@@ -60,6 +70,12 @@ if (TI_BUILD_TESTS)
include(cmake/TaichiTests.cmake)
endif()
+option(TI_BUILD_EXAMPLES "Build the CPP examples" ON)
+
+if (TI_BUILD_EXAMPLES)
+ include(cmake/TaichiExamples.cmake)
+endif()
+
include_directories(${PROJECT_SOURCE_DIR}/external/eigen)
message("C++ Flags: ${CMAKE_CXX_FLAGS}")
@@ -70,47 +86,66 @@ if (NOT TI_WITH_CUDA)
set(CUDA_TOOLKIT_ROOT_DIR "")
endif()
-message("python=${PYTHON_EXECUTABLE}")
-
-add_custom_target(
- generate_commit_hash
- COMMAND ${PYTHON_EXECUTABLE} "${CMAKE_CURRENT_LIST_DIR}/misc/generate_commit_hash.py"
-)
-add_dependencies(${CORE_LIBRARY_NAME} generate_commit_hash)
-
if (TI_WITH_CUDA)
set(CUDA_ARCH "cuda")
endif()
-find_program(CLANG_EXECUTABLE NAMES clang clang-7 clang-8 clang-9 clang-10)
+if (CLANG_EXECUTABLE)
+ message("Clang executable ${CLANG_EXECUTABLE}")
+elseif ("${CMAKE_CXX_COMPILER_ID}" MATCHES "Clang")
+ set (CLANG_EXECUTABLE ${CMAKE_CXX_COMPILER})
+ message("Clang executable using host compiler ${CLANG_EXECUTABLE}")
+else()
+ find_program(CLANG_EXECUTABLE NAMES clang clang-10 clang-11 clang-9 clang-8 clang-7)
+ message("Clang executable found at ${CLANG_EXECUTABLE}")
+endif()
+
if (NOT CLANG_EXECUTABLE)
- message(FATAL_ERROR "Cannot find any clang executable.")
+ message(FATAL_ERROR "Cannot find any clang executable.")
+endif()
+
+macro(check_clang_version)
+ execute_process(COMMAND ${CLANG_EXECUTABLE} --version OUTPUT_VARIABLE CLANG_VERSION_OUTPUT)
+ string(REGEX MATCH "([0-9]+)\\.[0-9]+(\\.[0-9]+)?" CLANG_VERSION "${CLANG_VERSION_OUTPUT}")
+ message("${CLANG_EXECUTABLE} --version: ${CLANG_VERSION}")
+
+ set(CLANG_VERSION_MAJOR "${CMAKE_MATCH_1}")
+endmacro()
+
+if (APPLE)
+ set(CLANG_OSX_FLAGS "-isysroot${CMAKE_OSX_SYSROOT}")
+ set(CLANG_HIGHEST_VERSION "13")
+else()
+ set(CLANG_HIGHEST_VERSION "11")
endif()
-find_program(LLVM_AS_EXECUTABLE NAMES llvm-as)
-if (NOT LLVM_AS_EXECUTABLE)
- message(FATAL_ERROR "Cannot find llvm-as executable.")
+check_clang_version()
+
+if (${CLANG_VERSION_MAJOR} VERSION_GREATER ${CLANG_HIGHEST_VERSION})
+ unset(CLANG_EXECUTABLE)
+ find_program(CLANG_EXECUTABLE NAMES clang-10 clang-11 clang-9 clang-8 clang-7)
+ if (NOT CLANG_EXECUTABLE)
+ message(FATAL_ERROR "${CLANG_EXECUTABLE} version: ${CLANG_VERSION}, required: <=${CLANG_HIGHEST_VERSION}. Condider passing -DCLANG_PATH=/path/to/clang to cmake to use a specific clang.")
+ else()
+ check_clang_version()
+ if (${CLANG_VERSION_MAJOR} VERSION_GREATER ${CLANG_HIGHEST_VERSION})
+ message(FATAL_ERROR "${CLANG_EXECUTABLE} version: ${CLANG_VERSION}, required: <=${CLANG_HIGHEST_VERSION}. Condider passing -DCLANG_PATH=/path/to/clang to cmake to use a specific clang.")
+ endif()
+ endif()
endif()
# Build llvm-runtime for host arch and cuda (if available)
foreach(arch IN LISTS HOST_ARCH CUDA_ARCH)
add_custom_target(
"generate_llvm_runtime_${arch}"
- COMMAND ${CLANG_EXECUTABLE} -S runtime.cpp -o "runtime_${arch}.ll" -fno-exceptions -emit-llvm -std=c++17 -D "ARCH_${arch}" -I ${PROJECT_SOURCE_DIR};
- COMMAND ${LLVM_AS_EXECUTABLE} "runtime_${arch}.ll" -o "runtime_${arch}.bc"
+ COMMAND ${CLANG_EXECUTABLE} ${CLANG_OSX_FLAGS} -c runtime.cpp -o "runtime_${arch}.bc" -fno-exceptions -emit-llvm -std=c++17 -D "ARCH_${arch}" -I ${PROJECT_SOURCE_DIR};
WORKING_DIRECTORY "${PROJECT_SOURCE_DIR}/taichi/runtime/llvm"
)
add_dependencies(${CORE_LIBRARY_NAME} "generate_llvm_runtime_${arch}")
endforeach()
-
-FILE(WRITE ${CMAKE_CURRENT_LIST_DIR}/taichi/common/version.h
- "#pragma once\n"
- "#define TI_VERSION_MAJOR \"${TI_VERSION_MAJOR}\"\n"
- "#define TI_VERSION_MINOR \"${TI_VERSION_MINOR}\"\n"
- "#define TI_VERSION_PATCH \"${TI_VERSION_PATCH}\"\n"
- "#define TI_CUDAVERSION \"${CUDA_VERSION}\"\n"
- )
+configure_file(taichi/common/version.h.in ${CMAKE_SOURCE_DIR}/taichi/common/version.h)
+configure_file(taichi/common/commit_hash.h.in ${CMAKE_SOURCE_DIR}/taichi/common/commit_hash.h)
option(TI_EXPORT_CORE "export taichi core" OFF)
diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md
new file mode 100644
index 0000000000000..1c6fa35ec9f4c
--- /dev/null
+++ b/CONTRIBUTING.md
@@ -0,0 +1,13 @@
+# Contributing Guide
+
+Thank you for your interest in contributing to Taichi! Please check out the [Contribution Guidelines](https://docs.taichi.graphics/lang/articles/contribution/contributor_guide) for how to make a contribution.
+
+## Developer installation
+
+Taichi is developed mainly in C++17 and Python3. Please check out the [Developer Installation](https://docs.taichi.graphics/lang/articles/contribution/dev_install) to build Taichi from source. Note that Taichi is LLVM-10.0.0 dependent and that we recommend installing [our pre-built LLVM libraries](https://docs.taichi.graphics/lang/articles/contribution/dev_install#installing-dependencies) for your platform.
+
+## Contribution opportunities
+
+Issues marked with ["welcome contribution"](https://github.com/taichi-dev/taichi/issues?q=is%3Aopen+is%3Aissue+label%3A%22welcome+contribution%22) are great places for starters. You can quickly get an idea of the entire workflow and how to join the community.
+
+**RFC**: We use the `RFC` (Request for Comments) mechanism to discuss and organize some of the more advanced and self-contained features. These are the projects that we would like to work on but still lack a concrete design or implementation roadmap for because of their complexity. We document these requests and the threaded proposals in the hope that we could provide the community with a good enough context and draw upon insights from the potentially passionate minds. You can find all the ongoing RFCs [here](https://github.com/taichi-dev/taichi/issues?q=is%3Aissue+is%3Aopen+label%3ARFC+), and you are also welcome to file new RFCs with us!
diff --git a/MANIFEST.in b/MANIFEST.in
index 0ca0ff3bd4272..cbc8a16b7dc89 100644
--- a/MANIFEST.in
+++ b/MANIFEST.in
@@ -1,14 +1,15 @@
include MANIFEST.in
+include version.txt
include python/*.txt
include python/*.py
include *.cfg
include python/taichi/*.md
include python/taichi/assets/*
recursive-include python/taichi/examples *.py
-include python/taichi/tests/*
-include python/taichi/lib/*.so
-include python/taichi/lib/*.pyd
-include python/taichi/lib/*.bc
+include python/taichi/_lib/core/*.so
+include python/taichi/_lib/core/*.pyd
+include python/taichi/_lib/runtime/*.bc
+include python/taichi/_lib/runtime/*.dylib
include python/taichi/shaders/*.spv
include python/taichi/shaders/*.vert
include python/taichi/shaders/*.frag
diff --git a/README.md b/README.md
index f6d03d805bd9d..d07cd9dc8345c 100644
--- a/README.md
+++ b/README.md
@@ -1,89 +1,178 @@
+
-
-
-
+
+**Taichi (太极)** is an open-source, imperative, parallel programming language for high-performance numerical computation. It is embedded in Python and uses just-in-time (JIT) compiler frameworks (e.g. LLVM) to offload compute-intensive Python code to the native GPU or CPU instructions.
-Advanced features of Taichi include [spatially sparse computing](https://docs.taichi.graphics/lang/articles/advanced/sparse), [differentiable programming](https://docs.taichi.graphics/lang/articles/advanced/differentiable_programming) [[examples]](https://github.com/yuanming-hu/difftaichi), and [quantized computation](https://github.com/taichi-dev/quantaichi).
+Advantages of Taichi:
-**Please check out our SIGGRAPH 2020 course on Taichi basics:** [YouTube](https://youtu.be/Y0-76n3aZFA), [Bilibili](https://www.bilibili.com/video/BV1kA411n7jk/), [slides (pdf)](https://yuanming.taichi.graphics/publication/2020-taichi-tutorial/taichi-tutorial.pdf).
+- Built around Python: Taichi shares almost the same syntax with Python, allowing you to write algorithms with minimal language barrier. It is also well integrated into the Python ecosystem, such as NumPy and PyTorch.
+- Flexibility: Taichi provides a set of generic data containers known as *SNode* (/ˈsnoʊd/), an effective mechanism for composing hierarchical, multi-dimensional fields. This can cover many use patterns in numerical simulation (e.g. [spatially sparse computing](https://docs.taichi.graphics/lang/articles/advanced/sparse)).
+- Performance: Through the `@ti.kernel` decorator, Taichi's JIT compiler automatically compiles your Python functions into efficient GPU or CPU machine code for parallel execution.
+- Portability: Write your code once and run it everywhere. Currently, Taichi supports most mainstream GPU APIs, such as CUDA and Vulkan.
+- ... and many more features! A cross-platform, Vulkan-based 3D visualizer, [differentiable programming](https://docs.taichi.graphics/lang/articles/advanced/differentiable_programming), [quantized computation](https://github.com/taichi-dev/quantaichi) (experimental), etc.
-**中文视频教程:** [[哔哩哔哩]](https://www.bilibili.com/video/BV1gA411j7H5), [[幻灯片]](https://yuanming.taichi.graphics/publication/2020-taichi-tutorial/taichi-tutorial.pdf)
+# Getting Started
-## Examples ([More...](misc/examples.md))
+## Installation
-
- 
-
-
+You can easily install Taichi with Python's package installer `pip`:
-
-
+```bash
+pip install taichi
+```
-## Installation [](https://pepy.tech/project/taichi)
+If you want to try out the latest features, we also provide a nightly package:
```bash
-python3 -m pip install taichi
+pip install -i https://test.pypi.org/simple/ taichi-nightly
```
-**Supported OS**: Windows, Linux, Mac OS X; **Python**: 3.6-3.9 (64-bit only); **Backends**: x64 CPUs, CUDA, Apple Metal, Vulkan, OpenGL Compute Shaders.
+*The nightly package can and will break from time to time!*
-Please build from source for other configurations (e.g., your CPU is ARM, or you want to try out our experimental C backend).
+**Supported environments**
-**Note:**
- - The PyPI package supports x64 CPU, CUDA 10/11, Metal, and OpenGL Compute Shader backends.
- - On Ubuntu 19.04+, please `sudo apt install libtinfo5`.
- - On Windows, please install [Microsoft Visual C++ Redistributable](https://aka.ms/vs/16/release/vc_redist.x64.exe) if you haven't.
- - [[All releases]](https://github.com/taichi-dev/taichi/releases)
+
+- Operating systems
+ - Windows[1](#win-note)
+ - Linux
+ - macOS
+- Python: 3.6 ~ 3.9 (64-bit only)
+- Compute backends
+ - x64/ARM CPUs
+ - CUDA
+ - Vulkan
+ - OpenGL (4.3+)
+ - Apple Metal
+ - WebAssembly (experiemental)
-|| **Linux (CUDA)** | **OS X (10.14+)** | **Windows** | **Documentation**|
-|:------|:-----|:-----|:-----|:-----|
-|**Build**|[](http://f11.csail.mit.edu:8080/job/taichi/)| [](https://travis-ci.com/taichi-dev/taichi) | [](https://ci.appveyor.com/project/yuanming-hu/taichi/branch/master)| [](https://app.netlify.com/sites/docs-taichi-graphics/deploys)|
-|**PyPI**|[](https://travis-ci.com/yuanming-hu/taichi-wheels-test)|[](https://travis-ci.com/yuanming-hu/taichi-wheels-test)|[](https://ci.appveyor.com/project/yuanming-hu/taichi-wheels-test) |
+1. On Windows, please install [Microsoft Visual C++ Redistributable](https://aka.ms/vs/16/release/vc_redist.x64.exe) first.
-## Developer Installation
+## Your first Taichi program
-Please follow [this doc](https://docs.taichi.graphics/lang/articles/contribution/dev_install) to learn how to build Taichi from source. Note that Taichi requires LLVM-10.0.0, and it is recommneded to use [our prebuilt LLVM libraries](https://docs.taichi.graphics/lang/articles/contribution/dev_install#installing-dependencies) for each platform.
+Here's how you can program a 2D fractal in Taichi:
-## Contributors
+```py
+# python/taichi/examples/simulation/fractal.py
-
+import taichi as ti
-*Note: contributor avatars above are randomly shuffled.*
+ti.init(arch=ti.gpu)
--------------------------------
+n = 320
+pixels = ti.field(dtype=float, shape=(n * 2, n))
-We welcome feedback and comments. If you would like to contribute to Taichi, please check out our [Contributor Guidelines](https://docs.taichi.graphics/lang/articles/contribution/contributor_guide).
-If you use Taichi in your research, please cite related papers:
+@ti.func
+def complex_sqr(z):
+ return ti.Vector([z[0]**2 - z[1]**2, z[1] * z[0] * 2])
+
+
+@ti.kernel
+def paint(t: float):
+ for i, j in pixels: # Parallelized over all pixels
+ c = ti.Vector([-0.8, ti.cos(t) * 0.2])
+ z = ti.Vector([i / n - 1, j / n - 0.5]) * 2
+ iterations = 0
+ while z.norm() < 20 and iterations < 50:
+ z = complex_sqr(z) + c
+ iterations += 1
+ pixels[i, j] = 1 - iterations * 0.02
+
+
+gui = ti.GUI("Julia Set", res=(n * 2, n))
+
+for i in range(1000000):
+ paint(i * 0.03)
+ gui.set_image(pixels)
+ gui.show()
+```
+
+If Taichi is properly installed, you should get the animation below 🎉:
+
+
-- [**(SIGGRAPH Asia 2019) Taichi: High-Performance Computation on Sparse Data Structures**](https://yuanming.taichi.graphics/publication/2019-taichi/taichi-lang.pdf) [[Video]](https://youtu.be/wKw8LMF3Djo) [[BibTex]](https://raw.githubusercontent.com/taichi-dev/taichi/master/misc/taichi_bibtex.txt) [[Code]](https://github.com/taichi-dev/taichi)
-- [**(ICLR 2020) DiffTaichi: Differentiable Programming for Physical Simulation**](https://arxiv.org/abs/1910.00935) [[Video]](https://www.youtube.com/watch?v=Z1xvAZve9aE) [[BibTex]](https://raw.githubusercontent.com/taichi-dev/taichi/master/misc/difftaichi_bibtex.txt) [[Code]](https://github.com/yuanming-hu/difftaichi)
-- [**(SIGGRAPH 2021) QuanTaichi: A Compiler for Quantized Simulations**](https://yuanming.taichi.graphics/publication/2021-quantaichi/quantaichi.pdf) [[Video]](https://www.youtube.com/watch?v=0jdrAQOxJlY) [[BibTex]](https://raw.githubusercontent.com/taichi-dev/taichi/master/misc/quantaichi_bibtex.txt) [[Code]](https://github.com/taichi-dev/quantaichi)
-## Links
-- [TaichiCon](https://github.com/taichi-dev/taichicon): Taichi developer conferences.
-- [GAMES 201 Lectures](https://github.com/taichi-dev/games201): (Chinese) A hands-on course on building advanced physics engines, based on Taichi.
-- [TaichiZoo](https://zoo.taichi.graphics): Running Taichi code in your browser [1](#zoo-disclaimer).
-- [加入太极图形](https://app.mokahr.com/apply/taichi/41024#/).
-- [太极图形课](https://github.com/taichiCourse01).
+# Documentation
+
+Taichi's documentation is available at https://docs.taichi.graphics.
+
+# Contacts
+
+We use these channels to report bugs, discuss design, show off demos and send announcements on a daily basis:
+
+- [GitHub Issues](https://github.com/taichi-dev/taichi/issues)
+- [GitHub Discussions](https://github.com/taichi-dev/taichi/discussions)
+- [Twitter](https://twitter.com/taichigraphics)
+- [Taichi 中文论坛](https://forum.taichi.graphics/)
+- Slack & Wechat groups: please send us a message at contact@taichi.graphics first, thanks!
+
+Should you spot any security issue, do not hesitate to report it by mailing to security@taichi.graphics.
+
+# Contributing
+
+If you would like to contribute to Taichi, please check out the [Contribution Guidelines](CONTRIBUTING.md).
+
+A huge thanks to all of our amazing contributors!
+
+
+
+*Contributor avatars are randomly shuffled.*
+
+# Resources
+
+## Demos
+
+- [Taichi examples](https://github.com/taichi-dev/taichi/tree/master/python/taichi/examples)
+- [Advanced Taichi examples](https://github.com/taichi-dev/advanced_examples)
+- [DiffTaichi](https://github.com/taichi-dev/difftaichi)
+- [Taichi elements](https://github.com/taichi-dev/taichi_elements)
+- [Taichi houdini](https://github.com/taichi-dev/taichi_houdini)
- [More...](misc/links.md)
-## Security
+
+
+
+
+
+
+
-Please disclose security issues responsibly to contact@taichi.graphics.
+## Lectures & talks
+- **SIGGRAPH 2020 course on Taichi basics**: [YouTube](https://youtu.be/Y0-76n3aZFA), [Bilibili](https://www.bilibili.com/video/BV1kA411n7jk/), [slides (pdf)](https://yuanming.taichi.graphics/publication/2020-taichi-tutorial/taichi-tutorial.pdf).
+- Chinagraph 2020 用太极编写物理引擎: [哔哩哔哩](https://www.bilibili.com/video/BV1gA411j7H5)
+- GAMES 201 高级物理引擎实战指南2020: [课件](https://github.com/taichi-dev/games201)
+- 太极图形课第一季：[课件](https://github.com/taichiCourse01)
+- [TaichiCon](https://github.com/taichi-dev/taichicon): Taichi developer conferences
+- More to come...
---
-1. TaichiZoo is still in its Beta version. If you've encountered any issue, please do not hesitate to [file a bug](https://github.com/taichi-dev/taichi-zoo-issue-tracker/issues/new/choose).
+If you use Taichi in your research, please cite related papers:
+
+- [**(SIGGRAPH Asia 2019) Taichi: High-Performance Computation on Sparse Data Structures**](https://yuanming.taichi.graphics/publication/2019-taichi/taichi-lang.pdf) [[Video]](https://youtu.be/wKw8LMF3Djo) [[BibTex]](https://raw.githubusercontent.com/taichi-dev/taichi/master/misc/taichi_bibtex.txt) [[Code]](https://github.com/taichi-dev/taichi)
+- [**(ICLR 2020) DiffTaichi: Differentiable Programming for Physical Simulation**](https://arxiv.org/abs/1910.00935) [[Video]](https://www.youtube.com/watch?v=Z1xvAZve9aE) [[BibTex]](https://raw.githubusercontent.com/taichi-dev/taichi/master/misc/difftaichi_bibtex.txt) [[Code]](https://github.com/yuanming-hu/difftaichi)
+- [**(SIGGRAPH 2021) QuanTaichi: A Compiler for Quantized Simulations**](https://yuanming.taichi.graphics/publication/2021-quantaichi/quantaichi.pdf) [[Video]](https://www.youtube.com/watch?v=0jdrAQOxJlY) [[BibTex]](https://raw.githubusercontent.com/taichi-dev/taichi/master/misc/quantaichi_bibtex.txt) [[Code]](https://github.com/taichi-dev/quantaichi)
diff --git a/appveyor.yml b/appveyor.yml
deleted file mode 100644
index cb426cd204ac6..0000000000000
--- a/appveyor.yml
+++ /dev/null
@@ -1,65 +0,0 @@
-#---------------------------------#
-# general configuration #
-#---------------------------------#
-
-# version format
-version: 0.0.{build}-{branch}
-
-#---------------------------------#
-# environment configuration #
-#---------------------------------#
-
-image: Visual Studio 2019
-clone_folder: C:\taichi
-
-#---------------------------------#
-# build configuration #
-#---------------------------------#
-
-platform: x64
-configuration: Release
-
-environment:
- matrix:
- - PYTHON: C:\Python36-x64\python.exe
- - PYTHON: C:\Python37-x64\python.exe
- - PYTHON: C:\Python38-x64\python.exe
- - PYTHON: C:\Python39-x64\python.exe
-
-skip_commits:
- files:
- - '*.md'
- - '*.rst'
- - docs
- - benchmarks
- - examples
- - misc
- - '.*'
-
-cache:
- - build -> CMakeLists.txt, cmake/*
-
-build_script:
- - set TAICHI_REPO_DIR=C:\taichi
- - "%PYTHON% %TAICHI_REPO_DIR%/misc/appveyor_filter.py || appveyor exit 0"
- - cd C:\
- - curl --retry 10 --retry-delay 5 https://github.com/taichi-dev/taichi_assets/releases/download/llvm10/taichi-llvm-10.0.0-msvc2019.zip -LO
- - 7z x taichi-llvm-10.0.0-msvc2019.zip -otaichi_llvm
- - curl --retry 10 --retry-delay 5 https://github.com/taichi-dev/taichi_assets/releases/download/llvm10/clang-10.0.0-win.zip -LO
- - "echo \"%APPVEYOR_REPO_COMMIT_MESSAGE%\" | grep '^\\[format\\]' && curl http://kun.csail.mit.edu:31415/%APPVEYOR_PULL_REQUEST_NUMBER% -LO || true"
- - 7z x clang-10.0.0-win.zip -otaichi_clang
- - set PATH=C:\taichi_llvm\bin;%PATH%;
- - set PATH=C:\taichi_clang\bin;%PATH%
- - clang --version
- - cd C:\taichi
- - set CI_SETUP_CMAKE_ARGS=-G "Visual Studio 16 2019" -A x64 -DLLVM_DIR=C:\taichi_llvm\lib\cmake\llvm
- # This is to fix the bug caused by execnet 1.2, the library xdist uses to schedule tests.
- # Reverting to v1.1 solves this bug(Python Fatal Error: IOError).
- - "%PYTHON% -m pip install -U execnet==1.1"
- - "%PYTHON% misc/ci_setup.py ci"
- - '%PYTHON% -c "import taichi"'
- - "%PYTHON% examples/algorithm/laplace.py"
- - "%PYTHON% bin/taichi diagnose"
- - "%PYTHON% bin/taichi test -Cvr2 -t2"
- - "cd python && %PYTHON% build.py try_upload"
- - "cd %TAICHI_REPO_DIR% && bash <(curl -s https://codecov.io/bash)"
diff --git a/benchmarks/README.md b/benchmarks/README.md
new file mode 100644
index 0000000000000..23af1a0f89dea
--- /dev/null
+++ b/benchmarks/README.md
@@ -0,0 +1,41 @@
+Install a few of extra requirements:
+```bash
+python3 -m pip install -r requirements.txt
+```
+
+## Run
+
+To run all benchmarks:
+```bash
+python3 run.py
+```
+
+## Result
+
+The benchmark results will be stored in the `results` folder in your current directory.
+If you wish to save the results as a single json file (`./results/results.json`):
+```bash
+python3 deserialize.py
+```
+Or you can specify the input and output path:
+```bash
+python3 deserialize.py --folder PATH_OF_RESULTS_FOLDER --output_path PATH_YOU_WIHS_TO_STORE
+```
+
+## Tools
+
+After getting benchmark results (`./results`), you can use a visualization tool to profile performance problems:
+```bash
+python3 visualization.py
+```
+
+You can specify the results file path:
+```bash
+python3 visualization.py --folder PATH_OF_RESULTS_FOLDER
+```
+
+The default host and port is `localhost:5006\visualization`.
+If you want to enable remote access, take the following steps:
+```bash
+python3 visualization.py --host YOUR_IP_ADDRESS --port PORT_YOU_WISH_TO_USE
+```
diff --git a/benchmarks/async_advection.py b/benchmarks/async_advection.py
deleted file mode 100644
index 916fd8314357a..0000000000000
--- a/benchmarks/async_advection.py
+++ /dev/null
@@ -1,126 +0,0 @@
-import math
-
-from utils import benchmark_async
-
-import taichi as ti
-
-# TODO: staggerred grid
-
-
-@benchmark_async
-def simple_advection(scale):
- n = 256 * 2**int((math.log(scale, 2)) // 2)
- x = ti.Vector.field(3, dtype=ti.f32, shape=(n, n))
- new_x = ti.Vector.field(3, dtype=ti.f32, shape=(n, n))
- v = ti.Vector.field(2, dtype=ti.f32, shape=(n, n))
- dx = 1 / n
- inv_dx = 1 / dx
- dt = 0.01
-
- stagger = ti.Vector([0.5, 0.5])
-
- @ti.func
- def Vector2(x, y):
- return ti.Vector([x, y])
-
- @ti.kernel
- def init():
- for i, j in v:
- v[i, j] = ti.Vector([j / n - 0.5, 0.5 - i / n])
-
- for i, j in ti.ndrange(n * 4, n * 4):
- ret = ti.taichi_logo(ti.Vector([i, j]) / (n * 4))
- x[i // 4, j // 4][0] += ret / 16
- x[i // 4, j // 4][1] += ret / 16
- x[i // 4, j // 4][2] += ret / 16
-
- @ti.func
- def vec(x, y):
- return ti.Vector([x, y])
-
- @ti.func
- def clamp(p):
- for d in ti.static(range(p.n)):
- p[d] = min(1 - 1e-4 - dx + stagger[d] * dx,
- max(p[d], stagger[d] * dx))
- return p
-
- @ti.func
- def sample_bilinear(x, p):
- p = clamp(p)
-
- p_grid = p * inv_dx - stagger
-
- I = ti.cast(ti.floor(p_grid), ti.i32)
- f = p_grid - I
- g = 1 - f
-
- return x[I] * (g[0] * g[1]) + x[I + vec(1, 0)] * (f[0] * g[1]) + x[
- I + vec(0, 1)] * (g[0] * f[1]) + x[I + vec(1, 1)] * (f[0] * f[1])
-
- @ti.func
- def velocity(p):
- return sample_bilinear(v, p)
-
- @ti.func
- def sample_min(x, p):
- p = clamp(p)
- p_grid = p * inv_dx - stagger
- I = ti.cast(ti.floor(p_grid), ti.i32)
-
- return min(x[I], x[I + vec(1, 0)], x[I + vec(0, 1)], x[I + vec(1, 1)])
-
- @ti.func
- def sample_max(x, p):
- p = clamp(p)
- p_grid = p * inv_dx - stagger
- I = ti.cast(ti.floor(p_grid), ti.i32)
-
- return max(x[I], x[I + vec(1, 0)], x[I + vec(0, 1)], x[I + vec(1, 1)])
-
- @ti.func
- def backtrace(I, dt): # RK3
- p = (I + stagger) * dx
- v1 = velocity(p)
- p1 = p - 0.5 * dt * v1
- v2 = velocity(p1)
- p2 = p - 0.75 * dt * v2
- v3 = velocity(p2)
- p -= dt * (2 / 9 * v1 + 1 / 3 * v2 + 4 / 9 * v3)
- return p
-
- @ti.func
- def semi_lagrangian(x, new_x, dt):
- for I in ti.grouped(x):
- new_x[I] = sample_bilinear(x, backtrace(I, dt))
-
- @ti.kernel
- def advect():
- semi_lagrangian(x(0), new_x(0), dt)
- semi_lagrangian(x(1), new_x(1), dt)
- semi_lagrangian(x(2), new_x(2), dt)
-
- for I in ti.grouped(x):
- x[I] = new_x[I]
-
- init()
-
- def task():
- for i in range(10):
- advect()
-
- ti.benchmark(task, repeat=100)
-
- visualize = False
-
- if visualize:
- gui = ti.GUI('Advection schemes', (n, n))
- for i in range(10):
- for _ in range(10):
- advect()
- gui.set_image(x.to_numpy())
- gui.show()
-
-
-if __name__ == '__main__':
- simple_advection()
diff --git a/benchmarks/async_cases.py b/benchmarks/async_cases.py
deleted file mode 100644
index af141f424ddb6..0000000000000
--- a/benchmarks/async_cases.py
+++ /dev/null
@@ -1,374 +0,0 @@
-import math
-import os
-import sys
-
-import taichi as ti
-
-sys.path.append(os.path.join(ti.core.get_repo_dir(), 'tests', 'python'))
-
-from fuse_test_template import (template_fuse_dense_x2y2z,
- template_fuse_reduction)
-from utils import *
-
-
-@benchmark_async
-def chain_copy(scale):
- template_fuse_dense_x2y2z(size=scale * 1024**2,
- repeat=1,
- benchmark_repeat=100,
- benchmark=True)
-
-
-@benchmark_async
-def increments(scale):
- template_fuse_reduction(size=scale * 1024**2,
- repeat=10,
- benchmark_repeat=10,
- benchmark=True)
-
-
-@benchmark_async
-def fill_array(scale):
- a = ti.field(dtype=ti.f32, shape=scale * 1024**2)
-
- @ti.kernel
- def fill():
- for i in a:
- a[i] = 1.0
-
- def repeated_fill():
- for _ in range(10):
- fill()
-
- ti.benchmark(repeated_fill, repeat=10)
-
-
-@benchmark_async
-def fill_scalar(scale):
- a = ti.field(dtype=ti.f32, shape=())
-
- @ti.kernel
- def fill():
- a[None] = 1.0
-
- def repeated_fill():
- for _ in range(1000):
- fill()
-
- ti.benchmark(repeated_fill, repeat=5)
-
-
-@benchmark_async
-def sparse_saxpy(scale):
- a = ti.field(dtype=ti.f32)
- b = ti.field(dtype=ti.f32)
-
- block_count = 2**int((math.log(scale, 2)) // 2) * 4
- block_size = 32
- # a, b always share the same sparsity
- ti.root.pointer(ti.ij, block_count).dense(ti.ij, block_size).place(a, b)
-
- @ti.kernel
- def initialize():
- for i, j in ti.ndrange(block_count * block_size,
- block_count * block_size):
- if (i // block_size + j // block_size) % 4 == 0:
- a[i, j] = i + j
-
- @ti.kernel
- def saxpy(x: ti.template(), y: ti.template(), alpha: ti.f32):
- for i, j in x:
- y[i, j] = alpha * x[i, j] + y[i, j]
-
- def task():
- initialize()
- saxpy(a, b, 2)
- saxpy(b, a, 1.1)
- saxpy(b, a, 1.1)
- saxpy(a, b, 1.1)
- saxpy(a, b, 1.1)
- saxpy(a, b, 1.1)
-
- ti.benchmark(task, repeat=10)
-
-
-@benchmark_async
-def autodiff(scale):
-
- n = 1024**2 * scale
-
- a = ti.field(dtype=ti.f32, shape=n, needs_grad=True)
- b = ti.field(dtype=ti.f32, shape=n)
- loss = ti.field(dtype=ti.f32, shape=(), needs_grad=True)
-
- @ti.kernel
- def compute_loss():
- for i in a:
- loss[None] += a[i]
-
- @ti.kernel
- def accumulate_grad():
- for i in a:
- b[i] += a.grad[i]
-
- def task():
- for i in range(10):
- with ti.Tape(loss=loss):
- # The forward kernel of compute_loss should be completely eliminated (except for the last one)
- compute_loss()
-
- accumulate_grad()
-
- ti.benchmark(task, repeat=10)
-
-
-@benchmark_async
-def stencil_reduction(scale):
- a = ti.field(dtype=ti.f32)
- b = ti.field(dtype=ti.f32)
- total = ti.field(dtype=ti.f32, shape=())
-
- block_count = scale * 64
- block_size = 1024
- # a, b always share the same sparsity
- ti.root.pointer(ti.i, block_count).dense(ti.i, block_size).place(a, b)
-
- @ti.kernel
- def initialize():
- for i in range(block_size, block_size * (block_count - 1)):
- a[i] = i
-
- @ti.kernel
- def stencil():
- for i in a:
- b[i] = a[i - 1] + a[i] + a[i + 1]
-
- @ti.kernel
- def reduce():
- for i in a:
- total[None] += b[i]
-
- @ti.kernel
- def clear_b():
- for i in a:
- b[i] = 0
-
- def task():
- initialize()
- for i in range(3):
- stencil()
- reduce()
- clear_b()
-
- ti.benchmark(task, repeat=5)
-
-
-@benchmark_async
-def mpm_splitted(scale):
- quality = int(3 * scale**(1 / 3))
- # Use a larger value for higher-res simulations
-
- n_particles, n_grid = 9000 * quality**2, 128 * quality
- dx, inv_dx = 1 / n_grid, float(n_grid)
- dt = 1e-4 / quality
- p_vol, p_rho = (dx * 0.5)**2, 1
- p_mass = p_vol * p_rho
- E, nu = 0.1e4, 0.2 # Young's modulus and Poisson's ratio
- mu_0, lambda_0 = E / (2 * (1 + nu)), E * nu / (
- (1 + nu) * (1 - 2 * nu)) # Lame parameters
- x = ti.Vector.field(2, dtype=float, shape=n_particles) # position
- v = ti.Vector.field(2, dtype=float, shape=n_particles) # velocity
- C = ti.Matrix.field(2, 2, dtype=float,
- shape=n_particles) # affine velocity field
- F = ti.Matrix.field(2, 2, dtype=float,
- shape=n_particles) # deformation gradient
- material = ti.field(dtype=int, shape=n_particles) # material id
- Jp = ti.field(dtype=float, shape=n_particles) # plastic deformation
- grid_v = ti.Vector.field(2, dtype=float,
- shape=(n_grid,
- n_grid)) # grid node momentum/velocity
- grid_m = ti.field(dtype=float, shape=(n_grid, n_grid)) # grid node mass
-
- @ti.kernel
- def substep():
- for i, j in grid_m:
- grid_v[i, j] = [0, 0]
- grid_m[i, j] = 0
- for p in x:
- F[p] = (ti.Matrix.identity(float, 2) +
- dt * C[p]) @ F[p] # deformation gradient update
- for p in x: # Particle state update and scatter to grid (P2G)
- base = (x[p] * inv_dx - 0.5).cast(int)
- fx = x[p] * inv_dx - base.cast(float)
- # Quadratic kernels [http://mpm.graphics Eqn. 123, with x=fx, fx-1,fx-2]
- w = [0.5 * (1.5 - fx)**2, 0.75 - (fx - 1)**2, 0.5 * (fx - 0.5)**2]
- h = ti.exp(
- 10 * (1.0 - Jp[p])
- ) # Hardening coefficient: snow gets harder when compressed
- if material[p] == 1: # jelly, make it softer
- h = 0.3
- mu, la = mu_0 * h, lambda_0 * h
- if material[p] == 0: # liquid
- mu = 0.0
- U, sig, V = ti.svd(F[p])
- J = 1.0
- for d in ti.static(range(2)):
- new_sig = sig[d, d]
- if material[p] == 2: # Snow
- new_sig = min(max(sig[d, d], 1 - 2.5e-2),
- 1 + 4.5e-3) # Plasticity
- Jp[p] *= sig[d, d] / new_sig
- sig[d, d] = new_sig
- J *= new_sig
- if material[
- p] == 0: # Reset deformation gradient to avoid numerical instability
- F[p] = ti.Matrix.identity(float, 2) * ti.sqrt(J)
- elif material[p] == 2:
- F[p] = U @ sig @ V.transpose(
- ) # Reconstruct elastic deformation gradient after plasticity
- stress = 2 * mu * (F[p] - U @ V.transpose()) @ F[p].transpose(
- ) + ti.Matrix.identity(float, 2) * la * J * (J - 1)
- stress = (-dt * p_vol * 4 * inv_dx * inv_dx) * stress
- affine = stress + p_mass * C[p]
- for i, j in ti.static(ti.ndrange(
- 3, 3)): # Loop over 3x3 grid node neighborhood
- offset = ti.Vector([i, j])
- dpos = (offset.cast(float) - fx) * dx
- weight = w[i][0] * w[j][1]
- grid_v[base +
- offset] += weight * (p_mass * v[p] + affine @ dpos)
- grid_m[base + offset] += weight * p_mass
- for i, j in grid_m:
- if grid_m[i, j] > 0: # No need for epsilon here
- grid_v[i, j] = (
- 1 / grid_m[i, j]) * grid_v[i, j] # Momentum to velocity
- grid_v[i, j][1] -= dt * 50 # gravity
- for i, j in grid_m:
- if grid_m[i, j] > 0: # No need for epsilon here
- if i < 3 and grid_v[i, j][0] < 0:
- grid_v[i, j][0] = 0 # Boundary conditions
- for i, j in grid_m:
- if grid_m[i, j] > 0: # No need for epsilon here
- if i > n_grid - 3 and grid_v[i, j][0] > 0: grid_v[i, j][0] = 0
- for i, j in grid_m:
- if grid_m[i, j] > 0: # No need for epsilon here
- if j < 3 and grid_v[i, j][1] < 0: grid_v[i, j][1] = 0
- for i, j in grid_m:
- if grid_m[i, j] > 0: # No need for epsilon here
- if j > n_grid - 3 and grid_v[i, j][1] > 0: grid_v[i, j][1] = 0
- for p in x: # grid to particle (G2P)
- base = (x[p] * inv_dx - 0.5).cast(int)
- fx = x[p] * inv_dx - base.cast(float)
- w = [
- 0.5 * (1.5 - fx)**2, 0.75 - (fx - 1.0)**2, 0.5 * (fx - 0.5)**2
- ]
- new_v = ti.Vector.zero(float, 2)
- new_C = ti.Matrix.zero(float, 2, 2)
- for i, j in ti.static(ti.ndrange(
- 3, 3)): # loop over 3x3 grid node neighborhood
- dpos = ti.Vector([i, j]).cast(float) - fx
- g_v = grid_v[base + ti.Vector([i, j])]
- weight = w[i][0] * w[j][1]
- new_v += weight * g_v
- new_C += 4 * inv_dx * weight * g_v.outer_product(dpos)
- v[p], C[p] = new_v, new_C
- for p in x:
- x[p] += dt * v[p] # advection
-
- group_size = n_particles // 3
-
- @ti.kernel
- def initialize():
- for i in range(n_particles):
- x[i] = [
- ti.random() * 0.2 + 0.3 + 0.10 * (i // group_size),
- ti.random() * 0.2 + 0.05 + 0.32 * (i // group_size)
- ]
- material[i] = i // group_size # 0: fluid 1: jelly 2: snow
- v[i] = ti.Matrix([0, 0])
- F[i] = ti.Matrix([[1, 0], [0, 1]])
- Jp[i] = 1
-
- initialize()
-
- def task():
- for s in range(int(2e-3 // dt)):
- substep()
-
- ti.benchmark(task, repeat=5)
-
-
-@benchmark_async
-def multires(scale):
- num_levels = 4
-
- x = []
- for i in range(num_levels):
- x.append(ti.field(dtype=ti.f32))
-
- # TODO: Using 1024 instead of 512 hangs the CUDA case. Need to figure out why.
- n = 512 * 1024 * scale
-
- block_size = 16
- assert n % block_size**2 == 0
-
- for i in range(num_levels):
- ti.root.pointer(ti.i, n // 2**i // block_size**2).pointer(
- ti.i, block_size).dense(ti.i, block_size).place(x[i])
-
- @ti.kernel
- def initialize():
- for i in range(n):
- x[0][i] = i
-
- @ti.kernel
- def downsample(l: ti.template()):
- for i in x[l]:
- if i % 2 == 0:
- x[l + 1][i // 2] = x[l][i]
-
- initialize()
-
- def task():
- for l in range(num_levels - 1):
- downsample(l)
-
- ti.benchmark(task, repeat=5)
-
-
-@benchmark_async
-def deep_hierarchy(scale):
- branching = 4
- num_levels = 8 + int(math.log(scale, branching))
-
- x = ti.field(dtype=ti.f32)
-
- n = 256 * 1024 * scale
-
- assert n % (branching**num_levels) == 0
-
- snode = ti.root
- for i in range(num_levels):
- snode = snode.pointer(ti.i, branching)
-
- snode.dense(ti.i, n // (branching**num_levels)).place(x)
-
- @ti.kernel
- def initialize():
- for i in range(n):
- x[i] = 0
-
- initialize()
-
- # Not fusible, but no modification to the mask/list of x either
- @ti.kernel
- def jitter():
- for i in x:
- if i % 2 == 0:
- x[i] += x[i + 1]
-
- def task():
- for i in range(5):
- jitter()
-
- ti.benchmark(task, repeat=5)
diff --git a/benchmarks/benchmark_async.py b/benchmarks/benchmark_async.py
deleted file mode 100644
index 9054631ff4fc2..0000000000000
--- a/benchmarks/benchmark_async.py
+++ /dev/null
@@ -1,36 +0,0 @@
-from async_advection import *
-from async_cases import *
-
-import taichi as ti
-
-rerun = True
-
-cases = [
- chain_copy, increments, fill_array, sparse_saxpy, autodiff,
- stencil_reduction, mpm_splitted, simple_advection, multires, deep_hierarchy
-]
-
-if rerun:
- for c in cases:
- print('*' * 30)
- print(f'* Running {c.__name__}')
- print('*' * 30)
- c()
-
-case_names = [c.__name__ for c in cases]
-
-ti.benchmark_plot(fn='benchmark.yml',
- cases=case_names,
- columns=[
- 'wall_clk_t', 'exec_t', 'launched_tasks',
- 'compiled_inst', 'compiled_tasks'
- ],
- column_titles=[
- 'Wall-clock time', 'Backend time', 'Tasks launched',
- 'Instructions emitted', 'Tasks compiled'
- ],
- archs=['cuda', 'x64'],
- title='Whole-Program Optimization Microbenchmarks',
- bars='sync_vs_async',
- left_margin=0.2,
- size=(11.5, 9))
diff --git a/benchmarks/deserialize.py b/benchmarks/deserialize.py
new file mode 100644
index 0000000000000..7c09d9dd72ffb
--- /dev/null
+++ b/benchmarks/deserialize.py
@@ -0,0 +1,99 @@
+import argparse
+import json
+import os
+from copy import deepcopy
+
+from utils import dump2json
+
+
+class ResultsBuilder():
+ def __init__(self, results_file_path: str):
+ self._suites_result = {}
+ self._file_path = results_file_path
+ self.load_suites_result()
+
+ def load_suites_result(self):
+ # benchmark info
+ info_path = os.path.join(self._file_path, '_info.json')
+ with open(info_path, 'r') as f:
+ info_dict = json.load(f)['suites']
+ # suite info
+ for suite_name, attrs in info_dict.items():
+ self._suites_result[suite_name] = {}
+ for arch in attrs['archs']:
+ self._suites_result[suite_name][arch] = {}
+ suite_info_path = os.path.join(self._file_path, suite_name,
+ arch, "_info.json")
+ with open(suite_info_path, 'r') as f:
+ suite_info_dict = json.load(f)
+ # case info
+ for case_name in suite_info_dict:
+ items = suite_info_dict[case_name]
+ items.pop('name')
+ items['metrics'] = items.pop('get_metric')
+ self._suites_result[suite_name][arch][case_name] = {
+ 'items': items
+ }
+ # cases result
+ for suite_name in self._suites_result:
+ for arch in self._suites_result[suite_name]:
+ for case_name in self._suites_result[suite_name][arch]:
+ case_info_path = os.path.join(self._file_path, suite_name,
+ arch, case_name + ".json")
+ with open(case_info_path, 'r') as f:
+ case_results = json.load(f)
+ remove_none_list = []
+ for name, data in case_results.items():
+ # remove case_name
+ data['tags'] = data['tags'][1:]
+ if data['result'] is None:
+ remove_none_list.append(name)
+ for name in remove_none_list:
+ case_results.pop(name)
+ self._suites_result[suite_name][arch][case_name][
+ 'results'] = case_results
+
+ def get_suites_result(self):
+ return self._suites_result
+
+ def save_results_as_json(self, costomized_dir=None):
+ file_path = os.path.join(self._file_path, 'results.json')
+ if costomized_dir != None:
+ file_path = os.path.join(costomized_dir, 'results.json')
+ with open(file_path, 'w') as f:
+ print(dump2json(self._suites_result), file=f)
+
+ def print_info(self):
+ # remove 'results' in self._suites_result, then print
+ info_dict = deepcopy(self._suites_result)
+ for suite_name in info_dict:
+ for arch in info_dict[suite_name]:
+ for case in info_dict[suite_name][arch]:
+ info_dict[suite_name][arch][case].pop('results')
+ print(dump2json(info_dict))
+
+
+if __name__ == '__main__':
+ parser = argparse.ArgumentParser()
+
+ parser.add_argument('-f',
+ '--folder',
+ default='./results',
+ dest='folder',
+ type=str,
+ help='Path of result folder. Defaults to ./results')
+
+ parser.add_argument('-o',
+ '--output_path',
+ default='./results',
+ dest='output_path',
+ type=str,
+ help='Path of result folder. Defaults to ./results')
+
+ args = parser.parse_args()
+ result_folder = args.folder
+ output_path = args.output_path
+
+ results = ResultsBuilder(result_folder)
+ results.save_results_as_json(output_path)
+ results.print_info()
diff --git a/benchmarks/fill_dense.py b/benchmarks/fill_dense.py
deleted file mode 100644
index 452c5075fe5d8..0000000000000
--- a/benchmarks/fill_dense.py
+++ /dev/null
@@ -1,124 +0,0 @@
-import taichi as ti
-
-# originally by @KLozes
-
-
-@ti.test()
-def benchmark_flat_struct():
- N = 4096
- a = ti.field(dtype=ti.f32, shape=(N, N))
-
- @ti.kernel
- def fill():
- for i, j in a:
- a[i, j] = 2.0
-
- return ti.benchmark(fill, repeat=500)
-
-
-@ti.test()
-def benchmark_flat_range():
- N = 4096
- a = ti.field(dtype=ti.f32, shape=(N, N))
-
- @ti.kernel
- def fill():
- for i, j in ti.ndrange(N, N):
- a[i, j] = 2.0
-
- return ti.benchmark(fill, repeat=700)
-
-
-@ti.test()
-def benchmark_nested_struct():
- a = ti.field(dtype=ti.f32)
- N = 512
-
- ti.root.dense(ti.ij, [N, N]).dense(ti.ij, [8, 8]).place(a)
-
- @ti.kernel
- def fill():
- for i, j in a:
- a[i, j] = 2.0
-
- return ti.benchmark(fill, repeat=700)
-
-
-@ti.test()
-def benchmark_nested_struct_listgen_8x8():
- a = ti.field(dtype=ti.f32)
- ti.cfg.demote_dense_struct_fors = False
- N = 512
-
- ti.root.dense(ti.ij, [N, N]).dense(ti.ij, [8, 8]).place(a)
-
- @ti.kernel
- def fill():
- for i, j in a:
- a[i, j] = 2.0
-
- return ti.benchmark(fill, repeat=1000)
-
-
-@ti.test()
-def benchmark_nested_struct_listgen_16x16():
- a = ti.field(dtype=ti.f32)
- ti.cfg.demote_dense_struct_fors = False
- N = 256
-
- ti.root.dense(ti.ij, [N, N]).dense(ti.ij, [16, 16]).place(a)
-
- @ti.kernel
- def fill():
- for i, j in a:
- a[i, j] = 2.0
-
- return ti.benchmark(fill, repeat=700)
-
-
-@ti.test()
-def benchmark_nested_range_blocked():
- a = ti.field(dtype=ti.f32)
- N = 512
-
- ti.root.dense(ti.ij, [N, N]).dense(ti.ij, [8, 8]).place(a)
-
- @ti.kernel
- def fill():
- for X in range(N * N):
- for Y in range(64):
- a[X // N * 8 + Y // 8, X % N * 8 + Y % 8] = 2.0
-
- return ti.benchmark(fill, repeat=800)
-
-
-@ti.test()
-def benchmark_nested_range():
- a = ti.field(dtype=ti.f32)
- N = 512
-
- ti.root.dense(ti.ij, [N, N]).dense(ti.ij, [8, 8]).place(a)
-
- @ti.kernel
- def fill():
- for j in range(N * 8):
- for i in range(N * 8):
- a[i, j] = 2.0
-
- return ti.benchmark(fill, repeat=1000)
-
-
-@ti.test()
-def benchmark_root_listgen():
- a = ti.field(dtype=ti.f32)
- ti.cfg.demote_dense_struct_fors = False
- N = 512
-
- ti.root.dense(ti.ij, [N, N]).dense(ti.ij, [8, 8]).place(a)
-
- @ti.kernel
- def fill():
- for i, j in a.parent():
- a[i, j] = 2.0
-
- return ti.benchmark(fill, repeat=800)
diff --git a/benchmarks/fill_sparse.py b/benchmarks/fill_sparse.py
deleted file mode 100644
index 64d15a65cb001..0000000000000
--- a/benchmarks/fill_sparse.py
+++ /dev/null
@@ -1,42 +0,0 @@
-import taichi as ti
-
-
-@ti.archs_support_sparse
-def benchmark_nested_struct():
- a = ti.field(dtype=ti.f32)
- N = 512
-
- ti.root.pointer(ti.ij, [N, N]).dense(ti.ij, [8, 8]).place(a)
-
- @ti.kernel
- def fill():
- for i, j in ti.ndrange(N * 8, N * 8):
- a[i, j] = 2.0
-
- fill()
-
- return ti.benchmark(fill)
-
-
-@ti.archs_support_sparse
-def benchmark_nested_struct_fill_and_clear():
- a = ti.field(dtype=ti.f32)
- N = 512
-
- ti.root.pointer(ti.ij, [N, N]).dense(ti.ij, [8, 8]).place(a)
-
- @ti.kernel
- def fill():
- for i, j in ti.ndrange(N * 8, N * 8):
- a[i, j] = 2.0
-
- @ti.kernel
- def clear():
- for i, j in a.parent():
- ti.deactivate(a.parent().parent(), [i, j])
-
- def task():
- fill()
- clear()
-
- return ti.benchmark(task, repeat=30)
diff --git a/benchmarks/memory_bound.py b/benchmarks/memory_bound.py
deleted file mode 100644
index 6bcd200099c12..0000000000000
--- a/benchmarks/memory_bound.py
+++ /dev/null
@@ -1,59 +0,0 @@
-import taichi as ti
-
-N = 1024**3 // 4 # 1 GB per buffer
-
-
-# 4 B/it
-@ti.test(exclude=ti.opengl)
-def benchmark_memset():
- a = ti.field(dtype=ti.f32, shape=N)
-
- @ti.kernel
- def memset():
- for i in a:
- a[i] = 1.0
-
- return ti.benchmark(memset, repeat=10)
-
-
-# 8 B/it
-@ti.test(exclude=ti.opengl)
-def benchmark_sscal():
- a = ti.field(dtype=ti.f32, shape=N)
-
- @ti.kernel
- def task():
- for i in a:
- a[i] = 0.5 * a[i]
-
- return ti.benchmark(task, repeat=10)
-
-
-# 8 B/it
-@ti.test(exclude=ti.opengl)
-def benchmark_memcpy():
- a = ti.field(dtype=ti.f32, shape=N)
- b = ti.field(dtype=ti.f32, shape=N)
-
- @ti.kernel
- def memcpy():
- for i in a:
- a[i] = b[i]
-
- return ti.benchmark(memcpy, repeat=10)
-
-
-# 12 B/it
-@ti.test(exclude=ti.opengl)
-def benchmark_saxpy():
- x = ti.field(dtype=ti.f32, shape=N)
- y = ti.field(dtype=ti.f32, shape=N)
- z = ti.field(dtype=ti.f32, shape=N)
-
- @ti.kernel
- def task():
- for i in x:
- a = 123
- z[i] = a * x[i] + y[i]
-
- return ti.benchmark(task, repeat=10)
diff --git a/benchmarks/microbenchmarks/__init__.py b/benchmarks/microbenchmarks/__init__.py
new file mode 100644
index 0000000000000..adbfa19ade092
--- /dev/null
+++ b/benchmarks/microbenchmarks/__init__.py
@@ -0,0 +1,12 @@
+from .atomic_ops import AtomicOpsPlan
+from .fill import FillPlan
+from .math_opts import MathOpsPlan
+from .matrix_ops import MatrixOpsPlan
+from .memcpy import MemcpyPlan
+from .saxpy import SaxpyPlan
+from .stencil2d import Stencil2DPlan
+
+benchmark_plan_list = [
+ AtomicOpsPlan, FillPlan, MathOpsPlan, MatrixOpsPlan, MemcpyPlan, SaxpyPlan,
+ Stencil2DPlan
+]
diff --git a/benchmarks/microbenchmarks/_items.py b/benchmarks/microbenchmarks/_items.py
new file mode 100644
index 0000000000000..af9b6747f46b6
--- /dev/null
+++ b/benchmarks/microbenchmarks/_items.py
@@ -0,0 +1,116 @@
+from microbenchmarks._utils import size2tag
+
+import taichi as ti
+
+
+class BenchmarkItem:
+ name = 'item'
+
+ def __init__(self):
+ self._items = {} # {'tag': impl, ...}
+
+ def get(self):
+ return self._items
+
+ def get_tags(self):
+ return list(self._items.keys())
+
+ def impl(self, tag: str):
+ return self._items[tag]
+
+ def remove(self, tags: list):
+ for tag in tags:
+ self._items.pop(tag)
+
+ def update(self, adict: dict):
+ self._items.update(adict)
+
+
+class DataType(BenchmarkItem):
+ name = 'dtype'
+ integer_list = ['i32', 'i64']
+
+ def __init__(self):
+ self._items = {
+ str(ti.i32): ti.i32,
+ str(ti.i64): ti.i64,
+ str(ti.f32): ti.f32,
+ str(ti.f64): ti.f64
+ }
+
+ def remove_integer(self):
+ self.remove(self.integer_list)
+
+ @staticmethod
+ def is_integer(dtype: str):
+ integer_list = ['i32', 'u32', 'i64', 'u64']
+ return True if dtype in integer_list else False
+
+
+class DataSize(BenchmarkItem):
+ name = 'dsize'
+
+ def __init__(self):
+ self._items = {}
+ for i in range(2, 10, 2): # [16KB,256KB,4MB,64MB]
+ size_bytes = (4**i) * 1024 # kibibytes(KiB) = 1024
+ self._items[size2tag(size_bytes)] = size_bytes
+
+
+class Container(BenchmarkItem):
+ name = 'container'
+
+ def __init__(self):
+ self._items = {'field': ti.field, 'ndarray': ti.ndarray}
+
+
+class MathOps(BenchmarkItem):
+ name = 'math_op'
+
+ #reference: https://docs.taichi.graphics/lang/articles/basic/operator
+ def __init__(self):
+ self._items = {
+ # Trigonometric
+ 'sin': ti.sin,
+ 'cos': ti.cos,
+ 'tan': ti.tan,
+ 'asin': ti.asin,
+ 'acos': ti.acos,
+ 'tanh': ti.tanh,
+ # Other arithmetic
+ 'sqrt': ti.sqrt,
+ 'rsqrt': ti.rsqrt, # A fast version for `1 / ti.sqrt(x)`.
+ 'exp': ti.exp,
+ 'log': ti.log,
+ 'round': ti.round,
+ 'floor': ti.floor,
+ 'ceil': ti.ceil,
+ 'abs': ti.abs,
+ }
+
+
+class AtomicOps(BenchmarkItem):
+ name = 'atomic_op'
+
+ def __init__(self):
+ self._items = {
+ 'atomic_add': ti.atomic_add,
+ 'atomic_sub': ti.atomic_sub,
+ 'atomic_and': ti.atomic_and,
+ 'atomic_or': ti.atomic_or,
+ 'atomic_xor': ti.atomic_xor,
+ 'atomic_max': ti.atomic_max,
+ 'atomic_min': ti.atomic_min
+ }
+
+ @staticmethod
+ def is_logical_op(op: str):
+ logical_op_list = ['atomic_and', 'atomic_or', 'atomic_xor']
+ return True if op in logical_op_list else False
+
+ @staticmethod
+ def is_supported_type(op: str, dtype: str):
+ if AtomicOps.is_logical_op(op) and not DataType.is_integer(dtype):
+ return False
+ else:
+ return True
diff --git a/benchmarks/microbenchmarks/_metric.py b/benchmarks/microbenchmarks/_metric.py
new file mode 100644
index 0000000000000..46c008c1e7949
--- /dev/null
+++ b/benchmarks/microbenchmarks/_metric.py
@@ -0,0 +1,47 @@
+from microbenchmarks._items import BenchmarkItem
+from microbenchmarks._utils import End2EndTimer, get_ti_arch
+
+import taichi as ti
+
+
+def end2end_executor(repeat, func, *args):
+ # compile & warmup
+ for i in range(repeat):
+ func(*args)
+
+ timer = End2EndTimer()
+ timer.tick()
+ for i in range(repeat):
+ func(*args)
+ time_in_s = timer.tock()
+ return time_in_s * 1000 / repeat #ms
+
+
+def kernel_executor(repeat, func, *args):
+ # compile & warmup
+ for i in range(repeat):
+ func(*args)
+ ti.clear_kernel_profile_info()
+ for i in range(repeat):
+ func(*args)
+ return ti.kernel_profiler_total_time() * 1000 / repeat #ms
+
+
+class MetricType(BenchmarkItem):
+ name = 'get_metric'
+
+ def __init__(self):
+ self._items = {
+ 'kernel_elapsed_time_ms': kernel_executor,
+ 'end2end_time_ms': end2end_executor
+ }
+
+ @staticmethod
+ def init_taichi(arch: str, tag_list: list):
+ if set(['kernel_elapsed_time_ms']).issubset(tag_list):
+ ti.init(kernel_profiler=True, arch=get_ti_arch(arch))
+ elif set(['end2end_time_ms']).issubset(tag_list):
+ ti.init(kernel_profiler=False, arch=get_ti_arch(arch))
+ else:
+ return False
+ return True
diff --git a/benchmarks/microbenchmarks/_plan.py b/benchmarks/microbenchmarks/_plan.py
new file mode 100644
index 0000000000000..52af1002c6530
--- /dev/null
+++ b/benchmarks/microbenchmarks/_plan.py
@@ -0,0 +1,89 @@
+import itertools
+
+from microbenchmarks._items import AtomicOps, DataType
+from microbenchmarks._metric import MetricType
+from microbenchmarks._utils import get_ti_arch, tags2name
+
+import taichi as ti
+
+
+class Funcs():
+ def __init__(self):
+ self._funcs = {}
+
+ def add_func(self, tag_list: list, func):
+ self._funcs[tags2name(tag_list)] = {'tags': tag_list, 'func': func}
+
+ def get_func(self, tags):
+ for name, item in self._funcs.items():
+ if set(item['tags']).issubset(tags):
+ return item['func']
+ return None
+
+
+class BenchmarkPlan:
+ def __init__(self, name='plan', arch='x64', basic_repeat_times=1):
+ self.name = name
+ self.arch = arch
+ self.basic_repeat_times = basic_repeat_times
+ self.info = {'name': self.name}
+ self.plan = {} # {'tags': [...], 'result': None}
+ self.items = {}
+ self.funcs = Funcs()
+
+ def create_plan(self, *items):
+ items_list = [[self.name]]
+ for item in items:
+ self.items[item.name] = item
+ items_list.append(item.get_tags())
+ self.info[item.name] = item.get_tags()
+ case_list = list(itertools.product(*items_list)) #items generate cases
+ for tags in case_list:
+ self.plan[tags2name(tags)] = {'tags': tags, 'result': None}
+ self._remove_conflict_items()
+
+ def add_func(self, tag_list, func):
+ self.funcs.add_func(tag_list, func)
+
+ def run(self):
+ for case, plan in self.plan.items():
+ tag_list = plan['tags']
+ MetricType.init_taichi(self.arch, tag_list)
+ _ms = self.funcs.get_func(tag_list)(self.arch,
+ self.basic_repeat_times,
+ **self._get_kwargs(tag_list))
+ plan['result'] = _ms
+ print(f'{tag_list}={_ms}')
+ ti.reset()
+ rdict = {'results': self.plan, 'info': self.info}
+ return rdict
+
+ def _get_kwargs(self, tags, impl=True):
+ kwargs = {}
+ tags = tags[1:] # tags = [case_name, item1_tag, item2_tag, ...]
+ for item, tag in zip(self.items.values(), tags):
+ kwargs[item.name] = item.impl(tag) if impl == True else tag
+ return kwargs
+
+ def _remove_conflict_items(self):
+ remove_list = []
+ #logical_atomic with float_type
+ if set([AtomicOps.name, DataType.name]).issubset(self.items.keys()):
+ for name, case in self.plan.items():
+ kwargs_tag = self._get_kwargs(case['tags'], impl=False)
+ atomic_tag = kwargs_tag[AtomicOps.name]
+ dtype_tag = kwargs_tag[DataType.name]
+ if not AtomicOps.is_supported_type(atomic_tag, dtype_tag):
+ remove_list.append(name)
+ #remove
+ for name in remove_list:
+ self.plan.pop(name)
+
+ def remove_cases_with_tags(self, tags: list):
+ remove_list = []
+ for case, plan in self.plan.items():
+ if set(tags).issubset(plan['tags']):
+ remove_list.append(case)
+ #remove
+ for case in remove_list:
+ self.plan.pop(case)
diff --git a/benchmarks/microbenchmarks/_utils.py b/benchmarks/microbenchmarks/_utils.py
new file mode 100644
index 0000000000000..8bea2f2821ea2
--- /dev/null
+++ b/benchmarks/microbenchmarks/_utils.py
@@ -0,0 +1,87 @@
+from time import perf_counter
+
+from taichi._lib import core as ti_core
+
+import taichi as ti
+
+
+class End2EndTimer:
+ def __init__(self):
+ self._ts1 = 0
+ self._ts2 = 0
+
+ def tick(self):
+ ti.sync()
+ self._ts1 = perf_counter()
+ return self._ts1
+
+ def tock(self):
+ ti.sync()
+ self._ts2 = perf_counter()
+ return self._ts2 - self._ts1
+
+
+def size2tag(size_in_byte):
+ size_subsection = [(0.0, 'B'), (1024.0, 'KB'), (1048576.0, 'MB'),
+ (1073741824.0, 'GB'),
+ (float('inf'), 'INF')] #B KB MB GB
+ for dsize, unit in reversed(size_subsection):
+ if size_in_byte >= dsize:
+ return str(int(size_in_byte / dsize)) + unit
+
+
+def tags2name(tag_list):
+ return '_'.join(tag_list)
+
+
+def dtype_size(ti_dtype):
+ dtype_size_dict = {ti.i32: 4, ti.i64: 8, ti.f32: 4, ti.f64: 8}
+ if ti_dtype not in dtype_size_dict:
+ raise RuntimeError('Unsupported ti.dtype: ' + str(type(ti_dtype)))
+ else:
+ return dtype_size_dict[ti_dtype]
+
+
+def get_ti_arch(arch: str):
+ arch_dict = {
+ 'cuda': ti.cuda,
+ 'vulkan': ti.vulkan,
+ 'opengl': ti.opengl,
+ 'metal': ti.metal,
+ 'x64': ti.x64,
+ 'cc': ti.cc
+ }
+ return arch_dict[arch]
+
+
+def scaled_repeat_times(arch: str, datasize, repeat=1):
+ if (arch == 'cuda') | (arch == 'vulkan') | (arch == 'opengl'):
+ repeat *= 10
+ if datasize <= 4 * 1024 * 1024:
+ repeat *= 10
+ return repeat
+
+
+def fill_random(dst, dtype, container):
+ @ti.kernel
+ def fill_template(dst: ti.template()):
+ for I in ti.grouped(dst):
+ dst[I] = ti.random(dtype)
+
+ @ti.kernel
+ def fill_1d_array(dst: ti.any_arr()):
+ for i in dst:
+ dst[i] = ti.random(dtype)
+
+ @ti.kernel
+ def fill_2d_array(dst: ti.any_arr()):
+ for i, j in dst:
+ dst[i, j] = ti.random(dtype)
+
+ if container == ti.ndarray:
+ if len(dst.shape) == 1:
+ fill_1d_array(dst)
+ elif len(dst.shape) == 2:
+ fill_2d_array(dst)
+ else:
+ fill_template(dst)
diff --git a/benchmarks/microbenchmarks/atomic_ops.py b/benchmarks/microbenchmarks/atomic_ops.py
new file mode 100644
index 0000000000000..9b78728bb4445
--- /dev/null
+++ b/benchmarks/microbenchmarks/atomic_ops.py
@@ -0,0 +1,42 @@
+from microbenchmarks._items import AtomicOps, Container, DataSize, DataType
+from microbenchmarks._metric import MetricType
+from microbenchmarks._plan import BenchmarkPlan
+from microbenchmarks._utils import dtype_size, fill_random, scaled_repeat_times
+
+import taichi as ti
+
+
+def reduction_default(arch, repeat, atomic_op, container, dtype, dsize,
+ get_metric):
+ repeat = scaled_repeat_times(arch, dsize, repeat)
+ num_elements = dsize // dtype_size(dtype)
+
+ x = container(dtype, shape=num_elements)
+ y = container(dtype, shape=())
+ y[None] = 0
+
+ @ti.kernel
+ def reduction_field(y: ti.template(), x: ti.template()):
+ for i in x:
+ atomic_op(y[None], x[i])
+
+ @ti.kernel
+ def reduction_array(y: ti.any_arr(), x: ti.any_arr()):
+ for i in x:
+ atomic_op(y[None], x[i])
+
+ fill_random(x, dtype, container)
+ func = reduction_field if container == ti.field else reduction_array
+ return get_metric(repeat, func, y, x)
+
+
+class AtomicOpsPlan(BenchmarkPlan):
+ def __init__(self, arch: str):
+ super().__init__('atomic_ops', arch, basic_repeat_times=10)
+ atomic_ops = AtomicOps()
+ atomic_ops.remove(
+ ['atomic_sub', 'atomic_and', 'atomic_xor', 'atomic_max'])
+ self.create_plan(atomic_ops, Container(), DataType(), DataSize(),
+ MetricType())
+ self.add_func(['field'], reduction_default)
+ self.add_func(['ndarray'], reduction_default)
diff --git a/benchmarks/microbenchmarks/fill.py b/benchmarks/microbenchmarks/fill.py
new file mode 100644
index 0000000000000..5d1dad91e8c82
--- /dev/null
+++ b/benchmarks/microbenchmarks/fill.py
@@ -0,0 +1,60 @@
+from microbenchmarks._items import BenchmarkItem, Container, DataSize, DataType
+from microbenchmarks._metric import MetricType
+from microbenchmarks._plan import BenchmarkPlan
+from microbenchmarks._utils import dtype_size, scaled_repeat_times
+
+import taichi as ti
+
+
+def fill_default(arch, repeat, container, dtype, dsize, get_metric):
+ @ti.kernel
+ def fill_field(dst: ti.template()):
+ for I in ti.grouped(dst):
+ dst[I] = ti.cast(0.7, dtype)
+
+ @ti.kernel
+ def fill_array(dst: ti.any_arr()):
+ for i in dst:
+ dst[i] = ti.cast(0.7, dtype)
+
+ repeat = scaled_repeat_times(arch, dsize, repeat)
+ num_elements = dsize // dtype_size(dtype)
+ x = container(dtype, num_elements)
+ func = fill_field if container == ti.field else fill_array
+ return get_metric(repeat, func, x)
+
+
+def fill_sparse(arch, repeat, container, dtype, dsize, get_metric):
+ repeat = scaled_repeat_times(arch, dsize, repeat=1)
+ # basic_repeat_time = 1: sparse-specific parameter
+ num_elements = dsize // dtype_size(dtype) // 8
+
+ block = ti.root.pointer(ti.i, num_elements)
+ x = ti.field(dtype)
+ block.dense(ti.i, 8).place(x)
+
+ @ti.kernel
+ def active_all():
+ for i in ti.ndrange(num_elements):
+ ti.activate(block, [i])
+
+ active_all()
+
+ @ti.kernel
+ def fill_const(dst: ti.template()):
+ for i in x:
+ dst[i] = ti.cast(0.7, dtype)
+
+ return get_metric(repeat, fill_const, x)
+
+
+class FillPlan(BenchmarkPlan):
+ def __init__(self, arch: str):
+ super().__init__('fill', arch, basic_repeat_times=10)
+ fill_container = Container()
+ fill_container.update({'sparse': None}) # None: implement by feature
+ self.create_plan(fill_container, DataType(), DataSize(), MetricType())
+ # use tag_list to label the customized implementation (funcs).
+ self.add_func(['field'], fill_default)
+ self.add_func(['ndarray'], fill_default)
+ self.add_func(['sparse'], fill_sparse)
diff --git a/benchmarks/microbenchmarks/math_opts.py b/benchmarks/microbenchmarks/math_opts.py
new file mode 100644
index 0000000000000..cf06d0bbcc500
--- /dev/null
+++ b/benchmarks/microbenchmarks/math_opts.py
@@ -0,0 +1,62 @@
+from microbenchmarks._items import BenchmarkItem, DataType, MathOps
+from microbenchmarks._metric import MetricType
+from microbenchmarks._plan import BenchmarkPlan
+from microbenchmarks._utils import dtype_size, scaled_repeat_times
+
+import taichi as ti
+
+
+def unary_ops_throughput_default(arch, repeat, math_op, dtype, element_num,
+ thread_for_loop, get_metric):
+ local_data_num = 16 #enough data to fill the instruction pipeline
+ global_vector = ti.Vector.field(local_data_num, dtype, element_num)
+
+ @ti.kernel
+ def op_throughput():
+ for e in global_vector:
+ local_vector = global_vector[e]
+ #loop
+ for j in range(thread_for_loop):
+ for k in ti.static(range(local_data_num)):
+ local_vector[k] = math_op(local_vector[k])
+ #epilogue
+ global_vector[e] = local_vector
+
+ @ti.kernel
+ def fill_random():
+ for e in global_vector:
+ for i in ti.static(range(local_data_num)):
+ global_vector[e][i] = ti.random()
+
+ fill_random()
+ return get_metric(repeat, op_throughput)
+
+
+class ElementNum(BenchmarkItem):
+ name = 'element_num'
+
+ def __init__(self):
+ self._items = {
+ 'element16384': 16384,
+ #enough threads for filling CUDA cores
+ }
+
+
+class ForLoopCycle(BenchmarkItem):
+ name = 'thread_for_loop'
+
+ def __init__(self):
+ self._items = {}
+ for i in range(1, 7):
+ cycles = 4 * pow(2, i) # [8 16 32 64 128 256]
+ self._items['threadloop' + str(cycles)] = cycles
+
+
+class MathOpsPlan(BenchmarkPlan):
+ def __init__(self, arch: str):
+ super().__init__('math_ops', arch, basic_repeat_times=10)
+ math_dtype = DataType()
+ math_dtype.remove_integer()
+ self.create_plan(MathOps(), math_dtype, ElementNum(), ForLoopCycle(),
+ MetricType())
+ self.add_func(['element16384'], unary_ops_throughput_default)
diff --git a/benchmarks/microbenchmarks/matrix_ops.py b/benchmarks/microbenchmarks/matrix_ops.py
new file mode 100644
index 0000000000000..df163b959917f
--- /dev/null
+++ b/benchmarks/microbenchmarks/matrix_ops.py
@@ -0,0 +1,109 @@
+from microbenchmarks._items import BenchmarkItem, DataType
+from microbenchmarks._metric import MetricType
+from microbenchmarks._plan import BenchmarkPlan
+
+import taichi as ti
+
+
+def matrix_operations_default(arch, repeat, matrix_op, block_mn, element_num,
+ dtype, get_metric):
+ m, n = block_mn
+ global_matrixA = ti.Matrix.field(m, n, dtype, shape=element_num)
+ global_matrixB = ti.Matrix.field(m, n, dtype, shape=element_num)
+ global_matrixC = ti.Matrix.field(m, n, dtype, shape=element_num)
+
+ @ti.kernel
+ def fill_matrixA():
+ for e in global_matrixA:
+ for i, j in ti.static(range(m, n)):
+ global_matrixA[e][i, j] = ti.random(dtype)
+
+ @ti.kernel
+ def fill_matrixB():
+ for e in global_matrixB:
+ for i, j in ti.static(range(m, n)):
+ global_matrixB[e][i, j] = ti.random(dtype)
+
+ @ti.kernel
+ def op_throughput():
+ for e in range(element_num):
+ #prelogue
+ A = global_matrixA[e]
+ B = global_matrixB[e]
+ C = ti.Matrix.zero(dtype, m, n)
+ C = matrix_op(C, A, B) #C += A@B
+ #loop
+ for i in range(2048):
+ for j in ti.static(range(4)): #16*4*4=256
+ A = matrix_op(A, C, B) #A += C@B
+ C = matrix_op(C, A, B) #C += A@B
+ B = matrix_op(B, A, C) #B += A@C
+ C = matrix_op(C, A, B) #C += A@B
+ #epilogue
+ global_matrixC[e] = C
+
+ fill_matrixA()
+ fill_matrixB()
+ return get_metric(repeat, op_throughput)
+
+
+@ti.func
+def matrix_add(C, A, B):
+ C = A + B
+ return C
+
+
+@ti.func
+def matrix_mul(C, A, B):
+ C = A @ B
+ return C
+
+
+@ti.func
+def matrix_mma(C, A, B):
+ """matrix multiply and add"""
+ C = A @ B + C
+ return C
+
+
+class MatrixOps(BenchmarkItem):
+ name = 'matrix_op'
+
+ def __init__(self):
+ self._items = {
+ 'mat_add': matrix_add,
+ 'mat_mul': matrix_mul,
+ 'mat_mma': matrix_mma,
+ }
+
+
+class BlockMN(BenchmarkItem):
+ name = 'block_mn'
+
+ def __init__(self):
+ self._items = {
+ 'block_mn11': (1, 1),
+ 'block_mn22': (2, 2),
+ 'block_mn33': (3, 3),
+ 'block_mn44': (4, 4),
+ }
+
+
+class ElementNum(BenchmarkItem):
+ name = 'element_num'
+
+ def __init__(self):
+ self._items = {
+ 'element16384': 16384,
+ #enough threads for filling CUDA cores
+ }
+
+
+class MatrixOpsPlan(BenchmarkPlan):
+ def __init__(self, arch: str):
+ super().__init__('matrix_ops', arch, basic_repeat_times=10)
+ dtype = DataType()
+ dtype.remove(['i64', 'f64'])
+ self.create_plan(MatrixOps(), BlockMN(), ElementNum(), dtype,
+ MetricType())
+ self.add_func(['element16384'], matrix_operations_default)
diff --git a/benchmarks/microbenchmarks/memcpy.py b/benchmarks/microbenchmarks/memcpy.py
new file mode 100644
index 0000000000000..ee7b654bf9c2a
--- /dev/null
+++ b/benchmarks/microbenchmarks/memcpy.py
@@ -0,0 +1,36 @@
+from microbenchmarks._items import Container, DataSize, DataType
+from microbenchmarks._metric import MetricType
+from microbenchmarks._plan import BenchmarkPlan
+from microbenchmarks._utils import dtype_size, fill_random, scaled_repeat_times
+
+import taichi as ti
+
+
+def memcpy_default(arch, repeat, container, dtype, dsize, get_metric):
+ @ti.kernel
+ def memcpy_field(dst: ti.template(), src: ti.template()):
+ for I in ti.grouped(dst):
+ dst[I] = src[I]
+
+ @ti.kernel
+ def memcpy_array(dst: ti.any_arr(), src: ti.any_arr()):
+ for I in ti.grouped(dst):
+ dst[I] = src[I]
+
+ repeat = scaled_repeat_times(arch, dsize, repeat)
+ num_elements = dsize // dtype_size(dtype) // 2 # y=x
+
+ x = container(dtype, num_elements)
+ y = container(dtype, num_elements)
+
+ func = memcpy_field if container == ti.field else memcpy_array
+ fill_random(x, dtype, container)
+ return get_metric(repeat, func, y, x)
+
+
+class MemcpyPlan(BenchmarkPlan):
+ def __init__(self, arch: str):
+ super().__init__('memcpy', arch, basic_repeat_times=10)
+ self.create_plan(Container(), DataType(), DataSize(), MetricType())
+ self.add_func(['field'], memcpy_default)
+ self.add_func(['ndarray'], memcpy_default)
diff --git a/benchmarks/microbenchmarks/saxpy.py b/benchmarks/microbenchmarks/saxpy.py
new file mode 100644
index 0000000000000..de8b32909af51
--- /dev/null
+++ b/benchmarks/microbenchmarks/saxpy.py
@@ -0,0 +1,39 @@
+from microbenchmarks._items import Container, DataSize, DataType
+from microbenchmarks._metric import MetricType
+from microbenchmarks._plan import BenchmarkPlan
+from microbenchmarks._utils import dtype_size, fill_random, scaled_repeat_times
+
+import taichi as ti
+
+
+def saxpy_default(arch, repeat, container, dtype, dsize, get_metric):
+
+ repeat = scaled_repeat_times(arch, dsize, repeat)
+ num_elements = dsize // dtype_size(dtype) // 3 #z=x+y
+
+ x = container(dtype, num_elements)
+ y = container(dtype, num_elements)
+ z = container(dtype, num_elements)
+
+ @ti.kernel
+ def saxpy_field(z: ti.template(), x: ti.template(), y: ti.template()):
+ for i in z:
+ z[i] = 17 * x[i] + y[i]
+
+ @ti.kernel
+ def saxpy_array(z: ti.any_arr(), x: ti.any_arr(), y: ti.any_arr()):
+ for i in z:
+ z[i] = 17 * x[i] + y[i]
+
+ fill_random(x, dtype, container)
+ fill_random(y, dtype, container)
+ func = saxpy_field if container == ti.field else saxpy_array
+ return get_metric(repeat, func, z, x, y)
+
+
+class SaxpyPlan(BenchmarkPlan):
+ def __init__(self, arch: str):
+ super().__init__('saxpy', arch, basic_repeat_times=10)
+ self.create_plan(Container(), DataType(), DataSize(), MetricType())
+ self.add_func(['field'], saxpy_default)
+ self.add_func(['ndarray'], saxpy_default)
diff --git a/benchmarks/microbenchmarks/stencil2d.py b/benchmarks/microbenchmarks/stencil2d.py
new file mode 100644
index 0000000000000..02cf21b2a3f66
--- /dev/null
+++ b/benchmarks/microbenchmarks/stencil2d.py
@@ -0,0 +1,135 @@
+from microbenchmarks._items import BenchmarkItem, Container, DataType
+from microbenchmarks._metric import MetricType
+from microbenchmarks._plan import BenchmarkPlan
+from microbenchmarks._utils import (dtype_size, fill_random,
+ scaled_repeat_times, size2tag)
+
+import taichi as ti
+
+stencil_common = [(0, 0), (0, -1), (0, 1), (1, 0)]
+
+
+def stencil_2d_default(arch, repeat, scatter, bls, container, dtype, dsize_2d,
+ get_metric):
+
+ dsize = dsize_2d[0] * dsize_2d[1]
+ repeat = scaled_repeat_times(arch, dsize, repeat)
+ num_elements_2d = (dsize_2d[0] // dtype_size(dtype), dsize_2d[1] // 2)
+
+ y = container(dtype, shape=num_elements_2d)
+ x = container(dtype, shape=num_elements_2d)
+
+ @ti.kernel
+ def stencil_2d_field(y: ti.template(), x: ti.template()):
+ for I in ti.grouped(x):
+ if ti.static(scatter):
+ for offset in ti.static(stencil_common):
+ y[I + ti.Vector(offset)] += x[I]
+ else: # gather
+ s = ti.cast(0.0, dtype)
+ for offset in ti.static(stencil_common):
+ s = s + x[I + ti.Vector(offset)]
+ y[I] = s
+
+ @ti.kernel
+ def stencil_2d_array(y: ti.any_arr(), x: ti.any_arr()):
+ for I in ti.grouped(x):
+ if ti.static(scatter):
+ for offset in ti.static(stencil_common):
+ y[I + ti.Vector(offset)] += x[I]
+ else: # gather
+ s = ti.cast(0.0, dtype)
+ for offset in ti.static(stencil_common):
+ s = s + x[I + ti.Vector(offset)]
+ y[I] = s
+
+ fill_random(x, dtype, container)
+ func = stencil_2d_field if container == ti.field else stencil_2d_array
+ return get_metric(repeat, func, y, x)
+
+
+def stencil_2d_sparse_bls(arch, repeat, scatter, bls, container, dtype,
+ dsize_2d, get_metric):
+
+ dsize = dsize_2d[0] * dsize_2d[1]
+ if dsize <= 4096 or dsize > 67108864: # 16KB <= dsize <= 64 MB: Sparse-specific parameters
+ return None
+ repeat = scaled_repeat_times(
+ arch, dsize, 1) # basic_repeat_time = 1: Sparse-specific parameters
+ block_elements_2d = (dsize_2d[0] // dtype_size(dtype) // 8,
+ dsize_2d[1] // 2 // 8)
+
+ block = ti.root.pointer(ti.ij, block_elements_2d)
+ y = ti.field(dtype)
+ x = ti.field(dtype)
+ block.dense(ti.ij, 8).place(y)
+ block.dense(ti.ij, 8).place(x)
+
+ @ti.kernel
+ def active_all():
+ for i, j in ti.ndrange(block_elements_2d[0], block_elements_2d[0]):
+ ti.activate(block, [i, j])
+
+ active_all()
+
+ @ti.kernel
+ def stencil_2d(y: ti.template(), x: ti.template()):
+ #reference: tests/python/bls_test_template.py
+ if ti.static(bls and not scatter):
+ ti.block_local(x)
+ if ti.static(bls and scatter):
+ ti.block_local(y)
+ ti.block_dim(64) # 8*8=64
+
+ for I in ti.grouped(x):
+ if ti.static(scatter):
+ for offset in ti.static(stencil_common):
+ y[I + ti.Vector(offset)] += x[I]
+ else: # gather
+ s = ti.cast(0.0, dtype)
+ for offset in ti.static(stencil_common):
+ s = s + x[I + ti.Vector(offset)]
+ y[I] = s
+
+ fill_random(x, dtype, container)
+ return get_metric(repeat, stencil_2d, y, x)
+
+
+class Scatter(BenchmarkItem):
+ name = 'scatter'
+
+ def __init__(self):
+ self._items = {'scatter': True, 'gether': False}
+
+
+class BloclLocalStorage(BenchmarkItem):
+ name = 'bls'
+
+ def __init__(self):
+ self._items = {'bls_on': True, 'bls_off': False}
+
+
+class DataSize2D(BenchmarkItem):
+ name = 'dsize_2d'
+
+ def __init__(self):
+ self._items = {}
+ for i in range(2, 10, 2): # [16KB,256KB,4MB,64MB]
+ size_bytes_2d = 32 * (2**i), 32 * (2**i)
+ size_bytes = size_bytes_2d[0] * size_bytes_2d[1]
+ self._items[size2tag(size_bytes)] = size_bytes_2d
+
+
+class Stencil2DPlan(BenchmarkPlan):
+ def __init__(self, arch: str):
+ super().__init__('stencil_2d', arch, basic_repeat_times=10)
+ container = Container()
+ container.update({'sparse': None}) # None: implement by feature
+ self.create_plan(Scatter(), BloclLocalStorage(), container, DataType(),
+ DataSize2D(), MetricType())
+ # no use for field & ndarray
+ self.remove_cases_with_tags(['field', 'bls1'])
+ self.remove_cases_with_tags(['ndarray', 'bls1'])
+ self.add_func(['field'], stencil_2d_default)
+ self.add_func(['ndarray'], stencil_2d_default)
+ self.add_func(['sparse'], stencil_2d_sparse_bls)
diff --git a/benchmarks/minimal.py b/benchmarks/minimal.py
deleted file mode 100644
index fa75cc6a9513c..0000000000000
--- a/benchmarks/minimal.py
+++ /dev/null
@@ -1,12 +0,0 @@
-import taichi as ti
-
-
-@ti.test()
-def benchmark_fill_scalar():
- a = ti.field(dtype=ti.f32, shape=())
-
- @ti.kernel
- def fill():
- a[None] = 1.0
-
- return ti.benchmark(fill, repeat=1000)
diff --git a/benchmarks/misc/membound.py b/benchmarks/misc/membound.py
deleted file mode 100644
index 55bfd3c9942e3..0000000000000
--- a/benchmarks/misc/membound.py
+++ /dev/null
@@ -1,107 +0,0 @@
-import time
-
-from membound_cases import fill, reduction, saxpy
-from utils import *
-
-import taichi as ti
-
-test_cases = [fill, saxpy, reduction]
-test_archs = [ti.cuda]
-test_dtype = [ti.i32, ti.i64, ti.f32, ti.f64]
-test_dsize = [(4**i) * kibibyte for i in range(1, 11)] #[4KB,16KB...1GB]
-test_repeat = 10
-results_evaluation = [geometric_mean]
-
-
-class BenchmarkResult:
- def __init__(self, name, arch, dtype, dsize, results_evaluation):
- self.test_name = name
- self.test_arch = arch
- self.data_type = dtype
- self.data_size = dsize
- self.min_time_in_us = []
- self.results_evaluation = results_evaluation
-
- def time2mdtableline(self):
- string = '|' + self.test_name + '.' + dtype2str[self.data_type] + '|'
- string += ''.join(
- str(round(time, 4)) + '|' for time in self.min_time_in_us)
- string += ''.join(
- str(round(item(self.min_time_in_us), 4)) + '|'
- for item in self.results_evaluation)
- return string
-
-
-class BenchmarkImpl:
- def __init__(self, func, archs, data_type, data_size):
- self.func = func
- self.name = func.__name__
- self.env = None
- self.device = None
- self.archs = archs
- self.data_type = data_type
- self.data_size = data_size
- self.benchmark_results = []
-
- def run(self):
- for arch in self.archs:
- for dtype in self.data_type:
- ti.init(kernel_profiler=True, arch=arch)
- print("TestCase[%s.%s.%s]" %
- (self.func.__name__, ti.core.arch_name(arch),
- dtype2str[dtype]))
- result = BenchmarkResult(self.name, arch, dtype,
- self.data_size, results_evaluation)
- for size in self.data_size:
- print("data_size = %s" % (size2str(size)))
- result.min_time_in_us.append(
- self.func(arch, dtype, size, test_repeat))
- time.sleep(0.2)
- self.benchmark_results.append(result)
-
- def print(self):
- i = 0
- for arch in self.archs:
- for dtype in self.data_type:
- for idx in range(len(self.data_size)):
- print(
- " test_case:[%s] arch:[%s] dtype:[%s] dsize:[%7s] >>> time:[%4.4f]"
- %
- (self.name, ti.core.arch_name(arch), dtype2str[dtype],
- size2str(self.benchmark_results[i].data_size[idx]),
- self.benchmark_results[i].min_time_in_us[idx]))
- i = i + 1
-
- def save2markdown(self, arch):
- header = '|kernel elapsed time(ms)' + ''.join(
- '|' for i in range(len(self.data_size) + len(results_evaluation)))
- lines = [header]
- for result in self.benchmark_results:
- if (result.test_arch == arch):
- lines.append(result.time2mdtableline())
- return lines
-
-
-class Membound:
- benchmark_imps = []
-
- def __init__(self):
- for case in test_cases:
- self.benchmark_imps.append(
- BenchmarkImpl(case, test_archs, test_dtype, test_dsize))
-
- def run(self):
- for case in self.benchmark_imps:
- case.run()
-
- def mdlines(self, arch):
- lines = []
- lines += md_table_header(self.__class__.__name__, arch, test_dsize,
- test_repeat, results_evaluation)
- for case in self.benchmark_imps:
- if arch in case.archs:
- lines += case.save2markdown(arch)
- else:
- continue
- lines.append('')
- return lines
diff --git a/benchmarks/misc/membound_cases.py b/benchmarks/misc/membound_cases.py
deleted file mode 100644
index b02b76382f008..0000000000000
--- a/benchmarks/misc/membound_cases.py
+++ /dev/null
@@ -1,76 +0,0 @@
-from utils import dtype_size, scale_repeat, size2str
-
-import taichi as ti
-
-
-def init_const(x, dtype, num_elements):
- @ti.kernel
- def init_const(x: ti.template(), n: ti.i32):
- for i in range(n):
- x[i] = ti.cast(0.7, dtype)
-
- init_const(x, num_elements)
-
-
-def membound_benchmark(func, num_elements, repeat):
- # compile the kernel first
- func(num_elements)
- ti.clear_kernel_profile_info()
- for i in range(repeat):
- func(num_elements)
- kernelname = func.__name__
- quering_result = ti.query_kernel_profile_info(kernelname)
- return quering_result.min
-
-
-def fill(arch, dtype, dsize, repeat=10):
-
- repeat = scale_repeat(arch, dsize, repeat)
- num_elements = dsize // dtype_size[dtype]
-
- x = ti.field(dtype, shape=num_elements)
-
- @ti.kernel
- def fill_const(n: ti.i32):
- for i in range(n):
- x[i] = ti.cast(0.7, dtype)
-
- return membound_benchmark(fill_const, num_elements, repeat)
-
-
-def saxpy(arch, dtype, dsize, repeat=10):
-
- repeat = scale_repeat(arch, dsize, repeat)
- num_elements = dsize // dtype_size[dtype] // 3 #z=x+y
-
- x = ti.field(dtype, shape=num_elements)
- y = ti.field(dtype, shape=num_elements)
- z = ti.field(dtype, shape=num_elements)
-
- @ti.kernel
- def saxpy(n: ti.i32):
- for i in range(n):
- z[i] = 17 * x[i] + y[i]
-
- init_const(x, dtype, num_elements)
- init_const(y, dtype, num_elements)
-
- return membound_benchmark(saxpy, num_elements, repeat)
-
-
-def reduction(arch, dtype, dsize, repeat=10):
-
- repeat = scale_repeat(arch, dsize, repeat)
- num_elements = dsize // dtype_size[dtype]
-
- x = ti.field(dtype, shape=num_elements)
- y = ti.field(dtype, shape=())
- y[None] = 0
-
- @ti.kernel
- def reduction(n: ti.i32):
- for i in range(n):
- y[None] += x[i]
-
- init_const(x, dtype, num_elements)
- return membound_benchmark(reduction, num_elements, repeat)
diff --git a/benchmarks/misc/run.py b/benchmarks/misc/run.py
deleted file mode 100644
index 69fff53d82a5c..0000000000000
--- a/benchmarks/misc/run.py
+++ /dev/null
@@ -1,33 +0,0 @@
-from membound import Membound
-
-import taichi as ti
-
-test_suites = [Membound]
-test_archs = [ti.cuda]
-
-
-class PerformanceMonitoring:
- suites = []
-
- def __init__(self):
- for s in test_suites:
- self.suites.append(s())
-
- def run(self):
- print("Running...")
- for s in self.suites:
- s.run()
-
- def write_md(self):
- filename = f'performance_result.md'
- with open(filename, 'w') as f:
- for arch in test_archs:
- for s in self.suites:
- lines = s.mdlines(arch)
- for line in lines:
- print(line, file=f)
-
-
-p = PerformanceMonitoring()
-p.run()
-p.write_md()
diff --git a/benchmarks/misc/utils.py b/benchmarks/misc/utils.py
deleted file mode 100644
index 5e6206116b34d..0000000000000
--- a/benchmarks/misc/utils.py
+++ /dev/null
@@ -1,59 +0,0 @@
-import taichi as ti
-
-kibibyte = 1024
-
-bls2str = {False: "BLS_off", True: "BLS_on"}
-dense2str = {False: "Struct_for", True: "Range_for"}
-
-dtype2str = {ti.i32: "i32", ti.i64: "i64", ti.f32: "f32", ti.f64: "f64"}
-dtype_size = {ti.i32: 4, ti.i64: 8, ti.f32: 4, ti.f64: 8}
-
-# for output string
-size_subsection = [(0.0, 'B'), (1024.0, 'KB'), (1048576.0, 'MB'),
- (1073741824.0, 'GB'), (float('inf'), 'INF')] #B KB MB GB
-
-
-def size2str(size_in_byte):
- for dsize, units in reversed(size_subsection):
- if size_in_byte >= dsize:
- return str(round(size_in_byte / dsize, 4)) + units
-
-
-def scale_repeat(arch, datasize, repeat=10):
- scaled = repeat
- if (arch == ti.gpu) | (arch == ti.opengl) | (arch == ti.cuda):
- scaled *= 10
- if datasize <= 4 * 1024 * 1024:
- scaled *= 10
- return scaled
-
-
-def geometric_mean(data_array):
- product = 1
- for data in data_array:
- product *= data
- return pow(product, 1.0 / len(data_array))
-
-
-def md_table_header(suite_name, arch, test_dsize, test_repeat,
- results_evaluation):
- header = '|' + suite_name + '.' + ti.core.arch_name(arch) + '|'
- header += ''.join('|' for i in range(len(test_dsize)))
- header += ''.join(item.__name__ + '|' for item in results_evaluation)
-
- layout = '|:--:|'
- layout += ''.join(
- ':--:|' for i in range(len(test_dsize) + len(results_evaluation)))
-
- size = '|**data size**|'
- size += ''.join(size2str(size) + '|' for size in test_dsize)
- size += ''.join('|' for i in range(len(results_evaluation)))
-
- repeat = '|**repeat**|'
- repeat += ''.join(
- str(scale_repeat(arch, size, test_repeat)) + '|'
- for size in test_dsize)
- repeat += ''.join('|' for i in range(len(results_evaluation)))
-
- lines = [header, layout, size, repeat]
- return lines
diff --git a/benchmarks/mpm2d.py b/benchmarks/mpm2d.py
deleted file mode 100644
index 4daa34d63e960..0000000000000
--- a/benchmarks/mpm2d.py
+++ /dev/null
@@ -1,241 +0,0 @@
-import time
-
-import numpy as np
-
-import taichi as ti
-
-
-@ti.test()
-def benchmark_range():
- quality = 1 # Use a larger value for higher-res simulations
- n_particles, n_grid = 9000 * quality**2, 128 * quality
- dx, inv_dx = 1 / n_grid, float(n_grid)
- dt = 1e-4 / quality
- p_vol, p_rho = (dx * 0.5)**2, 1
- p_mass = p_vol * p_rho
- E, nu = 0.1e4, 0.2 # Young's modulus and Poisson's ratio
- mu_0, lambda_0 = E / (2 * (1 + nu)), E * nu / (
- (1 + nu) * (1 - 2 * nu)) # Lame parameters
-
- x = ti.Vector.field(2, dtype=ti.f32, shape=n_particles) # position
- v = ti.Vector.field(2, dtype=ti.f32, shape=n_particles) # velocity
- C = ti.Matrix.field(2, 2, dtype=ti.f32,
- shape=n_particles) # affine velocity field
- F = ti.Matrix.field(2, 2, dtype=ti.f32,
- shape=n_particles) # deformation gradient
- material = ti.field(dtype=int, shape=n_particles) # material id
- Jp = ti.field(dtype=ti.f32, shape=n_particles) # plastic deformation
- grid_v = ti.Vector.field(2, dtype=ti.f32,
- shape=(n_grid,
- n_grid)) # grid node momemtum/velocity
- grid_m = ti.field(dtype=ti.f32, shape=(n_grid, n_grid)) # grid node mass
-
- @ti.kernel
- def substep():
- for i, j in ti.ndrange(n_grid, n_grid):
- grid_v[i, j] = [0, 0]
- grid_m[i, j] = 0
- for p in range(n_particles
- ): # Particle state update and scatter to grid (P2G)
- base = (x[p] * inv_dx - 0.5).cast(int)
- fx = x[p] * inv_dx - base.cast(float)
- # Quadratic kernels [http://mpm.graphics Eqn. 123, with x=fx, fx-1,fx-2]
- w = [0.5 * (1.5 - fx)**2, 0.75 - (fx - 1)**2, 0.5 * (fx - 0.5)**2]
- F[p] = (ti.Matrix.identity(ti.f32, 2) +
- dt * C[p]) @ F[p] # deformation gradient update
- h = ti.exp(
- 10 * (1.0 - Jp[p])
- ) # Hardening coefficient: snow gets harder when compressed
- if material[p] == 1: # jelly, make it softer
- h = 0.3
- mu, la = mu_0 * h, lambda_0 * h
- if material[p] == 0: # liquid
- mu = 0.0
- U, sig, V = ti.svd(F[p])
- J = 1.0
- for d in ti.static(range(2)):
- new_sig = sig[d, d]
- if material[p] == 2: # Snow
- new_sig = min(max(sig[d, d], 1 - 2.5e-2),
- 1 + 4.5e-3) # Plasticity
- Jp[p] *= sig[d, d] / new_sig
- sig[d, d] = new_sig
- J *= new_sig
- if material[
- p] == 0: # Reset deformation gradient to avoid numerical instability
- F[p] = ti.Matrix.identity(ti.f32, 2) * ti.sqrt(J)
- elif material[p] == 2:
- F[p] = U @ sig @ V.T(
- ) # Reconstruct elastic deformation gradient after plasticity
- stress = 2 * mu * (F[p] - U @ V.T()) @ F[p].T(
- ) + ti.Matrix.identity(ti.f32, 2) * la * J * (J - 1)
- stress = (-dt * p_vol * 4 * inv_dx * inv_dx) * stress
- affine = stress + p_mass * C[p]
- for i, j in ti.static(ti.ndrange(
- 3, 3)): # Loop over 3x3 grid node neighborhood
- offset = ti.Vector([i, j])
- dpos = (offset.cast(float) - fx) * dx
- weight = w[i][0] * w[j][1]
- grid_v[base +
- offset] += weight * (p_mass * v[p] + affine @ dpos)
- grid_m[base + offset] += weight * p_mass
- for i, j in ti.ndrange(n_grid, n_grid):
- if grid_m[i, j] > 0: # No need for epsilon here
- grid_v[i, j] = (
- 1 / grid_m[i, j]) * grid_v[i, j] # Momentum to velocity
- grid_v[i, j][1] -= dt * 50 # gravity
- if i < 3 and grid_v[i, j][0] < 0:
- grid_v[i, j][0] = 0 # Boundary conditions
- if i > n_grid - 3 and grid_v[i, j][0] > 0: grid_v[i, j][0] = 0
- if j < 3 and grid_v[i, j][1] < 0: grid_v[i, j][1] = 0
- if j > n_grid - 3 and grid_v[i, j][1] > 0: grid_v[i, j][1] = 0
- for p in range(n_particles): # grid to particle (G2P)
- base = (x[p] * inv_dx - 0.5).cast(int)
- fx = x[p] * inv_dx - base.cast(float)
- w = [
- 0.5 * (1.5 - fx)**2, 0.75 - (fx - 1.0)**2, 0.5 * (fx - 0.5)**2
- ]
- new_v = ti.Vector.zero(ti.f32, 2)
- new_C = ti.Matrix.zero(ti.f32, 2, 2)
- for i, j in ti.static(ti.ndrange(
- 3, 3)): # loop over 3x3 grid node neighborhood
- dpos = ti.Vector([i, j]).cast(float) - fx
- g_v = grid_v[base + ti.Vector([i, j])]
- weight = w[i][0] * w[j][1]
- new_v += weight * g_v
- new_C += 4 * inv_dx * weight * g_v.outer_product(dpos)
- v[p], C[p] = new_v, new_C
- x[p] += dt * v[p] # advection
-
- import random
- group_size = n_particles // 3
- for i in range(n_particles):
- x[i] = [
- random.random() * 0.2 + 0.3 + 0.10 * (i // group_size),
- random.random() * 0.2 + 0.05 + 0.32 * (i // group_size)
- ]
- material[i] = i // group_size # 0: fluid 1: jelly 2: snow
- v[i] = [0, 0]
- F[i] = [[1, 0], [0, 1]]
- Jp[i] = 1
-
- ti.benchmark(substep, repeat=4000)
-
-
-@ti.test(exclude=ti.opengl)
-def benchmark_struct():
- quality = 1 # Use a larger value for higher-res simulations
- n_particles, n_grid = 9000 * quality**2, 128 * quality
- dx, inv_dx = 1 / n_grid, float(n_grid)
- dt = 1e-4 / quality
- p_vol, p_rho = (dx * 0.5)**2, 1
- p_mass = p_vol * p_rho
- E, nu = 0.1e4, 0.2 # Young's modulus and Poisson's ratio
- mu_0, lambda_0 = E / (2 * (1 + nu)), E * nu / (
- (1 + nu) * (1 - 2 * nu)) # Lame parameters
-
- x = ti.Vector.field(2, dtype=ti.f32, shape=n_particles) # position
- v = ti.Vector.field(2, dtype=ti.f32, shape=n_particles) # velocity
- C = ti.Matrix.field(2, 2, dtype=ti.f32,
- shape=n_particles) # affine velocity field
- F = ti.Matrix.field(2, 2, dtype=ti.f32,
- shape=n_particles) # deformation gradient
- material = ti.field(dtype=int, shape=n_particles) # material id
- Jp = ti.field(dtype=ti.f32, shape=n_particles) # plastic deformation
- grid_v = ti.Vector.field(2, dtype=ti.f32,
- shape=(n_grid,
- n_grid)) # grid node momemtum/velocity
- grid_m = ti.field(dtype=ti.f32, shape=(n_grid, n_grid)) # grid node mass
-
- @ti.kernel
- def substep():
- for i, j in grid_m:
- grid_v[i, j] = [0, 0]
- grid_m[i, j] = 0
-
- for p in x: # Particle state update and scatter to grid (P2G)
- base = (x[p] * inv_dx - 0.5).cast(int)
- fx = x[p] * inv_dx - base.cast(float)
- # Quadratic kernels [http://mpm.graphics Eqn. 123, with x=fx, fx-1,fx-2]
- w = [0.5 * (1.5 - fx)**2, 0.75 - (fx - 1)**2, 0.5 * (fx - 0.5)**2]
- F[p] = (ti.Matrix.identity(ti.f32, 2) +
- dt * C[p]) @ F[p] # deformation gradient update
- h = ti.exp(
- 10 * (1.0 - Jp[p])
- ) # Hardening coefficient: snow gets harder when compressed
- if material[p] == 1: # jelly, make it softer
- h = 0.3
- mu, la = mu_0 * h, lambda_0 * h
- if material[p] == 0: # liquid
- mu = 0.0
- U, sig, V = ti.svd(F[p])
- J = 1.0
- for d in ti.static(range(2)):
- new_sig = sig[d, d]
- if material[p] == 2: # Snow
- new_sig = min(max(sig[d, d], 1 - 2.5e-2),
- 1 + 4.5e-3) # Plasticity
- Jp[p] *= sig[d, d] / new_sig
- sig[d, d] = new_sig
- J *= new_sig
- if material[
- p] == 0: # Reset deformation gradient to avoid numerical instability
- F[p] = ti.Matrix.identity(ti.f32, 2) * ti.sqrt(J)
- elif material[p] == 2:
- F[p] = U @ sig @ V.T(
- ) # Reconstruct elastic deformation gradient after plasticity
- stress = 2 * mu * (F[p] - U @ V.T()) @ F[p].T(
- ) + ti.Matrix.identity(ti.f32, 2) * la * J * (J - 1)
- stress = (-dt * p_vol * 4 * inv_dx * inv_dx) * stress
- affine = stress + p_mass * C[p]
- for i, j in ti.static(ti.ndrange(
- 3, 3)): # Loop over 3x3 grid node neighborhood
- offset = ti.Vector([i, j])
- dpos = (offset.cast(float) - fx) * dx
- weight = w[i][0] * w[j][1]
- grid_v[base +
- offset] += weight * (p_mass * v[p] + affine @ dpos)
- grid_m[base + offset] += weight * p_mass
-
- for i, j in grid_m:
- if grid_m[i, j] > 0: # No need for epsilon here
- grid_v[i, j] = (
- 1 / grid_m[i, j]) * grid_v[i, j] # Momentum to velocity
- grid_v[i, j][1] -= dt * 50 # gravity
- if i < 3 and grid_v[i, j][0] < 0:
- grid_v[i, j][0] = 0 # Boundary conditions
- if i > n_grid - 3 and grid_v[i, j][0] > 0: grid_v[i, j][0] = 0
- if j < 3 and grid_v[i, j][1] < 0: grid_v[i, j][1] = 0
- if j > n_grid - 3 and grid_v[i, j][1] > 0: grid_v[i, j][1] = 0
-
- for p in x: # grid to particle (G2P)
- base = (x[p] * inv_dx - 0.5).cast(int)
- fx = x[p] * inv_dx - base.cast(float)
- w = [
- 0.5 * (1.5 - fx)**2, 0.75 - (fx - 1.0)**2, 0.5 * (fx - 0.5)**2
- ]
- new_v = ti.Vector.zero(ti.f32, 2)
- new_C = ti.Matrix.zero(ti.f32, 2, 2)
- for i, j in ti.static(ti.ndrange(
- 3, 3)): # loop over 3x3 grid node neighborhood
- dpos = ti.Vector([i, j]).cast(float) - fx
- g_v = grid_v[base + ti.Vector([i, j])]
- weight = w[i][0] * w[j][1]
- new_v += weight * g_v
- new_C += 4 * inv_dx * weight * g_v.outer_product(dpos)
- v[p], C[p] = new_v, new_C
- x[p] += dt * v[p] # advection
-
- import random
- group_size = n_particles // 3
- for i in range(n_particles):
- x[i] = [
- random.random() * 0.2 + 0.3 + 0.10 * (i // group_size),
- random.random() * 0.2 + 0.05 + 0.32 * (i // group_size)
- ]
- material[i] = i // group_size # 0: fluid 1: jelly 2: snow
- v[i] = [0, 0]
- F[i] = [[1, 0], [0, 1]]
- Jp[i] = 1
-
- ti.benchmark(substep, repeat=4000)
diff --git a/benchmarks/requirements.txt b/benchmarks/requirements.txt
new file mode 100644
index 0000000000000..625298e5a2b64
--- /dev/null
+++ b/benchmarks/requirements.txt
@@ -0,0 +1,2 @@
+jsbeautifier
+bokeh
diff --git a/benchmarks/run.py b/benchmarks/run.py
index f7d7215d6cfd4..5c6574dd22cac 100644
--- a/benchmarks/run.py
+++ b/benchmarks/run.py
@@ -1,68 +1,62 @@
import os
+import warnings
-import taichi as ti
+from suite_microbenchmarks import MicroBenchmark
+from taichi._lib import core as ti_core
+from utils import datatime_with_format, dump2json
+benchmark_suites = [MicroBenchmark]
-def get_benchmark_dir():
- return os.path.join(ti.core.get_repo_dir(), 'benchmarks')
-
-class Case:
- def __init__(self, name, func):
- self.name = name
- self.func = func
- self.records = {}
-
- def __lt__(self, other):
- return self.name < other.name
-
- def __eq__(self, other):
- return self.name == other.name
-
- def run(self):
- print(f'==> {self.name}:')
- os.environ['TI_CURRENT_BENCHMARK'] = self.name
- self.func()
-
-
-class Suite:
- def __init__(self, filename):
- self.cases = []
- print(filename)
- self.name = filename[:-3]
- loc = {}
- exec(f'import {self.name} as suite', {}, loc)
- suite = loc['suite']
- case_keys = list(
- sorted(filter(lambda x: x.startswith('benchmark_'), dir(suite))))
- self.cases = [Case(k, getattr(suite, k)) for k in case_keys]
-
- def run(self):
- print(f'{self.name}:')
- for case in sorted(self.cases):
- case.run()
+class BenchmarkInfo:
+ def __init__(self):
+ """init with commit info"""
+ self.commit_hash = ti_core.get_commit_hash()
+ self.datetime = datatime_with_format()
+ self.suites = {}
+ print(f'commit_hash = {self.commit_hash}')
-class TaichiBenchmark:
+class BenchmarkSuites:
def __init__(self):
- self.suites = []
- benchmark_dir = get_benchmark_dir()
- for f in map(os.path.basename, sorted(os.listdir(benchmark_dir))):
- if f != 'run.py' and f.endswith('.py') and f[0] != '_':
- self.suites.append(Suite(f))
+ self._suites = []
+ for suite in benchmark_suites:
+ self._suites.append(suite())
def run(self):
- output_dir = os.environ.get('TI_BENCHMARK_OUTPUT_DIR', '.')
- filename = f'{output_dir}/benchmark.yml'
- try:
- with open(filename, 'r+') as f:
- f.truncate() # clear the previous result
- except FileNotFoundError:
- pass
- print("Running...")
- for s in self.suites:
- s.run()
-
-
-b = TaichiBenchmark()
-b.run()
+ for suite in self._suites:
+ suite.run()
+
+ def save(self, benchmark_dir='./'):
+ for suite in self._suites:
+ suite_dir = os.path.join(benchmark_dir, suite.suite_name)
+ os.makedirs(suite_dir, exist_ok=True)
+ suite.save_as_json(suite_dir)
+
+ def get_suites_info(self):
+ info_dict = {}
+ for suite in self._suites:
+ info_dict[suite.suite_name] = suite.get_benchmark_info()
+ return info_dict
+
+
+def main():
+
+ benchmark_dir = os.path.join(os.getcwd(), 'results')
+ os.makedirs(benchmark_dir, exist_ok=True)
+
+ #init & run
+ info = BenchmarkInfo()
+ suites = BenchmarkSuites()
+ suites.run()
+ #save benchmark results & info
+ suites.save(benchmark_dir)
+ info.suites = suites.get_suites_info()
+ info_path = os.path.join(benchmark_dir, '_info.json')
+ info_str = dump2json(info)
+ with open(info_path, 'w') as f:
+ print(info_str, file=f)
+
+
+if __name__ == '__main__':
+ main()
diff --git a/benchmarks/suite_microbenchmarks.py b/benchmarks/suite_microbenchmarks.py
new file mode 100644
index 0000000000000..dc8e8305404c1
--- /dev/null
+++ b/benchmarks/suite_microbenchmarks.py
@@ -0,0 +1,66 @@
+import os
+import time
+
+from microbenchmarks import benchmark_plan_list
+from utils import dump2json
+
+
+class MicroBenchmark:
+ suite_name = 'microbenchmarks'
+ config = {
+ 'cuda': {
+ 'enable': True
+ },
+ 'vulkan': {
+ 'enable': False
+ },
+ 'opengl': {
+ 'enable': False
+ }
+ }
+
+ def __init__(self):
+ self._results = {}
+ self._info = {}
+
+ def get_benchmark_info(self):
+ info_dict = {}
+ arch_list = []
+ for arch, item in self.config.items():
+ if item['enable'] == True:
+ arch_list.append(arch)
+ info_dict['archs'] = arch_list
+ return info_dict
+
+ def run(self):
+ for arch, item in self.config.items():
+ if item['enable'] == True:
+ arch_results = {}
+ self._info[arch] = {}
+ for plan in benchmark_plan_list:
+ plan_impl = plan(arch)
+ results = plan_impl.run()
+ self._info[arch][plan_impl.name] = results['info']
+ arch_results[plan_impl.name] = results['results']
+
+ self._results[arch] = arch_results
+
+ def save_as_json(self, suite_dir='./'):
+ for arch in self._results:
+ arch_dir = os.path.join(suite_dir, arch)
+ os.makedirs(arch_dir, exist_ok=True)
+ self._save_info_as_json(arch, arch_dir)
+ self._save_cases_as_json(arch, arch_dir)
+
+ def _save_info_as_json(self, arch, arch_dir='./'):
+ info_path = os.path.join(arch_dir, '_info.json')
+ with open(info_path, 'w') as f:
+ print(dump2json(self._info[arch]), file=f)
+
+ def _save_cases_as_json(self, arch, arch_dir='./'):
+ for case in self._info[arch]:
+ case_path = os.path.join(arch_dir, (case + '.json'))
+ case_results = self._results[arch][case]
+ with open(case_path, 'w') as f:
+ case_str = dump2json(case_results)
+ print(case_str, file=f)
diff --git a/benchmarks/summarize_async.py b/benchmarks/summarize_async.py
deleted file mode 100644
index e3ddf4ada499f..0000000000000
--- a/benchmarks/summarize_async.py
+++ /dev/null
@@ -1,31 +0,0 @@
-import math
-
-import yaml
-
-with open('benchmark.yml') as f:
- data = yaml.load(f)
-
-records = {}
-
-for case in data:
- for metrics in ['exec_t', 'launched_tasks']:
- for arch in ['x64', 'cuda']:
- key = metrics, arch
-
- d = []
- for scheme in ['sync', 'async']:
- d.append(data[case][metrics][arch][scheme])
-
- if key not in records:
- records[key] = []
-
- records[key].append(d[0] / d[1])
-
-for k in records:
- rec = records[k]
- # Compute geometric mean
- p = 1
- for r in rec:
- p *= r
- p = p**(1 / len(rec))
- print(f'{k}, {p:.3f}x')
diff --git a/benchmarks/utils.py b/benchmarks/utils.py
index a08a32eecc324..cc9d62371b1fe 100644
--- a/benchmarks/utils.py
+++ b/benchmarks/utils.py
@@ -1,25 +1,21 @@
+import datetime
import functools
+import json
import os
-import taichi as ti
+import jsbeautifier
-def benchmark_async(func):
- @functools.wraps(func)
- def body():
- for arch in [ti.cpu, ti.cuda]:
- for async_mode in [True, False]:
- os.environ['TI_CURRENT_BENCHMARK'] = func.__name__
- ti.init(arch=arch,
- async_mode=async_mode,
- kernel_profiler=True,
- verbose=False)
- if arch == ti.cpu:
- scale = 2
- else:
- # Use more data to hide compilation overhead
- # (since CUDA runs much faster than CPUs)
- scale = 64
- func(scale)
+def get_benchmark_dir():
+ return os.path.dirname(os.path.realpath(__file__))
- return body
+
+def dump2json(obj):
+ obj2dict = obj if type(obj) is dict else obj.__dict__
+ options = jsbeautifier.default_options()
+ options.indent_size = 4
+ return jsbeautifier.beautify(json.dumps(obj2dict), options)
+
+
+def datatime_with_format():
+ return datetime.datetime.now().isoformat()
diff --git a/build.ps1 b/build.ps1
new file mode 100644
index 0000000000000..9857d6d0928db
--- /dev/null
+++ b/build.ps1
@@ -0,0 +1,3 @@
+$stopwatch = [system.diagnostics.stopwatch]::startNew()
+python setup.py develop
+$stopwatch.Elapsed
diff --git a/ci/Dockerfile b/ci/Dockerfile
deleted file mode 100644
index bd8eb13fbfdd2..0000000000000
--- a/ci/Dockerfile
+++ /dev/null
@@ -1,82 +0,0 @@
-# Taichi Dockerfile for development
-ARG UBUNTU
-FROM nvidia/cuda:${UBUNTU}
-ARG PYTHON
-ARG TEST_OPTION
-ARG PYPI_PWD
-ARG COMMIT_SHA
-ENV PYPI_PWD=$PYPI_PWD
-ENV DEBIAN_FRONTEND=noninteractive
-LABEL maintainer="https://github.com/taichi-dev"
-
-RUN apt-get update && \
- apt-get install -y software-properties-common \
- $PYTHON \
- python3-pip \
- ${PYTHON}-dev\
- libtinfo-dev \
- clang-10 \
- wget \
- git \
- libx11-dev \
- libxrandr-dev \
- libxinerama-dev \
- libxcursor-dev \
- libxi-dev \
- libglu1-mesa-dev \
- freeglut3-dev \
- mesa-common-dev \
- build-essential \
- libssl-dev \
- libidn11-dev \
- libz-dev \
- unzip
-
-
-# Install the latest version of CMAKE v3.20.2 from source
-WORKDIR /
-RUN wget https://github.com/Kitware/CMake/releases/download/v3.20.5/cmake-3.20.5-linux-x86_64.tar.gz
-RUN tar xf cmake-3.20.5-linux-x86_64.tar.gz
-ENV PATH="/cmake-3.20.5-linux-x86_64/bin:$PATH"
-
-# Intall LLVM 10
-WORKDIR /
-# Make sure this URL gets updated each time there is a new prebuilt bin release
-RUN wget https://github.com/taichi-dev/taichi_assets/releases/download/llvm10_linux_patch2/taichi-llvm-10.0.0-linux.zip
-RUN unzip taichi-llvm-10.0.0-linux.zip
-ENV PATH="/taichi-llvm-10.0.0-linux/bin:$PATH"
-
-# Install Taichi from source
-ENV CC="clang-10"
-ENV CXX="clang++-10"
-WORKDIR /taichi-dev
-
-# Prevent docker caching when head changes
-ADD https://api.github.com/repos/taichi-dev/taichi/git/refs/heads/master version.json
-RUN git clone https://github.com/taichi-dev/taichi --branch=master
-WORKDIR ./taichi
-RUN git cat-file -t $COMMIT_SHA || exit 1
-RUN git checkout $COMMIT_SHA
-# Install Taichi's Python dependencies
-RUN $PYTHON -m pip install --user -r requirements_dev.txt
-# Build Taichi wheel from source
-RUN git submodule update --init --recursive --depth=1
-WORKDIR /taichi-dev/taichi
-WORKDIR python/
-ENV TAICHI_CMAKE_ARGS="-DTI_WITH_VULKAN:BOOL=OFF"
-RUN $PYTHON build.py build $TEST_OPTION
-WORKDIR ../
-RUN $PYTHON -m pip install dist/*.whl
-
-# Link Taichi source repo to Python Path
-ENV LANG="C.UTF-8"
-
-# Show ELF info
-RUN ldd build/libtaichi_core.so
-RUN strings build/libtaichi_core.so | grep GLIBC
-
-# Install twine and upload project to pypi.
-RUN $PYTHON -m pip install --user twine
-RUN ti test -vr2
-WORKDIR python/
-RUN $PYTHON build.py upload --skip_build $TEST_OPTION
diff --git a/ci/Dockerfile.manylinux2014.cpu b/ci/Dockerfile.manylinux2014.cpu
new file mode 100644
index 0000000000000..ceda86ee0901c
--- /dev/null
+++ b/ci/Dockerfile.manylinux2014.cpu
@@ -0,0 +1,55 @@
+# This file is generated by python Dockerfile_generator.py -o manylinux2014 -t cpu
+# Taichi Dockerfile (CPU only) for Manylinux2014 compliant
+FROM quay.io/pypa/manylinux2014_x86_64
+
+LABEL maintainer="https://github.com/taichi-dev"
+
+RUN yum check-update && \
+ yum install -y git \
+ cmake \
+ wget \
+ libXrandr
+
+# Build LLVM/Clang 10 from source
+WORKDIR /
+RUN wget https://github.com/llvm/llvm-project/releases/download/llvmorg-10.0.0/llvm-10.0.0.src.tar.xz
+RUN tar -xf llvm-10.0.0.src.tar.xz && rm llvm-10.0.0.src.tar.xz
+RUN wget https://github.com/llvm/llvm-project/releases/download/llvmorg-10.0.0/clang-10.0.0.src.tar.xz
+RUN tar -xf clang-10.0.0.src.tar.xz && rm clang-10.0.0.src.tar.xz
+RUN cp -r clang-10.0.0.src llvm-10.0.0.src/tools/clang
+
+WORKDIR /llvm-10.0.0.src/build
+RUN cmake .. -DLLVM_ENABLE_RTTI:BOOL=ON -DBUILD_SHARED_LIBS:BOOL=OFF -DCMAKE_BUILD_TYPE=Release -DLLVM_TARGETS_TO_BUILD="X86" -DLLVM_ENABLE_ASSERTIONS=ON -DLLVM_ENABLE_TERMINFO=OFF
+RUN make -j 8 && make install
+ENV CC="/usr/local/bin/clang"
+ENV CXX="/usr/local/bin/clang++"
+
+# Link gcc 10 to build Taichi
+WORKDIR /usr/lib/gcc/x86_64-redhat-linux/
+RUN ln -s /opt/rh/devtoolset-10/root/usr/lib/gcc/x86_64-redhat-linux/10 10
+# Check gcc-10 is used
+RUN clang++ -v
+
+# Create non-root user for running the container
+RUN useradd -ms /bin/bash dev
+WORKDIR /home/dev
+USER dev
+
+# Install miniconda
+RUN wget https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh && \
+ bash Miniconda3-latest-Linux-x86_64.sh -p /home/dev/miniconda -b
+ENV PATH="/home/dev/miniconda/bin:$PATH"
+
+# Set up multi-python environment
+RUN conda init bash
+RUN conda create -n py36 python=3.6 -y
+RUN conda create -n py37 python=3.7 -y
+RUN conda create -n py38 python=3.8 -y
+RUN conda create -n py39 python=3.9 -y
+
+# Load scripts for build and test
+WORKDIR /home/dev/scripts
+COPY ci/scripts/manylinux_build_wheel.sh manylinux_build_wheel.sh
+
+WORKDIR /home/dev
+ENV LANG="C.UTF-8"
diff --git a/ci/Dockerfile.manylinux2014.cuda b/ci/Dockerfile.manylinux2014.cuda
new file mode 100644
index 0000000000000..347c0f9968a81
--- /dev/null
+++ b/ci/Dockerfile.manylinux2014.cuda
@@ -0,0 +1,56 @@
+FROM nvidia/cudagl:11.2.2-devel-centos7
+
+LABEL maintainer="https://github.com/taichi-dev"
+
+RUN yum install -y git wget
+
+# Install cmake 3.x
+RUN yum install -y epel-release
+RUN yum install -y cmake3
+RUN ln -s /usr/bin/cmake3 /usr/bin/cmake
+
+# Install gcc 10 (https://git.centos.org/rpms/devtoolset-10-gcc)
+RUN yum install -y centos-release-scl
+RUN yum install -y devtoolset-10-gcc*
+ENV PATH="/opt/rh/devtoolset-10/root/usr/bin:$PATH"
+
+# Build LLVM/Clang 10 from source
+WORKDIR /
+RUN wget https://github.com/llvm/llvm-project/releases/download/llvmorg-10.0.0/llvm-10.0.0.src.tar.xz
+RUN tar -xf llvm-10.0.0.src.tar.xz && rm llvm-10.0.0.src.tar.xz
+RUN wget https://github.com/llvm/llvm-project/releases/download/llvmorg-10.0.0/clang-10.0.0.src.tar.xz
+RUN tar -xf clang-10.0.0.src.tar.xz && rm clang-10.0.0.src.tar.xz
+RUN cp -r clang-10.0.0.src llvm-10.0.0.src/tools/clang
+
+WORKDIR /llvm-10.0.0.src/build
+RUN cmake .. -DLLVM_ENABLE_RTTI:BOOL=ON -DBUILD_SHARED_LIBS:BOOL=OFF -DCMAKE_BUILD_TYPE=Release -DLLVM_TARGETS_TO_BUILD="X86;NVPTX" -DLLVM_ENABLE_ASSERTIONS=ON -DLLVM_ENABLE_TERMINFO=OFF
+RUN make -j 8 && make install
+ENV CC="/usr/local/bin/clang"
+ENV CXX="/usr/local/bin/clang++"
+
+# Link gcc 10 to build Taichi
+WORKDIR /usr/lib/gcc/x86_64-redhat-linux/
+RUN ln -s /opt/rh/devtoolset-10/root/usr/lib/gcc/x86_64-redhat-linux/10 10
+# Check gcc-10 is used
+RUN clang++ -v
+
+# Create non-root user for running the container
+RUN useradd -ms /bin/bash dev
+WORKDIR /home/dev
+USER dev
+
+# Install miniconda
+RUN wget https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh && \
+ bash Miniconda3-latest-Linux-x86_64.sh -p /home/dev/miniconda -b
+ENV PATH="/home/dev/miniconda/bin:$PATH"
+
+# Set up multi-python environment
+RUN conda init bash
+RUN conda create -n py36 python=3.6 -y
+RUN conda create -n py37 python=3.7 -y
+RUN conda create -n py38 python=3.8 -y
+RUN conda create -n py39 python=3.9 -y
+RUN conda create -n py310 python=3.10 -y
+
+WORKDIR /home/dev
+ENV LANG="C.UTF-8"
diff --git a/ci/Dockerfile.ubuntu.18.04 b/ci/Dockerfile.ubuntu.18.04
new file mode 100644
index 0000000000000..6addb38ebc664
--- /dev/null
+++ b/ci/Dockerfile.ubuntu.18.04
@@ -0,0 +1,101 @@
+# This file is generated by python Dockerfile_generator.py -o ubuntu -t gpu
+# Taichi Dockerfile for development
+FROM nvidia/cudagl:11.2.2-devel-ubuntu18.04
+# Use 11.2 instead of 11.4 to avoid forward compatibility issue on Nvidia driver 460
+
+ENV NVIDIA_DRIVER_CAPABILITIES compute,graphics,utility
+
+ENV DEBIAN_FRONTEND=noninteractive
+
+LABEL maintainer="https://github.com/taichi-dev"
+
+RUN apt-get update && \
+ apt-get install -y software-properties-common \
+ python3-pip \
+ libtinfo-dev \
+ clang-10 \
+ wget \
+ git \
+ unzip \
+ libxrandr-dev \
+ libxinerama-dev \
+ libxcursor-dev \
+ libxi-dev \
+ libglu1-mesa-dev \
+ freeglut3-dev \
+ mesa-common-dev \
+ libssl-dev \
+ libglm-dev \
+ libxcb-keysyms1-dev \
+ libxcb-dri3-dev \
+ libxcb-randr0-dev \
+ libxcb-ewmh-dev \
+ libpng-dev \
+ g++-multilib \
+ libmirclient-dev \
+ libwayland-dev \
+ bison \
+ libx11-xcb-dev \
+ liblz4-dev \
+ libzstd-dev \
+ qt5-default \
+ libglfw3 \
+ libglfw3-dev \
+ libjpeg-dev \
+ libvulkan-dev
+
+# Install the latest version of CMAKE v3.20.5 from source
+WORKDIR /
+RUN wget https://github.com/Kitware/CMake/releases/download/v3.20.5/cmake-3.20.5-linux-x86_64.tar.gz
+RUN tar xf cmake-3.20.5-linux-x86_64.tar.gz && \
+ rm cmake-3.20.5-linux-x86_64.tar.gz
+ENV PATH="/cmake-3.20.5-linux-x86_64/bin:$PATH"
+
+# Intall LLVM 10
+WORKDIR /
+# Make sure this URL gets updated each time there is a new prebuilt bin release
+RUN wget https://github.com/taichi-dev/taichi_assets/releases/download/llvm10_linux_patch2/taichi-llvm-10.0.0-linux.zip
+RUN unzip taichi-llvm-10.0.0-linux.zip && \
+ rm taichi-llvm-10.0.0-linux.zip
+ENV PATH="/taichi-llvm-10.0.0-linux/bin:$PATH"
+# Use Clang as the default compiler
+ENV CC="clang-10"
+ENV CXX="clang++-10"
+
+# Setting up Vulkan SDK
+# References
+# [1] https://github.com/edowson/docker-nvidia-vulkan
+# [2] https://gitlab.com/nvidia/container-images/vulkan/-/tree/master/docker
+WORKDIR /vulkan
+RUN wget https://sdk.lunarg.com/sdk/download/1.2.189.0/linux/vulkansdk-linux-x86_64-1.2.189.0.tar.gz
+RUN tar xf vulkansdk-linux-x86_64-1.2.189.0.tar.gz && \
+ rm vulkansdk-linux-x86_64-1.2.189.0.tar.gz
+# Locate Vulkan components
+ENV VULKAN_SDK="/vulkan/1.2.189.0/x86_64"
+ENV PATH="$VULKAN_SDK/bin:$PATH"
+ENV LD_LIBRARY_PATH="$VULKAN_SDK/lib${LD_LIBRARY_PATH:+:$LD_LIBRARY_PATH}"
+ENV VK_LAYER_PATH="$VULKAN_SDK/etc/vulkan/explicit_layer.d"
+WORKDIR /usr/share/vulkan/icd.d
+COPY vulkan/icd.d/nvidia_icd.json nvidia_icd.json
+
+# Create non-root user for running the container
+RUN useradd -ms /bin/bash dev
+WORKDIR /home/dev
+USER dev
+
+# Install miniconda
+RUN wget https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh && \
+ bash Miniconda3-latest-Linux-x86_64.sh -p /home/dev/miniconda -b
+ENV PATH="/home/dev/miniconda/bin:$PATH"
+
+# Set up multi-python environment
+RUN conda init bash
+RUN conda create -n py36 python=3.6 pytorch cudatoolkit=10.2 -c pytorch -y
+RUN conda create -n py37 python=3.7 pytorch cudatoolkit=10.2 -c pytorch -y
+RUN conda create -n py38 python=3.8 pytorch cudatoolkit=10.2 -c pytorch -y
+RUN conda create -n py39 python=3.9 pytorch cudatoolkit=10.2 -c pytorch -y
+# TODO add torch to 3.10 when supported
+RUN conda create -n py310 python=3.10 -y
+
+WORKDIR /home/dev
+ENV LANG="C.UTF-8"
diff --git a/ci/Dockerfile.ubuntu.18.04.cpu b/ci/Dockerfile.ubuntu.18.04.cpu
new file mode 100644
index 0000000000000..c7c60244b6837
--- /dev/null
+++ b/ci/Dockerfile.ubuntu.18.04.cpu
@@ -0,0 +1,57 @@
+# This file is generated by python Dockerfile_generator.py -o ubuntu -t cpu
+# Taichi Dockerfile for development
+FROM ubuntu:18.04
+
+ENV DEBIAN_FRONTEND=noninteractive
+
+LABEL maintainer="https://github.com/taichi-dev"
+
+RUN apt-get update && \
+ apt-get install -y software-properties-common \
+ python3-pip \
+ libtinfo-dev \
+ clang-10 \
+ wget \
+ git \
+ unzip \
+ libx11-xcb-dev \
+ zlib1g-dev
+# Install the latest version of CMAKE v3.20.5 from source
+WORKDIR /
+RUN wget https://github.com/Kitware/CMake/releases/download/v3.20.5/cmake-3.20.5-linux-x86_64.tar.gz
+RUN tar xf cmake-3.20.5-linux-x86_64.tar.gz && \
+ rm cmake-3.20.5-linux-x86_64.tar.gz
+ENV PATH="/cmake-3.20.5-linux-x86_64/bin:$PATH"
+
+# Intall LLVM 10
+WORKDIR /
+# Make sure this URL gets updated each time there is a new prebuilt bin release
+RUN wget https://github.com/taichi-dev/taichi_assets/releases/download/llvm10_linux_patch2/taichi-llvm-10.0.0-linux.zip
+RUN unzip taichi-llvm-10.0.0-linux.zip && \
+ rm taichi-llvm-10.0.0-linux.zip
+ENV PATH="/taichi-llvm-10.0.0-linux/bin:$PATH"
+# Use Clang as the default compiler
+ENV CC="clang-10"
+ENV CXX="clang++-10"
+
+# Create non-root user for running the container
+RUN useradd -ms /bin/bash dev
+WORKDIR /home/dev
+USER dev
+
+# Install miniconda
+RUN wget https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh && \
+ bash Miniconda3-latest-Linux-x86_64.sh -p /home/dev/miniconda -b
+ENV PATH="/home/dev/miniconda/bin:$PATH"
+
+# Set up multi-python environment
+RUN conda init bash
+RUN conda create -n py36 python=3.6 pytorch -y
+RUN conda create -n py37 python=3.7 pytorch -y
+RUN conda create -n py38 python=3.8 pytorch -y
+RUN conda create -n py39 python=3.9 pytorch -y
+# TODO add torch to 3.10 when supported
+RUN conda create -n py310 python=3.10 -y
+
+WORKDIR /home/dev
+ENV LANG="C.UTF-8"
diff --git a/Dockerfile b/ci/Dockerfile.ubuntu.20.04
similarity index 73%
rename from Dockerfile
rename to ci/Dockerfile.ubuntu.20.04
index 189f487dc870d..b38442b4bc486 100644
--- a/Dockerfile
+++ b/ci/Dockerfile.ubuntu.20.04
@@ -1,16 +1,17 @@
+# This file is generated by python Dockerfile_generator.py -o ubuntu -t gpu
# Taichi Dockerfile for development
FROM nvidia/cudagl:11.2.2-devel-ubuntu20.04
# Use 11.2 instead of 11.4 to avoid forward compatibility issue on Nvidia driver 460
ENV NVIDIA_DRIVER_CAPABILITIES compute,graphics,utility
+
ENV DEBIAN_FRONTEND=noninteractive
+
LABEL maintainer="https://github.com/taichi-dev"
-# Ubuntu 20.04 installs Python 3.8 by default
RUN apt-get update && \
apt-get install -y software-properties-common \
python3-pip \
- python-is-python3 \
libtinfo-dev \
clang-10 \
wget \
@@ -40,10 +41,10 @@ RUN apt-get update && \
qt5-default \
libglfw3 \
libglfw3-dev \
- vulkan-tools \
+ libjpeg-dev \
libvulkan-dev \
+ vulkan-tools \
vulkan-validationlayers-dev
-
# Install the latest version of CMAKE v3.20.5 from source
WORKDIR /
RUN wget https://github.com/Kitware/CMake/releases/download/v3.20.5/cmake-3.20.5-linux-x86_64.tar.gz
@@ -51,7 +52,6 @@ RUN tar xf cmake-3.20.5-linux-x86_64.tar.gz && \
rm cmake-3.20.5-linux-x86_64.tar.gz
ENV PATH="/cmake-3.20.5-linux-x86_64/bin:$PATH"
-
# Intall LLVM 10
WORKDIR /
# Make sure this URL gets updated each time there is a new prebuilt bin release
@@ -59,7 +59,9 @@ RUN wget https://github.com/taichi-dev/taichi_assets/releases/download/llvm10_li
RUN unzip taichi-llvm-10.0.0-linux.zip && \
rm taichi-llvm-10.0.0-linux.zip
ENV PATH="/taichi-llvm-10.0.0-linux/bin:$PATH"
-
+# Use Clang as the default compiler
+ENV CC="clang-10"
+ENV CXX="clang++-10"
# Setting up Vulkan SDK
# References
@@ -77,30 +79,26 @@ ENV VK_LAYER_PATH="$VULKAN_SDK/etc/vulkan/explicit_layer.d"
WORKDIR /usr/share/vulkan/icd.d
COPY ci/vulkan/icd.d/nvidia_icd.json nvidia_icd.json
+# Create non-root user for running the container
+RUN useradd -ms /bin/bash dev
+WORKDIR /home/dev
+USER dev
-# Install Taichi from source
-ENV CC="clang-10"
-ENV CXX="clang++-10"
-WORKDIR /taichi-dev
-# Prevent docker caching when head changes
-ADD https://api.github.com/repos/taichi-dev/taichi/git/refs/heads/master version.json
-RUN git clone --recursive https://github.com/taichi-dev/taichi --branch=master
-WORKDIR /taichi-dev/taichi
-RUN python3 -m pip install --user -r requirements_dev.txt
-# Update Torch version, otherwise cuda tests fail. See #2969.
-RUN python3 -m pip install torch==1.9.0+cu111 -f https://download.pytorch.org/whl/torch_stable.html
-RUN TAICHI_CMAKE_ARGS="-DTI_WITH_VULKAN:BOOL=ON -DTI_WITH_CUDA:BOOL=ON -DTI_WITH_OPENGL:BOOL=ON" python3 setup.py develop --user
-# Show ELF info
-RUN ldd build/libtaichi_core.so
-RUN strings build/libtaichi_core.so | grep GLIBC
+# Install miniconda
+RUN wget https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh && \
+ bash Miniconda3-latest-Linux-x86_64.sh -p /home/dev/miniconda -b
+ENV PATH="/home/dev/miniconda/bin:$PATH"
-# Link Taichi source repo to Python Path
-ENV PATH="/taichi-dev/taichi/bin:$PATH"
-ENV TAICHI_REPO_DIR="/taichi-dev/taichi/"
-ENV PYTHONPATH="$TAICHI_REPO_DIR/python:$PYTHONPATH"
-ENV LANG="C.UTF-8"
+# Set up multi-python environment
+RUN conda init bash
+RUN conda create -n py36 python=3.6 -y
+RUN conda create -n py37 python=3.7 -y
+RUN conda create -n py38 python=3.8 -y
+RUN conda create -n py39 python=3.9 -y
+
+# Load scripts for build and test
+WORKDIR /home/dev/scripts
+COPY ci/scripts/ubuntu_build_test.sh ubuntu_build_test.sh
-# Add Docker specific ENV
-ENV TI_IN_DOCKER=true
-WORKDIR /taichi-dev/taichi
-CMD /bin/bash
+WORKDIR /home/dev
+ENV LANG="C.UTF-8"
diff --git a/ci/Dockerfile.ubuntu.20.04.cpu b/ci/Dockerfile.ubuntu.20.04.cpu
new file mode 100644
index 0000000000000..23d228faa9500
--- /dev/null
+++ b/ci/Dockerfile.ubuntu.20.04.cpu
@@ -0,0 +1,59 @@
+# This file is generated by python Dockerfile_generator.py -o ubuntu -t cpu
+# Taichi Dockerfile for development
+FROM ubuntu:20.04
+
+ENV DEBIAN_FRONTEND=noninteractive
+
+LABEL maintainer="https://github.com/taichi-dev"
+
+RUN apt-get update && \
+ apt-get install -y software-properties-common \
+ python3-pip \
+ libtinfo-dev \
+ clang-10 \
+ wget \
+ git \
+ unzip \
+ libx11-xcb-dev
+
+# Install the latest version of CMAKE v3.20.5 from source
+WORKDIR /
+RUN wget https://github.com/Kitware/CMake/releases/download/v3.20.5/cmake-3.20.5-linux-x86_64.tar.gz
+RUN tar xf cmake-3.20.5-linux-x86_64.tar.gz && \
+ rm cmake-3.20.5-linux-x86_64.tar.gz
+ENV PATH="/cmake-3.20.5-linux-x86_64/bin:$PATH"
+
+# Intall LLVM 10
+WORKDIR /
+# Make sure this URL gets updated each time there is a new prebuilt bin release
+RUN wget https://github.com/taichi-dev/taichi_assets/releases/download/llvm10_linux_patch2/taichi-llvm-10.0.0-linux.zip
+RUN unzip taichi-llvm-10.0.0-linux.zip && \
+ rm taichi-llvm-10.0.0-linux.zip
+ENV PATH="/taichi-llvm-10.0.0-linux/bin:$PATH"
+# Use Clang as the default compiler
+ENV CC="clang-10"
+ENV CXX="clang++-10"
+
+# Create non-root user for running the container
+RUN useradd -ms /bin/bash dev
+WORKDIR /home/dev
+USER dev
+
+# Install miniconda
+RUN wget https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh && \
+ bash Miniconda3-latest-Linux-x86_64.sh -p /home/dev/miniconda -b
+ENV PATH="/home/dev/miniconda/bin:$PATH"
+
+# Set up multi-python environment
+RUN conda init bash
+RUN conda create -n py36 python=3.6 -y
+RUN conda create -n py37 python=3.7 -y
+RUN conda create -n py38 python=3.8 -y
+RUN conda create -n py39 python=3.9 -y
+
+# Load scripts for build and test
+WORKDIR /home/dev/scripts
+COPY ci/scripts/ubuntu_build_test_cpu.sh ubuntu_build_test_cpu.sh
+
+WORKDIR /home/dev
+ENV LANG="C.UTF-8"
diff --git a/ci/Dockerfile_generator.py b/ci/Dockerfile_generator.py
new file mode 100644
index 0000000000000..d75b2f45ab643
--- /dev/null
+++ b/ci/Dockerfile_generator.py
@@ -0,0 +1,351 @@
+import argparse
+import functools
+import sys
+from enum import Enum
+from functools import reduce
+from pathlib import Path
+
+OS = {
+ "windows": (),
+ "macos": (),
+ "manylinux2014": ("", ),
+ "ubuntu": (
+ "18.04",
+ "20.04",
+ )
+}
+HARDWARE = ("cpu", "gpu")
+
+HEAD_BLOCK = """# This file is generated by python Dockerfile_generator.py -o {os} -t {target}
+"""
+
+CPU_BASE_BLOCK = """# Taichi Dockerfile for development
+FROM {os}:{version}
+"""
+
+CPU_MANYLINUX_BASE_BLOCK = """# Taichi Dockerfile (CPU only) for Manylinux2014 compliant
+FROM quay.io/pypa/manylinux2014_x86_64
+"""
+
+GPU_BASE_BLOCK = """# Taichi Dockerfile for development
+FROM nvidia/cudagl:11.2.2-devel-ubuntu{version}
+# Use 11.2 instead of 11.4 to avoid forward compatibility issue on Nvidia driver 460
+"""
+
+CPU_YUM_INSTALL_BLOCK = """
+RUN yum check-update && \\
+ yum install -y git \\
+ cmake \\
+ wget \\
+ libXrandr
+"""
+
+CPU_APT_INSTALL_BLOCK = """
+RUN apt-get update && \\
+ apt-get install -y software-properties-common \\
+ python3-pip \\
+ libtinfo-dev \\
+ clang-10 \\
+ wget \\
+ git \\
+ unzip \\
+ libx11-xcb-dev
+"""
+
+GPU_APT_INSTALL_BLOCK = """
+RUN apt-get update && \\
+ apt-get install -y software-properties-common \\
+ python3-pip \\
+ libtinfo-dev \\
+ clang-10 \\
+ wget \\
+ git \\
+ unzip \\
+ libxrandr-dev \\
+ libxinerama-dev \\
+ libxcursor-dev \\
+ libxi-dev \\
+ libglu1-mesa-dev \\
+ freeglut3-dev \\
+ mesa-common-dev \\
+ libssl-dev \\
+ libglm-dev \\
+ libxcb-keysyms1-dev \\
+ libxcb-dri3-dev \\
+ libxcb-randr0-dev \\
+ libxcb-ewmh-dev \\
+ libpng-dev \\
+ g++-multilib \\
+ libmirclient-dev \\
+ libwayland-dev \\
+ bison \\
+ libx11-xcb-dev \\
+ liblz4-dev \\
+ libzstd-dev \\
+ qt5-default \\
+ libglfw3 \\
+ libglfw3-dev \\
+ libjpeg-dev \\
+ libvulkan-dev
+"""
+
+NVIDIA_DRIVER_CAPABILITIES_BLOCK = """
+ENV NVIDIA_DRIVER_CAPABILITIES compute,graphics,utility
+"""
+
+DEBIAN_NONINTERACTIVE_BLOCK = """
+ENV DEBIAN_FRONTEND=noninteractive
+"""
+
+MAINTAINER_BLOCK = """
+LABEL maintainer="https://github.com/taichi-dev"
+"""
+
+CMAKE_BLOCK = """
+# Install the latest version of CMAKE v3.20.5 from source
+WORKDIR /
+RUN wget https://github.com/Kitware/CMake/releases/download/v3.20.5/cmake-3.20.5-linux-x86_64.tar.gz
+RUN tar xf cmake-3.20.5-linux-x86_64.tar.gz && \\
+ rm cmake-3.20.5-linux-x86_64.tar.gz
+ENV PATH="/cmake-3.20.5-linux-x86_64/bin:$PATH"
+"""
+
+LLVM_BLOCK = """
+# Intall LLVM 10
+WORKDIR /
+# Make sure this URL gets updated each time there is a new prebuilt bin release
+RUN wget https://github.com/taichi-dev/taichi_assets/releases/download/llvm10_linux_patch2/taichi-llvm-10.0.0-linux.zip
+RUN unzip taichi-llvm-10.0.0-linux.zip && \\
+ rm taichi-llvm-10.0.0-linux.zip
+ENV PATH="/taichi-llvm-10.0.0-linux/bin:$PATH"
+# Use Clang as the default compiler
+ENV CC="clang-10"
+ENV CXX="clang++-10"
+"""
+
+LLVM_CLANG_FROM_SOURCE_BLOCK = """
+# Build LLVM/Clang 10 from source
+WORKDIR /
+RUN wget https://github.com/llvm/llvm-project/releases/download/llvmorg-10.0.0/llvm-10.0.0.src.tar.xz
+RUN tar -xf llvm-10.0.0.src.tar.xz && \
+ rm llvm-10.0.0.src.tar.xz
+RUN wget https://github.com/llvm/llvm-project/releases/download/llvmorg-10.0.0/clang-10.0.0.src.tar.xz
+RUN tar -xf clang-10.0.0.src.tar.xz && \
+ rm clang-10.0.0.src.tar.xz
+RUN cp -r clang-10.0.0.src llvm-10.0.0.src/tools/clang
+
+WORKDIR /llvm-10.0.0.src/build
+RUN cmake .. -DLLVM_ENABLE_RTTI:BOOL=ON -DBUILD_SHARED_LIBS:BOOL=OFF -DCMAKE_BUILD_TYPE=Release -DLLVM_TARGETS_TO_BUILD="X86" -DLLVM_ENABLE_ASSERTIONS=ON -DLLVM_ENABLE_TERMINFO=OFF
+RUN make -j 8 && \
+ make install
+ENV CC="/usr/local/bin/clang"
+ENV CXX="/usr/local/bin/clang++"
+"""
+
+GCC_LINK_BLOCK = """
+# Link gcc 10 to build Taichi
+WORKDIR /usr/lib/gcc/x86_64-redhat-linux/
+RUN ln -s /opt/rh/devtoolset-10/root/usr/lib/gcc/x86_64-redhat-linux/10 10
+# Check gcc-10 is used
+RUN clang++ -v
+"""
+
+USER_BLOCK = """
+# Create non-root user for running the container
+RUN useradd -ms /bin/bash dev
+WORKDIR /home/dev
+USER dev
+"""
+
+VULKAN_BLOCK = """
+# Setting up Vulkan SDK
+# References
+# [1] https://github.com/edowson/docker-nvidia-vulkan
+# [2] https://gitlab.com/nvidia/container-images/vulkan/-/tree/master/docker
+WORKDIR /vulkan
+RUN wget https://sdk.lunarg.com/sdk/download/1.2.189.0/linux/vulkansdk-linux-x86_64-1.2.189.0.tar.gz
+RUN tar xf vulkansdk-linux-x86_64-1.2.189.0.tar.gz && \\
+ rm vulkansdk-linux-x86_64-1.2.189.0.tar.gz
+# Locate Vulkan components
+ENV VULKAN_SDK="/vulkan/1.2.189.0/x86_64"
+ENV PATH="$VULKAN_SDK/bin:$PATH"
+ENV LD_LIBRARY_PATH="$VULKAN_SDK/lib${LD_LIBRARY_PATH:+:$LD_LIBRARY_PATH}"
+ENV VK_LAYER_PATH="$VULKAN_SDK/etc/vulkan/explicit_layer.d"
+WORKDIR /usr/share/vulkan/icd.d
+COPY ci/vulkan/icd.d/nvidia_icd.json nvidia_icd.json
+"""
+
+CONDA_BLOCK = """
+# Install miniconda
+RUN wget https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh && \\
+ bash Miniconda3-latest-Linux-x86_64.sh -p /home/dev/miniconda -b
+ENV PATH="/home/dev/miniconda/bin:$PATH"
+
+# Set up multi-python environment
+RUN conda init bash
+RUN conda create -n py36 python=3.6 -y
+RUN conda create -n py37 python=3.7 -y
+RUN conda create -n py38 python=3.8 -y
+RUN conda create -n py39 python=3.9 -y
+"""
+
+SCRIPTS_BLOCK = """
+# Load scripts for build and test
+WORKDIR /home/dev/scripts
+COPY ci/scripts/{script} {script}
+
+WORKDIR /home/dev
+ENV LANG="C.UTF-8"
+"""
+
+
+class Parser(argparse.ArgumentParser):
+ def error(self, message):
+ """Make it print help message by default."""
+ sys.stderr.write(f"error: {message}\n")
+ self.print_help()
+ sys.exit(2)
+
+
+class AvailableColors(Enum):
+ GRAY = 90
+ RED = 91
+ GREEN = 92
+ YELLOW = 93
+ BLUE = 94
+ PURPLE = 95
+ WHITE = 97
+ BLACK = 30
+ DEFAULT = 39
+
+
+def _apply_color(color: str, message: str) -> str:
+ """Dye message with color, fall back to default if it fails."""
+ color_code = AvailableColors["DEFAULT"].value
+ try:
+ color_code = AvailableColors[color.upper()].value
+ except KeyError:
+ pass
+ return f"\033[1;{color_code}m{message}\033[0m"
+
+
+def info(message: str, plain=False):
+ """Log the info to stdout"""
+ print(_apply_color("default", message) if not plain else message)
+
+
+def success(message: str):
+ """Log the success to stdout"""
+ print(_apply_color("green", f"[✔] {message}"))
+
+
+def error(message: str):
+ """Log the error to stderr"""
+ print(_apply_color("red", f"[✗] {message}"), file=sys.stderr)
+
+
+def warn(message: str):
+ """Log the warning to stdout"""
+ print(_apply_color("yellow", f"[!] {message}"))
+
+
+def main(arguments=None):
+ parser = Parser(description="""A CLI to generate Taichi CI Dockerfiles.
+ Example usage:
+ python3 Dockerfile_generator.py -o ubuntu -t cpu
+ """)
+ parser.add_argument(
+ "-o",
+ "--os",
+ help="The target os of the Dockerfile.",
+ required=True,
+ type=str,
+ choices=OS,
+ metavar="\b",
+ )
+ parser.add_argument(
+ "-t",
+ "--target",
+ help="The target hardware of the Dockerfile. [cpu/gpu]",
+ required=True,
+ type=str,
+ choices=HARDWARE,
+ metavar="\b",
+ )
+ args = parser.parse_args()
+ pwd = Path(__file__).resolve().parent
+
+ head_block = HEAD_BLOCK.format(os=args.os, target=args.target)
+
+ if args.target == "cpu":
+ info("Generating Dockerfile(s) for CPU.")
+
+ def f(os: str, version: str) -> str:
+ info(f"OS: {os}, version: {version}")
+
+ if os == "manylinux2014":
+ base_block = CPU_MANYLINUX_BASE_BLOCK
+ install_block = CPU_YUM_INSTALL_BLOCK
+ scripts_block = SCRIPTS_BLOCK.format(
+ script=f"manylinux_build_wheel.sh")
+
+ dockerfile = reduce(
+ lambda x, y: x + y,
+ (head_block, base_block, MAINTAINER_BLOCK, install_block,
+ LLVM_CLANG_FROM_SOURCE_BLOCK, GCC_LINK_BLOCK, USER_BLOCK,
+ CONDA_BLOCK, scripts_block))
+
+ filename = pwd / f"Dockerfile.{os}.cpu"
+ else:
+ base_block = CPU_BASE_BLOCK.format(os=os, version=version)
+ install_block = CPU_APT_INSTALL_BLOCK
+ scripts_block = SCRIPTS_BLOCK.format(
+ script=f"{os}_build_test_cpu.sh")
+ # ubuntu 18.04 needs special treatments
+ if os == "ubuntu" and version == "18.04":
+ install_block = install_block.rstrip() + """ \\
+ zlib1g-dev"""
+
+ dockerfile = reduce(
+ lambda x, y: x + y,
+ (head_block, base_block, DEBIAN_NONINTERACTIVE_BLOCK,
+ MAINTAINER_BLOCK, install_block, CMAKE_BLOCK, LLVM_BLOCK,
+ USER_BLOCK, CONDA_BLOCK, scripts_block))
+
+ filename = pwd / f"Dockerfile.{os}.{version}.cpu"
+
+ info(f"Storing at: {filename}")
+ with filename.open("w") as fp:
+ fp.write(dockerfile)
+ else:
+ info("Generating Dockerfile(s) for GPU.")
+
+ def f(os: str, version: str) -> str:
+ info(f"OS: {os}, version: {version}")
+ base_block = GPU_BASE_BLOCK.format(version=version)
+ scripts_block = SCRIPTS_BLOCK.format(script=f"{os}_build_test.sh")
+ install_block = GPU_APT_INSTALL_BLOCK
+
+ # ubuntu 20.04 needs special treatments
+ if os == "ubuntu" and version == "20.04":
+ install_block = install_block.rstrip() + """ \\
+ vulkan-tools \\
+ vulkan-validationlayers-dev"""
+
+ dockerfile = reduce(
+ lambda x, y: x + y,
+ (head_block, base_block, NVIDIA_DRIVER_CAPABILITIES_BLOCK,
+ DEBIAN_NONINTERACTIVE_BLOCK, MAINTAINER_BLOCK, install_block,
+ CMAKE_BLOCK, LLVM_BLOCK, VULKAN_BLOCK, USER_BLOCK,
+ CONDA_BLOCK, scripts_block))
+ filename = pwd / f"Dockerfile.{os}.{version}"
+ info(f"Storing at: {filename}")
+ with (filename).open("w") as fp:
+ fp.write(dockerfile)
+
+ list(map(functools.partial(f, args.os), OS[args.os]))
+ success("Dockerfile generation is complete.")
+
+
+if __name__ == "__main__":
+ main()
diff --git a/ci/scripts/manylinux_build_wheel.sh b/ci/scripts/manylinux_build_wheel.sh
new file mode 100755
index 0000000000000..e33157490eb8d
--- /dev/null
+++ b/ci/scripts/manylinux_build_wheel.sh
@@ -0,0 +1,31 @@
+#!/bin/bash
+
+set -ex
+
+# Parse ARGs
+for ARGUMENT in "$@"
+do
+ KEY=$(echo $ARGUMENT | cut -f1 -d=)
+ VALUE=$(echo $ARGUMENT | cut -f2 -d=)
+ case "$KEY" in
+ SHA) SHA=${VALUE} ;;
+ PY) PY=${VALUE} ;;
+ *)
+ esac
+done
+
+source /home/dev/miniconda/etc/profile.d/conda.sh
+conda activate $PY
+
+# Build Taichi from source
+git clone --recursive https://github.com/taichi-dev/taichi --branch=master
+cd taichi
+git checkout $SHA
+python3 -m pip install -r requirements_dev.txt -i http://repo.taichigraphics.com/repository/pypi/simple --trusted-host repo.taichigraphics.com
+# Add Docker specific ENV
+export TI_IN_DOCKER=true
+
+# TODO, unify this step with wheel build, check #3537
+TAICHI_CMAKE_ARGS="-DTI_WITH_VULKAN:BOOL=OFF -DTI_WITH_CUDA:BOOL=OFF -DTI_WITH_OPENGL:BOOL=OFF -DTI_WITH_CC:BOOL=OFF" python3 setup.py install
+# build.py is to be removed
+#cd python && python build.py build
diff --git a/ci/scripts/release_test.sh b/ci/scripts/release_test.sh
new file mode 100644
index 0000000000000..d64140b4f53bb
--- /dev/null
+++ b/ci/scripts/release_test.sh
@@ -0,0 +1,272 @@
+#!/usr/bin/env bash
+
+# Taichi release test suite
+
+# Usage: `bash release_test.sh`
+
+# This script is created mainly for eyeball-testing
+# that if all of the official examples are still working
+# with the latest version of Taichi.
+
+# Some of the test cases are fetched from external repositories
+# please reach out to us if you are the owner of those
+# repos and don't like us to do it.
+
+# You can add more tests into this script and plug-n-play
+# existing tests in the `taichi::test::main` function as
+# you need.
+
+function taichi::utils::set_debug {
+ if [ ${DEBUG} == "true" ]; then
+ set -x
+ fi
+ set -euo pipefail
+}
+
+function taichi::utils::logger {
+ # default: gray
+ if [ "$1" == "info" ]; then
+ printf '\e[1;90m%-6s\e[m\n' "$(date +"[%m-%d %H:%M:%S]") $2"
+ # error: red
+ elif [ "$1" == "error" ]; then
+ printf '\e[1;91m%-6s\e[m\n' "$(date +"[%m-%d %H:%M:%S]") $2"
+ # success: green
+ elif [ "$1" == "success" ]; then
+ printf '\e[1;92m%-6s\e[m\n' "$(date +"[%m-%d %H:%M:%S]") $2"
+ # warning: yellow
+ elif [ "$1" == "warning" ]; then
+ printf '\e[1;93m%-6s\e[m\n' "$(date +"[%m-%d %H:%M:%S]") $2"
+ # debug: gray
+ elif [ "$1" == "debug" ]; then
+ if [ "${DEBUG}" == "true" ]; then
+ printf '\e[1;90m%-6s\e[m\n' "$(date +"[%m-%d %H:%M:%S]") $2"
+ fi
+ else
+ printf "$1"
+ fi
+}
+
+function taichi::utils::logger::info {
+ taichi::utils::logger "info" "$1"
+}
+
+function taichi::utils::logger::error {
+ taichi::utils::logger "error" "$1"
+}
+
+function taichi::utils::logger::success {
+ taichi::utils::logger "success" "$1"
+}
+
+function taichi::utils::logger::warning {
+ taichi::utils::logger "warning" "$1"
+}
+
+function taichi::utils::logger::debug {
+ taichi::utils::logger "debug" "$1"
+}
+
+function taichi::utils::line {
+ printf '%.0s-' {1..20}; echo
+}
+
+function taichi::utils::git_clone {
+ local GIT_ORG=$1
+ local GIT_REPO=$2
+ git clone "git@github.com:${GIT_ORG}/${GIT_REPO}.git"
+}
+
+function taichi::utils::pause {
+ read -p "Press enter to continue"
+}
+
+function taichi::test::ggui {
+ local WORKDIR=${1}
+ local PATTERN="*_ggui.py"
+ local ORG="taichi-dev"
+ local REPO="taichi"
+
+ # divider
+ taichi::utils::line
+ taichi::utils::logger::info "Running GGUI examples"
+
+ # clone the repo
+ taichi::utils::git_clone "${ORG}" "${REPO}"
+ cd "${REPO}/python/taichi/examples/ggui_examples"
+
+ # run tests
+ for match in $(find ./ -name "${PATTERN}"); do
+ python "${match}"
+ taichi::utils::line
+ taichi::utils::pause
+ done
+
+ # go back to workdir
+ cd "${WORKDIR}"
+}
+
+function taichi::test::difftaichi {
+ local WORKDIR=${1}
+ local PATTERN="*.py"
+ local ORG="taichi-dev"
+ local REPO="difftaichi"
+
+ # divider
+ taichi::utils::line
+ taichi::utils::logger::info "Running DiffTaichi examples"
+
+ # clone the repo
+ taichi::utils::git_clone "${ORG}" "${REPO}"
+ cd "${REPO}/examples"
+
+ # run tests
+ for match in $(find ./ -name "${PATTERN}"); do
+ python "${match}"
+ taichi::utils::line
+ taichi::utils::pause
+ done
+
+ # go back to workdir
+ cd "${WORKDIR}"
+}
+
+function taichi::test::taichi_elements {
+ local WORKDIR=${1}
+ local PATTERN="demo_*.py"
+ local ORG="taichi-dev"
+ local REPO="taichi_elements"
+
+ # divider
+ taichi::utils::line
+ taichi::utils::logger::info "Running Taichi Elements examples"
+
+ # clone the repo
+ taichi::utils::git_clone "${ORG}" "${REPO}"
+ cd "${REPO}"
+
+ # install dependencies
+ python "download_ply.py"
+
+ # run tests
+ cd "${REPO}/demo"
+ for match in $(find ./ -name "${PATTERN}"); do
+ python "${match}"
+ taichi::utils::line
+ taichi::utils::pause
+ done
+
+ # run special tests
+ # FIXME: this does not work properly yet
+ # taichi::utils::logger::success $(ls)
+ # read -p "Please input the directory containing the generated particles, e.g. sim_2022-01-01_20-55-48" particles_dir
+ # python render_particles.py -i ./"${particles_dir}" \
+ # -b 0 -e 400 -s 1 \
+ # -o ./output \
+ # --gpu-memory 20 \
+ # -M 460 \
+ # --shutter-time 0.0 \
+ # -r 128
+
+ # go back to workdir
+ cd "${WORKDIR}"
+}
+
+function taichi::test::stannum {
+ local WORKDIR=${1}
+ local ORG="ifsheldon"
+ local REPO="stannum"
+
+ # divider
+ taichi::utils::line
+ taichi::utils::logger::info "Running Stannum examples"
+
+ # clone the repo
+ taichi::utils::git_clone "${ORG}" "${REPO}"
+ cd "${REPO}"
+
+ # run tests
+ pytest -v -s ./
+
+ # go back to workdir
+ cd "${WORKDIR}"
+}
+
+function taichi::test::sandyfluid {
+ local WORKDIR=${1}
+ local ORG="ethz-pbs21"
+ local REPO="SandyFluid"
+
+ # divider
+ taichi::utils::line
+ taichi::utils::logger::info "Running SandyFluid examples"
+
+ # clone the repo
+ taichi::utils::git_clone "${ORG}" "${REPO}"
+ cd "${REPO}"
+
+ # install dependencies
+ # remove the line contains pinned Taichi version for testing purposes
+ grep -v "taichi" requirements.txt > tmpfile && mv tmpfile requirements.txt
+ pip install -r requirements.txt
+
+ # run tests
+ python src/main.py
+
+ # go back to workdir
+ cd "${WORKDIR}"
+}
+
+function taichi::test::voxel_editor {
+ local WORKDIR=${1}
+ local ORG="taichi-dev"
+ local REPO="voxel_editor"
+
+ # divider
+ taichi::utils::line
+ taichi::utils::logger::info "Running Voxel Editor examples"
+
+ # clone the repo
+ taichi::utils::git_clone "${ORG}" "${REPO}"
+ cd "${REPO}"
+
+ # run tests
+ python voxel_editor.py
+
+ # go back to workdir
+ cd "${WORKDIR}"
+}
+
+function taichi::test::main {
+ # set debugging flag
+ DEBUG="false"
+
+ # create a temporary directory for testing
+ WORKDIR="$(mktemp -d)"
+ taichi::utils::logger::info "Running all tests within ${WORKDIR}"
+
+ # make sure to clean up the temp dir on exit
+ trap '{ rm -rf -- "$WORKDIR"; }' EXIT
+
+ # walk into the working dir
+ cd "${WORKDIR}"
+
+ # ggui examples
+ taichi::test::ggui "${WORKDIR}"
+
+ # difftaichi examples
+ taichi::test::difftaichi "${WORKDIR}"
+
+ # taichi_elements examples
+ taichi::test::taichi_elements "${WORKDIR}"
+
+ # stannum tests
+ taichi::test::stannum "${WORKDIR}"
+
+ # sandyfluid tests
+ taichi::test::sandyfluid "${WORKDIR}"
+
+ # voxel editor tests
+ taichi::test::voxel_editor "${WORKDIR}"
+}
+
+taichi::test::main
diff --git a/ci/scripts/ubuntu_build_test.sh b/ci/scripts/ubuntu_build_test.sh
new file mode 100755
index 0000000000000..ed5acc79c99a8
--- /dev/null
+++ b/ci/scripts/ubuntu_build_test.sh
@@ -0,0 +1,35 @@
+#!/bin/bash
+
+set -ex
+
+# Parse ARGs
+for ARGUMENT in "$@"
+do
+ KEY=$(echo $ARGUMENT | cut -f1 -d=)
+ VALUE=$(echo $ARGUMENT | cut -f2 -d=)
+ case "$KEY" in
+ SHA) SHA=${VALUE} ;;
+ PY) PY=${VALUE} ;;
+ *)
+ esac
+done
+
+source /home/dev/miniconda/etc/profile.d/conda.sh
+conda activate $PY
+
+# Build Taichi from source
+git clone --recursive https://github.com/taichi-dev/taichi --branch=master
+cd taichi
+git checkout $SHA
+python3 -m pip install -r requirements_dev.txt -i http://repo.taichigraphics.com/repository/pypi/simple --trusted-host repo.taichigraphics.com
+# Update Torch version, otherwise cuda tests fail. See #2969.
+python3 -m pip install torch==1.9.0+cu111 -f https://download.pytorch.org/whl/torch_stable.html -i http://repo.taichigraphics.com/repository/pypi/simple --trusted-host repo.taichigraphics.com
+TAICHI_CMAKE_ARGS="-DTI_WITH_VULKAN:BOOL=ON -DTI_WITH_CUDA:BOOL=ON -DTI_WITH_OPENGL:BOOL=ON" python3 setup.py install
+
+# Add Docker specific ENV
+export TI_IN_DOCKER=true
+
+# Run tests
+ti diagnose
+python tests/run_tests.py -vr2 -t2 -k "not ndarray and not torch"
+python tests/run_tests.py -vr2 -t1 -k "ndarray or torch"
diff --git a/ci/scripts/ubuntu_build_test_cpu.sh b/ci/scripts/ubuntu_build_test_cpu.sh
new file mode 100755
index 0000000000000..feba31b80e874
--- /dev/null
+++ b/ci/scripts/ubuntu_build_test_cpu.sh
@@ -0,0 +1,33 @@
+#!/bin/bash
+
+set -ex
+
+# Parse ARGs
+for ARGUMENT in "$@"
+do
+ KEY=$(echo $ARGUMENT | cut -f1 -d=)
+ VALUE=$(echo $ARGUMENT | cut -f2 -d=)
+ case "$KEY" in
+ SHA) SHA=${VALUE} ;;
+ PY) PY=${VALUE} ;;
+ *)
+ esac
+done
+
+source /home/dev/miniconda/etc/profile.d/conda.sh
+conda activate $PY
+
+# Build Taichi from source
+git clone --recursive https://github.com/taichi-dev/taichi --branch=master
+cd taichi
+git checkout $SHA
+python3 -m pip install -r requirements_dev.txt -i http://repo.taichigraphics.com/repository/pypi/simple --trusted-host repo.taichigraphics.com
+TAICHI_CMAKE_ARGS="-DTI_WITH_VULKAN:BOOL=OFF -DTI_WITH_CUDA:BOOL=OFF -DTI_WITH_OPENGL:BOOL=OFF" python3 setup.py install
+
+# Add Docker specific ENV
+export TI_IN_DOCKER=true
+
+# Run tests
+ti diagnose
+python tests/run_tests.py -vr2 -t2 -k "not ndarray and not torch"
+python tests/run_tests.py -vr2 -t1 -k "ndarray or torch"
diff --git a/cmake/PythonNumpyPybind11.cmake b/cmake/PythonNumpyPybind11.cmake
index 5957afc7999c3..311630dba74a8 100644
--- a/cmake/PythonNumpyPybind11.cmake
+++ b/cmake/PythonNumpyPybind11.cmake
@@ -51,39 +51,26 @@ execute_process(COMMAND ${PYTHON_EXECUTABLE} -c
sys.stdout.write(str(sys.version_info[1]))"
OUTPUT_VARIABLE PYTHON_MINOR_VERSION)
+
if (WIN32)
- link_directories(${PYTHON_LIBRARY_DIR}/../../libs)
- set(PYTHON_LIBRARIES ${PYTHON_LIBRARY_DIR}/../../libs/python3.lib)
- set(PYTHON_LIBRARIES ${PYTHON_LIBRARY_DIR}/../../libs/python3${PYTHON_MINOR_VERSION}.lib)
+ execute_process(COMMAND ${PYTHON_EXECUTABLE} -c
+ "import sys;sys.stdout.write(sys.base_prefix.replace('\\\\', '/'))"
+ OUTPUT_VARIABLE PYTHON_BASE_PREFIX)
+ link_directories(${PYTHON_BASE_PREFIX}/libs)
+ set(PYTHON_LIBRARIES ${PYTHON_BASE_PREFIX}/libs/python3.lib)
+ set(PYTHON_LIBRARIES ${PYTHON_BASE_PREFIX}/libs/python3${PYTHON_MINOR_VERSION}.lib)
else()
find_library(PYTHON_LIBRARY NAMES python${PYTHON_VERSION} python${PYTHON_VERSION}m PATHS ${PYTHON_LIBRARY_DIR}
NO_DEFAULT_PATH NO_SYSTEM_ENVIRONMENT_PATH PATH_SUFFIXES x86_64-linux-gnu)
set(PYTHON_LIBRARIES ${PYTHON_LIBRARY})
endif()
-# Creating python enters
-file(MAKE_DIRECTORY bin)
-file(WRITE ${CMAKE_SOURCE_DIR}/bin/ti "#!${PYTHON_EXECUTABLE_PATH}\nimport taichi\nexit(taichi.main())")
-execute_process(COMMAND chmod +x ${CMAKE_SOURCE_DIR}/bin/ti)
-execute_process(COMMAND cp ${CMAKE_SOURCE_DIR}/bin/ti ${CMAKE_SOURCE_DIR}/bin/taichi)
-
include_directories(${PYTHON_INCLUDE_DIRS})
message(" version: ${PYTHON_VERSION}")
message(" include: ${PYTHON_INCLUDE_DIRS}")
message(" library: ${PYTHON_LIBRARIES}")
-execute_process(COMMAND ${PYTHON_EXECUTABLE} -c
- "import git; from git import Repo; import sys;\
- sys.stdout.write(git.__version__)"
- OUTPUT_VARIABLE GITPYTHON_VERSION
- RESULT_VARIABLE GITPYTHON_IMPORT_RET)
-if (NOT GITPYTHON_IMPORT_RET)
- message(" gitpython version: ${GITPYTHON_VERSION}")
-else ()
- message(FATAL_ERROR "Cannot import git. Please install. ([sudo] pip3 install --user gitpython)")
-endif ()
-
execute_process(COMMAND ${PYTHON_EXECUTABLE} -c
"import numpy.distutils, sys;\
sys.stdout.write(':'.join(numpy.distutils.misc_util.get_numpy_include_dirs()))"
diff --git a/cmake/TaichiCXXFlags.cmake b/cmake/TaichiCXXFlags.cmake
index 030d58c6a7c29..79fe36770b646 100644
--- a/cmake/TaichiCXXFlags.cmake
+++ b/cmake/TaichiCXXFlags.cmake
@@ -17,10 +17,23 @@ if (MINGW)
set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -static-libstdc++")
endif ()
-if (MSVC)
+# Do not enable lto for APPLE since it made linking extremely slow.
+if (WIN32)
+ if (CMAKE_CXX_COMPILER_ID STREQUAL "Clang")
+ set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -flto=thin")
+ endif()
+endif()
+
+if (WIN32)
link_directories(${CMAKE_CURRENT_SOURCE_DIR}/external/lib)
- set(CMAKE_CXX_FLAGS
- "${CMAKE_CXX_FLAGS} /Zc:__cplusplus /std:c++17 /MP /Z7 /D \"_CRT_SECURE_NO_WARNINGS\" /D \"_ENABLE_EXTENDED_ALIGNED_STORAGE\"")
+ if (MSVC)
+ set(CMAKE_CXX_FLAGS
+ "${CMAKE_CXX_FLAGS} /Zc:__cplusplus /std:c++17 /bigobj /wd4244 /wd4267 /nologo /Zi /D \"_CRT_SECURE_NO_WARNINGS\" /D \"_ENABLE_EXTENDED_ALIGNED_STORAGE\"")
+ else()
+ set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -std=c++17 -fsized-deallocation -target x86_64-pc-windows-msvc")
+ set(CMAKE_CXX_FLAGS_DEBUG "${CMAKE_CXX_FLAGS_DEBUG} -gcodeview")
+ set(CMAKE_CXX_FLAGS_RELWITHDEBINFO "${CMAKE_CXX_FLAGS_RELWITHDEBINFO} -gcodeview")
+ endif()
else()
if ("${CMAKE_CXX_COMPILER_ID}" MATCHES "Clang")
message("Clang compiler detected. Using std=c++17.")
@@ -37,7 +50,7 @@ else()
endif ()
message("Building for processor ${CMAKE_SYSTEM_PROCESSOR}")
-if ("${CMAKE_SYSTEM_PROCESSOR}" STREQUAL "x86_64" OR "${CMAKE_SYSTEM_PROCESSOR}" STREQUAL "AMD64")
+if ("${CMAKE_SYSTEM_PROCESSOR}" STREQUAL "x86_64" OR "${CMAKE_SYSTEM_PROCESSOR}" STREQUAL "AMD64" OR "${CMAKE_SYSTEM_PROCESSOR}" STREQUAL "amd64")
if (MSVC)
set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} /D \"TI_ARCH_x64\"")
else()
@@ -48,6 +61,9 @@ if ("${CMAKE_SYSTEM_PROCESSOR}" STREQUAL "x86_64" OR "${CMAKE_SYSTEM_PROCESSOR}"
elseif ("${CMAKE_SYSTEM_PROCESSOR}" STREQUAL "aarch64" OR "${CMAKE_SYSTEM_PROCESSOR}" STREQUAL "arm64")
set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -DTI_ARCH_ARM")
set(ARCH "arm64")
+elseif ("${CMAKE_SYSTEM_PROCESSOR}" STREQUAL "x86")
+ set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -DTI_ARCH_x86")
+ set(ARCH "x86")
else()
message(FATAL_ERROR "Unknown processor type ${CMAKE_SYSTEM_PROCESSOR}")
endif()
diff --git a/cmake/TaichiCore.cmake b/cmake/TaichiCore.cmake
index c124205f1bf72..3545dc68b0fb2 100644
--- a/cmake/TaichiCore.cmake
+++ b/cmake/TaichiCore.cmake
@@ -1,9 +1,26 @@
option(USE_STDCPP "Use -stdlib=libc++" OFF)
+option(TI_WITH_LLVM "Build with LLVM backends" ON)
+option(TI_WITH_METAL "Build with the Metal backend" ON)
option(TI_WITH_CUDA "Build with the CUDA backend" ON)
option(TI_WITH_CUDA_TOOLKIT "Build with the CUDA toolkit" OFF)
option(TI_WITH_OPENGL "Build with the OpenGL backend" ON)
option(TI_WITH_CC "Build with the C backend" ON)
option(TI_WITH_VULKAN "Build with the Vulkan backend" OFF)
+option(TI_WITH_DX11 "Build with the DX11 backend" OFF)
+option(TI_EMSCRIPTENED "Build using emscripten" OFF)
+set(_TI_SYMBOL_VISIBILITY default)
+
+if(TI_EMSCRIPTENED)
+ set(TI_WITH_LLVM OFF)
+ set(TI_WITH_METAL OFF)
+ set(TI_WITH_CUDA OFF)
+ set(TI_WITH_OPENGL OFF)
+ set(TI_WITH_CC OFF)
+ set(TI_WITH_DX11 OFF)
+
+ set(TI_WITH_VULKAN ON)
+ set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -DTI_EMSCRIPTENED")
+endif()
if(UNIX AND NOT APPLE)
# Handy helper for Linux
@@ -34,16 +51,21 @@ if (WIN32)
endif()
set(TI_WITH_GGUI OFF)
-if(TI_WITH_CUDA AND TI_WITH_VULKAN)
+if(TI_WITH_VULKAN AND NOT TI_EMSCRIPTENED)
set(TI_WITH_GGUI ON)
endif()
-if (NOT EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/external/glad/src/glad.c")
+if (NOT EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/external/glad/src/gl.c")
set(TI_WITH_OPENGL OFF)
message(WARNING "external/glad submodule not detected. Settings TI_WITH_OPENGL to OFF.")
endif()
+if(NOT TI_WITH_LLVM)
+ set(TI_WITH_CUDA OFF)
+ set(TI_WITH_CUDA_TOOLKIT OFF)
+endif()
+
file(GLOB TAICHI_CORE_SOURCE
@@ -57,6 +79,7 @@ file(GLOB TAICHI_WASM_SOURCE "taichi/backends/wasm/*.cpp" "taichi/backends/wasm/
file(GLOB TAICHI_CUDA_SOURCE "taichi/backends/cuda/*.cpp" "taichi/backends/cuda/*.h")
file(GLOB TAICHI_METAL_SOURCE "taichi/backends/metal/*.h" "taichi/backends/metal/*.cpp" "taichi/backends/metal/shaders/*")
file(GLOB TAICHI_OPENGL_SOURCE "taichi/backends/opengl/*.h" "taichi/backends/opengl/*.cpp" "taichi/backends/opengl/shaders/*")
+file(GLOB TAICHI_DX11_SOURCE "taichi/backends/dx/*.h" "taichi/backends/dx/*.cpp")
file(GLOB TAICHI_CC_SOURCE "taichi/backends/cc/*.h" "taichi/backends/cc/*.cpp")
file(GLOB TAICHI_VULKAN_SOURCE "taichi/backends/vulkan/*.h" "taichi/backends/vulkan/*.cpp" "external/SPIRV-Reflect/spirv_reflect.c")
file(GLOB TAICHI_INTEROP_SOURCE "taichi/backends/interop/*.cpp" "taichi/backends/interop/*.h")
@@ -66,16 +89,22 @@ file(GLOB TAICHI_GGUI_SOURCE
"taichi/ui/*.cpp" "taichi/ui/*/*.cpp" "taichi/ui/*/*/*.cpp" "taichi/ui/*/*/*/*.cpp" "taichi/ui/*/*/*/*/*.cpp"
"taichi/ui/*.h" "taichi/ui/*/*.h" "taichi/ui/*/*/*.h" "taichi/ui/*/*/*/*.h" "taichi/ui/*/*/*/*/*.h"
)
+file(GLOB TAICHI_GGUI_GLFW_SOURCE
+ "taichi/ui/common/window_base.cpp"
+ "taichi/ui/backends/vulkan/window.cpp"
+)
list(REMOVE_ITEM TAICHI_CORE_SOURCE ${TAICHI_GGUI_SOURCE})
if(TI_WITH_GGUI)
add_definitions(-DTI_WITH_GGUI)
- list(APPEND TAICHI_CORE_SOURCE ${TAICHI_GGUI_SOURCE})
-
- include_directories(SYSTEM external/glm)
+ # Remove GLFW dependencies from the build for Android
+ if(ANDROID)
+ list(REMOVE_ITEM TAICHI_GGUI_SOURCE ${TAICHI_GGUI_GLFW_SOURCE})
+ endif()
+ list(APPEND TAICHI_CORE_SOURCE ${TAICHI_GGUI_SOURCE})
endif()
# These files are compiled into .bc and loaded as LLVM module dynamically. They should not be compiled into libtaichi. So they're removed here
@@ -95,14 +124,19 @@ file(GLOB TAICHI_OPENGL_REQUIRED_SOURCE
file(GLOB TAICHI_VULKAN_REQUIRED_SOURCE
"taichi/backends/vulkan/runtime.h"
"taichi/backends/vulkan/runtime.cpp"
- "taichi/backends/vulkan/snode_struct_compiler.cpp"
- "taichi/backends/vulkan/snode_struct_compiler.h"
)
list(REMOVE_ITEM TAICHI_CORE_SOURCE ${TAICHI_BACKEND_SOURCE})
-list(APPEND TAICHI_CORE_SOURCE ${TAICHI_CPU_SOURCE})
-list(APPEND TAICHI_CORE_SOURCE ${TAICHI_WASM_SOURCE})
+if(TI_WITH_LLVM)
+ set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -DTI_WITH_LLVM")
+ list(APPEND TAICHI_CORE_SOURCE ${TAICHI_CPU_SOURCE})
+ list(APPEND TAICHI_CORE_SOURCE ${TAICHI_WASM_SOURCE})
+else()
+ file(GLOB TAICHI_LLVM_SOURCE "taichi/llvm/*.cpp" "taichi/llvm/*.h")
+ list(REMOVE_ITEM TAICHI_CORE_SOURCE ${TAICHI_LLVM_SOURCE})
+endif()
+
list(APPEND TAICHI_CORE_SOURCE ${TAICHI_INTEROP_SOURCE})
@@ -115,14 +149,21 @@ if(NOT CUDA_VERSION)
set(CUDA_VERSION 10.0)
endif()
-# TODO(#529) include Metal source only on Apple MacOS, and OpenGL only when TI_WITH_OPENGL is ON
-list(APPEND TAICHI_CORE_SOURCE ${TAICHI_METAL_SOURCE})
+
+# By default, TI_WITH_METAL is ON for all platforms.
+# As of right now, on non-macOS platforms, the metal backend won't work at all.
+# We have future plans to allow metal AOT to run on non-macOS devices.
+if (TI_WITH_METAL)
+ set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -DTI_WITH_METAL")
+ list(APPEND TAICHI_CORE_SOURCE ${TAICHI_METAL_SOURCE})
+endif()
+
if (TI_WITH_OPENGL)
set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -DTI_WITH_OPENGL")
# Q: Why not external/glad/src/*.c?
# A: To ensure glad submodule exists when TI_WITH_OPENGL is ON.
- file(GLOB TAICHI_GLAD_SOURCE "external/glad/src/glad.c")
+ file(GLOB TAICHI_GLAD_SOURCE "external/glad/src/gl.c" "external/glad/src/egl.c")
list(APPEND TAICHI_CORE_SOURCE ${TAICHI_GLAD_SOURCE})
list(APPEND TAICHI_CORE_SOURCE ${TAICHI_OPENGL_SOURCE})
endif()
@@ -140,6 +181,11 @@ if (TI_WITH_VULKAN)
endif()
list(APPEND TAICHI_CORE_SOURCE ${TAICHI_VULKAN_REQUIRED_SOURCE})
+if (TI_WITH_DX11)
+ set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -DTI_WITH_DX11")
+ list(APPEND TAICHI_CORE_SOURCE ${TAICHI_DX11_SOURCE})
+endif()
+
# This compiles all the libraries with -fPIC, which is critical to link a static
# library into a shared lib.
set(CMAKE_POSITION_INDEPENDENT_CODE ON)
@@ -158,6 +204,14 @@ file(GLOB TAICHI_PYBIND_SOURCE
)
list(REMOVE_ITEM TAICHI_CORE_SOURCE ${TAICHI_PYBIND_SOURCE})
+file(GLOB TAICHI_EMBIND_SOURCE
+ "taichi/javascript/*.cpp"
+ "taichi/javascript/*.h"
+)
+if (TAICHI_EMBIND_SOURCE)
+ list(REMOVE_ITEM TAICHI_CORE_SOURCE ${TAICHI_EMBIND_SOURCE})
+endif()
+
# TODO(#2196): Rename these CMAKE variables:
# CORE_LIBRARY_NAME --> TAICHI_ISOLATED_CORE_LIB_NAME
# CORE_WITH_PYBIND_LIBRARY_NAME --> TAICHI_CORE_LIB_NAME
@@ -171,6 +225,7 @@ list(REMOVE_ITEM TAICHI_CORE_SOURCE ${TAICHI_PYBIND_SOURCE})
# everywhere in python.
set(CORE_LIBRARY_NAME taichi_isolated_core)
add_library(${CORE_LIBRARY_NAME} OBJECT ${TAICHI_CORE_SOURCE})
+set_target_properties(${CORE_LIBRARY_NAME} PROPERTIES CXX_VISIBILITY_PRESET ${_TI_SYMBOL_VISIBILITY})
if (APPLE)
# Ask OS X to minic Linux dynamic linking behavior
@@ -181,19 +236,26 @@ include_directories(${CMAKE_SOURCE_DIR})
include_directories(external/include)
include_directories(external/spdlog/include)
if (TI_WITH_OPENGL)
- include_directories(external/glad/include)
+ target_include_directories(${CORE_LIBRARY_NAME} PRIVATE external/glad/include)
endif()
+ target_include_directories(${CORE_LIBRARY_NAME} PRIVATE external/FP16/include)
set(LIBRARY_NAME ${CORE_LIBRARY_NAME})
-if (TI_WITH_OPENGL)
+# GLFW not available on Android
+if (TI_WITH_OPENGL OR TI_WITH_VULKAN AND NOT ANDROID AND NOT TI_EMSCRIPTENED)
set(GLFW_BUILD_DOCS OFF CACHE BOOL "" FORCE)
set(GLFW_BUILD_TESTS OFF CACHE BOOL "" FORCE)
set(GLFW_BUILD_EXAMPLES OFF CACHE BOOL "" FORCE)
+ if (APPLE)
+ set(GLFW_VULKAN_STATIC ON CACHE BOOL "" FORCE)
+ endif()
+
message("Building with GLFW")
add_subdirectory(external/glfw)
target_link_libraries(${LIBRARY_NAME} glfw)
+ target_include_directories(${CORE_LIBRARY_NAME} PRIVATE external/glfw/include)
endif()
if(DEFINED ENV{LLVM_DIR})
@@ -201,48 +263,50 @@ if(DEFINED ENV{LLVM_DIR})
message("Getting LLVM_DIR=${LLVM_DIR} from the environment variable")
endif()
-# http://llvm.org/docs/CMake.html#embedding-llvm-in-your-project
-find_package(LLVM REQUIRED CONFIG)
-message(STATUS "Found LLVM ${LLVM_PACKAGE_VERSION}")
-if(${LLVM_PACKAGE_VERSION} VERSION_LESS "10.0")
- message(FATAL_ERROR "LLVM version < 10 is not supported")
-endif()
-message(STATUS "Using LLVMConfig.cmake in: ${LLVM_DIR}")
-include_directories(${LLVM_INCLUDE_DIRS})
-message("LLVM include dirs ${LLVM_INCLUDE_DIRS}")
-message("LLVM library dirs ${LLVM_LIBRARY_DIRS}")
-add_definitions(${LLVM_DEFINITIONS})
-
-llvm_map_components_to_libnames(llvm_libs
- Core
- ExecutionEngine
- InstCombine
- OrcJIT
- RuntimeDyld
- TransformUtils
- BitReader
- BitWriter
- Object
- ScalarOpts
- Support
- native
- Linker
- Target
- MC
- Passes
- ipo
- Analysis
- )
-target_link_libraries(${LIBRARY_NAME} ${llvm_libs})
-
-if (APPLE AND "${CMAKE_SYSTEM_PROCESSOR}" STREQUAL "arm64")
- llvm_map_components_to_libnames(llvm_aarch64_libs AArch64)
- target_link_libraries(${LIBRARY_NAME} ${llvm_aarch64_libs})
-endif()
+if(TI_WITH_LLVM)
+ # http://llvm.org/docs/CMake.html#embedding-llvm-in-your-project
+ find_package(LLVM REQUIRED CONFIG)
+ message(STATUS "Found LLVM ${LLVM_PACKAGE_VERSION}")
+ if(${LLVM_PACKAGE_VERSION} VERSION_LESS "10.0")
+ message(FATAL_ERROR "LLVM version < 10 is not supported")
+ endif()
+ message(STATUS "Using LLVMConfig.cmake in: ${LLVM_DIR}")
+ include_directories(${LLVM_INCLUDE_DIRS})
+ message("LLVM include dirs ${LLVM_INCLUDE_DIRS}")
+ message("LLVM library dirs ${LLVM_LIBRARY_DIRS}")
+ add_definitions(${LLVM_DEFINITIONS})
+
+ llvm_map_components_to_libnames(llvm_libs
+ Core
+ ExecutionEngine
+ InstCombine
+ OrcJIT
+ RuntimeDyld
+ TransformUtils
+ BitReader
+ BitWriter
+ Object
+ ScalarOpts
+ Support
+ native
+ Linker
+ Target
+ MC
+ Passes
+ ipo
+ Analysis
+ )
+ target_link_libraries(${LIBRARY_NAME} ${llvm_libs})
+
+ if (APPLE AND "${CMAKE_SYSTEM_PROCESSOR}" STREQUAL "arm64")
+ llvm_map_components_to_libnames(llvm_aarch64_libs AArch64)
+ target_link_libraries(${LIBRARY_NAME} ${llvm_aarch64_libs})
+ endif()
-if (TI_WITH_CUDA)
- llvm_map_components_to_libnames(llvm_ptx_libs NVPTX)
- target_link_libraries(${LIBRARY_NAME} ${llvm_ptx_libs})
+ if (TI_WITH_CUDA)
+ llvm_map_components_to_libnames(llvm_ptx_libs NVPTX)
+ target_link_libraries(${LIBRARY_NAME} ${llvm_ptx_libs})
+ endif()
endif()
if (TI_WITH_CUDA_TOOLKIT)
@@ -261,38 +325,32 @@ else()
message(STATUS "TI_WITH_CUDA_TOOLKIT = OFF")
endif()
-add_subdirectory(external/SPIRV-Cross)
-target_include_directories(${CORE_LIBRARY_NAME} PRIVATE external/SPIRV-Cross)
-target_link_libraries(${CORE_LIBRARY_NAME} spirv-cross-glsl spirv-cross-core)
-
-if (TI_WITH_VULKAN)
- # Vulkan libs
- # https://cmake.org/cmake/help/latest/module/FindVulkan.html
- # https://github.com/PacktPublishing/Learning-Vulkan/blob/master/Chapter%2003/HandShake/CMakeLists.txt
- find_package(Vulkan REQUIRED)
-
- if(NOT Vulkan_FOUND)
- message(FATAL_ERROR "TI_WITH_VULKAN is ON but Vulkan could not be found")
- endif()
+if (TI_WITH_OPENGL)
+ set(SPIRV_CROSS_CLI false)
+ add_subdirectory(external/SPIRV-Cross)
+ target_include_directories(${CORE_LIBRARY_NAME} PRIVATE external/SPIRV-Cross)
+ target_link_libraries(${CORE_LIBRARY_NAME} spirv-cross-glsl spirv-cross-core)
+endif()
- message(STATUS "Vulkan_INCLUDE_DIR=${Vulkan_INCLUDE_DIR}")
- message(STATUS "Vulkan_LIBRARY=${Vulkan_LIBRARY}")
+if (TI_WITH_DX11)
+ set(SPIRV_CROSS_CLI false)
+ #target_include_directories(${CORE_LIBRARY_NAME} PRIVATE external/SPIRV-Cross)
+ target_link_libraries(${CORE_LIBRARY_NAME} spirv-cross-hlsl spirv-cross-core)
+endif()
- include_directories(external/SPIRV-Headers/include)
+# SPIR-V codegen is always there, regardless of Vulkan
+set(SPIRV_SKIP_EXECUTABLES true)
+set(SPIRV-Headers_SOURCE_DIR ${CMAKE_CURRENT_SOURCE_DIR}/external/SPIRV-Headers)
+add_subdirectory(external/SPIRV-Tools)
+# NOTE: SPIRV-Tools-opt must come before SPIRV-Tools
+# https://github.com/KhronosGroup/SPIRV-Tools/issues/1569#issuecomment-390250792
+target_link_libraries(${CORE_LIBRARY_NAME} SPIRV-Tools-opt ${SPIRV_TOOLS})
- set(SPIRV_SKIP_EXECUTABLES true)
- set(SPIRV-Headers_SOURCE_DIR ${CMAKE_CURRENT_SOURCE_DIR}/external/SPIRV-Headers)
- add_subdirectory(external/SPIRV-Tools)
- # NOTE: SPIRV-Tools-opt must come before SPIRV-Tools
- # https://github.com/KhronosGroup/SPIRV-Tools/issues/1569#issuecomment-390250792
- target_link_libraries(${CORE_LIBRARY_NAME} SPIRV-Tools-opt ${SPIRV_TOOLS})
+if (TI_WITH_VULKAN)
+ include_directories(SYSTEM external/Vulkan-Headers/include)
- # No longer link against vulkan, using volk instead
- #target_link_libraries(${CORE_LIBRARY_NAME} ${Vulkan_LIBRARY})
- include_directories(${Vulkan_INCLUDE_DIR})
- include_directories(external/volk)
+ include_directories(SYSTEM external/volk)
- # Is this the best way to include the SPIRV-Headers?
target_include_directories(${CORE_LIBRARY_NAME} PRIVATE external/SPIRV-Headers/include)
target_include_directories(${CORE_LIBRARY_NAME} PRIVATE external/SPIRV-Reflect)
target_include_directories(${CORE_LIBRARY_NAME} PRIVATE external/VulkanMemoryAllocator/include)
@@ -303,6 +361,12 @@ if (TI_WITH_VULKAN)
find_package(Threads REQUIRED)
target_link_libraries(${CORE_LIBRARY_NAME} Threads::Threads)
endif()
+
+ if (APPLE)
+ find_library(MOLTEN_VK libMoltenVK.dylib PATHS $HOMEBREW_CELLAR/molten-vk $VULKAN_SDK REQUIRED)
+ configure_file(${MOLTEN_VK} ${CMAKE_BINARY_DIR}/libMoltenVK.dylib COPYONLY)
+ message(STATUS "MoltenVK library ${MOLTEN_VK}")
+ endif()
endif ()
# Optional dependencies
@@ -312,17 +376,31 @@ if (APPLE)
endif ()
if (NOT WIN32)
- target_link_libraries(${CORE_LIBRARY_NAME} pthread stdc++)
- if (APPLE)
- # OS X
+ # Android has a custom toolchain so pthread is not available and should
+ # link against other libraries as well for logcat and internal features.
+ if (ANDROID)
+ target_link_libraries(${CORE_LIBRARY_NAME} android log)
+ else()
+ target_link_libraries(${CORE_LIBRARY_NAME} pthread stdc++)
+ endif()
+
+ if (UNIX AND NOT ${CMAKE_SYSTEM_NAME} MATCHES "Linux")
+ # OS X or BSD
else()
# Linux
target_link_libraries(${CORE_LIBRARY_NAME} stdc++fs X11)
target_link_libraries(${CORE_LIBRARY_NAME} -static-libgcc -static-libstdc++)
- if (NOT TI_EXPORT_CORE) # expose api for CHI IR Builder
+ if ((NOT TI_EXPORT_CORE) AND (NOT ${_TI_SYMBOL_VISIBILITY} STREQUAL hidden)) # expose api for CHI IR Builder
+ message(WARNING "Using linker.map to hide symbols!")
target_link_libraries(${CORE_LIBRARY_NAME} -Wl,--version-script,${CMAKE_CURRENT_SOURCE_DIR}/misc/linker.map)
endif ()
- target_link_libraries(${CORE_LIBRARY_NAME} -Wl,--wrap=log2f) # Avoid glibc dependencies
+ # Avoid glibc dependencies
+ if (TI_WITH_VULKAN)
+ target_link_libraries(${CORE_LIBRARY_NAME} -Wl,--wrap=log2f)
+ else()
+ # Enforce compatibility with manylinux2014
+ target_link_libraries(${CORE_LIBRARY_NAME} -Wl,--wrap=log2f -Wl,--wrap=exp2 -Wl,--wrap=log2 -Wl,--wrap=logf -Wl,--wrap=powf -Wl,--wrap=exp -Wl,--wrap=log -Wl,--wrap=pow)
+ endif()
endif()
else()
# windows
@@ -338,31 +416,59 @@ endforeach ()
message("PYTHON_LIBRARIES: " ${PYTHON_LIBRARIES})
-set(CORE_WITH_PYBIND_LIBRARY_NAME taichi_core)
-add_library(${CORE_WITH_PYBIND_LIBRARY_NAME} SHARED ${TAICHI_PYBIND_SOURCE})
-# It is actually possible to link with an OBJECT library
-# https://cmake.org/cmake/help/v3.13/command/target_link_libraries.html?highlight=target_link_libraries#linking-object-libraries
-target_link_libraries(${CORE_WITH_PYBIND_LIBRARY_NAME} PUBLIC ${CORE_LIBRARY_NAME})
-
-# These commands should apply to the DLL that is loaded from python, not the OBJECT library.
-if (MSVC)
- set_property(TARGET ${CORE_WITH_PYBIND_LIBRARY_NAME} APPEND PROPERTY LINK_FLAGS /DEBUG)
-endif ()
+if(NOT TI_EMSCRIPTENED)
+ set(CORE_WITH_PYBIND_LIBRARY_NAME taichi_core)
+ # Cannot compile Python source code with Android, but TI_EXPORT_CORE should be set and
+ # Android should only use the isolated library ignoring those source code.
+ if (NOT ANDROID)
+ add_library(${CORE_WITH_PYBIND_LIBRARY_NAME} SHARED ${TAICHI_PYBIND_SOURCE})
+ else()
+ add_library(${CORE_WITH_PYBIND_LIBRARY_NAME} SHARED)
+ endif ()
-if (WIN32)
- set_target_properties(${CORE_WITH_PYBIND_LIBRARY_NAME} PROPERTIES RUNTIME_OUTPUT_DIRECTORY
- "${CMAKE_CURRENT_SOURCE_DIR}/runtimes")
-endif ()
+ set_target_properties(${CORE_WITH_PYBIND_LIBRARY_NAME} PROPERTIES CXX_VISIBILITY_PRESET ${_TI_SYMBOL_VISIBILITY})
+ # Remove symbols from static libs: https://stackoverflow.com/a/14863432/12003165
+ if (LINUX)
+ target_link_options(${CORE_WITH_PYBIND_LIBRARY_NAME} PUBLIC -Wl,--exclude-libs=ALL)
+ endif()
+ # It is actually possible to link with an OBJECT library
+ # https://cmake.org/cmake/help/v3.13/command/target_link_libraries.html?highlight=target_link_libraries#linking-object-libraries
+ target_link_libraries(${CORE_WITH_PYBIND_LIBRARY_NAME} PUBLIC ${CORE_LIBRARY_NAME})
+
+ # These commands should apply to the DLL that is loaded from python, not the OBJECT library.
+ if (MSVC)
+ set_property(TARGET ${CORE_WITH_PYBIND_LIBRARY_NAME} APPEND PROPERTY LINK_FLAGS /DEBUG)
+ endif ()
+
+ if (WIN32)
+ set_target_properties(${CORE_WITH_PYBIND_LIBRARY_NAME} PROPERTIES RUNTIME_OUTPUT_DIRECTORY
+ "${CMAKE_CURRENT_SOURCE_DIR}/runtimes")
+ endif ()
+endif()
+if(TI_EMSCRIPTENED)
+ set(CORE_WITH_EMBIND_LIBRARY_NAME taichi)
+ add_executable(${CORE_WITH_EMBIND_LIBRARY_NAME} ${TAICHI_EMBIND_SOURCE})
+ target_link_libraries(${CORE_WITH_EMBIND_LIBRARY_NAME} PUBLIC ${CORE_LIBRARY_NAME})
+ target_compile_options(${CORE_WITH_EMBIND_LIBRARY_NAME} PRIVATE "-Oz")
+ # target_compile_options(${CORE_LIBRARY_NAME} PRIVATE "-Oz")
+ set_target_properties(${CORE_LIBRARY_NAME} PROPERTIES LINK_FLAGS "-s ERROR_ON_UNDEFINED_SYMBOLS=0 -s ASSERTIONS=1")
+ set_target_properties(${CORE_WITH_EMBIND_LIBRARY_NAME} PROPERTIES LINK_FLAGS "--bind -s MODULARIZE=1 -s EXPORT_NAME=createTaichiModule -s WASM=0 --memory-init-file 0 -Oz --closure 1 -s ERROR_ON_UNDEFINED_SYMBOLS=0 -s ASSERTIONS=1 -s NO_DISABLE_EXCEPTION_CATCHING")
+endif()
if(TI_WITH_GGUI)
+ include_directories(SYSTEM PRIVATE external/glm)
# Dear ImGui
add_definitions(-DIMGUI_IMPL_VULKAN_NO_PROTOTYPES)
set(IMGUI_DIR external/imgui)
- include_directories(external/glfw/include)
include_directories(SYSTEM ${IMGUI_DIR} ${IMGUI_DIR}/backends ..)
+if(ANDROID)
+ add_library(imgui ${IMGUI_DIR}/backends/imgui_impl_android.cpp ${IMGUI_DIR}/backends/imgui_impl_vulkan.cpp ${IMGUI_DIR}/imgui.cpp ${IMGUI_DIR}/imgui_draw.cpp ${IMGUI_DIR}/imgui_tables.cpp ${IMGUI_DIR}/imgui_widgets.cpp)
+else()
+ include_directories(external/glfw/include)
add_library(imgui ${IMGUI_DIR}/backends/imgui_impl_glfw.cpp ${IMGUI_DIR}/backends/imgui_impl_vulkan.cpp ${IMGUI_DIR}/imgui.cpp ${IMGUI_DIR}/imgui_draw.cpp ${IMGUI_DIR}/imgui_tables.cpp ${IMGUI_DIR}/imgui_widgets.cpp)
+endif()
target_link_libraries(${CORE_LIBRARY_NAME} imgui)
endif()
diff --git a/cmake/TaichiExamples.cmake b/cmake/TaichiExamples.cmake
new file mode 100644
index 0000000000000..742119013d623
--- /dev/null
+++ b/cmake/TaichiExamples.cmake
@@ -0,0 +1,30 @@
+cmake_minimum_required(VERSION 3.0)
+
+if(NOT TI_EMSCRIPTENED)
+
+set(EXAMPLES_NAME taichi_cpp_examples)
+
+file(GLOB_RECURSE TAICHI_EXAMPLES_SOURCE
+"cpp_examples/main.cpp"
+"cpp_examples/run_snode.cpp"
+"cpp_examples/autograd.cpp"
+"cpp_examples/aot_save.cpp"
+)
+
+include_directories(
+ ${PROJECT_SOURCE_DIR},
+)
+
+add_executable(${EXAMPLES_NAME} ${TAICHI_EXAMPLES_SOURCE})
+if (WIN32)
+ # Output the executable to bin/ instead of build/Debug/...
+ set(EXAMPLES_OUTPUT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/bin")
+ set_target_properties(${EXAMPLES_NAME} PROPERTIES RUNTIME_OUTPUT_DIRECTORY ${EXAMPLES_OUTPUT_DIR})
+ set_target_properties(${EXAMPLES_NAME} PROPERTIES RUNTIME_OUTPUT_DIRECTORY_DEBUG ${EXAMPLES_OUTPUT_DIR})
+ set_target_properties(${EXAMPLES_NAME} PROPERTIES RUNTIME_OUTPUT_DIRECTORY_RELEASE ${EXAMPLES_OUTPUT_DIR})
+ set_target_properties(${EXAMPLES_NAME} PROPERTIES RUNTIME_OUTPUT_DIRECTORY_MINSIZEREL ${EXAMPLES_OUTPUT_DIR})
+ set_target_properties(${EXAMPLES_NAME} PROPERTIES RUNTIME_OUTPUT_DIRECTORY_RELWITHDEBINFO ${EXAMPLES_OUTPUT_DIR})
+endif()
+target_link_libraries(${EXAMPLES_NAME} taichi_isolated_core)
+
+endif()
diff --git a/cmake/TaichiMain.cmake b/cmake/TaichiMain.cmake
deleted file mode 100644
index 3504c0049ed38..0000000000000
--- a/cmake/TaichiMain.cmake
+++ /dev/null
@@ -1,10 +0,0 @@
-if (WIN32)
- message("Added executable 'ti' for Windows")
- # On Windows, generate a ti.exe as entry
- add_executable(ti python/ti.cpp)
- set_target_properties(ti PROPERTIES RUNTIME_OUTPUT_DIRECTORY "${CMAKE_CURRENT_SOURCE_DIR}/bin")
- set_target_properties(ti PROPERTIES RUNTIME_OUTPUT_DIRECTORY_DEBUG "${CMAKE_CURRENT_SOURCE_DIR}/bin")
- set_target_properties(ti PROPERTIES RUNTIME_OUTPUT_DIRECTORY_RELEASE "${CMAKE_CURRENT_SOURCE_DIR}/bin")
- set_target_properties(ti PROPERTIES RUNTIME_OUTPUT_DIRECTORY_MINSIZEREL "${CMAKE_CURRENT_SOURCE_DIR}/bin")
- set_target_properties(ti PROPERTIES RUNTIME_OUTPUT_DIRECTORY_RELWITHDEBINFO "${CMAKE_CURRENT_SOURCE_DIR}/bin")
-endif ()
diff --git a/cmake/TaichiTests.cmake b/cmake/TaichiTests.cmake
index 2946b758e56cc..2b6da207c82ff 100644
--- a/cmake/TaichiTests.cmake
+++ b/cmake/TaichiTests.cmake
@@ -12,6 +12,8 @@ endif()
# 2. Re-implement the legacy CPP tests using googletest
file(GLOB_RECURSE TAICHI_TESTS_SOURCE
"tests/cpp/analysis/*.cpp"
+ "tests/cpp/aot/*.cpp"
+ "tests/cpp/backends/*.cpp"
"tests/cpp/codegen/*.cpp"
"tests/cpp/common/*.cpp"
"tests/cpp/ir/*.cpp"
diff --git a/cpp_examples/aot_save.cpp b/cpp_examples/aot_save.cpp
new file mode 100644
index 0000000000000..cc88ebbb53403
--- /dev/null
+++ b/cpp_examples/aot_save.cpp
@@ -0,0 +1,78 @@
+#include "taichi/ir/ir_builder.h"
+#include "taichi/ir/statements.h"
+#include "taichi/program/program.h"
+
+void aot_save() {
+ using namespace taichi;
+ using namespace lang;
+ auto program = Program(Arch::vulkan);
+
+ program.config.advanced_optimization = false;
+
+ int n = 10;
+
+ // program.materialize_runtime();
+ auto *root = new SNode(0, SNodeType::root);
+ auto *pointer = &root->dense(Axis(0), n, false);
+ auto *place = &pointer->insert_children(SNodeType::place);
+ place->dt = PrimitiveType::i32;
+ program.add_snode_tree(std::unique_ptr(root), /*compile_only=*/true);
+
+ auto aot_builder = program.make_aot_module_builder(Arch::vulkan);
+
+ std::unique_ptr kernel_init, kernel_ret;
+
+ {
+ /*
+ @ti.kernel
+ def init():
+ for index in range(n):
+ place[index] = index
+ */
+ IRBuilder builder;
+ auto *zero = builder.get_int32(0);
+ auto *n_stmt = builder.get_int32(n);
+ auto *loop = builder.create_range_for(zero, n_stmt, 0, 4);
+ {
+ auto _ = builder.get_loop_guard(loop);
+ auto *index = builder.get_loop_index(loop);
+ auto *ptr = builder.create_global_ptr(place, {index});
+ builder.create_global_store(ptr, index);
+ }
+
+ kernel_init =
+ std::make_unique(program, builder.extract_ir(), "init");
+ }
+
+ {
+ /*
+ @ti.kernel
+ def ret():
+ sum = 0
+ for index in place:
+ sum = sum + place[index];
+ return sum
+ */
+ IRBuilder builder;
+ auto *sum = builder.create_local_var(PrimitiveType::i32);
+ auto *loop = builder.create_struct_for(pointer, 0, 4);
+ {
+ auto _ = builder.get_loop_guard(loop);
+ auto *index = builder.get_loop_index(loop);
+ auto *sum_old = builder.create_local_load(sum);
+ auto *place_index =
+ builder.create_global_load(builder.create_global_ptr(place, {index}));
+ builder.create_local_store(sum, builder.create_add(sum_old, place_index));
+ }
+ builder.create_return(builder.create_local_load(sum));
+
+ kernel_ret = std::make_unique(program, builder.extract_ir(), "ret");
+ kernel_ret->insert_ret(PrimitiveType::i32);
+ }
+
+ aot_builder->add_field("place", place, true, place->dt, {n}, 1, 1);
+ aot_builder->add("init", kernel_init.get());
+ aot_builder->add("ret", kernel_ret.get());
+ aot_builder->dump(".", "aot.tcb");
+ std::cout << "done" << std::endl;
+}
diff --git a/examples/chi_examples/main.cpp b/cpp_examples/autograd.cpp
similarity index 57%
rename from examples/chi_examples/main.cpp
rename to cpp_examples/autograd.cpp
index 7424b5389676b..0b3bc8f43ed95 100644
--- a/examples/chi_examples/main.cpp
+++ b/cpp_examples/autograd.cpp
@@ -1,147 +1,7 @@
-#include
-
#include "taichi/ir/ir_builder.h"
#include "taichi/ir/statements.h"
#include "taichi/program/program.h"
-void run_snode() {
- /*
- import taichi as ti, numpy as np
- ti.init()
- #ti.init(print_ir = True)
-
- n = 10
- place = ti.field(dtype = ti.i32)
- ti.root.pointer(ti.i, n).place(place)
-
- @ti.kernel
- def init():
- for index in range(n):
- place[index] = index
-
- @ti.kernel
- def ret() -> ti.i32:
- sum = 0
- for index in place:
- sum = sum + place[index]
- return sum
-
- @ti.kernel
- def ext(ext_arr: ti.ext_arr()):
- for index in place:
- ext_arr[index] = place[index]
-
- init()
- print(ret())
- ext_arr = np.zeros(n, np.int32)
- ext(ext_arr)
- #ext_arr = place.to_numpy()
- print(ext_arr)
- */
-
- using namespace taichi;
- using namespace lang;
- auto program = Program(Arch::x64);
- /*CompileConfig config_print_ir;
- config_print_ir.print_ir = true;
- prog_.config = config_print_ir;*/ // print_ir = True
-
- int n = 10;
- program.materialize_runtime();
- auto *root = new SNode(0, SNodeType::root);
- auto *pointer = &root->pointer(Index(0), n);
- auto *place = &pointer->insert_children(SNodeType::place);
- place->dt = PrimitiveType::i32;
- program.add_snode_tree(std::unique_ptr(root));
-
- std::unique_ptr kernel_init, kernel_ret, kernel_ext;
-
- {
- /*
- @ti.kernel
- def init():
- for index in range(n):
- place[index] = index
- */
- IRBuilder builder;
- auto *zero = builder.get_int32(0);
- auto *n_stmt = builder.get_int32(n);
- auto *loop = builder.create_range_for(zero, n_stmt, 1, 0, 4);
- {
- auto _ = builder.get_loop_guard(loop);
- auto *index = builder.get_loop_index(loop);
- auto *ptr = builder.create_global_ptr(place, {index});
- builder.create_global_store(ptr, index);
- }
-
- kernel_init =
- std::make_unique(program, builder.extract_ir(), "init");
- }
-
- {
- /*
- @ti.kernel
- def ret():
- sum = 0
- for index in place:
- sum = sum + place[index];
- return sum
- */
- IRBuilder builder;
- auto *sum = builder.create_local_var(PrimitiveType::i32);
- auto *loop = builder.create_struct_for(pointer, 1, 0, 4);
- {
- auto _ = builder.get_loop_guard(loop);
- auto *index = builder.get_loop_index(loop);
- auto *sum_old = builder.create_local_load(sum);
- auto *place_index =
- builder.create_global_load(builder.create_global_ptr(place, {index}));
- builder.create_local_store(sum, builder.create_add(sum_old, place_index));
- }
- builder.create_return(builder.create_local_load(sum));
-
- kernel_ret = std::make_unique(program, builder.extract_ir(), "ret");
- }
-
- {
- /*
- @ti.kernel
- def ext(ext: ti.ext_arr()):
- for index in place:
- ext[index] = place[index];
- # ext = place.to_numpy()
- */
- IRBuilder builder;
- auto *loop = builder.create_struct_for(pointer, 1, 0, 4);
- {
- auto _ = builder.get_loop_guard(loop);
- auto *index = builder.get_loop_index(loop);
- auto *ext = builder.create_external_ptr(
- builder.create_arg_load(0, PrimitiveType::i32, true), {index});
- auto *place_index =
- builder.create_global_load(builder.create_global_ptr(place, {index}));
- builder.create_global_store(ext, place_index);
- }
-
- kernel_ext = std::make_unique(program, builder.extract_ir(), "ext");
- kernel_ext->insert_arg(get_data_type(), true);
- }
-
- auto ctx_init = kernel_init->make_launch_context();
- auto ctx_ret = kernel_ret->make_launch_context();
- auto ctx_ext = kernel_ext->make_launch_context();
- std::vector ext_arr(n);
- ctx_ext.set_arg_external_array(0, taichi::uint64(ext_arr.data()), n);
-
- (*kernel_init)(ctx_init);
- (*kernel_ret)(ctx_ret);
- std::cout << program.fetch_result(0) << std::endl;
- (*kernel_ext)(ctx_ext);
- for (int i = 0; i < n; i++)
- std::cout << ext_arr[i] << " ";
- std::cout << std::endl;
-}
-
void autograd() {
/*
import taichi as ti, numpy as np
@@ -211,16 +71,17 @@ void autograd() {
}
};
- auto *snode = &root->dense(0, n).insert_children(SNodeType::place);
+ auto *snode =
+ &root->dense(Axis(0), n, false).insert_children(SNodeType::place);
snode->dt = PrimitiveType::f32;
snode->grad_info = std::make_unique(
- &root->dense(0, n).insert_children(SNodeType::place));
+ &root->dense(Axis(0), n, false).insert_children(SNodeType::place));
snode->get_grad()->dt = PrimitiveType::f32;
snode->get_grad()->grad_info = std::make_unique();
return snode;
};
auto *a = get_snode_grad(), *b = get_snode_grad(), *c = get_snode_grad();
- program.add_snode_tree(std::unique_ptr(root));
+ program.add_snode_tree(std::unique_ptr(root), /*compile_only=*/false);
std::unique_ptr kernel_init, kernel_forward, kernel_backward,
kernel_ext;
@@ -230,7 +91,7 @@ void autograd() {
auto *zero = builder.get_int32(0);
auto *one = builder.get_int32(1);
auto *n_stmt = builder.get_int32(n);
- auto *loop = builder.create_range_for(zero, n_stmt, 1, 0, 4);
+ auto *loop = builder.create_range_for(zero, n_stmt, 0, 4);
{
auto _ = builder.get_loop_guard(loop);
auto *i = builder.get_loop_index(loop);
@@ -253,7 +114,7 @@ void autograd() {
auto get_kernel_cal = [&](bool grad) -> Kernel * {
IRBuilder builder;
- auto *loop = builder.create_struct_for(a, 1, 0, 4);
+ auto *loop = builder.create_struct_for(a, 0, 4);
{
auto _ = builder.get_loop_guard(loop);
auto *i = builder.get_loop_index(loop);
@@ -272,7 +133,7 @@ void autograd() {
{
IRBuilder builder;
- auto *loop = builder.create_struct_for(a, 1, 0, 4);
+ auto *loop = builder.create_struct_for(a, 0, 4);
{
auto _ = builder.get_loop_guard(loop);
auto *i = builder.get_loop_index(loop);
@@ -306,9 +167,12 @@ void autograd() {
auto ctx_backward = kernel_backward->make_launch_context();
auto ctx_ext = kernel_ext->make_launch_context();
std::vector ext_a(n), ext_b(n), ext_c(n);
- ctx_ext.set_arg_external_array(0, taichi::uint64(ext_a.data()), n);
- ctx_ext.set_arg_external_array(1, taichi::uint64(ext_b.data()), n);
- ctx_ext.set_arg_external_array(2, taichi::uint64(ext_c.data()), n);
+ ctx_ext.set_arg_external_array(0, taichi::uint64(ext_a.data()), n,
+ /*is_device_allocation=*/false);
+ ctx_ext.set_arg_external_array(1, taichi::uint64(ext_b.data()), n,
+ /*is_device_allocation=*/false);
+ ctx_ext.set_arg_external_array(2, taichi::uint64(ext_c.data()), n,
+ /*is_device_allocation=*/false);
(*kernel_init)(ctx_init);
(*kernel_forward)(ctx_forward);
@@ -324,9 +188,3 @@ void autograd() {
std::cout << ext_c[i] << " ";
std::cout << std::endl;
}
-
-int main() {
- run_snode();
- autograd();
- return 0;
-}
diff --git a/cpp_examples/main.cpp b/cpp_examples/main.cpp
new file mode 100644
index 0000000000000..af8656a7d5465
--- /dev/null
+++ b/cpp_examples/main.cpp
@@ -0,0 +1,14 @@
+#include "taichi/ir/ir_builder.h"
+#include "taichi/ir/statements.h"
+#include "taichi/program/program.h"
+
+void run_snode();
+void autograd();
+void aot_save();
+
+int main() {
+ run_snode();
+ autograd();
+ aot_save();
+ return 0;
+}
diff --git a/cpp_examples/run_snode.cpp b/cpp_examples/run_snode.cpp
new file mode 100644
index 0000000000000..992f6ae1d79f2
--- /dev/null
+++ b/cpp_examples/run_snode.cpp
@@ -0,0 +1,141 @@
+#include "taichi/ir/ir_builder.h"
+#include "taichi/ir/statements.h"
+#include "taichi/program/program.h"
+
+void run_snode() {
+ /*
+ import taichi as ti, numpy as np
+ ti.init()
+ #ti.init(print_ir = True)
+
+ n = 10
+ place = ti.field(dtype = ti.i32)
+ ti.root.pointer(ti.i, n).place(place)
+
+ @ti.kernel
+ def init():
+ for index in range(n):
+ place[index] = index
+
+ @ti.kernel
+ def ret() -> ti.i32:
+ sum = 0
+ for index in place:
+ sum = sum + place[index]
+ return sum
+
+ @ti.kernel
+ def ext(ext_arr: ti.ext_arr()):
+ for index in place:
+ ext_arr[index] = place[index]
+
+ init()
+ print(ret())
+ ext_arr = np.zeros(n, np.int32)
+ ext(ext_arr)
+ #ext_arr = place.to_numpy()
+ print(ext_arr)
+ */
+
+ using namespace taichi;
+ using namespace lang;
+ auto program = Program(Arch::x64);
+ /*CompileConfig config_print_ir;
+ config_print_ir.print_ir = true;
+ prog_.config = config_print_ir;*/ // print_ir = True
+
+ int n = 10;
+ program.materialize_runtime();
+ auto *root = new SNode(0, SNodeType::root);
+ auto *pointer = &root->pointer(Axis(0), n, false);
+ auto *place = &pointer->insert_children(SNodeType::place);
+ place->dt = PrimitiveType::i32;
+ program.add_snode_tree(std::unique_ptr(root), /*compile_only=*/false);
+
+ std::unique_ptr kernel_init, kernel_ret, kernel_ext;
+
+ {
+ /*
+ @ti.kernel
+ def init():
+ for index in range(n):
+ place[index] = index
+ */
+ IRBuilder builder;
+ auto *zero = builder.get_int32(0);
+ auto *n_stmt = builder.get_int32(n);
+ auto *loop = builder.create_range_for(zero, n_stmt, 0, 4);
+ {
+ auto _ = builder.get_loop_guard(loop);
+ auto *index = builder.get_loop_index(loop);
+ auto *ptr = builder.create_global_ptr(place, {index});
+ builder.create_global_store(ptr, index);
+ }
+
+ kernel_init =
+ std::make_unique(program, builder.extract_ir(), "init");
+ }
+
+ {
+ /*
+ @ti.kernel
+ def ret():
+ sum = 0
+ for index in place:
+ sum = sum + place[index];
+ return sum
+ */
+ IRBuilder builder;
+ auto *sum = builder.create_local_var(PrimitiveType::i32);
+ auto *loop = builder.create_struct_for(pointer, 0, 4);
+ {
+ auto _ = builder.get_loop_guard(loop);
+ auto *index = builder.get_loop_index(loop);
+ auto *sum_old = builder.create_local_load(sum);
+ auto *place_index =
+ builder.create_global_load(builder.create_global_ptr(place, {index}));
+ builder.create_local_store(sum, builder.create_add(sum_old, place_index));
+ }
+ builder.create_return(builder.create_local_load(sum));
+
+ kernel_ret = std::make_unique(program, builder.extract_ir(), "ret");
+ }
+
+ {
+ /*
+ @ti.kernel
+ def ext(ext: ti.ext_arr()):
+ for index in place:
+ ext[index] = place[index];
+ # ext = place.to_numpy()
+ */
+ IRBuilder builder;
+ auto *loop = builder.create_struct_for(pointer, 0, 4);
+ {
+ auto _ = builder.get_loop_guard(loop);
+ auto *index = builder.get_loop_index(loop);
+ auto *ext = builder.create_external_ptr(
+ builder.create_arg_load(0, PrimitiveType::i32, true), {index});
+ auto *place_index =
+ builder.create_global_load(builder.create_global_ptr(place, {index}));
+ builder.create_global_store(ext, place_index);
+ }
+
+ kernel_ext = std::make_unique(program, builder.extract_ir(), "ext");
+ kernel_ext->insert_arg(get_data_type(), true);
+ }
+
+ auto ctx_init = kernel_init->make_launch_context();
+ auto ctx_ret = kernel_ret->make_launch_context();
+ auto ctx_ext = kernel_ext->make_launch_context();
+ std::vector ext_arr(n);
+ ctx_ext.set_arg_external_array(0, taichi::uint64(ext_arr.data()), n, false);
+
+ (*kernel_init)(ctx_init);
+ (*kernel_ret)(ctx_ret);
+ std::cout << program.fetch_result(0) << std::endl;
+ (*kernel_ext)(ctx_ext);
+ for (int i = 0; i < n; i++)
+ std::cout << ext_arr[i] << " ";
+ std::cout << std::endl;
+}
diff --git a/docs/fragments/.gitkeep b/docs/fragments/.gitkeep
new file mode 100644
index 0000000000000..e69de29bb2d1d
diff --git a/docs/lang/api/index.md b/docs/lang/api/index.md
deleted file mode 100644
index 3973c95ba54bc..0000000000000
--- a/docs/lang/api/index.md
+++ /dev/null
@@ -1,9 +0,0 @@
----
-sidebar_position: 1
----
-
-# API Docs
-
-:::info Taichi language API reference
-Please refer to https://api-docs.taichi.graphics for the most up-to-date Taichi language's API reference while we are working on improving this page, thanks for your understanding!
-:::
diff --git a/docs/lang/articles/advanced/_category_.json b/docs/lang/articles/advanced/_category_.json
index 84fc16f9348e5..321360c5fbf3d 100644
--- a/docs/lang/articles/advanced/_category_.json
+++ b/docs/lang/articles/advanced/_category_.json
@@ -1,4 +1,4 @@
{
- "label": "Advanced Programming",
+ "label": "Advanced Topics",
"position": 3
}
diff --git a/docs/lang/articles/advanced/differentiable_programming.md b/docs/lang/articles/advanced/differentiable_programming.md
index 7bdd46768cc2b..40377b5974f7d 100644
--- a/docs/lang/articles/advanced/differentiable_programming.md
+++ b/docs/lang/articles/advanced/differentiable_programming.md
@@ -79,14 +79,14 @@ print('dy/dx =', x.grad[None], ' at x =', x[None])
A common problem in physical simulation is that it's usually easy to compute
energy but hard to compute force on every particle,
e.g [Bond bending (and torsion) in molecular dynamics](https://github.com/victoriacity/taichimd/blob/5a44841cc8dfe5eb97de51f1d46f1bede1cc9936/taichimd/interaction.py#L190-L220)
-and [FEM with hyperelastic energy functions](https://github.com/taichi-dev/taichi/blob/master/examples/simulation/fem128.py).
+and [FEM with hyperelastic energy functions](https://github.com/taichi-dev/taichi/blob/master/python/taichi/examples/simulation/fem128.py).
Recall that we can differentiate
(negative) potential energy to get forces: `F_i = -dU / dx_i`. So once you have coded
a kernel that computes the potential energy, you may use Taichi's autodiff
system to obtain the derivatives and then `F_i` on each particle.
Take
-[examples/simulation/ad_gravity.py](https://github.com/taichi-dev/taichi/blob/master/examples/simulation/ad_gravity.py)
+[examples/simulation/ad_gravity.py](https://github.com/taichi-dev/taichi/blob/master/python/taichi/examples/simulation/ad_gravity.py)
as an example:
```python
@@ -158,9 +158,9 @@ start up.
:::tip
See
-[examples/simulation/mpm_lagrangian_forces.py](https://github.com/taichi-dev/taichi/blob/master/examples/simulation/mpm_lagrangian_forces.py)
+[examples/simulation/mpm_lagrangian_forces.py](https://github.com/taichi-dev/taichi/blob/master/python/taichi/examples/simulation/mpm_lagrangian_forces.py)
and
-[examples/simulation/fem99.py](https://github.com/taichi-dev/taichi/blob/master/examples/simulation/fem99.py)
+[examples/simulation/fem99.py](https://github.com/taichi-dev/taichi/blob/master/python/taichi/examples/simulation/fem99.py)
for examples on using autodiff-based force evaluation MPM and FEM.
:::
@@ -169,6 +169,10 @@ for examples on using autodiff-based force evaluation MPM and FEM.
As mentioned above, `ti.Tape()` can only track a 0D field as the output variable.
If there're multiple output variables that you want to back-propagate
gradients to inputs, `kernel.grad()` should be used instead of `ti.Tape()`.
+Different from using `ti.Tape()`, you need to set the `grad` of the output variables themselves to `1` manually
+before calling the `kernel.grad()`. The reason is that the `grad` of the output variables themselves
+will always be multiplied to the `grad` with respect to the inputs at the end of the back-propagation.
+If using `ti.Tape()`, the program will help you do this under the hood.
```python {13-14}
import taichi as ti
@@ -188,6 +192,8 @@ def func():
for i in range(N):
x[i] = i
+
+# Set the `grad` of the output variables to `1` before calling `func.grad()`.
loss.grad[None] = 1
loss2.grad[None] = 1
@@ -296,41 +302,22 @@ func_break_rule_2.grad()
assert x.grad[1] == 4.0
assert x.grad[2] == 3.0
```
-### Kernel Simplicity Rule
-:::note Kernel Simplicity Rule
-Kernel body must consist of multiple simply nested for-loops. For example, each for-loop can either contain exactly one (nested) for-loop (and no other statements), or a group of statements without loops.
+### Avoid mixed usage of parallel for-loop and non-for statements
+
+Mixed usage of parallel for-loops and non-for statements are not supported in the autodiff system.
+Please split the two kinds of statements into different kernels.
+
+:::note
+Kernel body must only consist of either multiple for-loops or non-for statements.
:::
Example:
```python
@ti.kernel
-def differentiable_task1():
- # Good: simple for loop
- for i in x:
- x[i] = y[i]
-
-@ti.kernel
-def differentiable_task2():
- # Good: one nested for loop
- for i in range(10):
- for j in range(20):
- for k in range(300):
- ... do whatever you want, as long as there are no loops
-
-@ti.kernel
-def differentiable_task3():
- # Bad: the outer for loop contains two for loops.
- for i in range(10):
- for j in range(20):
- ...
- for j in range(20):
- ...
-
-@ti.kernel
-def differentiable_task4():
- # Bad: mixed usage of for-loop and a statement without looping. Please split them into two kernels.
+def differentiable_task():
+ # Bad: mixed usage of a parallel for-loop and a statement without looping. Please split them into two kernels.
loss[None] += x[0]
for i in range(10):
...
@@ -345,27 +332,50 @@ to open a [github issue](https://github.com/taichi-dev/taichi/issues/new?assigne
if you see any silent wrong results.
:::
-### Workaround kernel simplicity rule
+### Write differentiable code inside Taichi kernel
-:::tip
-**static for-loops** (e.g. `for i in ti.static(range(4))`) will get
-unrolled by the Python frontend preprocessor and therefore does not
-count as a level of loop.
-:::
+Taichi compiler only captures the code in the Taichi scope when performing the source code transformation for autodiff. Therefore, only the code written in Taichi scope is auto-differentiated. Although you can modify the `grad` of a field in python scope manually, the code is not auto-differentiated.
+
+Example:
+
+```python
+import taichi as ti
+
+ti.init()
+x = ti.field(dtype=float, shape=(), needs_grad=True)
+loss = ti.field(dtype=float, shape=(), needs_grad=True)
-For instance, we can rewrite `differentiable_task3` listed above using `ti.static`:
-``` python
@ti.kernel
-def differentiable_task3():
- # Good: ti.static unrolls the inner loops so that it now only has one simple for loop.
- for i in range(10):
- for j in ti.static(range(20)):
- ...
- for j in ti.static(range(20)):
- ...
+def differentiable_task():
+ for l in range(3):
+ loss[None] += ti.sin(x[None]) + 1.0
+
+@ti.kernel
+def manipulation_in_kernel():
+ loss[None] += ti.sin(x[None]) + 1.0
+
+
+x[None] = 0.0
+with ti.Tape(loss=loss):
+ # The line below in python scope only contribute to the forward pass
+ # but not the backward pass i.e., not auto-differentiated.
+ loss[None] += ti.sin(x[None]) + 1.0
+
+ # Code in Taichi scope i.e. inside Taichi kernels, is auto-differentiated.
+ manipulation_in_kernel()
+ differentiable_task()
+
+# The outputs are 5.0 and 4.0
+print(loss[None], x.grad[None])
+
+# You can modify the grad of a field manually in python scope, e.g., clear the grad.
+x.grad[None] = 0.0
+# The output is 0.0
+print(x.grad[None])
```
+
## Extending Taichi Autodiff system
diff --git a/docs/lang/articles/advanced/layout.md b/docs/lang/articles/advanced/layout.md
index be7088838e14a..dc41ca0ce38e1 100644
--- a/docs/lang/articles/advanced/layout.md
+++ b/docs/lang/articles/advanced/layout.md
@@ -3,58 +3,33 @@ sidebar_position: 2
---
# Fields (advanced)
+Modern processor cores compute orders of magnitude faster than their equipped memory systems. To shrink this performance gap, multi-level cache systems and high-bandwidth multi-channel memories are built into computer architectures.
-This section introduces some advanced features of Taichi fields.
-Make sure you have gone through [Fields](../basic/field).
+After familiar yourself with the basics of Taichi [Fields](../basic/field), this article helps you one step further by explaining the underlying memory layout that is essential to write high-performance Taichi programs. In particular, we present how to organize an efficient data layout and how to manage memory occupancy.
-## Packed mode
+## How to organize an efficient data layout
-By default, all non-power-of-two dimensions of a field are automatically
-padded to a power of two. For instance, a field of shape `(18, 65)` will
-have an internal shape of `(32, 128)`. Although the padding has many benefits
-such as allowing fast and convenient bitwise operations for coordinate handling,
-it will consume potentially much more memory than expected.
+In this section, we introduce how to organize data layouts in Taichi fields. The central principle of efficient data layout is _locality_. Generally speaking, a program with desirable locality has at least one of the following features:
+* Dense data structure
+* Loop over data in small-range (within 32KB is good for most processors)
+* Sequential load/store
-If you would like to reduce memory usage, you can use the optional packed
-mode. In packed mode, no padding is applied such that a field does not have a
-larger internal shape than the defined shape when some of its dimensions
-are not powers of two. The downside is that the runtime performance will
-regress slightly.
-
-A switch named `packed` for `ti.init()` decides whether to use packed mode:
-
-```python
-ti.init() # default: packed=False
-a = ti.field(ti.i32, shape=(18, 65)) # padded to (32, 128)
-```
-
-```python
-ti.init(packed=True)
-a = ti.field(ti.i32, shape=(18, 65)) # no padding
-```
-
-## Advanced data layouts
+:::note
-Apart from shape and data type, you can also specify the data layout of a
-field in a recursive manner. This may allow you to achieve better performance.
-Normally, you don't have to worry about the performance nuances between
-different layouts, and you can just use the default one (simply by specifying
-`shape` when creating fields) as a start.
+Be aware that data are always fetched from memory in blocks (pages). The hardware has no knowledge about how a specific data element is used in the block. The processor blindly fetch the entire block according to the requested memory address. Therefore, the memory bandwidth is wasted when data are not fully utilized.
-However, when a field gets large, a proper data layout may be critical to
-performance, especially for memory-bound applications. A carefully designed
-data layout has much better spatial locality, which will significantly
-improve cache/TLB-hit rates and cache line utilization.
+For sparse fields, refer to [Sparse computation](./sparse.md).
-Taichi decouples computation from data structures, and the Taichi compiler
-automatically optimizes data accesses on a specific data layout. This allows
-you to quickly experiment with different data layouts and figure out the most
-efficient one on a specific task and computer architecture.
+:::
### Layout 101: from `shape` to `ti.root.X`
-The following declares a 0-D field:
+
+In basic usages, we use the `shape` descriptor to construct a field. Taichi provides flexible statements to describe more advanced data organizations, the `ti.root.X`.
+Let's get some familiarity with examples:
+
+* Declare a 0-D field:
```python {1-2}
x = ti.field(ti.f32)
@@ -63,7 +38,7 @@ ti.root.place(x)
x = ti.field(ti.f32, shape=())
```
-The following declares a 1D field of shape `3`:
+* Declare a 1-D field of shape `3`:
```python {1-2}
x = ti.field(ti.f32)
@@ -72,7 +47,7 @@ ti.root.dense(ti.i, 3).place(x)
x = ti.field(ti.f32, shape=3)
```
-The following declares a 2D field of shape `(3, 4)`:
+* Declare a 2-D field of shape `(3, 4)`:
```python {1-2}
x = ti.field(ti.f32)
@@ -80,46 +55,48 @@ ti.root.dense(ti.ij, (3, 4)).place(x)
# is equivalent to:
x = ti.field(ti.f32, shape=(3, 4))
```
+You can also nest two 1D `dense` statements to describe the same 2D array.
+```python {1-2}
+x = ti.field(ti.f32)
+ti.root.dense(ti.i, 3).dense(ti.j, 4).place(x)
+```
-After being comfortable with these equivalent definitions, you can move forward
-and see how to change the data layout.
+In a nutshell, the `ti.root.X` statement progressively binds a shape to the corresponding axis.
+By nesting multiple statements, we can construct a field with higher dimensions.
+
-### Row-major versus column-major
+In order to traverse the nested statements, we can use `struct-for`:
+```python {1}
+for i, j in A:
+ A[i, j] += 1
+```
+The order to access `A`, namely the order to iterate `i` and `j`, affects the program performance subtly. The Taichi compiler is capable to automatically deduce the underlying data layout and apply a proper access order. This is an advantage over most general-purpose programming languages where the access order has to be optimized manually.
-As you might have learned in a computer architecture course,
-address spaces are linear in modern computers. To
-simplify the discussions, data type size will not be considered and will always
-be treated as 1. Assume the starting address of a field is `base`. Then for 1D
-Taichi fields, the address of the `i`-th element is simply `base + i`.
+### Row-major versus column-major
-However, a multi-dimensional field has to be flattened in order to fit into the
-1D address space. For example, there are two ways to store a 2D field of size `(3, 2)`:
+Memory address space is linear as you might have learnt from a computer architecture course. Without loss of generality, we omit the differences in data types and assume each data element has size 1. Moreover, we denote the starting memory address of a field as `base`, and the indexing formula for 1D Taichi fields is `base + i` for the `i`-th element.
-- Row-major: let the address of the `(i, j)`-th element be `base + i * 2 + j`;
-- Column-major: let the address of the `(i, j)`-th element be
- `base + j * 3 + i`.
+For multi-dimensional fields, we can flatten the high-dimension index into the linear memory address space in two ways: Take a 2D field of shape `(M, N)` as an instance, we can either store `M` rows with `N`-length 1D buffers, say the _row-major_ way, or store `N` columns, say the _column-major_ way. The index flatten formula for the `(i, j)`-th element is `base + i * N + j` for row-major and `base + j * M + i` for column-major, respectively.
-To specify which layout to use (default layout is row-major):
+We can easily derive that elements in the same row are close in memory for row-major fields. The selection of the optimal layout is based on how the elements are accessed, namely, the access patterns. Patterns such as frequently accessing elements of the same row in a column-major field typically lead to performance degradation.
+The default Taichi field layout is row-major. With the `ti.root` statements, fields can be defined as follows:
```python
-ti.root.dense(ti.i, 3).dense(ti.j, 2).place(x) # row-major
-ti.root.dense(ti.j, 2).dense(ti.i, 3).place(y) # column-major
+ti.root.dense(ti.i, M).dense(ti.j, N).place(x) # row-major
+ti.root.dense(ti.j, N).dense(ti.i, M).place(y) # column-major
```
-
-Both `x` and `y` have shape `(3, 2)`, and they can be accessed in the same
-manner with `x[i, j]` and `y[i, j]`, where `0 <= i < 3 && 0 <= j < 2`. However,
-they have different memory layouts:
-
+In the code above, the axis denotation in the rightmost `dense` statement indicates the continuous axis. For the `x` field, elements in the same row (with same `i` and different `j`) are close in memory, hence it's row-major; For the `y` field, elements in the same column (same `j` and different `i`) are close, hence it's column-major. With an example of (2, 3), we visualize the memory layouts of `x` and `y` as follows:
```
# address: low ........................................... high
# x: x[0, 0] x[0, 1] x[1, 0] x[1, 1] x[2, 0] x[2, 1]
# y: y[0, 0] y[1, 0] y[2, 0] y[0, 1] y[1, 1] y[2, 1]
```
-:::note
+It is worth noting that the accessor is unified for Taichi fields: the `(i, j)`-th element in the field is accessed with the identical 2D index `x[i, j]` and `y[i, j]`. Taichi handles the layout variants and applies proper indexing equations internally. Thanks to this feature, users can specify their desired layout at definition, and use the fields without concerning about the underlying memory organizations. To change the layout, it's sufficient to just swap the order of `dense` statements, and leave rest of the code intact.
-For those who are familiar with C/C++, here is what they look like in C code:
+:::note
+For readers who are familiar with C/C++, below is an example C code snippet that demonstrates data access in 2D arrays:
```c
int x[3][2]; // row-major
int y[2][3]; // column-major
@@ -132,49 +109,60 @@ for (int i = 0; i < 3; i++) {
}
```
+The accessors of `x` and `y` are in reverse order between row-major arrays and column-major arrays, respectively. Compared with Taichi fields, there are much more code to revise when you change the memory layout.
+
:::
-### Array of Structures (AoS) versus Structure of Arrays (SoA)
+
+### AoS versus SoA
-Fields of same shape can be placed together.
+AoS means _array of structures_ and SoA means _structure of arrays_. Consider an RGB image with 4 pixels and 3 color channels, an AoS layout stores `RGBRGBRGBRGB` while an SoA layout stores `RRRRGGGGBBBB`.
-For example, the following places two 1D fields of shape `3` together, which
-is called Array of Structures (AoS):
+The selection of AoS or SoA layout largely depends on the access pattern to the field. Let's discuss a scenario to process large RGB images. The two layouts have the following arrangements in memory:
+```
+# address: low ...................... high
+# AoS: RGBRGBRGBRGBRGBRGB.............
+# SoA: RRRRR...RGGGGGGG...GBBBBBBB...B
+```
+To calculate grey scale of each pixel, we need all color channels but do not require the value of other pixels. In this case, the AoS layout has a better memory access pattern: Since color channels are stored continuously, and adjacent channels can be fetched instantly. The SoA layout is not a good option because the color channels of a pixel are stored far apart in the memory space.
+We describe how to construct AoS and SoA fields with our `ti.root.X` statements. The SoA fields are trivial:
```python
-ti.root.dense(ti.i, 3).place(x, y)
+x = ti.field(ti.f32)
+y = ti.field(ti.f32)
+ti.root.dense(ti.i, M).place(x)
+ti.root.dense(ti.i, M).place(y)
```
-
-Their memory layout is:
-
+where M is the length of `x` and `y`.
+The data elements in `x` and `y` are continuous in memory:
```
-# address: low ......................... high
-# x[0] y[0] x[1] y[1] x[2] y[2]
+# address: low ................................. high
+# x[0] x[1] x[2] ... y[0] y[1] y[2] ...
```
-By contrast, the following places these two fields separately, which is called
-Structure of Arrays (SoA):
+For AoS fields, we construct the field with
```python
-ti.root.dense(ti.i, 3).place(x)
-ti.root.dense(ti.i, 3).place(y)
+x = ti.field(ti.f32)
+y = ti.field(ti.f32)
+ti.root.dense(ti.i, M).place(x, y)
```
-
-Now their memory layout is:
-
+The memroy layout then becomes
```
-# address: low ......................... high
-# x[0] x[1] x[2] y[0] y[1] y[2]
+# address: low .............................. high
+# x[0] y[0] x[1] y[1] x[2] y[2] ...
```
+Here, `place` interleaves the elements of Taichi fields `x` and `y`.
-**To improve spatial locality of memory accesses, it may be helpful to
-place data elements that are often accessed together within
-relatively close addresses.** Take a simple 1D wave equation solver as an example:
+As previously introduced, the access methods to `x` and `y` remain the same for both AoS and SoA. Therefore, the data layout can be changed flexibly without revising the application logic.
+
+For better illustration, let's see an example of an 1D wave equation solver:
```python
N = 200000
pos = ti.field(ti.f32)
vel = ti.field(ti.f32)
+# SoA placement
ti.root.dense(ti.i, N).place(pos)
ti.root.dense(ti.i, N).place(vel)
@@ -182,127 +170,122 @@ ti.root.dense(ti.i, N).place(vel)
def step():
pos[i] += vel[i] * dt
vel[i] += -k * pos[i] * dt
+
+...
```
+The above code snippet defines SoA fields and a `step` kernel that sequentially accesses each element.
+The kernel fetches an element from `pos` and `vel` for every iteration, respectively.
+For SoA fields, the closest distance of any two elements in memory is `N`, which is unlikely to be efficient for large `N`.
-Here, `pos` and `vel` are placed separately, so the distance in address
-space between `pos[i]` and `vel[i]` is `200000`. This results in poor spatial
-locality and poor performance. A better way is to place them together:
+We hereby switch the layout to AoS as follows:
```python
+N = 200000
+pos = ti.field(ti.f32)
+vel = ti.field(ti.f32)
+# AoS placement
ti.root.dense(ti.i, N).place(pos, vel)
+
+@ti.kernel
+def step():
+ pos[i] += vel[i] * dt
+ vel[i] += -k * pos[i] * dt
```
+Merely revising the place statement is sufficient to change the layout. With this optimization, the instant elements `pos[i]` and `vel[i]` are now adjacent in memory, which is more efficient.
-Then `vel[i]` is placed right next to `pos[i]`, which can increase spatial
-locality and therefore improve performance.
+
-Struct-fors on nested dense data structures will automatically follow their
-layout in memory. For example, if 2D scalar field `A` is defined in row-major,
+
-As you may notice, only dense data layouts are covered in this section. For sparse
-data layouts, see [Sparse computation](./sparse.md).
+
-2D field, column-major:
-```python
-A = ti.field(ti.f32)
-ti.root.dense(ti.j, 256).dense(ti.i, 256).place(A)
-```
-_8x8_-blocked 2D field of size _1024x1024_:
+### AoS extension: hierarchical fields
+
+Sometimes we want to access memory in a complex but fixed pattern, like traversing an image in 8x8 blocks. The apparent best practice is to flatten each 8x8 block and concatenate them together. From a Taichi user's perspective, however, the field is no longer a flat buffer. It now has a hierarchy with two levels: The image level and the block level. Equivalently, the field is an array of implicit 8x8 block structures.
+
+We demonstrate the statements as follows:
```python
-density = ti.field(ti.f32)
-ti.root.dense(ti.ij, (128, 128)).dense(ti.ij, (8, 8)).place(density)
+# Flat field
+val = ti.field(ti.f32)
+ti.root.dense(ti.ij, (M, N)).place(val)
```
-3D particle positions and velocities, AoS:
-
```python
-pos = ti.Vector.field(3, dtype=ti.f32)
-vel = ti.Vector.field(3, dtype=ti.f32)
-ti.root.dense(ti.i, 1024).place(pos, vel)
-# equivalent to
-ti.root.dense(ti.i, 1024).place(pos.get_scalar_field(0),
- pos.get_scalar_field(1),
- pos.get_scalar_field(2),
- vel.get_scalar_field(0),
- vel.get_scalar_field(1),
- vel.get_scalar_field(2))
+# Hierarchical field
+val = ti.field(ti.f32)
+ti.root.dense(ti.ij, (M // 8, N // 8)).dense(ti.ij, (8, 8)).place(val)
```
+where `M` and `N` are multiples of 8. We encourage you to try this out! The performance difference can be significant!
-3D particle positions and velocities, SoA:
+## How to manage memory occupancy
-```python
-pos = ti.Vector.field(3, dtype=ti.f32)
-vel = ti.Vector.field(3, dtype=ti.f32)
-for i in range(3):
- ti.root.dense(ti.i, 1024).place(pos.get_scalar_field(i))
-for i in range(3):
- ti.root.dense(ti.i, 1024).place(vel.get_scalar_field(i))
-```
+### Manual field allocation and destruction
-## Dynamic field allocation and destruction
+Generally Taichi manages memory allocation and destruction without disturbing the users. However, there are times that users want explicit control over their memory allocations.
-You can use the `FieldsBuilder` class for dynamic field allocation and destruction.
-`FieldsBuilder` has the same data structure declaration APIs as `ti.root`,
-including `dense()`. After declaration, you need to call the `finalize()`
-method to compile it to an `SNodeTree` object.
+In this scenario, Taichi provides the `FieldsBuilder` for manual field memory allocation and destruction. `FieldsBuilder` features identical declaration APIs as `ti.root`. The extra step is to invoke `finalize()` at the end of all declarations. The `finalize()` returns an `SNodeTree` object to handle subsequent destructions.
-A simple example is:
+Let's see a simple example:
```python
import taichi as ti
@@ -318,53 +301,34 @@ x = ti.field(dtype=ti.f32)
fb1.dense(ti.ij, (5, 5)).place(x)
fb1_snode_tree = fb1.finalize() # Finalizes the FieldsBuilder and returns a SNodeTree
func(x)
+...
+fb1_snode_tree.destroy() # Destruction
fb2 = ti.FieldsBuilder()
y = ti.field(dtype=ti.f32)
fb2.dense(ti.i, 5).place(y)
fb2_snode_tree = fb2.finalize() # Finalizes the FieldsBuilder and returns a SNodeTree
func(y)
+...
+fb2_snode_tree.destroy() # Destruction
```
+Actually, the above demonstrated `ti.root` statements are implemented with `FieldsBuilder`, despite that `ti.root` has the capability to automatically manage memory allocations and recycling.
-In fact, `ti.root` is implemented by `FieldsBuilder` implicitly, so you can
-allocate the fields directly under `ti.root`:
-```python
-import taichi as ti
-ti.init() # Implicitly: ti.root = ti.FieldsBuilder()
+### Packed mode
-@ti.kernel
-def func(v: ti.template()):
- for I in ti.grouped(v):
- v[I] += 1
+By default, Taichi implicitly fits a field in a larger buffer with power-of-two dimensions. We take the power-of-two padding convention because it is widely adopted in computer graphics. The design enables fast indexing with bitwise arithmetic and better memory address alignment, while trading off memory occupations.
-x = ti.field(dtype=ti.f32)
-ti.root.dense(ti.ij, (5, 5)).place(x)
-func(x) # Automatically calls ti.root.finalize()
-# Implicitly: ti.root = ti.FieldsBuilder()
+For example, a `(18, 65)` field is materialized with a `(32, 128)` buffer, which is acceptable. As field size grows, the padding strategy can be exaggeratedly unbearable: `(129, 6553600)` will be expanded to `(256, 6335600)`, which allocates considerable unsed blank memory. Therefore, Taichi provides the optional packed mode to allocate buffer that tightly fits the requested field shape. It is especially useful when memory usage is a major concern.
-y = ti.field(dtype=ti.f32)
-ti.root.dense(ti.i, 5).place(y)
-func(y) # Automatically calls ti.root.finalize()
+To leverage the packed mode, spcifify `packed` in `ti.init()` argument:
+```python
+ti.init() # default: packed=False
+a = ti.field(ti.i32, shape=(18, 65)) # padded to (32, 128)
```
-Furthermore, if you don't want to use the fields under a certain `SNodeTree`
-anymore, you could call the `destroy()` method on the finalized `SNodeTree`
-object, which will recycle its memory into the memory pool:
-
-```py
-import taichi as ti
-ti.init()
-
-@ti.kernel
-def func(v: ti.template()):
- for I in ti.grouped(v):
- v[I] += 1
-
-fb = ti.FieldsBuilder()
-x = ti.field(dtype=ti.f32)
-fb.dense(ti.ij, (5, 5)).place(x)
-fb_snode_tree = fb.finalize() # Finalizes the FieldsBuilder and returns a SNodeTree
-func(x)
-
-fb_snode_tree.destroy() # x cannot be used anymore
+```python
+ti.init(packed=True)
+a = ti.field(ti.i32, shape=(18, 65)) # no padding
```
+
+You might observe mild performance regression with the packed mode due to more complex adressing and memory alignment. Therefore, the packed mode should be specified only when memory capacity is a major concern.
diff --git a/docs/lang/articles/advanced/meta.md b/docs/lang/articles/advanced/meta.md
index 9b9f7b54d8d61..82947bed60728 100644
--- a/docs/lang/articles/advanced/meta.md
+++ b/docs/lang/articles/advanced/meta.md
@@ -18,7 +18,7 @@ Every kernel in Taichi is a template kernel, even if it has no template argument
## Template metaprogramming
-By using `ti.template()` as a argument type hint, a Taichi field can be passed into a kernel. Template programming also enables the code to be reused for fields with different shapes:
+By using `ti.template()` as an argument type hint, a Taichi field or a python object can be passed into a kernel. Template programming also enables the code to be reused for fields with different shapes:
```python {2}
@ti.kernel
@@ -38,6 +38,10 @@ copy_1D(a, b)
copy_1D(c, d)
```
+:::note
+If a template parameter is not a Taichi object, it cannot be reassigned inside Taichi kernel.
+:::
+
:::note
The template parameters are inlined into the generated kernel after compilation.
:::
@@ -122,6 +126,12 @@ def foo():
Using compile-time evaluation allows for some computation to be executed when kernels are instantiated. This helps the compiler to conduct optimization and reduce
computational overhead at runtime:
+### Static Scope
+`ti.static` is a function which receives one argument. It is a hint for the compiler to evaluate the argument at compile time.
+The scope of the argument of `ti.static` is called static-scope.
+
+### Compile-time branching
+
- Use `ti.static` for compile-time branching (for those who are familiar with
C++17, this is similar to [if
constexpr](https://en.cppreference.com/w/cpp/language/if).):
@@ -139,6 +149,8 @@ def static():
One of the two branches of the `static if` will be discarded after compilation.
:::
+### Loop unrolling
+
- Use `ti.static` for forced loop unrolling:
```python {3}
@@ -173,3 +185,28 @@ def reset():
# The inner loop must be unrolled since j is an index for accessing a vector
x[i][j] = 0
```
+
+## Compile-time recursion of `ti.func`
+
+A compile-time recursive function is a function with recursion that can be recursively inlined at compile time. The condition which determines whether to recurse is evaluated at compile time.
+
+You can combine [compile-time branching](#compile-time-evaluations) and [template](#template-metaprogramming) to write compile-time recursive functions.
+
+For example, `sum_from_one_to` is a compile-time recursive function that calculates the sum of numbers from `1` to `n`.
+
+```python {1-6}
+@ti.func
+def sum_from_one_to(n: ti.template()) -> ti.i32:
+ ret = 0
+ if ti.static(n > 0):
+ ret = n + sum_from_one_to(n - 1)
+ return ret
+
+@ti.kernel
+def sum_from_one_to_ten():
+ print(sum_from_one_to(10)) # prints 55
+```
+
+:::caution WARNING
+When the recursion is too deep, it is not recommended to use compile-time recursion because deeper compile-time recursion expands to longer code during compilation, resulting in increased compilation time.
+:::
diff --git a/docs/lang/articles/advanced/performance.md b/docs/lang/articles/advanced/performance.md
index eaed9c0d5f295..d2de92ffec58b 100644
--- a/docs/lang/articles/advanced/performance.md
+++ b/docs/lang/articles/advanced/performance.md
@@ -12,7 +12,7 @@ the target architecture. Nevertheless, for Ninjas who strive for the last few %
of performance, we also provide some APIs to allow developers fine-tune their
applications. For example, specifying a suitable `ti.block_dim` could yield an almost
3x performance boost in
-[examples/mpm3d.py](https://github.com/taichi-dev/taichi/blob/master/examples/mpm3d.py).
+[examples/mpm3d.py](https://github.com/taichi-dev/taichi/blob/master/python/taichi/examples/mpm3d.py).
:::note
For **performance profiling** utilities, please see [Profiler section of the Contribution Guide](../misc/profiler.md).
diff --git a/docs/lang/articles/advanced/sparse.md b/docs/lang/articles/advanced/sparse.md
index 7c17b41631e2c..fd0528d688c80 100644
--- a/docs/lang/articles/advanced/sparse.md
+++ b/docs/lang/articles/advanced/sparse.md
@@ -2,50 +2,29 @@
sidebar_position: 3
---
-# Sparse computation
-
-Compiler-level support for spatially sparse computation is a unique feature of Taichi.
-
-
-
-Figure: A swinging "Taichi" pattern represented with a 512x512 sparse grid. The sparse grid has a multi-level *tree* structure.
-White stands for inactive tree nodes, and active tree nodes are darker.
-
-The sparse grid above has the following structure:
-- The grid is divided into 8x8 `block1` containers;
-- Each `block1` container has 4x4 `block2` cells;
-- Each `block2` container has 4x4 `block3` cells;
-- Each `block3` container has 4x4 pixel cells;
-- Each pixel contains an `i32` value `x[i, j]`.
+# Sparse spatial data structures
:::note
-For more information about *cells* and *containers*, see [**Data structure organization**](../misc/internal.md#data-structure-organization).
-In this article, you can assume *containers* and *cells* are the same.
+Prerequisite: please read the [Fields](lang/articles/basic/field.md), [Fields (advanced)](lang/articles/advanced/layout.md), and [SNodes](lang/articles/misc/internal.md#data-structure-organization) first.
:::
-Taichi allows you to define sparse data structures effortlessly. For example, you can define the grid above as
-
-```python
-x = ti.field(dtype=ti.i32)
-
-block1 = ti.root.pointer(ti.ij, 8)
-block2 = block1.pointer(ti.ij, 4)
-block3 = block2.pointer(ti.ij, 4)
-block3.dense(ti.ij, 4).place(x)
-```
-[[Full source code of this animation]](https://github.com/taichi-dev/taichi/blob/master/examples/features/sparse/taichi_sparse.py)
-
-Intuitively, a sparse grid in Taichi allows you to use memory space more wisely, since only tree nodes involved in computation are allocated.
-Now, let's take a step back and think about *why we need sparse grids, how to define them in Taichi, and how to compute on these data structures*.
+
+Figure: A 3D fluid simulation that uses both particles and grids. Left to right: particles, 1x1x1 voxels, 4x4x4 blocks, 16x16x16 blocks.
## Motivation
High-resolution 2D/3D grids are often needed in large-scale spatial computation, such as physical simulation, rendering, and 3D reconstruction.
-However, these grids tend to consume a huge amount of memory space and computation.
+However, these grids tend to consume a huge amount of memory space and computation if we use dense data structures (see [field](lang/articles/basic/field.md) and [field advanced](lang/articles/advanced/layout.md)).
While a programmer may allocate large dense grids to store spatial data (especially physical quantities such as a density or velocity field),
oftentimes, they only care about a small fraction of this dense grid since the rest may be empty space (vacuum or air).
-In short, the regions of interest in sparse grids may only occupy a small fraction of the whole bounding box.
+
+
+
+
+
+
+For example, the regions of interest in sparse grids shown above may only occupy a small fraction of the whole bounding box.
If we can leverage such "spatial sparsity" and focus computation on the regions we care about,
we will significantly save storage and computing power.
@@ -53,265 +32,258 @@ we will significantly save storage and computing power.
The key to leverage spatial sparsity is to replace *dense* grids with *sparse* grids.
:::
-On a sparse data structure, we consider a pixel, voxel, or a grid node to be *active*,
-if it is allocated and involved in computation.
-The rest of the grid is simply *inactive*.
-The *activity* of a leaf or intermediate cell is a boolean value. The activity value of a cell is `True` if and only if the cell is *active*.
-
-Below is a 2D multi-physics simulation (material point method) with 256x256 grid cells.
-Since the simulated objects do not fully occupy the whole domain, we would like to *adaptively* allocate the underlying simulation grid.
-We subdivide the whole simulation domain into 16x16 *blocks*,
-and each block has 16x16 *grid cells*.
-Memory allocation can then happen at *block* granularity,
-and we only consume memory space of blocks that are actually in the simulation.
+The traditional sparse spatial data stuctures are [Quadtrees](https://en.wikipedia.org/wiki/Quadtree) (2D) and
+[Octrees](https://en.wikipedia.org/wiki/Octree) (3D). Since dereferencing pointers is relatively costly on modern computer architectures, compared to quadtrees and octrees, it is more performance-friendly to use shallower trees with larger branching factors.
+[VDB](https://www.openvdb.org/) and [SPGrid](http://pages.cs.wisc.edu/~sifakis/papers/SPGrid.pdf) are such examples.
+In Taichi, programmers can compose data structures similar to VDB and SPGrid with SNodes. The advantages of Taichi sparse spatial data structures include
+1. Access with indices, which just like accessing a dense data structure.
+2. Automatic parallelization when iterating.
+3. Automatic memory access optimization.
-
-(Note the changing distribution of active blocks throughout the simulation.)
:::note
-**Backend compatibility**: The LLVM backends (CPU/CUDA) and the Metal backend offer the full functionality of sparse computation.
-Other backends provide no or limited support of sparse computation.
+**Backend compatibility**: The LLVM backends (CPU/CUDA) and the Metal backend offer the full functionality of computation on sparse spatial data structures.
:::
+
:::note
-Sparse matrices are usually **not** implemented in Taichi via (spatially-) sparse data structures. Use `ti.SparseMatrixBuilder` instead.
+Sparse matrices are usually **not** implemented in Taichi via sparse spatial data structures. See [sparse matrix](lang/articles/advanced/sparse_matrix.md) instead.
:::
-## Defining sparse data structures in Taichi
+## Sparse spatial data structures in Taichi
-Ideally, it would be nice to have a sparse voxel data structure that consumes space or computation only when the voxels are active.
-Practically, Taichi programmers use hierarchical data structures (trees) to organize sparse voxel data.
+Sparse spatial data structures in Taichi are usually composed of `pointer`, `bitmasked`, `dynamic`, and `dense` SNodes. A SNode tree merely composed of `dense` SNodes is not a sparse spatial data structure.
-### Data structure hierarchy
-
-Traditionally, [Quadtrees](https://en.wikipedia.org/wiki/Quadtree) (2D) and
-[Octrees](https://en.wikipedia.org/wiki/Octree) (3D) are often adopted.
-Since dereferencing pointers is relatively costly on modern computer architectures,
-compared to quadtrees and octrees, it is more performance-friendly to use shallower trees with larger branching factors.
-[VDB](https://www.openvdb.org/) and [SPGrid](http://pages.cs.wisc.edu/~sifakis/papers/SPGrid.pdf) are such examples.
-In Taichi, programmers can compose data structures similar to VDB and SPGrid with SNodes.
-
-
-Figure: A 3D fluid simulation that uses both particles and grids. Left to right: particles, 1x1x1 voxels, 4x4x4 blocks, 16x16x16 blocks.
+On a sparse spatial data structure, we consider a pixel, voxel, or a grid node to be *active*,
+if it is allocated and involved in the computation.
+The rest of the grid is simply *inactive*.
+In SNode terms, the *activity* of a leaf or intermediate cell is a boolean value. The activity value of a cell is `True` if and only if the cell is *active*. When writing to an inactive cell, Taichi automatically activates it. Taichi also provides manual manipulation of the activity of a cell, see [Explicitly manipulating and querying sparsity](#explicitly-manipulating-and-querying-sparsity).
-#### Blocked leaf cells and bitmasks
+:::note
+Reading an inactive pixel returns zero.
+:::
-While a null pointer can effectively represent an empty sub-tree, at the leaf level using 64 bits to represent the activity
-of a single voxel can consume too much space.
-For example, if each voxel contains a single `f32` value (4 bytes),
-the 64-bit pointer pointing to the value would take 8 bytes.
-The fact that storage costs of pointers are higher than the space to store the value themselves
-goes against our goal to use sparse data structures to save space.
+### Pointer SNode
-To amortize the storage cost of pointers, programmers usually organize voxels in a *blocked* manner
-and let the pointers directly point to the blocks (instead of voxels).
+```python {2} title=pointer.py
+x = ti.field(ti.f32)
+block = ti.root.pointer(ti.ij, (4,4))
+pixel = block.dense(ti.ij, (2,2))
+pixel.place(x)
-One caveat of this design is that voxels in the same `dense` block can no longer change their activity flexibly.
-Instead, they share a single activity flag. To address this issue,
-the `bitmasked` SNode additionally allocates 1-bit per voxel data to represent the voxel activity.
+@ti.kernel
+def activate():
+ x[2,3] = 1.0
+ x[2,4] = 2.0
-### A typical sparse data structure
+@ti.kernel
+def print_active():
+ for i, j in block:
+ print("Active block", i, j)
+ # output: Active block 1 1
+ # Active block 1 2
+ for i, j in x:
+ print('field x[{}, {}] = {}'.format(i, j, x[i, j]))
+ # output: field x[2, 2] = 0.000000
+ # field x[2, 3] = 1.000000
+ # field x[3, 2] = 0.000000
+ # field x[3, 3] = 0.000000
+ # field x[2, 4] = 2.000000
+ # field x[2, 5] = 0.000000
+ # field x[3, 4] = 0.000000
+ # field x[3, 5] = 0.000000
+```
+The code snippet above creates an 8x8 sparse grid, with the top-level being a 4x4 pointer array (line 2 of `pointer.py`),
+and each pointer pointing to a 2x2 dense block.
+You can write and read the sparse field like a dense field using indices. The below figure shows the active blocks and pixels in green.
-Sparse data structures in Taichi are usually composed of `pointer`, `dense`, and `bitmasked` SNodes.
-The code snippet below creates an 8x8 sparse grid, with the top level being 4x4 pointer arrays,
-and each pointer points to a 2x2 dense block.
+
-```python
-x = ti.field(dtype=ti.i32)
+
-block = ti.root.pointer(ti.ij, (4, 4))
-pixel = block.dense(ti.ij, (2, 2))
-pixel.place(x)
-```
+
-
+Executing the `activate()` function automatically activates `block[1,1]`, which includes `x[2,3]`, and `block[1,2]`, which includes `x[2,4]`. Other pixels of `block[1,1]` (`x[2,2], x[3,2], x[3,3]`) and `block[1,2]` (`x[2,5], x[3,4], x[3,5]`) are also implicitly activated because all pixels in the dense block share the same activity value.
-## Computation on sparse data structures
+In fact, the sparse field is a SNode tree shown in the following figure. You could use the struct-for loop to loop over the different levels of the SNode tree like the `print_active()` function in `pointer.py`. `for i, j in block` would loop over all active `pointer` SNodes. `for i, j in pixel` would loop over all active `dense` SNodes.
-### Activation on write
+
-When writing to an inactive cell on a sparse data structure, Taichi automatically populates the data structure.
+
-For example, when executing `x[2, 3] = 2` on the aforementioned sparse grid `x`,
-Taichi automatically activates `block[1, 1]` so that `pixel[2, 3]` is allocated.
+
-:::note
-Reading an inactive voxel returns zero.
-:::
-### Sparse struct-fors
-Efficiently looping over sparse grid cells that distribute irregularly can be challenging, especially on parallel devices such as GPUs.
-In Taichi, *struct-for's* natively support sparse data structures and only loops over currently active voxels.
-The Taichi system ensures efficient parallelization.
-You can loop over different levels of the tree.
-The code below demonstrates the creation and manipulation of a sparse grid:
+### Bitmasked SNode
-```python
-import taichi as ti
+While a null pointer can effectively represent an empty sub-tree, at the leaf level using 64 bits to represent the activity
+of a single pixel can consume too much space.
+For example, if each pixel contains a single `f32` value (4 bytes),
+the 64-bit pointer pointing to the value would take 8 bytes.
+The fact that storage costs of pointers are higher than the space to store the value themselves
+goes against our goal to use sparse spatial data structures to save space.
-use_bitmask = True
+To amortize the storage cost of pointers, you could organize pixels in a *blocked* manner
+and let the pointers directly point to the blocks like the data structure defined in `pointer.py`.
-ti.init()
+One caveat of this design is that pixels in the same `dense` block can no longer change their activity flexibly.
+Instead, they share a single activity flag. To address this issue,
+the `bitmasked` SNode additionally allocates 1-bit per pixel data to represent the pixel activity.
-x = ti.field(dtype=ti.i32)
-block = ti.root.pointer(ti.ij, (4, 4))
-if use_bitmask:
- pixel = block.bitmasked(ti.ij, (2, 2))
-else:
- pixel = block.dense(ti.ij, (2, 2))
+```python {3} title=bitmasked.py
+x = ti.field(ti.f32)
+block = ti.root.pointer(ti.ij, (4,4))
+pixel = block.bitmasked(ti.ij, (2,2))
pixel.place(x)
@ti.kernel
-def sparse_struct_for():
- x[2, 3] = 2
- x[5, 6] = 3
+def activate():
+ x[2,3] = 1.0
+ x[2,4] = 2.0
+@ti.kernel
+def print_active():
+ for i, j in block:
+ print("Active block", i, j)
for i, j in x:
print('field x[{}, {}] = {}'.format(i, j, x[i, j]))
-
- for i, j in block:
- print('Active block: [{}, {}]'.format(i, j))
-
-print('use_bitmask = {}'.format(use_bitmask))
-sparse_struct_for()
```
-When `bitmask = True`, the program above outputs
-```
-field x[2, 3] = 2
-field x[5, 6] = 3
-Active block: [1, 1]
-Active block: [2, 3]
-```
+The code snippet above also creates an 8x8 sparse grid. The only difference between `bitmasked.py` and `pointer.py` is that the bitmasked SNode replaces the dense SNode (line 3). As shown in the figure below, the active blocks are the same as `pointer.py`. However, the bitmasked pixels in the block are not all activated, because each of them has an activity value.
-When `bitmask = False`, we get
-```
-field x[2, 2] = 0
-field x[2, 3] = 2
-field x[3, 2] = 0
-field x[3, 3] = 0
-field x[4, 6] = 0
-field x[4, 7] = 0
-field x[5, 6] = 3
-field x[5, 7] = 0
-Active block: [1, 1]
-Active block: [2, 3]
-```
+
-When using a `dense` SNode as the leaf block,
-activating `x[2, 3]` also implicitly activates other pixels in `block[1, 1]`, i.e., `x[2, 2]`, `x[3, 2]`, and `x[3, 3]`.
-Without a bitmask, these pixels in the same `block` share the same activity.
+
-### Explicitly manipulating and querying sparsity
+
-Taichi also provides APIs that explicitly manipulates data structure sparsity.
-- Use `ti.is_active(snode, [i, j, ...])` to query if `snode[i, j, ...]` is active or not.
-- `ti.activate/deactivate(snode, [i, j, ...])` to explicitly activate or deactivate a cell of `snode[i, j, ...]`.
-- Use `snode.deactivate_all()` to deactivate all cells of SNode `snode`. This operation also recursively deactivates all its children.
-- Use `ti.deactivate_all_snodes()` to deactivate all cells of all SNodes with sparsity.
-- Use `ti.rescale_index(descendant_snode/field, ancestor_snode, index)` to compute the ancestor index given a descendant index.
+The bitmasked SNodes are like dense SNodes with auxiliary activity values.
+
-Below is an example of these APIs:
+
-```python
-import taichi as ti
+
-ti.init()
+### Dynamic SNode
-x = ti.field(dtype=ti.i32)
-block1 = ti.root.pointer(ti.ij, (4, 4))
-block2 = block1.pointer(ti.ij, (2, 2))
-pixel = block2.dense(ti.ij, (2, 2))
+To support variable-length fields, Taichi provides dynamic SNodes. The code snippet below first creates a 5x1 dense block (line 2). Then each cell of the dense block contains a variable-length dynamic container (line 3). The maximum length of the dynamic container is 5. In the `make_lists()` function, you can use `ti.append()` to add a value to the end of a dynamic SNode. `x.parent()` is the same as `pixel`. The dense field `l` stores the length of each dynamic SNode.
+
+```python {3} title=dynamic.py
+x = ti.field(ti.i32)
+block = ti.root.dense(ti.i, 5)
+pixel = block.dynamic(ti.j, 5)
pixel.place(x)
+l = ti.field(ti.i32)
+ti.root.dense(ti.i, 5).place(l)
@ti.kernel
-def sparse_api_demo():
- ti.activate(block1, [0, 1])
- ti.activate(block2, [1, 2])
+def make_lists():
+ for i in range(5):
+ for j in range(i):
+ ti.append(x.parent(), i, j * j) # ti.append(pixel, i, j * j)
+ l[i] = ti.length(x.parent(), i) # [0, 1, 2, 3, 4]
+```
- for i, j in x:
- print('field x[{}, {}] = {}'.format(i, j, x[i, j]))
- # outputs:
- # field x[2, 4] = 0
- # field x[2, 5] = 0
- # field x[3, 4] = 0
- # field x[3, 5] = 0
- for i, j in block2:
- print('Active block2: [{}, {}]'.format(i, j))
- # output: Active block2: [1, 2]
+
- for i, j in block1:
- print('Active block1: [{}, {}]'.format(i, j))
- # output: Active block1: [0, 1]
+
- for j in range(4):
- print('Activity of block2[2, {}] = {}'.format(j, ti.is_active(block2, [1, j])))
+
- ti.deactivate(block2, [1, 2])
+## Computation on sparse spatial data structures
- for i, j in block2:
- print('Active block2: [{}, {}]'.format(i, j))
- # output: nothing
+### Sparse struct-fors
- for i, j in block1:
- print('Active block1: [{}, {}]'.format(i, j))
- # output: Active block1: [0, 1]
+Efficiently looping over sparse grid cells that distribute irregularly can be challenging, especially on parallel devices such as GPUs.
+In Taichi, *struct-for*s natively support sparse spatial data structures and only loop over currently active pixels with automatic efficient parallelization.
- print(ti.rescale_index(x, block1, ti.Vector([9, 17])))
- # output = [2, 4]
+### Explicitly manipulating and querying sparsity
- # Note: ti.Vector is optional in ti.rescale_index.
- print(ti.rescale_index(x, block1, [9, 17]))
- # output = [2, 4]
+Taichi also provides APIs that explicitly manipulate data structure sparsity. You can manually **check** the activity of a SNode, **activate** a SNode, or **deactivate** a SNode. We now illustrate these functions based on the field defined below.
- ti.activate(block2, [1, 2])
+```python
+x = ti.field(dtype=ti.i32)
+block1 = ti.root.pointer(ti.ij, (3, 3))
+block2 = block1.pointer(ti.ij, (2, 2))
+pixel = block2.bitmasked(ti.ij, (2, 2))
+pixel.place(x)
+```
-sparse_api_demo()
+#### 1. Activity checking
+You can use `ti.is_active(snode, [i, j, ...])` to explicitly query if `snode[i, j, ...]` is active or not.
+```python
@ti.kernel
-def check_activity(snode: ti.template(), i: ti.i32, j: ti.i32):
+def activity_checking(snode: ti.template(), i: ti.i32, j: ti.i32):
print(ti.is_active(snode, [i, j]))
-check_activity(block2, 1, 2) # output = 1
-block2.deactivate_all()
-check_activity(block2, 1, 2) # output = 0
-check_activity(block1, 0, 1) # output = 1
-ti.deactivate_all_snodes()
-check_activity(block1, 0, 1) # output = 0
+for i in range(3):
+ for j in range(3):
+ activity_checking(block1, i, j)
+for i in range(6):
+ for j in range(6):
+ activity_checking(block2, i, j)
+for i in range(12):
+ for j in range(12):
+ activity_checking(pixel, i, j)
+```
+#### 2. Activation
+You can use `ti.activate(snode, [i, j, ...])` to explicitly activate a cell of `snode[i, j, ...]`.
+```python
+@ti.kernel
+def activate_snodes()
+ ti.activate(block1, [1, 0])
+ ti.activate(block2, [3, 1])
+ ti.activate(pixel, [7, 3])
+
+activity_checking(block1, [1, 0]) # output: 1
+activity_checking(block2, [3, 1]) # output: 1
+activity_checking(pixel, [7, 3]) # output: 1
```
+
+
+
+
+
+
+#### 3. Deactivation
+- Use `ti.deactivate(snode, [i, j, ...])` to explicitly deactivate a cell of `snode[i, j, ...]`.
+- Use `snode.deactivate_all()` to deactivate all cells of SNode `snode`. This operation also recursively deactivates all its children.
+- Use `ti.deactivate_all_snodes()` to deactivate all cells of all SNodes with sparsity.
+
+When deactivation happens, the Taichi runtime automatically recycles and zero-fills memory of the deactivated containers.
:::note
For performance reasons, `ti.activate(snode, index)` only activates `snode[index]`.
-The programmer must ensure all ancestor containers of `snode[index]` are already active.
+The programmer must ensure all ancestor containers of `snode[index]` is already active.
Otherwise, this operation results in undefined behavior.
Similarly, `ti.deactivate` ...
- does **not** recursively deactivate all the descendants of a cell.
-- does **not** trigger an deactivation of its parent container, even if all the children of the parent container are deactivated.
+- does **not** trigger deactivation of its parent container, even if all the children of the parent container are deactivated.
:::
-:::note
-When deactivation happens, the Taichi runtime automatically recycles and zero-fills memory of the deactivated containers.
-:::
+#### 4. Ancestor index query
+You can use `ti.rescale_index(descendant_snode/field, ancestor_snode, index)` to compute the ancestor index given a descendant index.
-:::note
-While it is possible to directly use `[i // 2, j // 2]` to compute the `block` index given `pixel` index,
-doing so couples computation code with the internal configuration of data structures (in this case, the size of `block` containers).
+```python
+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]
+```
-Use `ti.rescale_index` to avoid hard-coding internal information of data structures.
-:::
+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.
## Further reading
-Please read our [paper](https://yuanming.taichi.graphics/publication/2019-taichi/taichi-lang.pdf),
-watch the [introduction video](https://www.youtube.com/watch?v=wKw8LMF3Djo), or check out
-the SIGGRAPH Asia 2019 [slides](https://yuanming.taichi.graphics/publication/2019-taichi/taichi-lang-slides.pdf)
-for more details on sparse computation.
+Please read the SIGGRAPH Asia 2019 [paper](https://yuanming.taichi.graphics/publication/2019-taichi/taichi-lang.pdf) or watch the associated
+[introduction video](https://www.youtube.com/watch?v=wKw8LMF3Djo) with [slides](https://yuanming.taichi.graphics/publication/2019-taichi/taichi-lang-slides.pdf)
+for more details on computation of sparse spatial data structures.
-[Taichi elements](https://github.com/taichi-dev/taichi_elements) implement a high-performance
-MLS-MPM solver on Taichi sparse grids.
+[Taichi elements](https://github.com/taichi-dev/taichi_elements) implement a high-performance MLS-MPM solver on Taichi sparse grids.
diff --git a/docs/lang/articles/advanced/sparse_matrix.md b/docs/lang/articles/advanced/sparse_matrix.md
index 9a619991f6e14..5e69c1d7a457b 100644
--- a/docs/lang/articles/advanced/sparse_matrix.md
+++ b/docs/lang/articles/advanced/sparse_matrix.md
@@ -22,7 +22,7 @@ n = 4
K = ti.linalg.SparseMatrixBuilder(n, n, max_num_triplets=100)
@ti.kernel
-def fill(A: ti.linalg.sparse_matrix_builder()):
+def fill(A: ti.types.sparse_matrix_builder()):
for i in range(n):
A[i, i] += 1 # Only += and -= operators are supported for now.
@@ -146,7 +146,7 @@ K = ti.linalg.SparseMatrixBuilder(n, n, max_num_triplets=100)
b = ti.field(ti.f32, shape=n)
@ti.kernel
-def fill(A: ti.linalg.sparse_matrix_builder(), b: ti.template(), interval: ti.i32):
+def fill(A: ti.types.sparse_matrix_builder(), b: ti.template(), interval: ti.i32):
for i in range(n):
A[i, i] += 2.0
@@ -184,5 +184,5 @@ print(f">>>> Computation was successful?: {isSuccess}")
## Examples
Please have a look at our two demos for more information:
-+ [Stable fluid](https://github.com/taichi-dev/taichi/blob/master/examples/simulation/stable_fluid.py): A 2D fluid simulation using a sparse Laplacian matrix to solve Poisson's pressure equation.
-+ [Implicit mass spring](https://github.com/taichi-dev/taichi/blob/master/examples/simulation/implicit_mass_spring.py): A 2D cloth simulation demo using sparse matrices to solve the linear systems.
++ [Stable fluid](https://github.com/taichi-dev/taichi/blob/master/python/taichi/examples/simulation/stable_fluid.py): A 2D fluid simulation using a sparse Laplacian matrix to solve Poisson's pressure equation.
++ [Implicit mass spring](https://github.com/taichi-dev/taichi/blob/master/python/taichi/examples/simulation/implicit_mass_spring.py): A 2D cloth simulation demo using sparse matrices to solve the linear systems.
diff --git a/docs/lang/articles/basic/_category_.json b/docs/lang/articles/basic/_category_.json
index 60b7ef558624d..245518655465e 100644
--- a/docs/lang/articles/basic/_category_.json
+++ b/docs/lang/articles/basic/_category_.json
@@ -1,4 +1,4 @@
{
- "label": "Taichi Language Basic Concepts",
+ "label": "Basic Concepts",
"position": 2
}
diff --git a/docs/lang/articles/basic/differences_between_taichi_and_python_programs.md b/docs/lang/articles/basic/differences_between_taichi_and_python_programs.md
new file mode 100644
index 0000000000000..89de2268c05f5
--- /dev/null
+++ b/docs/lang/articles/basic/differences_between_taichi_and_python_programs.md
@@ -0,0 +1,162 @@
+---
+sidebar_position: 2
+---
+
+# Differences between Taichi and Python programs
+
+Although Taichi uses Python as the frontend, it follows a different set of rules in many aspects, including:
+
+1. [Taichi only supports return statement outside non-static `if`/`for`/`while` scope in the program](#return-statement)
+2. [Variables defined inside an `if`/`for`/`while` block cannot be accessed outside the block.](#variable-scoping)
+3. [Taichi does not fully support some language features of Python.](#unsupportedpartially-supported-python-language-features)
+ - [Set, list, dictionary and operator `in`](#set-list-dictionary-and-operator-in)
+ - [Comprehensions](#comprehensions)
+ - [Operator `is`](#operator-is)
+
+## Return statement and return type annotation
+
+- If a Taichi kernel/function does not have a return statement, it must not have return type annotation.
+- If a Taichi kernel has a return statement, it must have return type annotation.
+- If a Taichi function has a return statement, return type annotation is recommended, and it will be mandatory in the future.
+
+```python {3,7,10,14}
+@ti.kernel
+def error_kernel_no_return_annotation():
+ return 0 # Error: Have return statement but have no return type annotation
+
+@ti.kernel
+def error_kernel_no_return() -> ti.i31: # Error: Have return type annotation but have no return statement
+ pass
+
+@ti.func
+def error_func_no_return() -> ti.i31: # Error: Have return type annotation but have no return statement
+ pass
+```
+
+- The return statement can not be in a scope of non-static `if`/`for`/`while`.
+
+```python {4}
+@ti.kernel
+def error_return_inside_non_static_if(a: ti.i32) -> ti.i32:
+ if a:
+ return 1 # Error: Return statement inside if scope
+```
+
+- The compiler discards code after the first return statement.
+
+```python {4-5}
+@ti.kernel
+def discarded_after_first_return(a: ti.i32) -> ti.i32:
+ return 1
+ if a: # Discarded
+ return 1 # Discarded
+
+discarded_after_first_return(0) # OK: returns 1
+```
+- If there are [compile-time evaluations](/lang/articles/advanced/meta#compile-time-evaluations) in the code, make sure there is a return statement under all circumstances.
+Otherwise, error occurs when a branch is chosen which does not have return statement.
+```python {7-8,15-16,21,23-24}
+@ti.kernel
+def return_inside_static_if(a: ti.template()) -> ti.i32:
+ if ti.static(a):
+ return 1
+ return 0
+
+return_inside_static_if(1) # OK: Returns 1
+return_inside_static_if(0) # OK: Returns 0
+
+@ti.kernel
+def return_inside_static_if_no_return_outside(a: ti.template()) -> ti.i32:
+ if ti.static(a):
+ return 1
+
+return_inside_static_if_no_return_outside(1) # OK: Returns 1
+return_inside_static_if_no_return_outside(0) # Error: No return statement
+
+@ti.kernel
+def ok_return_inside_static_for() -> ti.i32:
+ a = 0
+ for i in ti.static(range(10)): # Static for
+ a += i
+ if ti.static(i == 8): # Static if
+ return a # OK: Returns 36
+```
+
+## Variable scoping
+
+In Python, a variable defined inside an `if`/`for`/`while` block can be accessed outside the block.
+**However**, in Taichi, the variables can only be accessed **within the block it is defined**.
+
+```python {5,13,17,22}
+@ti.kernel
+def error_access_var_outside_for() -> ti.i32:
+ for i in range(10):
+ a = i
+ return a # Error: variable "a" not found
+
+@ti.kernel
+def error_access_var_outside_if(a: ti.i32) -> ti.i32:
+ if a:
+ b = 1
+ else:
+ b = 2
+ return b # Error: variable "b" not found
+
+@ti.kernel
+def ok_define_var_before_if(a: ti.i32) -> ti.i32:
+ b = 0
+ if a:
+ b = 1
+ else:
+ b = 2
+ return b # OK: "b" is defined before "if"
+
+ok_define_var_before_if(0) # Returns 2
+```
+
+## Unsupported/partially supported Python language features
+
+### Set, list, dictionary and operator `in`
+
+Currently, Taichi does not support `set`.
+
+List and dictionary before assigning to a variable works as the python list and dictionary.
+However, after assigning to a variable, the content of the list and the values (not keys) of the dictionary are converted to Taichi variables.
+
+Taichi does not have a runtime implementation of `in` currently. Therefore, operator `in` and `not in` only works in [static scope](/lang/articles/advanced/meta#static-scope) (inside `ti.static()`).
+
+```python {3,11-12,20}
+@ti.kernel
+def list_without_assign() -> ti.i32:
+ if ti.static(1 in [1, 2]): # [1, 2]
+ return 1
+ return 0
+
+list_without_assign() # Returns 1
+
+@ti.kernel
+def list_assigned() -> ti.i32:
+ a = [1, 2] # a: [Variable(1), Variable(2)]
+ if ti.static(1 in a): # 1 is not in [Variable(1), Variable(2)]
+ return 1
+ return 0
+
+list_assigned() # Returns 0
+
+@ti.kernel
+def error_non_static_in():
+ if i in [1, 2]: # Error: Cannot use `in` outside static scope
+ pass
+```
+
+### Comprehensions
+
+Taichi partially supports list comprehension and dictionary comprehension,
+but does not support set comprehension.
+
+For list comprehensions and dictionary comprehensions, the `if`s and `for`s in them are evaluated at compile time.
+The iterators and conditions are implicitly in [static scope](/lang/articles/advanced/meta#static-scope).
+
+### Operator `is`
+
+Currently, Taichi does not support operator `is` and `is not`.
diff --git a/docs/lang/articles/basic/external.md b/docs/lang/articles/basic/external.md
index 2904f40b60577..3e4e70fc07464 100644
--- a/docs/lang/articles/basic/external.md
+++ b/docs/lang/articles/basic/external.md
@@ -1,5 +1,5 @@
---
-sidebar_position: 4
+sidebar_position: 5
---
# Interacting with external arrays
@@ -19,8 +19,8 @@ support NumPy, e.g. `matplotlib`.
```python {8}
@ti.kernel
def my_kernel():
- for i in x:
- x[i] = i * 2
+ for i in x:
+ x[i] = i * 2
x = ti.field(ti.f32, 4)
my_kernel()
@@ -41,12 +41,38 @@ print(x[2]) # 3
print(x[3]) # 5
```
+Likewise, Taichi fields can be **imported from and exported to PyTorch tensors**:
+```python
+@ti.kernel
+def my_kernel():
+ for i in x:
+ x[i] = i * 2
+
+x = ti.field(ti.f32, 4)
+my_kernel()
+x_torch = x.to_torch()
+print(x_torch) # torch.tensor([0, 2, 4, 6])
+
+x.from_numpy(torch.tensor([1, 7, 3, 5]))
+print(x[0]) # 1
+print(x[1]) # 7
+print(x[2]) # 3
+print(x[3]) # 5
+```
+When calling `to_torch()`, specify the PyTorch device where the Taichi field is exported using the `device` argument:
+```python
+x = ti.field(ti.f32, 4)
+x.fill(3.0)
+x_torch = x.to_torch(device="cuda:0")
+print(x_torch.device) # device(type='cuda', index=0)
+```
+
## External array shapes
-Shapes of Taichi fields and those of corresponding NumPy arrays are closely
+Shapes of Taichi fields and those of corresponding NumPy arrays or PyTorch tensors are closely
connected via the following rules:
-- For scalar fields, **the shape of NumPy array is exactly the same as
+- For scalar fields, **the shape of NumPy array or PyTorch tensor equals the shape of
the Taichi field**:
```python
@@ -60,7 +86,7 @@ field.from_numpy(array) # the input array must be of shape (256, 512)
```
- For vector fields, if the vector is `n`-D, then **the shape of NumPy
- array should be** `(*field_shape, vector_n)`:
+ array or Pytorch tensor should be** `(*field_shape, vector_n)`:
```python
field = ti.Vector.field(3, ti.i32, shape=(256, 512))
@@ -74,7 +100,7 @@ field.from_numpy(array) # the input array must be of shape (256, 512, 3)
```
- For matrix fields, if the matrix is `n`-by-`m` (`n x m`), then **the shape of NumPy
-array should be** `(*field_shape, matrix_n, matrix_m)`:
+array or Pytorch Tensor should be** `(*field_shape, matrix_n, matrix_m)`:
```python
field = ti.Matrix.field(3, 4, ti.i32, shape=(256, 512))
@@ -88,7 +114,8 @@ array.shape # (256, 512, 3, 4)
field.from_numpy(array) # the input array must be of shape (256, 512, 3, 4)
```
-- For struct fields, the external array will be exported as **a dictionary of arrays** with the keys being struct member names and values being struct member arrays. Nested structs will be exported as nested dictionaries:
+- For struct fields, the external array will be exported as **a dictionary of NumPy arrays or PyTorch tensors** with keys
+being struct member names and values being struct member arrays. Nested structs will be exported as nested dictionaries:
```python
field = ti.Struct.field({'a': ti.i32, 'b': ti.types.vector(float, 3)} shape=(256, 512))
@@ -104,7 +131,7 @@ field.from_numpy(array_dict) # the input array must have the same keys as the fi
## Using external arrays as Taichi kernel arguments
-Use the type hint `ti.ext_arr()` for passing external arrays as kernel
+Use type hint `ti.ext_arr()` or `ti.any_arr()` to pass external arrays as kernel
arguments. For example:
```python {10}
@@ -135,3 +162,31 @@ for i in range(n):
for j in range(m):
assert a[i, j] == i * j + i + j
```
+
+Note that the elements in an external array must be indexed using a single square bracket.
+This contrasts with a Taichi vector or matrix field where field and matrix indices are indexed separately:
+```python
+@ti.kernel
+def copy_vector(x: ti.template(), y: ti.ext_arr()):
+ for i, j in ti.ndrange(n, m):
+ for k in ti.static(range(3)):
+ y[i, j, k] = x[i, j][k] # correct
+ # y[i][j][k] = x[i, j][k] incorrect
+ # y[i, j][k] = x[i, j][k] incorrect
+```
+Also, external arrays in a Taichi kernel are indexed using its **physical memory layout**. For PyTorch users,
+this implies that the PyTorch tensor [needs to be made contiguous](https://pytorch.org/docs/stable/generated/torch.Tensor.contiguous.html)
+before being passed into a Taichi kernel:
+```python
+@ti.kernel
+def copy_scalar(x: ti.template(), y: ti.ext_arr()):
+ for i, j in x:
+ y[i, j] = x[i, j]
+
+x = ti.field(dtype=int, shape=(3, 3))
+y = torch.Tensor([[1, 2, 3], [4, 5, 6], [7, 8, 9]])
+y = y.T # Transposing the tensor returns a view of the tensor which is not contiguous
+copy(x, y) # error!
+copy(x, y.clone()) # correct
+copy(x, y.contiguous()) # correct
+```
diff --git a/docs/lang/articles/basic/field.md b/docs/lang/articles/basic/field.md
index 592ca50aa3a45..d0779f3592d0c 100644
--- a/docs/lang/articles/basic/field.md
+++ b/docs/lang/articles/basic/field.md
@@ -1,102 +1,168 @@
---
-sidebar_position: 3
+sidebar_position: 4
---
# Fields
+Taichi fields are used to store data.
+In general, fields are global data containers that can be read and written from both the Python scope and the Taichi scope.
-Fields are **global** variables provided by Taichi. **Global** indicates that fields can be read/written from both the Python scope and the Taichi scope. A field can be considered as a multi-dimensional array of elements, and it can be either **dense** or **sparse**. Similar to a NumPy `ndarray` object, a field has a data type and a shape. Moreover, an element of a field can be a scalar, a **vector**, a **matrix**, or a **struct**.
+A field has its own data type and shape and can be considered as a multi-dimensional array of elements.
+An element of a field can be a **scalar**, a **vector**, a **matrix**, or a **struct**.
+The sparsity of a field element is **dense** by default, but it can also be **sparse**, as detailed described in [Sparse spatial data structures](/lang/articles/advanced/sparse).
-The term **field** is borrowed from mathematics and physics. If you
-have already known [scalar field](https://en.wikipedia.org/wiki/Scalar_field) (e.g., heat field) or vector field (e.g., [gravitational field](https://en.wikipedia.org/wiki/Gravitational_field)) in mathematics and physics, it will be straightforward to understand the fields in Taichi.
-
-To be noticed:
-* Fields are always accessed by indices.
-* Field values are initially zero.
-* Sparse fields are initially inactive.
-
-:::tip
-In earlier versions of Taichi, you could not allocate new fields after executing the first kernel. Since Taichi v0.8.0, you can use a new class `FieldsBuilder` for dynamic field allocation and destruction. For more details, please see [Field (advanced)](/lang/articles/advanced/layout).
+:::note
+The term **field** is borrowed from mathematics and physics.
+If you have already known [scalar field](https://en.wikipedia.org/wiki/Scalar_field) (e.g., heat field) or vector field (e.g., [gravitational field](https://en.wikipedia.org/wiki/Gravitational_field)) in mathematics and physics,
+it will be straightforward to understand the fields in Taichi.
:::
## Scalar fields
+We start introducing fields from this very basic type, the elements of scalar fields are simply scalars.
+* A 0D scalar field is a single scalar.
+* A 1D scalar field is an array.
+* A 2D scalar field can be used to represent a 2D regular grid of values.
+* A 3D scalar field can be used for volumetric data.
-A simple example might help you understand scalar fields. Assume you have a rectangular wok on the top of a fire. At each point of the wok, there would be a temperature. The surface of the wok forms a heat field. The width and height of the wok are similar to the `shape` of the Taichi scalar field. The temperature (0-D scalar) is like the element of the Taichi scalar field. We could use the following field to represent the
-heat field on the wok:
-
+### Declaration
``` python
-heat_field = ti.field(dtype=ti.f32, shape=(width_wok, height_wok))
+import taichi as ti
+ti.init(arch=ti.cpu)
+
+energy = ti.field(ti.f32, shape=()) # 0-D
+linear_array = ti.field(ti.i32, shape=128) # 1-D
+gray_scale_image = ti.field(ti.u8, shape=(640, 480)) # 2-D
+volumetric_data = ti.field(ti.f32, shape=(32, 32, 32)) # 3-D
```
### Access elements of scalar fields
-- If `x` is a 3D scalar field (`ti.field(dtype=ti.f32, shape=(10, 20, 30)`), access its element with `x[i, j, k]` (`0 <= i < 10, 0 <= j < 20, 0 <= k < 30`).
-- When accessing 0-D field `x`, use `x[None] = 0` instead of `x = 0`. A 0-D field looks like `energy = ti.field(dtype=ti.f32, shape=())`.
+``` python
+energy[None] = 10.0
+linear_array[0] = 1
+gray_scale_image[1,2] = 255
+volumetric_data[3,3,3] = 2.0
+```
-:::caution
-Please **always** use indexing to access entries in fields.
+### Meta data
+``` python
+linear_array.shape # (128,)
+volumetric_data.dtype # f32
+```
+
+:::note
+* Field values are initially zero.
+* Fields are **always** accessed by indices. When accessing 0-D field `x`, use `x[None] = 0` instead of `x = 0`.
+:::
+
+### Example
+An example might help you understand scalar fields.
+Assume you have a gray-scale image. At each point in the image, there would be a pixel value. The width and height of the image are similar to the `shape` of the Taichi scalar field. The pixel value (0-D scalar) is like the element of the Taichi scalar field. We could use the following code to generate a gray-scale image with random pixel values:
+
+``` python {5}
+import taichi as ti
+
+ti.init(arch=ti.cpu)
+width, height = 640,480
+gray_scale_image = ti.field(dtype=ti.f32, shape=(width, height))
+
+@ti.kernel
+def fill_image():
+ for i,j in gray_scale_image:
+ gray_scale_image[i,j] = ti.random()
+
+fill_image()
+
+gui = ti.GUI('gray-scale image with random values', (width, height))
+while gui.running:
+ gui.set_image(gray_scale_image)
+ gui.show()
+```
+
+:::tip
+In earlier versions of Taichi, you could not allocate new fields after executing the first kernel. Since Taichi v0.8.0, you can use a new class `FieldsBuilder` for dynamic field allocation and destruction. For more details, please see [Field (advanced)](/lang/articles/advanced/layout).
:::
## Vector fields
-We are all live in a gravitational field which is a vector field. At each position of the 3D space, there is a gravity force vector. The gravitational field could be represented with:
+We are all living in a gravitational field, which is a vector field. At each position in 3D space, there is a gravity force vector. The gravitational field could be represented by:
```python
gravitational_field = ti.Vector.field(n=3, dtype=ti.f32, shape=(x, y, z))
```
`x, y, z` are the sizes of each dimension of the 3D space respectively. `n` is the number of elements of the gravity force vector.
### Access elements of vector fields
+There are **two** indexing operators `[]` when you access a member of a vector field: the first is for field indexing, and the second is for vector indexing.
- The gravity force vector could be accessed by `gravitational_field[i, j, k]` (`0 <= i < x, 0 <= j < y, 0 <= k < z`).
- The `p`-th member of the gravity force vector could be accessed by `gravitational_field[i, j, k][p]` (`0 <= p < n`).
- The 0-D vector field `x = ti.Vector.field(n=3, dtype=ti.f32, shape=())` should be accessed by `x[None][p]` (`0 <= p < n`).
-:::note
-As you may have noticed, there are **two** indexing operators `[]` when you access a member of a vector from a vector field: the first is for field indexing, and the second is for vector indexing.
-:::
+### Example
+This example helps you understand how to access vector fields:
+``` python
+import taichi as ti
+ti.init(arch=ti.cpu)
-## Matrix fields
+n,w,h = 3,128,64
+vec_field = ti.Vector.field(n, dtype=ti.f32, shape=(w,h))
+
+@ti.kernel
+def fill_vector():
+ for i,j in vec_field:
+ for k in ti.static(range(n)):
+ #ti.static unrolls the inner loops
+ vec_field[i,j][k] = ti.random()
+
+fill_vector()
+print(vec_field[w-1,h-1][n-1])
+```
+## Matrix fields
Field elements can also be matrices. In continuum mechanics, each
infinitesimal point in a material exists a strain and a stress tensor. The strain and stress tensor is a 3 by 3 matrix in the 3D space. To represent this tensor field we could use:
```python
strain_tensor_field = ti.Matrix.field(n=3, m=3, dtype=ti.f32, shape=(x, y, z))
```
-
`x, y, z` are the sizes of each dimension of the 3D material respectively. `n, m` are the dimensions of the strain tensor.
In a general case, suppose you have a `128 x 64` field called `A`, and each element is
a `3 x 2` matrix, you can define it with `A = ti.Matrix.field(3, 2, dtype=ti.f32, shape=(128, 64))`.
### Access elements of matrix fields
-- If you want to get the matrix of grid node `i, j`, please use
- `mat = A[i, j]`. `mat` is simply a `3 x 2` matrix.
-- To get the element on the first row and second column of that
- matrix, use `mat[0, 1]` or `A[i, j][0, 1]`.
+There are **two** indexing operators `[]` when you access a member of a matrix from a matrix field:
+the first is for field indexing, and the second is for matrix indexing.
+- If you want to get the element `i, j` of the matrix field, please use `mat = A[i, j]`. `mat` is simply a `3 x 2` matrix.
+- To get the member on the first row and second column of that element `mat`, use `mat[0, 1]` or `A[i, j][0, 1]`.
- The 0-D matrix field `x = ti.Matrix.field(n=3, m=4, dtype=ti.f32, shape=())` should be accessed by `x[None][p, q]` (`0 <= p < n, 0 <= q < m`).
-
-:::note
-- As you may have noticed, there are **two** indexing operators `[]`
- when you access a member of a matrix from a matrix field: the
- first is for field indexing, and the second is for matrix indexing.
- `ti.Vector` is simply an alias of `ti.Matrix`.
-:::
-### Matrix size
+### Example
+This example helps you understand element and member in matrix fields:
+``` python
+matrix_field = ti.Matrix.field(n = 2, m = 3, dtype = ti.f32, shape = (2, 2))
+Element = matrix_field[0, 0]
+Member = matrix_field[0, 1][1,1]
+```
+
-For performance reasons matrix operations will be unrolled during the compile stage, therefore we
-suggest using only small matrices. For example, `2x1`, `3x3`, `4x4`
+### Matrix size
+For performance reasons, matrix operations will be unrolled during the compile stage.
+Therefore we suggest using only small matrices. For example, `2x1`, `3x3`, `4x4`
matrices are fine, yet `32x6` is probably too big as a matrix size.
-:::caution
-Due to the unrolling mechanisms, operating on large matrices (e.g.
-`32x128`) can lead to a very long compilation time and low performance.
-:::
-
If you have a dimension that is too large (e.g. `64`), it's better to
declare a field of size `64`. E.g., instead of declaring
`ti.Matrix.field(64, 32, dtype=ti.f32, shape=(3, 2))`, declare
`ti.Matrix.field(3, 2, dtype=ti.f32, shape=(64, 32))`. Try to put large
dimensions to fields instead of matrices.
+:::caution
+Due to the unrolling mechanism, operating on large matrices (e.g.
+`32x128`) can lead to a very long compilation time and low performance.
+:::
+
## Struct fields
-In addition to vectors and matrices, field elements can be user-defined structs. A struct variable may contain scalars, vectors/matrices, or other structs as its members. A struct field is created by providing a dictionary of the name and data type of each member. For example, a 1D field of particles with position, velocity, acceleration, and mass for each particle can be represented as:
+Field elements can be user-defined structs.
+Struct fields are created by providing the name and data type of each member variable in a dictionary format.
+Member variables of struct fields might be scalars, vectors, matrices, or other struct fields.
+For example, a 1-D field of particles with position, velocity, acceleration, and mass can be declared as:
```python
particle_field = ti.Struct.field({
"pos": ti.types.vector(3, ti.f32),
@@ -105,7 +171,8 @@ particle_field = ti.Struct.field({
"mass": ti.f32,
}, shape=(n,))
```
-[Compound types](type.md#compound-types) (`ti.types.vector`, `ti.types.matrix`, and `ti.types.struct`) need to be used to create vectors, matrices, or structs as field members. Apart from using `ti.Struct.field`, the above particle field can be alternatively created using field creation from compound types as:
+
+[Compound types](type.md#compound-types) (`ti.types.vector`, `ti.types.matrix`, and `ti.types.struct`) are used to declare vectors, matrices, or structs as field members. Apart from using `ti.Struct.field`, the above particle field can also be declared by using the field of compound types:
```python
vec3f = ti.types.vector(3, ti.f32)
particle = ti.types.struct(
@@ -113,6 +180,7 @@ particle = ti.types.struct(
)
particle_field = particle.field(shape=(n,))
```
+
Members of a struct field can be accessed either locally (i.e., member of a struct field element) or globally (i.e., member field of a struct field):
```python
# set the position of the first particle to origin
diff --git a/docs/lang/articles/basic/operator.md b/docs/lang/articles/basic/operator.md
new file mode 100644
index 0000000000000..1fb675cb86535
--- /dev/null
+++ b/docs/lang/articles/basic/operator.md
@@ -0,0 +1,325 @@
+---
+sidebar_position: 4
+---
+
+# Operators
+Here we present the supported operators in Taichi for both primitive types and
+compound types such as matrices.
+
+## Supported operators for primitive types
+### Arithmetic operators
+
+| Operation | Result |
+| --------- | ------------------------------- |
+| `-a` | `a` negated |
+| `+a` | `a` unchanged |
+| `a + b` | sum of `a` and `b` |
+| `a - b` | difference of `a` and `b` |
+| `a * b` | product of `a` and `b` |
+| `a / b` | quotient of `a` and `b` |
+| `a // b` | floored quotient of `a` and `b` |
+| `a % b` | remainder of `a / b` |
+| `a ** b` | `a` to the power of `b` |
+
+:::note
+
+The `%` operator in Taichi follows the Python style instead of C style,
+e.g.,
+
+```python
+# In Taichi-scope or Python-scope:
+print(2 % 3) # 2
+print(-2 % 3) # 1
+```
+
+For C-style mod (`%`), please use `ti.raw_mod`. This function also receives floating points as arguments.
+
+`ti.raw_mod(a, b)` returns `a - b * int(float(a) / b)`.
+
+```python
+print(ti.raw_mod(2, 3)) # 2
+print(ti.raw_mod(-2, 3)) # -2
+print(ti.raw_mod(3.5, 1.5)) # 0.5
+```
+:::
+
+:::note
+
+Python3 distinguishes `/` (true division) and `//` (floor division), e.g., `1.0 / 2.0 = 0.5`, `1 / 2 = 0.5`, `1 // 2 = 0`,
+`4.2 // 2 = 2`. Taichi follows the same design:
+
+- **True divisions** on integral types first cast their
+ operands to the default floating point type.
+- **Floor divisions** on floating point types first cast their
+ operands to the default integral type.
+
+To avoid such implicit casting, you can manually cast your operands to
+desired types, using `ti.cast`. Please see
+[Default precisions](#default-precisions) for more details on
+default numerical types.
+
+Taichi also provides `ti.raw_div` function which performs true division if one of the operands is floating point type
+and performs floor division if both operands are integral types.
+
+```python
+print(ti.raw_div(5, 2)) # 2
+print(ti.raw_div(5, 2.0)) # 2.5
+```
+
+:::
+
+
+### Comparison operators
+
+| Operation | Result |
+| ------------------ | ------------------------------------------------------------- |
+| `a == b` | if `a` is equal to `b`, then True, else False |
+| `a != b` | if `a` is not equal to `b`, then True, else False |
+| `a > b` | if `a` is strictly greater than `b`, then True, else False |
+| `a < b` | if `a` is strictly less than `b`, then True, else False |
+| `a >= b` | if `a` is greater than or equal to `b`, then True, else False |
+| `a <= b` | if `a` is less than or equal to `b`, then True, else False |
+
+### Logical operators
+
+| Operation | Result |
+| ------------------ | ------------------------------------------------------------- |
+| `not a` | if `a` is False, then True, else False |
+| `a or b` | if `a` is False, then `b`, else `a` |
+| `a and b` | if `a` is False, then `a`, else `b` |
+
+### Conditional operations
+
+The result of conditional expression `a if cond else b` is `a` if `cond` is True, or `b` otherwise.
+`a` and `b` must have a same type.
+
+The conditional expression does short-circuit evaluation, which means the branch not chosen is not evaluated.
+
+```python
+a = ti.field(ti.i32, shape=(10,))
+for i in range(10):
+ a[i] = i
+
+@ti.kernel
+def cond_expr(ind: ti.i32) -> ti.i32:
+ return a[ind] if ind < 10 else 0
+
+cond_expr(3) # returns 3
+cond_expr(10) # returns 0, a[10] is not evaluated
+```
+
+
+For element-wise conditional operations on Taichi vectors and matrices,
+Taichi provides `ti.select(cond, a, b)` which **does not** do short-circuit evaluation.
+```python {4}
+cond = ti.Vector([1, 0])
+a = ti.Vector([2, 3])
+b = ti.Vector([4, 5])
+ti.select(cond, a, b) # ti.Vector([2, 5])
+```
+
+### Bitwise operators
+
+| Operation | Result |
+| ----------------------- | ----------------------------------- |
+| `~a` | the bits of `a` inverted |
+| `a & b` | bitwise and of `a` and `b` |
+| `a ^ b` | bitwise exclusive or of `a` and `b` |
+| a | b | bitwise or of `a` and `b` |
+| `a << b` | left-shift `a` by `b` bits |
+| `a >> b` | right-shift `a` by `b` bits |
+
+:::note
+
+The `>>` operation denotes the
+[Shift Arithmetic](https://en.wikipedia.org/wiki/Arithmetic_shift) Right (SAR) operation.
+For the [Shift Logical](https://en.wikipedia.org/wiki/Logical_shift) Right (SHR) operation,
+consider using `ti.bit_shr()`. For left shift operations, SAL and SHL are the
+same.
+
+
+:::
+
+### Trigonometric functions
+
+```python
+ti.sin(x)
+ti.cos(x)
+ti.tan(x)
+ti.asin(x)
+ti.acos(x)
+ti.atan2(x, y)
+ti.tanh(x)
+```
+
+### Other arithmetic functions
+
+```python
+ti.sqrt(x)
+ti.rsqrt(x) # A fast version for `1 / ti.sqrt(x)`.
+ti.exp(x)
+ti.log(x)
+ti.round(x)
+ti.floor(x)
+ti.ceil(x)
+ti.sum(x)
+ti.max(x, y, ...)
+ti.min(x, y, ...)
+ti.abs(x) # Same as `abs(x)`
+ti.pow(x, y) # Same as `pow(x, y)` and `x ** y`
+```
+
+### Builtin-alike functions
+
+```python
+abs(x) # Same as `ti.abs(x, y)`
+pow(x, y) # Same as `ti.pow(x, y)` and `x ** y`.
+```
+
+### Random number generator
+
+```python
+ti.random(dtype=float)
+```
+
+:::note
+
+`ti.random` supports `u32`, `i32`, `u64`, `i64`, and all floating point types.
+The range of the returned value is type-specific.
+
+| Type | Range |
+| --- | --- |
+| i32 | -2,147,483,648 to 2,147,483,647 |
+| u32 | 0 to 4,294,967,295 |
+| i64 | -9,223,372,036,854,775,808 to 9,223,372,036,854,775,807 |
+| u64 | 0 to 18,446,744,073,709,551,615 |
+| floating point | 0.0 to 1.0 |
+
+:::
+
+### Supported atomic operations
+
+In Taichi, augmented assignments (e.g., `x[i] += 1`) are automatically
+[atomic](https://en.wikipedia.org/wiki/Fetch-and-add).
+
+:::caution
+
+When modifying global variables in parallel, make sure you use atomic
+operations. For example, to sum up all the elements in `x`,
+
+```python
+@ti.kernel
+def sum():
+ for i in x:
+ # Approach 1: OK
+ total[None] += x[i]
+
+ # Approach 2: OK
+ ti.atomic_add(total[None], x[i])
+
+ # Approach 3: Wrong result since the operation is not atomic.
+ total[None] = total[None] + x[i]
+```
+:::
+
+:::note
+
+When atomic operations are applied to local values, the Taichi compiler
+will try to demote these operations into their non-atomic counterparts.
+:::
+
+Apart from the augmented assignments, explicit atomic operations, such
+as `ti.atomic_add`, also do read-modify-write atomically. These
+operations additionally return the **old value** of the first argument.
+For example,
+
+```python
+x[i] = 3
+y[i] = 4
+z[i] = ti.atomic_add(x[i], y[i])
+# now x[i] = 7, y[i] = 4, z[i] = 3
+```
+
+Below is a list of all explicit atomic operations:
+
+| Operation | Behavior |
+| --------------------- | ---------------------------------------------------------------------------------------------------- |
+| `ti.atomic_add(x, y)` | atomically compute `x + y`, store the result in `x`, and return the old value of `x` |
+| `ti.atomic_sub(x, y)` | atomically compute `x - y`, store the result in `x`, and return the old value of `x` |
+| `ti.atomic_and(x, y)` | atomically compute `x & y`, store the result in `x`, and return the old value of `x` |
+| `ti.atomic_or(x, y)` | atomically compute x | y, store the result in `x`, and return the old value of `x` |
+| `ti.atomic_xor(x, y)` | atomically compute `x ^ y`, store the result in `x`, and return the old value of `x` |
+| `ti.atomic_max(x, y)` | atomically compute `max(x, y)`, store the result in `x`, and return the old value of `x` |
+| `ti.atomic_min(x, y)` | atomically compute `min(x, y)`, store the result in `x`, and return the old value of `x` |
+
+:::note
+
+Supported atomic operations on each backend:
+
+| type | CPU | CUDA | OpenGL | Metal | C source |
+| ---- | ---- | ---- | ------ | ----- | -------- |
+| i32 |:heavy_check_mark:|:heavy_check_mark:|:heavy_check_mark:|:heavy_check_mark:|:heavy_check_mark:|
+| f32 |:heavy_check_mark:|:heavy_check_mark:|:heavy_check_mark:|:heavy_check_mark:|:heavy_check_mark:|
+| i64 |:heavy_check_mark:|:heavy_check_mark:|:large_orange_diamond:|:x:|:heavy_check_mark:|
+| f64 |:heavy_check_mark:|:heavy_check_mark:|:large_orange_diamond:|:x:|:heavy_check_mark:|
+
+(:large_orange_diamond: requires extension)
+:::
+
+
+## Supported operators for matrices
+
+The previously mentioned operations on primitive types can also be applied on
+compound types such as matrices.
+In these cases, they are applied in an element-wise manner. For example:
+
+```python
+B = ti.Matrix([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]])
+C = ti.Matrix([[3.0, 4.0, 5.0], [6.0, 7.0, 8.0]])
+
+A = ti.sin(B)
+# is equivalent to
+for i in ti.static(range(2)):
+ for j in ti.static(range(3)):
+ A[i, j] = ti.sin(B[i, j])
+
+A = B ** 2
+# is equivalent to
+for i in ti.static(range(2)):
+ for j in ti.static(range(3)):
+ A[i, j] = B[i, j] ** 2
+
+A = B ** C
+# is equivalent to
+for i in ti.static(range(2)):
+ for j in ti.static(range(3)):
+ A[i, j] = B[i, j] ** C[i, j]
+
+A += 2
+# is equivalent to
+for i in ti.static(range(2)):
+ for j in ti.static(range(3)):
+ A[i, j] += 2
+
+A += B
+# is equivalent to
+for i in ti.static(range(2)):
+ for j in ti.static(range(3)):
+ A[i, j] += B[i, j]
+```
+
+In addition, the following methods are supported matrices operations:
+
+```python
+a = ti.Matrix([[2, 3], [4, 5]])
+a.transpose() # the transposed matrix of `a`, will not effect the data in `a`.
+a.trace() # the trace of matrix `a`, the returned scalar value can be computed as `a[0, 0] + a[1, 1] + ...`.
+a.determinant() # the determinant of matrix `a`.
+a.inverse() # (ti.Matrix) the inverse of matrix `a`.
+a@a # @ denotes matrix multiplication
+```
+
+:::note
+For now, determinant() and inverse() only works in Taichi-scope, and the
+size of the matrix must be 1x1, 2x2, 3x3 or 4x4.
+:::
diff --git a/docs/lang/articles/basic/overview.md b/docs/lang/articles/basic/overview.md
index 80bb9d7c9abd1..c558e06e0ae22 100644
--- a/docs/lang/articles/basic/overview.md
+++ b/docs/lang/articles/basic/overview.md
@@ -2,22 +2,51 @@
sidebar_position: 0
---
-# Why new programming language
+# Why a new programming language
-Taichi is a high-performance programming language for computer graphics
-applications. The design goals are
+Imagine you'd like to write a new particle-based fluid algorithm. You started simple, didn't spend much time before finding a reference C++/CUDA work online (or derived the work from your labmate, unfortunately). `cmake .. && make`, you typed. Oops, cmake threw out an error due to a random incompatible third party library. Installed and rebuilt, now it passed. Then you ran it, which immediately segfaulted (without any stacktrace, of course). Then you started gazing at the code, placed the necessary asset files at the right place, fixed a few dangling pointers and reran. It... actually worked, until you plugged in your revised algorithm. Now another big fight with the GPU or CPU code. More often than not, you get lost in the language details.
-- Productivity
-- Performance
-- Portability
-- Spatially sparse computation
-- Differentiable programming
-- Metaprogramming
+If all these sound too familiar to you, congratulations, you are probably looking at the right solution.
-## Design decisions
+Born from the MIT CSAIL lab, Taichi was designed to facilitate computer graphics researchers' everyday life, by helping them quickly implement visual computing and physics simulation algorithms that are executable on GPU. The path Taichi took was an innovative one: Taichi is embedded in Python and uses modern just-in-time (JIT) frameworks (for example LLVM, SPIR-V) to offload the Python source code to native GPU or CPU instructions, offering the performance at both development time and runtime.
-- Decouple computation from data structures
-- Domain-specific compiler optimizations
-- Megakernels
-- Two-scale automatic differentiation
-- Embedding in Python
+To be fair, a domain-specific language (DSL) with a Python frontend is not something new. In the past few years, frameworks like Halide, PyTorch, and TVM have matured into the de facto standards in areas such as image processing and deep learning (DL). What distinguishes Taichi the most from these frameworks is its imperative programming paradigm. As a DSL, Taichi is not so specialized in a particular computing pattern. This provides better flexibility. While one may argue that flexibility usually comes at the cost of not being fully optimized, we often find this not the case for a few reasons:
+
+* Taichi's workload typically does *not* exhibit an exploitable pattern (e.g., element-wise operations), meaning that the arithmetic intensity is bounded anyway. By simply switching to the GPU backend, one can already enjoy a nice performance gain.
+* Unlike the traditional DL frameworks, where operators are simple math expressions and have to be fused at the graph level to achieve higher arithmetic intensity, Taichi's imperative paradigm makes it quite easy to write a large amount of computation in a single kernel. We call it *mega-kernel*.
+* Taichi heavily optimizes the source code using various compiler technologies: common subexpression elimination, dead code elimination, control flow graph analysis, etc. These optimizations are backend neutral, because Taichi hosts its own intermediate representation (IR) layer.
+* JIT compilation provides additional optimization opportunities.
+
+That said, Taichi goes beyond a Python JIT transpiler. One of the initial design goals is to *decouple the computation from the data structures*. The mechanism that Taichi provides is a set of generic data containers, called *SNode* (/ˈsnoʊd/). SNodes can be used to compose hierarchical, dense or sparse, multi-dimensional fields conveniently. Switching between array-of-structures and structure-of-arrays layouts is usually a matter of ≤10 lines of code. This has sparked many use cases in numerical simulation. If you are interested to learn them, please check out [Fields (advanced)](https://docs.taichi.graphics/lang/articles/advanced/layout), [Sparse spatial data structures](https://docs.taichi.graphics/lang/articles/advanced/sparse), or [the original Taichi paper](https://yuanming.taichi.graphics/publication/2019-taichi/taichi-lang.pdf).
+
+The concept of decoupling is further extended to the type system. With GPU memory capacity and bandwidth becoming the major bottlenecks nowadays, it is vital to be able to pack more data per memory unit. Since 2021, Taichi has introduced customizable quantized types, allowing for the definition of fixed point or floating point numbers with arbitrary bits (still needs to be under 64). This has allowed an MPM simulation of over 400 million particles on a single GPU device. Learn more details in [the QuanTaichi paper](https://yuanming.taichi.graphics/publication/2021-quantaichi/quantaichi.pdf).
+
+Taichi is intuitive. If you know Python, you know Taichi. If you write Taichi, you awaken your GPU (or CPU as a fallback). Ever since its debut, this simple idea has gained so much popularity, that many were attracted to contribute new backends, including Vulkan, OpenGL and DirectX (working in progress). Without our strong and dedicated community, Taichi would never have been where it is now.
+
+Going forward, we see many new opportunities lying ahead, and would like to share some of our vision with you.
+
+**Academia**
+
+90% of the research code will be trashed due to the nature of research where assumptions keep being broken and ideas keep being iterated. Swiftly coding without thinking too much about performance may lead to incorrect conclusions, while pre-matured code optimization can be a waste of time and often produces a tangled mess. The high performance and productivity are, therefore, extremely helpful for research projects.
+
+Taichi will keep embracing the academia. The key features we have (or plan to have) for high-performance computing research projects include small-scale linear algebra (inside kernels), large-scale sparse systems, and efficient neighbor accessing for both structured and unstructured data.
+
+Taichi also provides an automatic differentiation module via source code transformation (at IR level), making it a sweet differentiable simulation tool for machine learning projects.
+
+**Apps & game engine integration**
+
+One huge advantange of Taichi lies in its portability, thanks to the support for a wide variety of backends. During the development process, we have also recognized the increasing demands from our industry users for multi-platform packaging and deployment. Below shows an experimental demo of integrating Taichi with Unity. By exporting Taichi kernels as SPIR-V shaders, we can easily import them into a Unity project.
+
+
+
+**General-purpose computing**
+
+While originally designed for physics simulation, Taichi has found its application in many other areas that can be boosted by GPU general-purpose computing.
+
+* [taichimd](https://github.com/victoriacity/taichimd): Interactive, GPU-accelerated Molecular (& Macroscopic) Dynamics using the Taichi programming language
+* [TaichiSLAM](https://github.com/xuhao1/TaichiSLAM): a 3D Dense mapping backend library of SLAM based Taichi-Lang, designed for the aerial swarm.
+* [Stannum](https://github.com/ifsheldon/stannum): Fusing Taichi into PyTorch.
+
+**Maybe a new frontend?**
+
+The benefit of adopting the compiler approach is that you can decouple the frontend from the backend. Taichi is *currently* embedded in Python, but who says it needs to stay that way? Stay tuned :-)
diff --git a/docs/lang/articles/basic/syntax.md b/docs/lang/articles/basic/syntax.md
index c61311e10d781..afbd33300e466 100644
--- a/docs/lang/articles/basic/syntax.md
+++ b/docs/lang/articles/basic/syntax.md
@@ -4,6 +4,29 @@ sidebar_position: 1
# Kernels and functions
+Taichi has two types of functions: Taichi kernels, and Taichi functions.
+
+Scope inside Taichi kernels and Taichi functions is called Taichi scope, and scope outside them is called Python scope.
+
+A Taichi kernel is the entrypoint of a Taichi program, and it is similar to a `__global__` function in CUDA. It can only be called inside Python scope.
+
+A Taichi function can only be called inside Taichi scope, and it is similar to a `__device__` function in CUDA.
+
+Major differences between Taichi kernels and Taichi functions are listed in the table below.
+
+| | Taichi kernels | Taichi functions |
+| :--- | :--- | :--- |
+| Can be called in | Python scope | Taichi scope |
+| Argument type annotation | Mandatory | Recommended |
+| Return type annotation | Mandatory| Recommended |
+| Return value | Scalar/Vector/Matrix | Arbitrary |
+| Max number of total elements in arguments | 8 (for OpenGL and CC) or 64 (other) | Unlimited |
+| Max number of return values in a return statement | 1 | Unlimited |
+| Max number of total elements in return values | 30 | Unlimited |
+
+
+
+
## Taichi-scope vs Python-scope
Code decorated by `@ti.kernel` or `@ti.func` is in the **Taichi-scope**.
@@ -45,7 +68,6 @@ For people from CUDA, Taichi kernels are similar to `__global__` functions.
Kernels can have multiple arguments, which support passing values from Python-scope to Taichi-scope conveniently.
:::note
-For kernels executed on OpenGL and CC backends, the number of arguments is limited to 8.
:::
Kernel arguments must be type hinted:
@@ -59,18 +81,22 @@ my_kernel(24, 3.2) # prints: 27.2
```
:::note
+Taichi supports scalars, `ti.Matrix` and`ti.Vector` as kernel arguments.
+The total number of elements in kernel arguments must not exceed 8 on OpenGL and CC backends, or 64 on other backends.
+The number of elements in a scalar argument is 1, and the number of elements in a `ti.Matrix` or`ti.Vector` is the number of elements inside it.
-For now, Taichi supports scalars as kernel arguments. Specifying `ti.Matrix` or
-`ti.Vector` as an argument is not supported yet:
-
-```python {2,7}
+```python {2,7,11}
@ti.kernel
-def valid_kernel(vx: ti.f32, vy: ti.f32):
+def valid_scalar_argument(vx: ti.f32, vy: ti.f32):
v = ti.Vector([vx, vy])
...
@ti.kernel
-def error_kernel(v: ti.Vector): # Error: Invalid type annotation
+def valid_matrix_argument(u: ti.i32, v: ti.types.matrix(2, 2, ti.i32)): # OK: has 5 elements in total
+ ...
+
+@ti.kernel
+def error_too_many_arguments(u: ti.i32, v: ti.i64, w: ti.types.matrix(7, 9, ti.i64)): # Error: has 65 elements in total
...
```
@@ -78,7 +104,7 @@ def error_kernel(v: ti.Vector): # Error: Invalid type annotation
### Return value
-It is optional for a kernel to have a return value. If specified, it must be a type hinted **scalar** value:
+It is optional for a kernel to have a return value. If specified, it must be a type hinted **scalar/vector/matrix** value:
```python {2}
@ti.kernel
@@ -100,21 +126,19 @@ print(my_kernel()) # 128, cast into ti.i32
:::note
-For now, a kernel can only have one scalar return value. Returning
-`ti.Matrix`, `ti.Vector` or Python-style tuple is not supported:
-
-```python {3,9}
+For now, a kernel can only have one return value, and the number of elements in the return value must not exceed 30.
+```python {2,6,10}
@ti.kernel
-def valid_kernel() -> ti.f32:
+def valid_scalar_return() -> ti.f32:
return 128.0 # Return 128.0
@ti.kernel
-def error_kernel() -> ti.Matrix:
- return ti.Matrix([[1, 0], [0, 1]]) # Compilation error
+def valid_matrix_return() -> ti.types.matrix(2, 2, ti.i32):
+ return ti.Matrix([[1, 0], [0, 1]])
@ti.kernel
-def error_kernel() -> (ti.i32, ti.f32):
+def error_multiple_return() -> (ti.i32, ti.f32):
x = 1
y = 0.5
return x, y # Compilation error
@@ -141,7 +165,7 @@ differentiation. Instead, it is recommended to store the result into a global va
`loss[None]`).
:::
-### Functions
+## Functions
A Python function decorated by `@ti.func` is a **Taichi function**:
@@ -170,8 +194,9 @@ Taichi functions can be nested.
:::
:::caution
-Currently, all functions are force-inlined. Therefore, no recursion is
-allowed.
+Currently, all functions are force-inlined. Therefore, no runtime recursion is allowed.
+
+Compile-time recursion is an advanced metaprogramming feature for experienced programmers. See [Metaprogramming](/lang/articles/advanced/meta#compile-time-recursion-of-tifunc) for more information.
:::
### Arguments and return values
@@ -230,27 +255,6 @@ def my_kernel():
...
```
-:::note
-
-Unlike kernels, functions **do support vectors or matrices as arguments
-and return values**:
-
-```python {2,6}
-@ti.func
-def sdf(u): # functions support matrices and vectors as arguments. No type-hints needed.
- return u.norm() - 1
-
-@ti.kernel
-def render(d_x: ti.f32, d_y: ti.f32): # Kernels do not support vector/matrix arguments yet.
- d = ti.Vector([d_x, d_y])
- p = ti.Vector([0.0, 0.0])
- t = sdf(p)
- p += d * t
- ...
-```
-
-:::
-
:::caution
Functions with multiple `return` statements are not supported for now.
diff --git a/docs/lang/articles/basic/type.md b/docs/lang/articles/basic/type.md
index 28161a96bac7b..a67d2935544ab 100644
--- a/docs/lang/articles/basic/type.md
+++ b/docs/lang/articles/basic/type.md
@@ -1,106 +1,52 @@
---
-sidebar_position: 2
+sidebar_position: 3
---
# Type system
-Data types in Taichi consist of Primitive Types and Compound Types. Primitive Types are the numerical data types used by backends, while Compound Types are user-defined types of data records composed of multiple members.
+Data types in Taichi consist of _primitive types_ and _compound types_. Primitive types are the numerical data types used by different backends, while compound types are user-defined types of data records composed of multiple members.
## Primitive types
-Taichi supports common numerical data types. Each type is denoted as a
-character indicating its _category_ and a number of _precision bits_,
-e.g., `i32` and `f64`.
+Taichi supports common numerical data types as its primitive types. Each type is denoted as a
+character indicating its _category_ followed by a number indicating its _precision bits_. The
+_category_ can be either `i` (for signed integers), `u` (for unsigned integers), or `f` (for floating-point numbers). The _precision bits_ can be either `8`, `16`, `32`, or `64`,
+which represents the number of **bits** for storing the data. For example, the two most commonly used types:
-The _category_ can be one of:
+- `i32` represents a 32-bit signed integer;
+- `f32` represents a 32-bit floating-point number.
-- `i` for signed integers, e.g. 24, -32
-- `u` for unsigned integers, e.g. 128, 256
-- `f` for floating point numbers, e.g. 3.14, 1.0, 1e-4
+### Supported primitive types on each backend
-The _digital number_ can be one of:
+| type | CPU | CUDA | OpenGL | Metal | Vulkan |
+| ---- | ---------------- | ---------------- | -------------------- | ---------------- | -------------------- |
+| i8 |:heavy_check_mark:|:heavy_check_mark:|:x: |:heavy_check_mark:|:large_orange_diamond:|
+| i16 |:heavy_check_mark:|:heavy_check_mark:|:x: |:heavy_check_mark:|:large_orange_diamond:|
+| i32 |:heavy_check_mark:|:heavy_check_mark:|:heavy_check_mark: |:heavy_check_mark:|:heavy_check_mark: |
+| i64 |:heavy_check_mark:|:heavy_check_mark:|:large_orange_diamond:|:x: |:large_orange_diamond:|
+| u8 |:heavy_check_mark:|:heavy_check_mark:|:x: |:heavy_check_mark:|:large_orange_diamond:|
+| u16 |:heavy_check_mark:|:heavy_check_mark:|:x: |:heavy_check_mark:|:large_orange_diamond:|
+| u32 |:heavy_check_mark:|:heavy_check_mark:|:x: |:heavy_check_mark:|:heavy_check_mark: |
+| u64 |:heavy_check_mark:|:heavy_check_mark:|:x: |:x: |:large_orange_diamond:|
+| f16 |:heavy_check_mark:|:heavy_check_mark:|:x: |:x: |:heavy_check_mark: |
+| f32 |:heavy_check_mark:|:heavy_check_mark:|:heavy_check_mark: |:heavy_check_mark:|:heavy_check_mark: |
+| f64 |:heavy_check_mark:|:heavy_check_mark:|:heavy_check_mark: |:x: |:large_orange_diamond:|
-- `8`
-- `16`
-- `32`
-- `64`
+(:large_orange_diamond: Requiring extensions of the backend)
-It represents how many **bits** are used in storing the data. The larger
-the bit number, the higher the precision is.
+### Default types for integers and floating-point numbers
-For example, the two most commonly used types:
-
-- `i32` represents a 32-bit signed integer.
-- `f32` represents a 32-bit floating point number.
-
-## Supported primitive types
-
-Currently, supported primitive types in Taichi are
-
-- int8 `ti.i8`
-- int16 `ti.i16`
-- int32 `ti.i32`
-- int64 `ti.i64`
-- uint8 `ti.u8`
-- uint16 `ti.u16`
-- uint32 `ti.u32`
-- uint64 `ti.u64`
-- float32 `ti.f32`
-- float64 `ti.f64`
-
-:::note
-
-Supported types on each backend:
-
-| type | CPU/CUDA | OpenGL | Metal | Vulkan |
-| ---- | -------- | ------- | ----- | -------- |
-| i8 | > OK | > N/A | > OK | > EXT |
-| i16 | > OK | > N/A | > OK | > EXT |
-| i32 | > OK | > OK | > OK | > OK |
-| i64 | > OK | > EXT | > N/A | > EXT |
-| u8 | > OK | > N/A | > OK | > EXT |
-| u16 | > OK | > N/A | > OK | > EXT |
-| u32 | > OK | > N/A | > OK | > OK |
-| u64 | > OK | > N/A | > N/A | > EXT |
-| f32 | > OK | > OK | > OK | > OK |
-| f64 | > OK | > OK | > N/A | > EXT |
-
-(OK: supported, EXT: require extension, N/A: not available)
-:::
-
-:::note
-Boolean types are represented using `ti.i32`.
-:::
-
-## Type promotion
-
-Binary operations on different types will give you a promoted type,
-following the C programming language convention, e.g.:
-
-- `i32 + f32 = f32` (integer + float = float)
-- `i32 + i64 = i64` (less-bits + more-bits = more-bits)
-
-Basically it will try to choose the more precise type to contain the
-result value.
-
-## Default precisions
-
-By default, all numerical literals have 32-bit precisions. For example,
-`42` has type `ti.i32` and `3.14` has type `ti.f32`.
-
-Default integer and float-point precisions (`default_ip` and
-`default_fp`) can be specified when initializing Taichi:
+An integer literal, e.g., `42`, has default type `ti.i32`, while a floating-point literal,
+e.g., `3.14`, has default type `ti.f32`. This behavior can be changed by explicitly specifying
+default types when initializing Taichi:
```python
-ti.init(default_fp=ti.f32)
-ti.init(default_fp=ti.f64)
-
-ti.init(default_ip=ti.i32)
-ti.init(default_ip=ti.i64)
+ti.init(default_ip=ti.i64) # set default integer type to ti.i64
+ti.init(default_fp=ti.f64) # set default floating-point type to ti.f64
```
-Also note that you may use `float` or `int` in type definitions as
-aliases for default precisions, e.g.:
+In addition, you can use `int` as an alias for the default integer type, and `float` as an alias
+for the default floating-point type:
```python
ti.init(default_ip=ti.i64, default_fp=ti.f32)
@@ -113,102 +59,114 @@ y = ti.field(ti.i64, 5)
def func(a: float) -> int:
...
-
# is equivalent to:
def func(a: ti.f32) -> ti.i64:
...
```
-## Type casts
+### Explicit type casting
-All data types are static in the **Taichi scope**. Therefore, casts are needed when you want to assign a certain type of data to another one.
-
-### Implicit casts
+Just like programming in other languages, you may encounter situations where you have a certain
+type of data, but it is not feasible for the assignment or calculation you want to perform. In this
+case, you can do *explicit type casting*. There are two kinds of explicit type casting in Taichi,
+namely *normal casting* and *bit casting*.
:::caution
-The type of a variable is **determined on its initialization**.
+In Taichi-scope, the type of a variable is **static** and **determined on its initialization**.
+That is, you can never change the type of a variable. The compiler relies on this compile-time
+information to check the validity of expressions in Taichi programs.
:::
-When a _low-precision_ variable is assigned to a _high-precision_
-variable, it will be implicitly promoted to the _high-precision_ type
-and no warning will be raised:
+#### Normal casting
-```python {4}
-@ti.kernel
-def foo():
- a = 3.14
- a = 1
- print(a) # 1.0
-```
+`ti.cast()` is used for normal type casting as in other programming languages:
-When a _high-precision_ variable is assigned to a _low-precision_ type,
-it will be implicitly down-cast into the _low-precision_ type and Taichi
-will raise a warning:
-
-```python {4}
+```python {4-5}
@ti.kernel
def foo():
- a = 1
a = 3.14
- print(a) # 3
+ b = ti.cast(a, ti.i32) # 3
+ c = ti.cast(b, ti.f32) # 3.0
```
-### Explicit casts
-
-You may use `ti.cast` to explicitly cast scalar values between different
+You can also use `int()` and `float()` to convert values to default integer and floating-point
types:
```python {4-5}
@ti.kernel
def foo():
a = 3.14
- b = ti.cast(a, ti.i32) # 3
- c = ti.cast(b, ti.f32) # 3.0
+ b = int(a) # 3
+ c = float(b) # 3.0
```
-Equivalently, use `int()` and `float()` to convert values to float-point
-or integer types of default precisions:
+#### Bit casting
+
+Use `ti.bit_cast()` to cast a value into another type **with its underlying bits preserved**:
```python {4-5}
@ti.kernel
def foo():
a = 3.14
- b = int(a) # 3
- c = float(b) # 3.0
+ b = ti.bit_cast(a, ti.i32) # 1078523331
+ c = ti.bit_cast(b, ti.f32) # 3.14
```
-### Casting vectors and matrices
+Note that the new type must have the same precision bits as the old type (`i32`->`f64` is not
+allowed). Use this operation with caution.
-Type casts applied to vectors/matrices are element-wise:
+:::note
+`ti.bit_cast` is equivalent to `reinterpret_cast` in C++.
+:::
-```python {4,6}
+### Implicit type casting
+
+When you accidentally use a value in a place where a different type is expected, implicit type
+casting is triggered for the following cases.
+
+:::caution
+Relying on implicit type casting is bad practice and one major source of bugs.
+:::
+
+#### In binary operations
+
+Following the [implicit conversion rules](https://en.cppreference.com/w/c/language/conversion) of the C programming language, Taichi implicitly casts
+binary operation operands into a *common type* if they have different types. Some simple but most
+commonly used rules to determine the common type of two types are listed below:
+
+- `i32 + f32 = f32` (int + float = float)
+- `i32 + i64 = i64` (low precision bits + high precision bits = high precision bits)
+
+#### In assignments
+
+When a value is assigned to a variable with a different type, the value is implicitly cast into that
+type. If the type of the variable differs from the common type of the variable and the value, a
+warning about losing precisions is raised.
+
+In the following example, variable `a` is initialized with type `float`. On the next line, the
+assignment casts `1` from `int` to `float` implicitly without any warning because the type of the
+variable is the same as the common type `float`:
+
+```python {4}
@ti.kernel
def foo():
- u = ti.Vector([2.3, 4.7])
- v = int(u) # ti.Vector([2, 4])
- # If you are using ti.i32 as default_ip, this is equivalent to:
- v = ti.cast(u, ti.i32) # ti.Vector([2, 4])
+ a = 3.14
+ a = 1
+ print(a) # 1.0
```
-### Bit-casts
-
-Use `ti.bit_cast` to bit-cast a value into another data type. The
-underlying bits will be preserved in this cast. The new type must have
-the same width as the the old type. For example, bit-casting `i32` to
-`f64` is not allowed. Use this operation with caution.
+In the following example, variable `a` is initialized with type `int`. On the next line, the
+assignment casts `3.14` from `float` to `int` implicitly with a warning because the type of the
+variable differs from the common type `float`:
-```python {4-5}
+```python {4}
@ti.kernel
def foo():
+ a = 1
a = 3.14
- b = ti.bit_cast(a, ti.i32) # 1078523331
- c = ti.bit_cast(b, ti.f32) # 3.14
+ print(a) # 3
```
-:::note
-For people from C++, `ti.bit_cast` is equivalent to `reinterpret_cast`.
-:::
-
## Compound types
User-defined compound types can be created using the `ti.types` module. Supported compound types include vectors, matrices, and structs:
@@ -219,6 +177,7 @@ my_vec3f = ti.types.vector(3, float)
my_mat2f = ti.types.matrix(2, 2, float)
my_ray3f = ti.types.struct(ro=my_vec3f, rd=my_vec3f, l=ti.f32)
```
+In this example, we define four compound types for creating fields and local variables.
### Creating fields
@@ -234,8 +193,10 @@ vec1 = ti.Vector.field(2, dtype=ti.i32, shape=(128, 128, 128))
mat2 = ti.Matrix.field(2, 2, dtype=ti.i32, shape=(24, 32))
ray3 = ti.Struct.field({'ro': my_vec3f, 'rd': my_vec3f, 'l': ti.f32}, shape=(1024, 768))
```
+In this example, we define three fields in two different ways but of exactly the same effect.
### Creating local variables
+
Compound types can be directly called to create vector, matrix or struct instances. Vectors, matrices and structs can be created using GLSL-like broadcast syntax since their shapes are already known:
```python
ray1 = my_ray3f(0.0) # ti.Struct(ro=[0.0, 0.0, 0.0], rd=[0.0, 0.0, 0.0], l=0.0)
@@ -244,3 +205,17 @@ mat1 = my_mat2f(1.0) # ti.Matrix([[1.0, 1.0], [1.0, 1.0]])
vec2 = my_vec3f(my_vec2i(0), 1) # ti.Vector([0.0, 0.0, 1.0]), will perform implicit cast
ray2 = my_ray3f(ro=vec1, rd=vec2, l=1.0)
```
+In this example, we define five local variables, each of a different type. In the definition statement of `vec2`, `my_vec3f()` performs an implicit cast operation when combining `my_vec2i(0)` with `1`.
+
+### Type casting on vectors and matrices
+
+Type casting on vectors/matrices is element-wise:
+
+```python {4,6}
+@ti.kernel
+def foo():
+ u = ti.Vector([2.3, 4.7])
+ v = int(u) # ti.Vector([2, 4])
+ # If you are using ti.i32 as default_ip, this is equivalent to:
+ v = ti.cast(u, ti.i32) # ti.Vector([2, 4])
+```
diff --git a/docs/lang/articles/contribution/_category_.json b/docs/lang/articles/contribution/_category_.json
index 74136a441f97d..e5adcd33e7752 100644
--- a/docs/lang/articles/contribution/_category_.json
+++ b/docs/lang/articles/contribution/_category_.json
@@ -1,4 +1,4 @@
{
"label": "Contribution Guide",
- "position": 5
+ "position": 6
}
diff --git a/docs/lang/articles/contribution/contributor_guide.md b/docs/lang/articles/contribution/contributor_guide.md
index 492ab454ecad6..8f2599cf1f5f4 100644
--- a/docs/lang/articles/contribution/contributor_guide.md
+++ b/docs/lang/articles/contribution/contributor_guide.md
@@ -4,295 +4,300 @@ sidebar_position: 1
# Contribution guidelines
-First of all, thank you for contributing! We welcome all kinds of contributions, including but not limited to
-
-- Bug fixes
-- New feature proposals and implementations
-- Documentation improvements and translations
-- More user-friendly error messages
-- New test cases and examples
-- Compiler performance enhancements
-- High-quality blog posts and tutorials
-- Participation in the [Taichi forum](https://forum.taichi.graphics/)
-- Introducing Taichi to your friends or simply staring [the
- project on GitHub](https://github.com/taichi-dev/taichi)
-- Typo fixes in the documentation, code or comments (please go ahead and
- make a pull request for minor issues like these)
-
-:::tip reminder
-Please take some time to familiarize yourself with this contribution guide before opening a pull request.
-For more details regarding development of the Taichi compiler, read the [development tips](./development_tips).
-:::
-## Where to find contribution opportunities
-
-- Issues marked with ["good first
-issue"](https://github.com/taichi-dev/taichi/issues?q=is%3Aopen+is%3Aissue+label%3A%22good+first+issue%22)
-are great chances for starters.
-- Issues marked with ["welcome
- contribution"](https://github.com/taichi-dev/taichi/issues?q=is%3Aopen+is%3Aissue+label%3A%22welcome+contribution%22)
- are slightly more challenging but still friendly to beginners.
-
-## How to take over an issue
-
-- Please first leave a comment (e.g. _I know how to fix this and would
- like to help!_) on the issue, so that people know someone is already
- working on it. This helps prevent redundant work.
-- If no core developer has commented and described a potential
- solution on the issue, please briefly describe your plan, and wait
- for a core developer to reply before you start. This helps keep
- implementations simple and effective.
-
-## High-level guidelines
-
-- Be pragmatic: practically solving problems is our ultimate goal.
-- No overkills: always use _easy_ solutions to solve easy problems, so
- that you have time and energy for real hard ones.
-- Almost every design decision has pros and cons. A decision is
- *good* if its pros outweigh its cons. Always think about
- both sides.
-- Debugging is hard. Changesets should be small so that sources of
- bugs can be easily pinpointed.
-- Unit tests and integration tests are our friends.
-:::note
-"There are two ways of constructing a software design: One way is to
-make it so simple that there are obviously no deficiencies, and the
-other way is to make it so complicated that there are no obvious
-deficiencies. _The first method is far more difficult_."
-— [C.A.R. Hoare](https://en.wikipedia.org/wiki/Tony_Hoare)
-:::
+Thank you for your interest in contributing to Taichi. Taichi was born as an academic research project. Though we are working hard to improve its code quality, Taichi has a long way to go to become a mature, large-scale engineering project. This is also why we decided to open source Taichi from the very beginning: We rely on our community to help Taichi evolve and thrive. From document updates, bug fix, to feature implementation, wherever you spot an issue, you are very welcome to file a PR (pull request) with us!:-)
-One thing to keep in mind is that, Taichi was originally born as an
-academic research project. This usually means that some parts did not
-have the luxury to go through a solid design. While we are always trying
-to improve the code quality, it doesn't mean that the project is free
-from technical debts. Some places may be confusing or overly
-complicated. Whenever you spot one, you are more than welcome to shoot
-us a PR! :-)
-
-## Effective communication
-
-A few tips for effective communication in the Taichi community:
-
-- How much information one effectively conveys, is way more important
- than how many words one typed.
-- Be constructive. Be polite. Be organized. Be concise.
-- Bulleted lists are our friends.
-- Proofread before you post: if you are the reader, can you understand
- what you typed?
-- If you are not a native speaker, consider using a spell checker such
- as [Grammarly](https://app.grammarly.com/).
-
-Please base your discussion and feedback on facts, and not personal
-feelings. It is very important for all of us to maintain a friendly and
-blame-free community. Some examples:
-
-:::tip Acceptable :-)
-This design could be confusing to new Taichi users.
-:::
+Centered around the common process of taking on an issue, testing, and making a corresponding PR, this document provides guidelines, tips, and major considerations for Taichi's contributors. We highly recommend that you spend some time familiarizing yourself with this contribution guide before contributing to Taichi.
-:::danger Not Acceptable
-This design is terrible.
-:::
+## General guidelines and tips
-## Making good pull requests (PRs)
-
-- PRs with **small** changesets are preferred.
- - A PR should ideally address **only one issue**.
- - It is fine to include off-topic **trivial** refactoring such as
- typo fixes;
- - The reviewers reserve the right to ask PR authors to remove
- off-topic **non-trivial** changes.
- - When implementing a complex feature, consider breaking it down into
- small PRs to keep a more detailed development history and to
- interact with core developers more frequently.
-- PR titles should be short sentences describing the changes and following
- [certain format](./contributor_guide#pr-title-format-and-tags).
-- In the description of a PR, it will be nice to link relevant GitHub issues
- (e.g. `fixes #issue_number`) or provide a little context on the motivation.
- Some important implementation decisions you made in the PR is also helpful.
-- If you want early feedback from core developers,
- - Open a PR in
- [Draft](https://github.blog/2019-02-14-introducing-draft-pull-requests/)
- state on GitHub to share your progress;
- - Make sure you @ the corresponding developers in the comments or
- request reviews from them.
-- All PRs should ideally come with corresponding **tests**. See [Testing](./contributor_guide#testing).
-- All PRs should come with **documentation updates**, except for
- internal compiler implementations. See [Documentation](./contributor_guide#documentation).
-- All PRs must pass **continuous integration tests** before they get
- merged. See [Using continuous integration](./contributor_guide#using-continuous-integration).
-- All PRs must pass **code format checks**. See [Enforcing code style](./contributor_guide#enforcing-code-style).
-- Read a great article from Google on [how to have your PR merged
- quickly](https://testing.googleblog.com/2017/06/code-health-too-many-comments-on-your.html).
- [\[PDF\]](https://github.com/yuanming-hu/public_files/blob/master/graphics/taichi/google_review_comments.pdf)
-- All commits in a PR will always be **squashed and merged into master
- as a single commit**. However, PR authors **should not squash commits on their own**.
-- If you are making multiple PRs,
- - Independent PRs should be based on **different** branches
- forking from `master`;
- - PRs with dependencies should be raised only after all
- prerequisite PRs are merged into `master`.
-
-
-## PR reviewing & merging
-
-- Please try to follow these tips from Google:
- - [Code Health: Understanding Code In
- Review](https://testing.googleblog.com/2018/05/code-health-understanding-code-in-review.html);
- [\[PDF\]](https://github.com/yuanming-hu/public_files/blob/master/graphics/taichi/google_understanding_code.pdf)
- - [Code Health: Respectful Reviews == Useful
- Reviews](https://testing.googleblog.com/2019/11/code-health-respectful-reviews-useful.html).
- [\[PDF\]](https://github.com/yuanming-hu/public_files/blob/master/graphics/taichi/google_respectful_reviews.pdf)
-- The merger should always **squash and merge** PRs into the master
- branch.
-- The master branch is required to have a **linear history**.
-- Make sure the PR passes **continuous integration tests**, except for
- cases like documentation updates. See [Using continuous integration](./contributor_guide#using-continuous-integration).
-- Make sure the title follows [PR tag rules](./contributor_guide#pr-title-format-and-tags).
-
-## Using continuous integration
-
-- Continuous Integration (CI) will **build** and **test** your
- commits in a PR in multiple environments.
-- Currently, Taichi uses [Github Actions](https://github.com/features/actions).
-- CI will be triggered every time you push commits to an open PR.
-- You can prepend `[skip ci]` to your commit message to avoid
- triggering CI. For example, a commit with the message
- `[skip ci] This commit will not trigger CI` will not trigger CI.
-- A tick on the left-hand side of a commit hash means CI passed, while
- a cross means CI failed.
-
-## Enforcing code style
-
-- Locally, you can run `ti format` in the command line to re-format
- code style. Note that you have to install `clang-format-10` and
- `yapf v0.31.0` locally before you use `ti format`.
-
-- If you don't have these formatting tools locally, feel free to
- leverage GitHub Actions: simply comment `/format` in a PR
- (e.g., [#2481](https://github.com/taichi-dev/taichi/pull/2481#issuecomment-872226701))
- and then [Taichi Gardener](https://github.com/taichi-gardener)
- will automatically push a commit to your branch that formats the code.
- Note if you want to make more changes afterwards, you'll need to
- `git pull` first.
-
-- For your C++ code, please also follow [C++ style](./cpp_style).
-
-## PR title format and tags
-
-PR titles will be part of the commit history reflected in the `master`
-branch, therefore it is important to keep PR titles readable.
-
-- Please always prepend **at least one tag** such as `[Lang]` to PR
- titles:
- - When using multiple tags, make sure there is exactly one
- space between tags;
- - For example, `[Lang][refactor]` (no space) should be replaced
- by `[Lang] [refactor]`.
-- The first letter of the PR title body should be capitalized:
- - For example, `[Doc] improve documentation` should be replaced by
- `[Doc] Improve documentation`;
- - `[Lang] "ti.sqr(x)" is now deprecated` is fine because `"`
- is a symbol.
-- Please do not include back quotes ("`") in PR titles.
-- Good examples include `[Metal] Support bitmasked SNode`, `[Vulkan]
- ti.atomic_min/max support`, or `[Opt] [ir] Enhanced intra-function optimizations`.
-
-Frequently used tags:
-
-- `[CPU]`, `[CUDA]`, `[Metal]`, `[Vulkan]`, `[OpenGL]`: backends;
-- `[LLVM]`: the LLVM backend shared by CPUs and CUDA;
-- `[Lang]`: frontend language features, including syntax sugars;
-- `[Std]`: standard library, e.g., `ti.Matrix` and `ti.Vector`;
-- `[Sparse]`: sparse computation;
-- `[IR]`: intermediate representation;
-- `[Opt]`: IR optimization passes;
-- `[GUI]`: the built-in GUI system;
-- `[Refactor]`: code refactoring;
-- `[CLI]`: commandline interfaces, e.g., the `ti` command;
-- `[Doc]`: documentation under [docs/](https://github.com/taichi-dev/taichi/blob/master/docs/);
-- `[Example]`: examples under [examples/](https://github.com/taichi-dev/taichi/blob/master/examples/);
-- `[Test]`: tests under [tests/](https://github.com/taichi-dev/taichi/blob/master/tests/);
-- `[Linux]`: Linux platform;
-- `[Mac]`: macOS platform;
-- `[Windows]`: Windows platform;
-- `[Perf]`: performance improvements;
-- `[CI]`: CI/CD workflow;
-- `[Misc]`: something that doesn't belong to any category, such as
- version bump, reformatting;
-- `[Bug]`: bug fixes.
-
-Check out more tags in
- [misc/prtags.json](https://github.com/taichi-dev/taichi/blob/master/misc/prtags.json). When introducing a new tag, please update
- [misc/prtags.json](https://github.com/taichi-dev/taichi/blob/master/misc/prtags.json) in the first PR with that tag, so that people can
- follow.
+This section provides some general guidelines for the Taichi community and tips that we find practically useful.
-:::note
+### Be pragmatic & no overkills
-We do appreciate all kinds of contributions, yet we should not expose
-the title of every PR to end-users. Therefore the changelog will
-distinguish *what the user should know* from *what the
-developers are doing*. This is done by **capitalizing PR
-tags**:
-
-- PRs with visible or notable features to the users should be marked
- with tags starting with **the first letter capitalized**, e.g.,
- `[Metal]`, `[Vulkan]`, `[IR]`, `[Lang]`, `[CLI]`. These PRs will be
- [highlighted in the release note](https://github.com/taichi-dev/taichi/blob/master/misc/make_changelog.py)
- for end-users, therefore it is important to make sure your PR title is
- effective in telling what your PR does.
-- Other PRs (underlying development or intermediate implementation)
- should use tags with **everything in lowercase letters**, e.g.,
- `[metal]`, `[vulkan]`, `[ir]`, `[lang]`, `[cli]`.
-- Because of the way the release changelog is generated, there
- should be **at most one capitalized tag** in a PR title to prevent
- duplicate PR highlights. For example,
- `[GUI] [Mac] Support modifier keys` ([#1189](https://github.com/taichi-dev/taichi/pull/1189))
- is an improper tag choice, and we
- should have used `[gui] [Mac] Support modifier keys in GUI` instead.
- Please capitalize the tag that is the *most* relevant to the PR.
-:::
+Always use straightforward (sometimes even brute-force) solutions: Complicated code usually suggests a lack of design or over-engineering.
+
+> - "There are two ways of constructing a software design: One way is to make it so simple that there are obviously no deficiencies, and the other way is to make it so complicated that there are no obvious deficiencies. *The first method is far more difficult*." — [C.A.R. Hoare](https://en.wikipedia.org/wiki/Tony_Hoare)
+> - "Perfection (in design) is achieved not when there is nothing more to add, but rather when there is nothing more to take away." — [Antoine de Saint-Exupéry](https://en.wikipedia.org/wiki/The_Cathedral_and_the_Bazaar)
+
+### Juxtapose pros and cons
+
+When it comes to making a design decision, weigh up its pros and cons. A design is *good to go* so long as its advantages outweigh its disadvantages.
+
+### Communicate effectively
+
+Our ultimate goal is to build a sustainable, prosperous Taichi community, and effective communication is the cornerstone of that goal. Following are tips that may contribute to effective communication:
+
+- Concise:
+ - The message behind your words outweighs the number of your words. Use as few words as possible to drive your point home.
+ - Use tables, figures, and lists where possible.
+
+- Professional:
+ - Read twice before you post: Would your point get across with your words?
+ - Use a spell checker, such as [Grammarly](https://app.grammarly.com/), to improve your writing in terms of grammar, style, and tone.
+
+- Constructive and courteous: Base your feedback and discussions on facts, *NOT* on personal feelings.
+ - Acceptable😃: *"This design could be confusing to new Taichi users. If it were designed this way, it could..."*
+ - Undesirable😞: ~~*"This design is terrible."*~~
+
+## What you can contribute
+
+
+
+We welcome all kinds of contributions, including but not limited to:
+
+- Fixing a bug
+- Proposing and implementing new features
+- Improving or refactoring an existing document
+- Suggesting more friendly error messages
+- Adding new test cases and examples (demos)
+- Posting blog articles and tutorials
+- Enhancing compiler performance
+- Minor updates to documentation, codes, or annotations.
+
+## Take over an issue
+
+Except for minor updates, most PRs start from a developer taking over an issue. This section provides some corresponding tips and best practices.
+
+### Where to find issues for starters
+
+| Issue Tag | Description | Target developer |
+| ------------------------------------------------------------ | ------------------------- | ---------------------------------------------- |
+| [good first issue](https://github.com/taichi-dev/taichi/issues?q=is:open+is:issue+label:"good+first+issue") | Issues that are easy to start with | Developers new to Taichi |
+| [welcome contribution](https://github.com/taichi-dev/taichi/issues?q=is:open+is:issue+label:"welcome+contribution") | Issues *slightly* more challenging | Developers who wish to dive deeper into Taichi |
+
+### Best practices
+
+- When you plan to take over an issue:
+ - **Best practice**: Leave a message claiming that you are working on it.
+ - **Goal**: Avoid unnecessary repeated work.
+ - **Example**: *"I know how to fix this and would like to help."*
+- After you take over an issue:
+ - **Best practice**:
+ 1. Briefly describe how you plan to handle it (if no solution has been provided).
+ 2. Hold off until a core developer responds to your action plan.
+ - **Goal**: Keep your implementation neat and effective.
+ - **Example**: See [#2610](https://github.com/taichi-dev/taichi/issues/2610).
+
+## References for documentation updates
+
+As part of the effort to increase visibility of the community and to improve developer experience, we highly recommend including documentation updates in your PR if applicable. Here are some of the documentation-specific references and tips:
+
+- Documentation source files are hosted under [docs/](https://github.com/taichi-dev/taichi/blob/master/docs/).
+- We use GitHub Flavored Markdown (GFM) and [Docusaurus](https://docusaurus.io/) to build our documentation site. For information on the supported Markdown syntax, see the [Documentation Writing Guide](./doc_writing).
+- When it comes to writing, we adhere to the [Google Developer Documentation Style Guide](https://developers.google.com/style/).
+- For instructions on setting up a local server and previewing your updated documentation in real-time, see the [Local Development](https://github.com/taichi-dev/docs.taichi.graphics#local-development).
+
+## Add test cases for your local changes
+
+If your PR is to implement a new feature, we recommend that you write your own test cases to cover corner cases for your codes before filing a PR.
+
+- To write a Python test case, see the [Workflow for writing a Python test](./write_test).
+- To write a C++ test case, see the [Workflow for writing a C++ test](./writing_cpp_tests).
+
+## Conduct style checks and integration tests locally
+
+We highly recommend that you complete code style checks and integration tests on your local computer before filing a PR.
+
+### Enforce code style
+
+1. Ensure that you have installed `clang-format-10`.
+2. Ensure that you have installed `yapf v0.31.0`.
+3. Re-format your code style:
+
+```
+python misc/code_format.py
+```
+
+ How to install clang-format-10 on M1 Mac
+
+1. Download and extract [Clang + LLVM 10.0.0 pre-built binary for macOS](https://github.com/llvm/llvm-project/releases/download/llvmorg-10.0.0/clang+llvm-10.0.0-x86_64-apple-darwin.tar.xz)
+
+2. Copy the `clang-format` binary to `~/.local/bin` and add `~/.local/bin` to `PATH`
+
+```shell
+mkdir -p ~/.local/bin
+cp clang+llvm-10.0.0-x86_64-apple-darwin/bin/clang-format ~/.local/bin/clang-format-10
+echo "export PATH=$HOME/.local/bin:\$PATH" >> ~/.zshrc
+source ~/.zshrc
+```
+
+Please refer to [this](./dev_install#llvm-as-cannot-be-opened-on-macos) if you get an error message like `clang-format-10 can’t be opened because Apple cannot check it for malicious software on macOS`.
+
+
+
+
+ What if I didn't format my code style locally?
+
+1. Have your reviewer leave a comment `/format` in your PR to enable GitHub Actions. See [#2481](https://github.com/taichi-dev/taichi/pull/2481).
+ *[Taichi Gardener](https://github.com/taichi-gardener)* *automatically pushes a commit to your branch to format your code.*
-## Testing
+2. If you wish to submit more changes after someone leaves the `/format` comment, ensure that your branch is up to date with your remote counterpart.
-Tests should be added to [tests/](https://github.com/taichi-dev/taichi/blob/master/tests/). We
-have both Python tests and C++ tests.
+
-### Python tests
+
-- Use `ti test` to run all the tests.
-- Use `ti test -v` for verbose outputs.
-- Use `ti test -s` for original output from the tests.
-- Use `ti test -a ` to test only specified backends, e.g.,
- `ti test -a cuda,metal`.
-- Use `ti test -na ` to test all backends excluding specified ones,
- e.g., `ti test -na opengl,x64`.
-- Use `ti test ` to run tests in specified files. For example,
- `ti test numpy_io` will run all tests in [tests/python/test_numpy_io.py](https://github.com/taichi-dev/taichi/blob/master/tests/python/test_numpy_io.py).
-- Use `ti test -k ` to run tests that match the specified key. For
- example, `ti test linalg -k "cross or diag"` will run `test_cross`
- and `test_diag` in [tests/python/test_linalg.py](https://github.com/taichi-dev/taichi/blob/master/tests/python/test_linalg.py).
-- Use `ti test -t ` to set custom number of threads for parallel testing.
+> For more style information for your C++ code, see [our C++ style](./cpp_style).
+
+### Run integration tests
+
+To run all the C++ and Python tests:
+`python tests/run_tests.py`
+
+- **Example 1:**
+`python tests/run_tests.py -v -t3 -a cpu,metal -s`
+ - `-v`: Verbose output.
+ - `-t `: Set a custom number of threads for parallel testing.
+ - `-a `: Test only the specified backends (separated by comma).
+ - `-s`: Original output from the tests.
+
+- **Example 2:**
+`python tests/run_tests.py numpy_io`
+ - ``: Run test cases in specified files only (separated by comma).
+ - This command runs all tests in [tests/python/test_numpy_io.py](https://github.com/taichi-dev/taichi/blob/master/tests/python/test_numpy_io.py).
+
+- **Example 3:**
+`python tests/run_tests.py linalg -k "cross or diag"`
+ - `-k `: Run only the tests that match the specified keys (supports expression in a key string).
+ - This command runs `test_cross()` and `test_diag()` in [tests/python/test_linalg.py](https://github.com/taichi-dev/taichi/blob/master/tests/python/test_linalg.py).
+
+- **To show all available options**
+`python tests/run_tests.py -h`
+
+> We have both Python and C++ test cases, but C++ test cases are disabled by default. To enable C++ test cases:
+>
+> 1. Build Taichi from source using the `python setup.py develop` command.
+> 2. Set `TAICHI_CMAKE_ARGS="-DTI_BUILD_TESTS:BOOL=ON"`.
+
+## File a pull request (PR)
+
+Now you get to the point where you need to get your hands dirty with your PRs. This section provides the following:
+- [Considerations when you create PRs](#considerations)
+- [PR naming conventions](#pr-naming-conventions)
+- [PR review & merging checklist](#pr-review-merging-checklist)
+
+### Considerations
+
+
+
+- **When implementing a complex feature:**
+
+ - Consider breaking it down to multiple separate, self-contained PRs to provide the community with a clearer context and keep a more traceable development history.
+
+- **When creating a PR:**
+
+ - Have your PR address only one issue:
+ - In this way, you keep your changesets small so that potential issues can be readily identified.
+ - If you include in your PR irrevelant implementations, ensure that they are minor.
+ - Your reviewers have the right to request you to remove massive, irrevelant changes from your PR.
+ - If your PR is to implement a new feature, ensure that you have designed test cases for it. See [Add test cases for your local changes](#add-test-cases-for-your-local-changes).
+ - You are required to conduct code style checks and integration tests locally for your PR. See [Conduct style checks and integration tests locally](#conduct-style-checks-and-integration-tests-locally)
+
+- **When describing your PR:**
+ - Provide sufficient information in the description of your PR to provide the community with clearer context:
+ - Link to a specific GitHub issue if applicable, for example `fixes #`.
+ - Share important design decisions in your description.
+
+- **If you create a PR still in progress:**
+
+ - Click **Convert to draft** on your PR page to convert the PR to draft, indicating that you are still working on it.
+ - Click **Ready for review** when you are all set and up for a review.
+ - See [Draft](https://github.blog/2019-02-14-introducing-draft-pull-requests/) for more information.
+
+
+### PR naming conventions
+
+Your PR will make it into the commit history in the the master branch or even Taichi's release notes, therefore it is important to keep your PR title self-explanatory. This section describes our PR naming conventions:
+
+```Gherkin
+[tag1] [tag2]...[tagN] Your PR title must be short but carry necessary info
+
+^----^ ^----^...^----^ ^--------------------------------------------------^
+
+| | | |
+
+| | | +---> Capitalize the initial of your title.
+
+| | +---> Adjacent tags are separated with precisely one space.
+
+| +---> Frequently used tags: [cuda], [lang], [ci], [ir], [refactor].
+
++---> Prepend at least one tag to your PR title.
+```
+
+- **Tag naming conventions:**
+ - Prepend at least one tag, such as `[lang]`, to your PR title.
+ - If you have multiple tags, separate adjacent tags with one space.
+ - See [misc/prtags.json](https://github.com/taichi-dev/taichi/blob/master/misc/prtags.json) for a full list of available tags.
+ - We differentiate PRs for end-users from PRs for developers by *capitalizing tag initial*.
+ - If a PR deals with a feature visible to the end-users, initialize the most relevant tag and the PR will [make it into the release notes](https://github.com/taichi-dev/taichi/blob/master/misc/make_changelog.py). For example, `[Metal]`, `[Vulkan]`, `[IR]`, `[Lang]`, or `[CUDA]`. Ensure that your PR title has *AT MOST* one tag dealt this way.
+ - If a PR deals with the underlying or intermediate implementation, then it is for the developers and you need to ensure that all its tags are *in lowercase*. For example, `[metal]`, `[vulkan]`, `[ir]`, `[lang]`, or `[cuda]`.
+
+ :::danger INCORRECT
+ `[Lang][refactor]` (sans space)
+ :::
+
+ :::tip CORRECT
+ `[Lang] [refactor]`
+ :::
+
+ :::danger INCORRECT
+ `[GUI] [Mac] Support modifier keys` (both tags have their initial capitalized)
+ :::
+
+ :::tip CORRECT
+ `[gui] [Mac] Support modifier keys` (only one tag has its initial capitalized)
+ :::
+
+- **Title naming conventions:**
+ - Keep your PR title short enough but ensure that it carries necessary information.
+ - Do not include back quotes ("\`") in your PR title.
+ - Capitalize the initial letter of your title, which is the word immediately after your tag(s).
+
+ :::danger INCORRECT
+ `[Doc] improve documentation` (the initial of the title is not capitalized)
+ :::
+
+ :::tip CORRECT
+ `[Doc] Improve documentation`
+ :::
+
+:::note
+
+Following are some frequently used tags:
+
+- `[cuda]`: Backend-specific changes.
+- `[lang]`: Frontend language features, including syntax sugars.
+- `[ir]`: Intermediate representation-specific changes.
+- `[refactor]`: Code refactoring changes.
+- `[ci]`: CI/CD workflow-specific changes.
+- `[Doc]`: Documentation updates.
+
+When introducing a new tag, ensure that you add it to [misc/prtags.json](https://github.com/taichi-dev/taichi/blob/master/misc/prtags.json) so that others can follow.
+
+:::
-For more options, see `ti test -h`.
+### PR review & merging checklist
-For more details on how to write a Python test case, see
-[Workflow for writing a Python test](./write_test).
+Follow this checklist during PR review or merging:
-### C++ tests
+1. Ensure that your PR title follows our [naming conventions](#pr-naming-conventions).
+2. Ensure that Taichi's master branch has a *linear history*. See [Linear vs Non-Linear History](https://idiv-biodiversity.github.io/git-knowledge-base/linear-vs-nonlinear.html) for more information.
+3. Ensure that your PR passes all Continuous Integration (CI) tests before merging it.
-For more details on C++ tests, see
-[Workflow for writing a CPP test](./writing_cpp_tests).
+ CI is triggered each time you push a commit to an open PR. It builds and tests all commits in your PR in multiple environments. Keep an eye on the CI test results:
+ - A ✔️ on the left-hand side of a commit hash: CI has passed,
+ - A ❌ on the left-hand side of a commit hash: CI has failed.
-## Documentation
+Here, we do not want to repeat some best practices summarized in the following Google blog articles. But please spare a couple of minutes reading them if your PR is being reviewed or if you are reviewing a PR. They have our recommendation!
+ - [Code Health: Understanding Code In Review](https://testing.googleblog.com/2018/05/code-health-understanding-code-in-review.html)
+ - [Code Health: Respectful Reviews == Useful Reviews](https://testing.googleblog.com/2019/11/code-health-respectful-reviews-useful.html)
+ - [How to have your PR merged quickly](https://testing.googleblog.com/2017/06/code-health-too-many-comments-on-your.html)
-Documentation source files are under [docs/](https://github.com/taichi-dev/taichi/blob/master/docs/) of [**the main Taichi repo**](https://github.com/taichi-dev/taichi).
-An automatic service syncs the updated content with our [documentation repo](https://github.com/taichi-dev/docs.taichi.graphics) and deploys the documentation at [the Taichi documentation site](https://docs.taichi.graphics).
+## Still have issues?
-We use [Markdown](https://www.markdownguide.org/getting-started/) (.md) to write documentation.
-Please see [the documentation writing guide](./doc_writing) for more tips.
+If you encounter any issue that is not covered here, feel free to report it by asking us on GitHub discussions or by [opening an issue on GitHub](https://github.com/taichi-dev/taichi/issues/new?labels=potential+bug&template=bug_report.md) and including the details. We are always there to help!
-To set up a local server and preview your documentation edits in real time,
-see instructions for [Local Development](https://github.com/taichi-dev/docs.taichi.graphics#local-development).
+Finally, thanks again for your interest in contributing to Taichi. We look forward to seeing your contributions!
diff --git a/docs/lang/articles/contribution/dev_install.md b/docs/lang/articles/contribution/dev_install.md
index 366a191975319..852030a0361a4 100644
--- a/docs/lang/articles/contribution/dev_install.md
+++ b/docs/lang/articles/contribution/dev_install.md
@@ -4,306 +4,560 @@ sidebar_position: 2
# Developer installation
-:::note
-End users should use the [pip packages](../get-started.md) instead of building from source.
-:::
-
-This section documents how to configure the Taichi development environment and build Taichi from source for compiler developers. The installation instructions might vary among different operating systems. We also provide a Dockerfile which helps setup a containerized development environment with CUDA support based on the Ubuntu docker image.
+## Target audience
-## Installing dependencies
-1. Python: Currently, 3.6/3.7/3.8/3.9 are supported.
- - If you are on an Apple M1 machine, you might want to install `conda` from [Miniforge](https://github.com/conda-forge/miniforge/#download).
+Developers who are interested in the compiler, computer graphics, or high-performance computing, and would like to contribute new features or bug fixes to the [Taichi programming language](https://github.com/taichi-dev/taichi).
-2. Clang: Make sure you have `clang-8` (or later) on Linux, or download `clang-10` on Windows:
- - On OSX: Normally, you don’t need to do anything.
- - On Ubuntu: Execute `sudo apt install libtinfo-dev clang-8`.
- - On Arch Linux, download `llvm == 10.0.0` prebuilt binary for `ubuntu 18.04` from [here](https://releases.llvm.org/download.html#10.0.1). Then update environment variables `TAICHI_CMAKE_ARGS` and `PATH`:
+:::danger IMPORTANT
- ```bash
- export TAICHI_CMAKE_ARGS="-DCMAKE_CXX_COMPILER=/bin/clang++:$TAICHI_CMAKE_ARGS"
- export PATH=/bin:$PATH
- ```
+This installation guide is *NOT* intended for end users who only wish to do simulation or high performance numerical computation. We recommend that end users install Taichi via `pip install taichi` and that there is no need for you to build Taichi from source. Doing both at the same time may cause unnecessary conflicts.
- - On other Linux distributions, please search [this site](https://pkgs.org) for clang version \>= 8.
- - On Windows: Please download [clang-10](https://github.com/taichi-dev/taichi_assets/releases/download/llvm10/clang-10.0.0-win.zip). Make sure you add the `bin` folder containing `clang.exe` to the `PATH` environment variable.
+See the [Get Started](https://docs.taichi.graphics/) for more information on quickly setting up Taichi for end users.
-:::note
-On Linux, `clang` is the **only** supported compiler for compiling the Taichi package.
:::
-:::note
-On Linux, some additional packages might be required to build Taichi. E.g., on Ubuntu 20.04, you may need `libxi-dev` `libxcursor-dev` `libxinerama-dev` `libxrandr-dev` `libx11-dev` `libgl-dev` `libtinfo5`. please check the output of of CMake when building from source.
-:::
+## Introduction
-3. LLVM: Make sure you have version 10.0.0 installed. Taichi uses a **customized LLVM**, which we provided as binaries depending on your system environment. Note that the pre-built binaries from the LLVM official website or other sources may not work.
- - [LLVM 10.0.0 for Linux](https://github.com/taichi-dev/taichi_assets/releases/download/llvm10_linux_patch2/taichi-llvm-10.0.0-linux.zip)
- - [LLVM 10.0.0 for macOS](https://github.com/taichi-dev/taichi_assets/releases/download/llvm10/taichi-llvm-10.0.0-macos.zip)
- - [LLVM 10.0.0 for Windows MSVC 2019](https://github.com/taichi-dev/taichi_assets/releases/download/llvm10/taichi-llvm-10.0.0-msvc2019.zip)
+ This installation guide covers the following:
+
+ - [Prerequisites for building Taichi from source](#prerequisites)
+ - [Installing optional dependencies](#install-optional-dependencies)
+ - [Building Taichi from source](#build-taichi-from-source)
+ - [Troubleshooting and debugging](#troubleshooting-and-debugging)
+ - [Frequently asked questions](#frequently-asked-questions)
:::note
-When using the above pre-built LLVM for Taichi, please add `$LLVM_FOLDER/bin` to `PATH`, e.g., `export PATH=/bin:$PATH` on Linux.
+
+Installation instructions vary depending on which operating system (OS) you are using. Choose the right OS or platform before you proceed.
+
:::
- - If the previous LLVM binaries do not work, please build from source:
- - For Linux & Mac OSX:
-
- ```bash
- wget https://github.com/llvm/llvm-project/releases/download/llvmorg-10.0.0/llvm-10.0.0.src.tar.xz
- tar xvJf llvm-10.0.0.src.tar.xz
- cd llvm-10.0.0.src
- mkdir build
- cd build
- cmake .. -DLLVM_ENABLE_RTTI:BOOL=ON -DBUILD_SHARED_LIBS:BOOL=OFF -DCMAKE_BUILD_TYPE=Release -DLLVM_TARGETS_TO_BUILD="X86;NVPTX" -DLLVM_ENABLE_ASSERTIONS=ON -DLLVM_ENABLE_TERMINFO=OFF
- # If you are building on Apple M1, use -DLLVM_TARGETS_TO_BUILD="AArch64".
- # If you are building on NVIDIA Jetson TX2, use -DLLVM_TARGETS_TO_BUILD="ARM;NVPTX"
- # If you are building for a PyPI release, add -DLLVM_ENABLE_Z3_SOLVER=OFF to reduce the library dependency.
- make -j 8
- sudo make install
- # Check your LLVM installation
- llvm-config --version # You should get 10.0.0
- ```
-
- - For Windows:
-
- ```bash
- # For Windows
- # LLVM 10.0.0 + MSVC 2019
- cmake .. -G "Visual Studio 16 2019" -A x64 -DLLVM_ENABLE_RTTI:BOOL=ON -DBUILD_SHARED_LIBS:BOOL=OFF -DCMAKE_BUILD_TYPE=Release -DLLVM_TARGETS_TO_BUILD="X86;NVPTX" -DLLVM_ENABLE_ASSERTIONS=ON -Thost=x64 -DLLVM_BUILD_TESTS:BOOL=OFF -DCMAKE_INSTALL_PREFIX=installed
- ```
-
- - Then open `LLVM.sln` and use Visual Studio 2017+ to build.
- - Please make sure you are using the `Release` configuration.
- After building the `INSTALL` project (under folder
- `CMakePredefinedTargets` in the Solution Explorer window).
- - If you use MSVC 2019, **make sure you use C++17** for the
- `INSTALL` project.
- - After the build is complete, find your LLVM binaries and
- headers in `build/installed`.
- - Please add `build/installed/bin` to `PATH`. Later, when you build Taichi, please use `cmake -DLLVM_DIR=/build/installed/lib/cmake/llvm`.
-
-
-### Setting up CUDA (optional)
+## Prerequisites
+
+
+
+
+
+| Category | Prerequisites |
+|:----------------------------:|------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|
+| OS | macOS / Ubuntu / Arch Linux / Other Linux distributions |
+| Python | 3.6/3.7/3.8/3.9 We recommend installing Python from [Miniforge](https://github.com/conda-forge/miniforge/#download) conda if you are on a MacBook with M1 chip. |
+| Clang++ | 8≤ Clang++ <12 |
+| LLVM | 10.0.0 (Taichi customized version) |
+| Command line tools for Xcode | For macOS users only: `xcode-select --install ` |
+
+
+
+
+
+| Category | Prerequisites |
+|:-------------:|------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|
+| OS | Windows 7/8/10/11 |
+| Python | 3.6/3.7/3.8/3.9 |
+| Clang++ | 8≤ Clang++ <12 (We provide pre-built versions in the clang section) |
+| LLVM | 10.0.0 (Taichi customized version) |
+| Visual Studio | Visual Studio 2019/2022 with "Desktop Development with C++" component. If you want to use Clang++ as the compiler, also install "C++ Clang Compiler for Windows" component |
+
+
+
+
+
+### Install Clang
+
+
+This Clang compiler is used to compile the Taichi device runtime. It is **not required** to use this compiler for the C++ compiler.
+
+
+
+
+
+
+1. Ensure that the Clang that ships with your MacBook has a version ≥8 and <12:
+
+ ```
+ clang --version
+ ```
+
+2. If your Clang version is ≥12, install Clang 11:
+
+ ```
+ brew install llvm@11
+ export CXX=/opt/homebrew/opt/llvm@11/bin/clang++
+ ```
+
+
+
+
+
+Download and extract [Clang 10.0.0 pre-built binary for windows](https://github.com/taichi-dev/taichi_assets/releases/download/llvm10/clang-10.0.0-win.zip).
+
+
+
+
+
+```
+sudo apt install clang-10
+```
+
+:::tip NOTE
+
+- Some Linux distributions may require additional packages to build Taichi. For example, you may need `libxi-dev` `libxcursor-dev` `libxinerama-dev` `libxrandr-dev` `libx11-dev` `libgl-dev` for Ubuntu 20.04. Keep an eye on the output of CMake when building from source.
+- If this installation fails, you may want to `apt-get` the corresponding Clang package for your distribution following [this page](https://apt.llvm.org/).
-:::note
-To build with NVIDIA GPU support, CUDA 10.0+ is needed. This installation guide works for Ubuntu 16.04+.
:::
-If you don't have CUDA, go to [this website](https://developer.nvidia.com/cuda-downloads) and download the installer.
+
-- To check if CUDA is installed, run `nvcc --version` or
- `cat /usr/local/cuda/version.txt`.
-- On **Ubuntu** we recommend choosing `deb (local)` as **Installer
- Type**.
-- On **Arch Linux**, you can easily install CUDA via `pacman -S cuda`
- without downloading the installer manually.
+
+
+1. Download [Clang + LLVM 10.0.0 pre-built binary for Ubuntu 18.04](https://github.com/llvm/llvm-project/releases/download/llvmorg-10.0.0/clang+llvm-10.0.0-x86_64-linux-gnu-ubuntu-18.04.tar.xz).
+2. Update the environment variables `TAICHI_CMAKE_ARGS` and `PATH`:
+
+ ```shell
+ export TAICHI_CMAKE_ARGS="-DCMAKE_CXX_COMPILER=/bin/clang++ $TAICHI_CMAKE_ARGS"
+
+ export PATH=/bin:$PATH
+ ```
+
+ :::tip NOTE
+
+ Some Linux distributions may require additional packages to build Taichi. Keep an eye on the output of CMake when building from source.
+
+ :::
+
+
+
+
+
+Search [this site](https://pkgs.org/) for a Clang version that Taichi supports.
+
+:::tip NOTE
+
+Some Linux distributions may require additional packages to build Taichi. Keep an eye on the output of CMake when building from source.
-:::note
-If you are using a machine with an earlier CUDA version and/or old generation GPUs. We suggest to consult the [Compatibility Document](https://docs.nvidia.com/deploy/cuda-compatibility/) and the [CUDA Installation Guide](https://docs.nvidia.com/cuda/cuda-installation-guide-linux/index.html) first.
:::
-### Setting up Vulkan (optional)
+
+
+
+
+### Install LLVM
+
+#### Install pre-built, customized LLVM binaries
+
+We provide pre-built, customized LLVM binaries. For now, Taichi supports LLVM 10.0.0 only.
+
+1. Download and install customized binaries from the following list per your system environment:
+
+
+
+
+ LLVM 10.0.0 for Linux
+
+
+ LLVM 10.0.0 for macOS (without M1 chip)
+
+
+ LLVM 10.0.0 for macOS (with M1 chip)
+
+
+ LLVM 10.0.0 for Windows MSVC 2019
+
+
+
+2. Configure environment variable:
+
+
+
+
+
+1. Add LLVM to your PATH variable:
+ ```
+ echo "export PATH=/bin:\$PATH" >> ~/.bashrc
+ ```
+2. Update your path for the remainder of the session:
+
+ ```shell
+ source ~/.bashrc
+ ```
+
+
+
+
+
+Add an environment variable `LLVM_DIR` with value ``
+
+
+
+
+
+
+
+
+Build LLVM 10.0.0 from source
+
+We provide instructions here if you need to build LLVM 10.0.0 from source.
+
+
+
+
+
+```shell
+wget https://github.com/llvm/llvm-project/releases/download/llvmorg-10.0.0/llvm-10.0.0.src.tar.xz
+
+tar xvJf llvm-10.0.0.src.tar.xz
+
+cd llvm-10.0.0.src
+
+mkdir build
+
+cd build
+
+cmake .. -DLLVM_ENABLE_RTTI:BOOL=ON -DBUILD_SHARED_LIBS:BOOL=OFF -DCMAKE_BUILD_TYPE=Release -DLLVM_TARGETS_TO_BUILD="X86;NVPTX" -DLLVM_ENABLE_ASSERTIONS=ON -DLLVM_ENABLE_TERMINFO=OFF
+
+# If you are building on Apple M1, use -DLLVM_TARGETS_TO_BUILD="AArch64".
+
+# If you are building on NVIDIA Jetson TX2, use -DLLVM_TARGETS_TO_BUILD="ARM;NVPTX"
+
+# If you are building for a PyPI release, add -DLLVM_ENABLE_Z3_SOLVER=OFF to reduce the library dependency.
+
+make -j 8
+
+sudo make install
+
+# Check your LLVM installation
+
+llvm-config --version # You should get 10.0.0
+```
+
+
+
+
+
+```shell
+# For Windows
+
+# LLVM 10.0.0 + MSVC 2019
+
+cmake .. -G "Visual Studio 16 2019" -A x64 -DLLVM_ENABLE_RTTI:BOOL=ON -DBUILD_SHARED_LIBS:BOOL=OFF -DCMAKE_BUILD_TYPE=Release -DLLVM_TARGETS_TO_BUILD="X86;NVPTX" -DLLVM_ENABLE_ASSERTIONS=ON -Thost=x64 -DLLVM_BUILD_TESTS:BOOL=OFF -DCMAKE_INSTALL_PREFIX=installed
+```
+
+1. Use Visual Studio 2017+ to build **LLVM.sln**.
+2. Ensure that you use the **Release** configuration. After building the `INSTALL` project (under folde **CMakePredefinedTargets** in the Solution Explorer window).
+3. If you use MSVC 2019, ensure that you use **C++17** for the `INSTALL` project.
+4. When the build completes, add an environment variable `LLVM_DIR` with value `/build/installed/lib/cmake/llvm`.
+
+
+
+
+
+
+
+## Install optional dependencies
+
+[CUDA](https://en.wikipedia.org/wiki/CUDA) is NVIDIA's answer to high-performance computing. Taichi has implemented a backend based on CUDA 10.0.0+. Vulkan is a next-generation, cross-platform API, open standard for 3D graphics and computing. Taichi has added a Vulkan backend as of v0.8.0.
+
+This section provides instructions on installing these two optional dependencies.
+
+
+Install CUDA
+
+This section works for you if you have a Nvidia GPU supporting CUDA. Note that the required CUDA version is 10.0+.
+
+To install CUDA:
+
+import Tabs from '@theme/Tabs';
+import TabItem from '@theme/TabItem';
+
+
+
+
+
+1. Go to [the official site](https://developer.nvidia.com/cuda-downloads) to download the installer.
+2. Choose **deb (local)** as **Installer Type**.
+3. Check if CUDA is properly installed:
+
+ ```
+ nvidia-smi
+ ```
+
+
+
+
+
+1. `pacman -S cuda`
+2. Check if CUDA is properly installed:
+
+ ```
+ nvidia-smi
+ ```
+
+
+
+
+
+1. Go to [the official site](https://developer.nvidia.com/cuda-downloads) and download the installer.
+2. Choose **exe (local)** as **Installer Type**.
+3. Check if CUDA is properly installed:
+
+ ```
+ nvidia-smi
+ ```
+
+
+
+
+
+
+
+Install Vulkan
+
+You must install the Vulkan SDK in order to debug Taichi's Vulkan backend. To proceed:
+
+
+
+
+
+1. Go to [Vulkan's SDK download page](https://vulkan.lunarg.com/sdk/home) and follow the instructions for your OS.
+2. Check if environment variables `VULKAN_SDK`, `PATH`, `LD_LIBRARY_PATH`, and `VK_LAYER_PATH` are updated.
+
+ > The SDK for Ubuntu provides a `setup-env.sh` for updating these variables.
+
+3. Ensure that you have a Vulkan driver from a GPU vendor properly installed.
+
+ > On Ubuntu, check if a JSON file with a name corresponding to your GPU vendor is in: `/etc/vulkan/icd.d/` or `/usr/share/vulkan/icd.d/`.
+
+4. Check if the SDK is properly installed: `vulkaninfo`.
+
+5. If the SDK is properly installed, add an environment variable `TAICHI_CMAKE_ARGS` with the value `-DTI_WITH_VULKAN:BOOL=ON` to enable the Vulkan backend: (Otherwise Vulkan backend is disabled by default when compiling from source.)
+
+ ```shell
+ export TAICHI_CMAKE_ARGS="$TAICHI_CMAKE_ARGS -DTI_WITH_VULKAN:BOOL=ON"
+ ```
+
+
+
+
+
+1. Go to [Vulkan's SDK download page](https://vulkan.lunarg.com/sdk/home) and follow the instructions for your OS.
+2. Set the environment variable `VULKAN_SDK` to `C:/VulkanSDK/${YOUR_VULKAN_VERSION}`.
+3. If the SDK is properly installed, add an environment variable `TAICHI_CMAKE_ARGS` with the value `-DTI_WITH_VULKAN:BOOL=ON` to enable the Vulkan backend:
+
+ ```shell
+ $env:TAICHI_CMAKE_ARGS += " -DTI_WITH_VULKAN:BOOL=ON"
+ ```
+
+
+
+
+
-If you wish to build taichi with Vulkan. You will need to install the Vulkan SDK. Please visit [this website](https://vulkan.lunarg.com/sdk/home) and follow the instructions for your OS.
-- If you are working on Windows, please set the environment variable `VULKAN_SDK` to `C:/VulkanSDK/${YOUR_VULKAN_VERSION}`. (For example, when using Vulkan 1.2.189.0, set `VULKAN_SDK` to `C:/VulkanSDK/1.2.189.0`).
-- On Linux, also make sure the environment variable `VULKAN_SDK` `PATH` `LD_LIBRARY_PATH` and `VK_LAYER_PATH` are updated. On Ubuntu, the downloaded SDK provides a `setup-env.sh` that can be sourced.
-- Make sure you have a Vulkan driver from a GPU vendor installed. On Ubuntu, you
- can verify there is a JSON file in one of these two locations: `/etc/vulkan/icd.d/` or `/usr/share/vulkan/icd.d`.
-- You can verify the installation of the Vulkan SDK by running `vkvia`, `vulkaninfo`, and/or `vkcube`.
+## Build Taichi from source
-After Vulkan is successfully installed. You can build Taichi with Vulkan by adding an environment variable `TAICHI_CMAKE_ARGS` with the value `-DTI_WITH_VULKAN:BOOL=ON`.
+
-### Setting up Taichi for development
+
-1. Clone the Taichi repo **recursively**, and build:
+1. Clone the Taichi repo *recursively* and build[^1]:
- ```bash
+ ```shell
git clone --recursive https://github.com/taichi-dev/taichi
+
cd taichi
+
python3 -m pip install --user -r requirements_dev.txt
- # export CXX=/path/to/clang # Uncomment if clang is not system default compiler.
- python3 setup.py develop --user # Optionally add DEBUG=1 to keep debug information.
+
+ # Exports CXX=/path/to/clang++ # Uncomment if clang++ is not default compiler of the system. Note that clang is not acceptable due to requirements of some submodules.
+
+ # export DEBUG=1 #Uncomment it if you wish to keep debug information.
+
+ python3 setup.py develop --user
```
-:::note
-We use `MSBUILD.exe` to build the generated project on Windows. Please note that Windows
-could have multiple instances of `MSBUILD.exe` shipped with different
-products. Please make sure you add the path for `MSBUILD.exe` within your
-MSVS directory and make it a higher priority (for instance than the one
-shipped with .NET).
-:::
+2. Try out some of the demos in the **examples/** folder to see if Taichi is properly installed. For example:
+
+ ```shell
+ python3 examples/simulation/mpm128.py
+ ```
:::note
-`python setup.py develop` command (recommended for developers) works very similarly to
-`setup.py install` command (recommended for users) except
-that it doesn't actually install anything. It fits developer need better since edits
-on python file take effect immediately without rebuilding. You only need to rerun `develop`
-commands when you change a project’s C extensions or similarly compiled files. See
-[development mode](https://setuptools.pypa.io/en/stable/userguide/development_mode.html) for more details.
+
+[^1]Although the two commands work similarly, `python setup.py develop` is recommended for you as a developer and `python setup.py install`more for end users. The difference is:
+
+- The `develop` command does not actually install anything but only symbolically links the source code to the deployment directory.
+- The `install` command deep copies the source code so that end users need to rerun the command every time they modify the source code.
+
+The `develop` command serves the developers' needs better because edits to the Python files take effect immediately without the need to rerun the command. A rerun is needed only if you have modified the project's C extension or compiled files. See the [Development Mode](https://setuptools.pypa.io/en/stable/userguide/development_mode.html) for more information.
+
:::
-2. Check out the `examples` folder for runnable examples. Run them with commands
- like `python3 examples/simulation/mpm128.py`.
+
-3. Execute `python3 -m taichi test` to run all the tests. It may take
- up to 5 minutes to run all tests.
+
-4. Execute `python3 setup.py clean` to clean up the local information of your
- previous builds. This allows a fresh build without any cache from the previous
- builds. Note that to uninstall the Taichi package from your Python
- environment, please use `pip uninstall taichi`.
+1. Set-up the environment variable `TAICHI_CMAKE_ARGS` with value `-DCLANG_EXECUTABLE=;/bin/clang.exe -DLLVM_AS_EXECUTABLE=/bin/llvm-as.exe`
+2. Open the "x64 Native Tools Command Prompt" for VS2019 or VS2022. Please make sure you opened the x64 version. (Or load the Visual Studio environment yourself)
+3. Clone the Taichi repo *recursively* & install python dependencies
-## Conda
-To avoid directly installing Taichi's dependencies into your existing
-Python environment, we have provided a pre-defined `conda` environment.
-You can find the instructions [here](https://github.com/taichi-dev/taichi/blob/master/conda/README.md).
+ ```shell
+ git clone --recursive https://github.com/taichi-dev/taichi
-:::note
-This step only helps you setup the development environment,
-you would still need to run `python3 setup.py develop` to re-build
-Taichi.
-:::
+ cd taichi
-## Docker
+ python -m pip install --user -r requirements_dev.txt
+ ```
-For those who prefer to use Docker, we also provide a Dockerfile which
-helps setup the Taichi development environment with CUDA support based
-on Ubuntu docker image.
+4. Build taichi by using `python setup.py develop`
:::note
-In order to follow the instructions in this section, please make sure
-you have the [Docker Desktop (or Engine for
-Linux)](https://www.docker.com/products/docker-desktop) installed and
-set up properly.
-:::
-### Build the Docker image
+[^1]Although the two commands work similarly, `python setup.py develop` is recommended for you as a developer and `python setup.py install`more for end users. The difference is:
+
+- The `develop` command does not actually install anything but only symbolically links the source code to the deployment directory.
+- The `install` command deep copies the source code so that end users need to rerun the command every time they modify the source code.
+
+The `develop` command serves the developers' needs better because edits to the Python files take effect immediately without the need to rerun the command. A rerun is needed only if you have modified the project's C extension or compiled files. See the [Development Mode](https://setuptools.pypa.io/en/stable/userguide/development_mode.html) for more information.
-From within the root directory of the taichi Git repository, execute
-`docker build -t taichi:latest .` to build a Docker image based off the
-local master branch tagged with _latest_. Since this builds the image
-from source, please expect up to 40 mins build time if you don't have
-cached Docker image layers.
+:::
:::note
-In order to save the time on building Docker images, you could always
-visit our [Docker Hub
-repository](https://hub.docker.com/r/taichidev/taichi) and pull the
-versions of pre-built images you would like to use.
+If you want to build Taichi with Clang or maybe utilize `ccache` to cache and speed-up builds, add the following to the end of environment variable `TAICHI_CMAKE_ARGS`: ` -DCMAKE_CXX_COMPILER=clang++ -DCMAKE_C_COMPILER=clang`.
-For example, to pull a image built from release v0.6.17, run
-`docker pull taichidev/taichi:v0.6.17`
:::
-:::caution
+
-The nature of Docker container determines that no changes to the file
-system on the container could be preserved once you exit from the
-container. If you want to use Docker as a persistent development
-environment, we recommend you [mount the taichi Git repository to the
-container as a volume](https://docs.docker.com/storage/volumes/) and set
-the Python path to the mounted directory.
-:::
+
-### Use Docker image on macOS (CPU only)
+## Troubleshooting and debugging
-1. Make sure `XQuartz` and `socat` are installed:
+### `llvm-as` cannot be opened on macOS
-```bash
-brew cask install xquartz
-brew install socat
-```
+**Description**
-2. Temporally disable the xhost access-control: `xhost +`.
-3. Start the Docker container with
- `docker run -it -e DISPLAY=$(ipconfig getifaddr en0):0 taichidev/taichi:v0.6.17`.
-4. Do whatever you want within the container, e.g. you could run tests
- or an example, try: `ti test` or `ti example mpm88`.
-5. Exit from the container with `exit` or `ctrl+D`.
-6. **[To keep your xhost safe]** Re-enable the xhost access-control:
- `xhost -`.
-
-### Use Docker image on Ubuntu (with CUDA support)
-
-1. Make sure your host machine has CUDA properly installed and
- configured. Usually you could verify it by running `nvidia-smi`.
-2. Make sure [NVIDIA Container Toolkit](https://github.com/NVIDIA/nvidia-docker) is properly
- installed:
-
-```bash
-distribution=$(. /etc/os-release;echo $ID$VERSION_ID)
-curl -s -L https://nvidia.github.io/nvidia-docker/gpgkey | sudo apt-key add -
-curl -s -L https://nvidia.github.io/nvidia-docker/$distribution/nvidia-docker.list | sudo tee /etc/apt/sources.list.d/nvidia-docker.list
-
-sudo apt-get update && sudo apt-get install -y nvidia-container-toolkit
-sudo systemctl restart docker
-```
+Gets an error message `llvm-as can’t be opened because Apple cannot check it for malicious software on macOS`.
-3. Make sure `xorg` is installed: `sudo apt-get install xorg`.
-4. Temporally disable the xhost access-control: `xhost +`.
-5. Start the Docker container with
- `sudo docker run -it --gpus all -e DISPLAY=$DISPLAY -v /tmp/.X11-unix:/tmp/.X11-unix taichidev/taichi:v0.6.17`.
-6. Do whatever you want within the container, e.g. you could run tests
- or an example, try: `ti test` or `ti example mpm88`.
-7. Exit from the container with `exit` or `ctrl+D`.
-8. **[To keep your xhost safe]** Re-enable the xhost access-control:
- `xhost -`.
+**Workaround**
+One-off: **System Preferences > Security & Privacy > General > Allow anyway**.
-## Troubleshooting developer installation
+### Permission denied
-- If `python3 setup.py develop`(or `python3 setup.py install`) gives `permission denied` error, it means you're
- installing into system python without write permission. You can work around this by:
- - `python3 setup.py develop --user` or `python3 setup.py install --user`.
- - Install conda and use python from conda enviroments.
+**Description**
-- If `make` fails to compile and reports
- `fatal error: 'spdlog/XXX.h' file not found`, please try runing
- `git submodule update --init --recursive --depth=1`.
+Gets a `permission denied` after `python3 setup.py develop` or `python3 setup.py install`.
-- If the build succeeded but running any Taichi code results in errors
- like
+**Root cause**
- ```
- Bitcode file (/tmp/taichi-tero94pl/runtime//runtime_x64.bc) not found
- ```
+You were trying to install packages into the Python environment without write permission.
- please double check `clang` is in your `PATH`:
+**Workaround**
- ```bash
- clang --version
- # version should be >= 7
- ```
+1. `python3 setup.py develop --user` or `python3 setup.py install --user`.
+2. Install Conda and use python from within the conda environment.
- and our **Taichi configured** `llvm-as`:
+### `make` fails to compile
- ```bash
- llvm-as --version
- # version should be >= 8
- which llvm-as
- # should be /usr/local/bin/llvm-as or /opt/XXX/bin/llvm-as, which is our configured installation
- ```
+**Description**
+
+`make` fails to compile and reports `fatal error: 'spdlog/XXX.h' file not found`.
+
+**Root cause**
+
+You did not use the `--recursive` flag when cloning the Taichi repository.
+
+**Workaround**
+
+Run `git submodule update --init --recursive --depth=1`.
+
+### `which python` still returns the system's Python location
- If not, please install `clang` and **build LLVM from source** with
- instructions above in [dev_install](#installing-dependencies-1),
- then add their path to environment variable `PATH`.
+**Description**
-- If you don't have `wget` on OSX, try installing [homebrew](https://brew.sh/) and then run `brew install wget`.
+`which python` still returns the system's Python location after Conda is installed.
-- If you get a new Apple machine, you might need to run `xcode-select --install` first.
+**Workaround**
-- If you installed `conda` but `which python` still points to the system `python` location, run the following commands to enable it:
+Run the following commands to activate Conda:
+
+```shell
+source /bin/activate
+
+conda init
+```
+
+## Frequently asked questions
+
+### How can I get a fresh Taichi build?
+
+1. Clean up cache from your previous builds:
```
- source /bin/activate
- conda init
+ python3 setup.py clean
```
-- See also [Installation Troubleshooting](../misc/install.md) for issues
- that may share with end-user installation.
+2. Uninstall the Taichi package from your Python environment:
+
+- `python setup.py develop --uninstall`, if you build Taichi using `python setup.py develop`.
+- `pip uninstall taichi`, if you build Taichi using `python setup.py install`.
+
+### What if I don't have `wget` on my macOS?
+
+1. Install [Homebrew](https://brew.sh/).
+2. Use Homebrew to install `wget`:
+
+ `brew install wget`
+
+## Still have issues?
+
+- See [Installation Troubleshooting](../misc/install.md) for issues that may share with the end-user installation.
-- If you encounter other issues, feel free to report (please include the details) by [opening an
- issue on
- GitHub](https://github.com/taichi-dev/taichi/issues/new?labels=potential+bug&template=bug_report.md).
- We are willing to help!
+- If you encounter any issue that is not covered here, feel free to report it by [opening an issue on GitHub](https://github.com/taichi-dev/taichi/issues/new?labels=potential+bug&template=bug_report.md) and including the details. We are always there to help!
diff --git a/docs/lang/articles/contribution/development_tips.md b/docs/lang/articles/contribution/development_tips.md
index fc1f825a221f2..065206bcff030 100644
--- a/docs/lang/articles/contribution/development_tips.md
+++ b/docs/lang/articles/contribution/development_tips.md
@@ -71,9 +71,6 @@ When creating a Taichi program using
`ti.init(arch=desired_arch, **kwargs)`, pass in the following parameters
to make the Taichi compiler print out IRs in different stages:
-- `print_preprocessed=True`: print results of the frontend Python
- AST transform. The resulting scripts will generate a Taichi Frontend
- AST when executed.
- `print_ir=True`: print the Taichi IR transformation process of
kernel (excluding accessors) compilation.
- `print_accessor_ir=True`: print the IR transformation process of
diff --git a/docs/lang/articles/contribution/style_guide_en.md b/docs/lang/articles/contribution/style_guide_en.md
new file mode 100644
index 0000000000000..15f282a4fd367
--- /dev/null
+++ b/docs/lang/articles/contribution/style_guide_en.md
@@ -0,0 +1,374 @@
+---
+sidebar_position: 12
+---
+
+# Taichi Technical Documentation Style Guide
+
+This is a reference for the developers and users at Taichi community to improve their writing and the consistency of Taichi's documentation. You can find detailed style, usage, and grammar in the following sections.
+
+## General principles
+
+### Style and tone
+
+- Use active voice when possible.
+- Write for scannability; use bullet points, short paragraphs, and sections/headers to break up your content.
+- Oxford comma: In a list of three or more items, add a comma before the conjunction (for example: "Android, iOS, and Windows").
+- Spell check your content.
+- Be friendly by using "you".
+- Review your content. Edit out any information your reader does not need to know.
+- Remove ambiguity by choosing words with clear meaning.
+
+Avoid the following:
+
+- Exclamation marks, except in code snippets.
+- Using phrases such as "simply" or "it is that simple" or "it is easy" in a procedure.
+- Do not use dangling modifiers. A modifier "dangles" when the sentence is not clear about what is being modified.
+
+### Write for a global audience
+
+- Use simple verbs. For example, do not use words like utilize when the simpler word use conveys the same information.
+- Do not use idiomatic or colloquial expressions.
+- Avoid making negative voice constructions.
+- Do not use double negative.
+- Keep paragraphs short. Dense pages full of text are intimidating for readers.
+- Address the reader directly by using “you” instead of “the developer”.
+- Be inclusive in your writing. Use gender-neutral pronouns.
+- Be consistent in your word usage. Do not use the same word to mean different things, and vice versa.
+
+## Language and grammar
+
+### Abbreviations and acronyms
+
+Abbreviations are the shortened version of a word or phrase used to represent the whole. Examples include "s" for "seconds,” "approx." for "approximately," and "e.g." for "exempli gratia" (meaning "for example").
+Abbreviations and acronyms can affect the clarity of Taichi content for the reader. While many are understood by our readers and do not need to be spelled out, for new or novel terms always spell out the first mention of an abbreviated term in the text, followed immediately by the abbreviation in parentheses . Use the abbreviated form for all subsequent references of the abbreviation on the same page.
+
+#### Latin abbreviations
+
+Do not use Latin abbreviations in your technical writing.
+Many abbreviations derived from Latin are used in written English. Examples include "e.g." for "exempli gratia" (meaning "for example"), "i.e." for "id est" (meaning "in other words"), and "etc." for "et cetera" (meaning "and the other things").
+Plain language principles suggest avoiding these abbreviated terms.
+
+#### Contractions
+
+Do not use contractions except in FAQs.
+A contraction is a word or phrase that is shortened by dropping one or more letters. Examples include "aren't" for "are not", "let's" for "let us", and "can’t” for “cannot”. While any native English reader understands them, they add unnecessary complexity for non-native readers. For example, contractions that end with the letter "s" can be mistaken for possessive nouns. In business communication, the use of contractions is frowned upon as they make the tone of the writing too informal.
+The only exception to this rule is when you are writing content for an FAQ. The more conversational tone of an FAQ allows for the use of contractions in titles and headers .
+
+### Articles (a, an, the)
+
+"A" and "an" are indefinite articles and are used before a singular countable noun. They refer to any member of a group. "The" is a definite article. It is used before singular and plural nouns and refers to one or more particular members of a group.
+Sound rule for using "a" and "an"
+The way a word or acronym is spoken determines whether "a" or "an" precedes it. Use "a" before a word that starts with a consonant sound, and "an" for words that begin with a vowel sound. For example "a URL", and "an SDK".
+
+### Capitalization
+
+- Use an uppercase letter to begin the first word of the text immediately following a colon.
+- Use sentence case for captions and other figure-related text.
+- Use sentence case for items in all types of lists.
+- Use sentence case for all the elements in a table: contents, headings, labels, and captions.
+- Refer to web page titles with the same casing used on the page.
+
+### Ornamental words
+
+An expletive is an added word or phrase that does not contribute meaning to the sentence. The most common expletives are "there are" and "there is".
+- Not recommended: There are 10 users in the workspace.
+- Recommended: The workspace has 10 users.
+
+### Direct address or imperative mood
+Use the imperative mood for task steps or a command line, a shell command for example.
+Use third person singular for a description of an API method.
+The imperative mood keeps the content concise. The direct address is more personal.
+- Not recommended: Developers can download the SDK from here.
+- Better: You can download the SDK from here.
+- Recommended: Download the SDK from here.
+
+### Gender-neutral
+
+- Avoid using "his" or "her", or "he/she".
+- Use the second person, "you", or the collective noun.
+
+### Present tense
+
+- Use the present tense as it creates concise sentences and provides a tone of immediacy. An exception to this rule is the release date of an SDK or other product. Always frame the release date in the past tense, as that tense will only be correct on the day of release. For example, use "v0.8.0 was released on September 23, 2021", NOT "v0.8.0 is released on September 23, 2021".
+- Avoid using "will" unless you want to stress that something happens at a later point.
+- Use future tense if there is a significant time delay that matters in the context.
+
+### Second person
+
+- In general, use second person "you" (sometimes implicit) in your docs.
+- In glossary terms, avoid using person where possible. Use "developer" to refer to the reader if necessary.
+
+### Clause order
+- Put the most important information at the beginning of a sentence, followed by what the user can do with that information.
+- Provide the context before you provide the instruction; that way, the reader can skip the additional information if it does not apply to their circumstance.
+
+### Punctuations
+
+#### Colons
+
+- The first word after the colon should be in uppercase.
+- When a colon introduces a list, the phrase before the colon must be a complete sentence.
+
+#### Ampersands
+
+- Do not use ampersands ("&") unless they are part of a name, UI string, or in a piece of code.
+
+#### Hyphens
+
+- All words following a hyphen are in lowercase, even if the word is at the beginning of a sentence. For example, "32-bit", or "Multi-threaded".
+- Use a hyphen to form a compound adjective which is an adjective made up of more than one word. Examples include, "A 25-minute interval", "16-bit color", "a six-figure price", and more.
+- Use a hyphen to indicate a common second element. For example, "a 10- to 11-digit number", "a six- or seven-hour delay", "a two-, three-, or fourfold increase".
+- Many common prefixes, such as "co", "de", "pre", "pro", "re", and "sub" do not need a hyphen.
+- Do not use a hyphen when forming a compound adjective with an adverb that ends in "ly".
+
+#### Spaces
+
+- Add a space before an opening parenthesis. Example: operating system (OS)
+- Use only one space after full stops. Example: Add a space. One after the full stop.
+- Use one space between numbers and units. Example: 256 Kbps.
+- Use spaces around mathematical symbols. Example: V = off, width x height, x < y. Use spaces around dimensions. Example: 3.2 x 3.6 x 0.6 mm.
+Note that million is not a unit, so there is no need to add a space between a number and M. For example, 10M is the right Taichi style.
+
+## Plagiarism
+
+Plagiarism puts the firm in a questionable position. Ensure that you do not copy and paste anything that you find from an online search to our technical documentation. As a tip, you can paraphrase contents that you find online.
+
+## Formatting
+
+### Headings and titles
+
+Headings assist the reader in scanning content, helping them discover exactly what they are seeking. They provide structure and are visual points of reference for the reader.
+Use headers to help outline your draft content. Some other points for consideration:
+- Capitalize all words in a document title, except for articles and prepositions.
+- Use sentence case for section titles.
+- Be descriptive and concise.
+- Focus on what the reader needs to know and what they can accomplish.
+- Use ampersands or other symbols only if they appear in a UI or product name.
+- Do NOT conclude a heading with a period or colon. (An exception are FAQs whose titles are often phrased as a conversation with the reader).
+
+### Table headings
+
+When referencing something specific (such as a unit of measure) in a table header, do not repeat it in the cells in that column. For example, if a table column header uses “Kbps”, then there is no need to repeat it in the cells for that column.
+
+### Information
+
+#### Note
+Provides supplemental information that may not apply to all readers, but is important for those specific readers to know.
+Wrap the notes in:
+:::note
+This is a note.
+:::
+
+#### Warning
+Suggests proceeding with caution.
+Wrap the notes in
+:::caution WARNING
+This is a warning.
+:::
+
+#### DANGER
+Designed to guide the reader away from a circumstance that poses some form of problem or hazard.
+Stronger than a Caution; it means "Don't do this."
+Wrap the notes in:
+:::danger DANGER
+This is a danger!
+:::
+
+## Writing examples
+
+### Example 1
+
+- Not recommended: Taichi Zoo uses cookies for security, improvement and analytics purposes.
+- Recommended: Taichi Zoo uses cookies for security, improvement, and analytics purposes.
+
+**Comments:**
+In a list of three or more items, add a comma before the conjunction。
+
+### Example 2
+
+- Not recommended: Two of the most commonly used types:
+ - f32 represents a 32-bit floating point number.
+ - i32 represents a 32-bit signed integer.
+- Recommended: Two of the most commonly used types:
+ - f32: 32-bit floating point number.
+ - i32: 32-bit signed integer.
+
+**Comments:**
+
+Avoid repetitive information in a bullet list.
+
+### Example 3
+
+- Not recommended: If you run into this situation, Taichi's handy automatic differentiation (autodiff) system comes to the rescue!
+- Recommended: Taichi's automatic differentiation (autodiff) system addresses this situation.
+
+**Comments:**
+
+- Avoid subjective descriptions, such as "handy" and "very", and dramatic expressions, for example "come to the rescue" in a technical document.
+
+### Example 4
+
+- Not recommended: ScopedProfileris used to analyze the performance of the Taichi JIT compiler (host).
+- Recommended: ScopedProfiler analyzes the performance of the Taichi JIT compiler (host).
+
+**Comments:**
+
+- Use third person singular when describing a function, a method, or a callback.
+- Use active voice as much as possible in a technical document.
+
+### Example 5
+
+- Not recommended: The easiest way is to make use of ti.GUI.
+- Recommended: The easiest way is to use ti.GUI.
+
+**Comments:**
+
+Use simple verbs. A noun phrase, for example "make use of", is usually wordier than its original verb form, in this case "use".
+
+### Example 6
+
+- Not recommended: Use ti video -f40for creating a video with 40 FPS.
+- Recommended: Use ti video -f40to create a video with a frame rate of 40 FPS.
+
+### Example 7
+
+- Not recommended: Write less bugs.
+- Recommended: Write fewer bugs.
+
+**Comments:**
+
+"Less" describes uncountable noun; "fewer" describes countable noun.
+
+### Example 8
+
+- Not recommended: Sometimes user may want to override the gradients provided by the Taichi autodiff system.
+- Recommended: Sometimes you may want to override the gradients provided by the Taichi autodiff system.
+
+**Comments:**
+
+Address your audience directly by using "you".
+
+### Example 9
+
+- Not recommended: Compared to FLAT , query speed is much faster. Compared with IVFFLAT , less disk and CPU/GPU memory is required for IVF_SQ8.
+- Recommended: IVF_SQ8 has a much higher query speed than FLAT, and requires less disk space and CPU/GPU memory than IVFFLAT.
+
+**Comments:**
+
+- IVF_SQ8 has a much faster query speed than FLAT (has). The second instance of "has" here can be omitted.
+- "Compared to" and "Compared with" are usually wordy.
+
+### Example 10
+
+- Not recommended: Different from IVF_SQ8 , IVF_SQ8H uses a GPU-based coarse quantizer that greatly reduces the quantization time .
+- Recommended: Unlike IVF_SQ8 , IVF_SQ8H uses a GPU-based coarse quantizer , which greatly reduces the time to quantize.
+
+**Comments:**
+
+- In technical writing, one word is always better than two.
+- Which is used in a non-restrictive attributive clause; that is used in a restrictive attributive clause. Always precede a which-clause with a comma.
+
+
+### Example 11
+
+- Not recommended: When you use a client to update the following parameters, the updates take effect immediately.
+- Recommended: Updates to the following parameters from a client take effect immediately:
+
+**Comments:**
+
+- The original is wordy.
+
+### Example 12
+
+- Not recommended: Vectors are quantized to 8-bit floats , which may cause accuracy loss.
+- Recommended: Vectors are quantized to 8-bit floats. This may cause accuracy loss.
+- Not recommended: However, the process to build a search graph requires a lot of computations for distances between vectors, which results in high computation costs.
+- Recommended: However, the process of building a search graph requires a lot of computations for distances between vectors, resulting in high computation costs.
+
+**Comments:**
+
+- You cannot use which to refer to an entire preceding clause. Which only modifies the noun or noun phrase ( noun + prep. + noun) immediately preceding it. Use this to refer to the entire preceding clause .
+
+### Example 13
+
+- Not recommended: Make sure the Memory available to Docker Engine exceeds the sum of insert_buffer_size and cpu_cache_capacity you set in the config.yaml file.
+- Recommended: Ensure that the Memory available to Docker Engine exceeds the sum of insert_buffer_size and cpu_cache_capacity , both of which are defined in config.yaml.
+
+**Comments:**
+
+- When it comes to technical writing, do not use more than one word when one word can convey the same information.
+- Always use that to lead an appositive clause.
+- If you have already spelt out the file name, you do not need to emphasize it is a file. Your readers can tell for themselves.
+
+### Example 14
+
+- Not recommended: Start the Prometheus server, with the --config.file flag pointing to the configuration file:
+$ ./prometheus --config.file=prometheus.yml
+
+- Recommended: Start the Prometheus server and specify the configuration file:
+$ ./prometheus --config.file=prometheus.yml
+
+**Comments:**
+
+- Misuse of with. With modifies the subject of the main clause.
+- The original is horribly specific. The revised version speaks for itself.
+
+### Example 15
+
+- Not recommended: This document talks about the following topics:
+- Recommended: This document covers the following topics:
+
+**Comments:**
+
+- Anthropomorphism is not accepted in technical documents.
+
+### Example 16
+
+- Not recommended:
+ - True: Enables the debug mode.
+ - False: Disables the debug mode.
+- Recommended:
+ - True: Enable debug mode.
+ - False: Disable debug mode.
+
+
+**Comments:**
+
+- Use imperative mood when desbribing a binary parameter; use third person singular when describing a function, a method, or a callback.
+- Do not use the definite article before the word mode.
+
+### Example 17
+
+- Not recommended: This parameter is used to enable or disable Write Ahead Log (WAL).
+- Recommended: This parameter enables or disables Write Ahead Log (WAL).
+
+**Comments:**
+
+- Clean, clear, and straight to the point!
+
+### Example 18
+
+- Not recommended: Active monitoring helps you identify problems early. But it is also essential to create alerting rules that promptly send notifications when there are events that require investigation or intervention.
+- Recommended: Proactively monitoring metrics helps identify issues when they emerge. Creating alerting rules for events that require immediate intervention is essential as well.
+
+**Comments:**
+
+- Do not use "but" to lead a separate sentence.
+- Way too many that-clauses!
+- The "there be" construction is always awkward. An expletive is an added word or phrase that does not contribute meaning to the sentence. The most common expletives are "there are" and "there is".
+
+### Example 19
+
+- Not recommended: However, for delete operations, the operation speed is faster when write ahead log is enabled.
+- Recommended: Delete operations are faster when write ahead log is enabled.
+
+**Comments:**
+
+- You cannot say faster speed. You can say higher speed or greater speed . You can also say an operation is faster.
+- The original is wordy.
+
+## English style guide references
+
+- Microsoft Writing Style Guide
+- The Chicago Manual of Style
+- Merriam-Webster's Dictionary
diff --git a/docs/lang/articles/contribution/utilities.md b/docs/lang/articles/contribution/utilities.md
index d537ec7df20c3..2fd7dfc5fa62c 100644
--- a/docs/lang/articles/contribution/utilities.md
+++ b/docs/lang/articles/contribution/utilities.md
@@ -86,7 +86,7 @@ int func(void *p) {
## Benchmarking and regression tests
- Run `ti benchmark` to run tests in benchmark mode. This will record
- the performance of `ti test`, and save it in `benchmarks/output`.
+ the performance of `python tests/run_tests.py`, and save it in `benchmarks/output`.
- Run `ti regression` to show the difference between the previous
result in `benchmarks/baseline`. And you can see if the performance
is increasing or decreasing after your commits. This is really
@@ -150,7 +150,7 @@ when the program crashes.
```python
# Python
-ti.set_gdb_trigger(True)
+ti.init(gdb_trigger=True)
```
```cpp
@@ -188,7 +188,7 @@ in is executed in test.
not C++ yet.
```bash
-ti test -C # run tests and save results to .coverage
+python tests/run_tests.py -C # run tests and save results to .coverage
coverage report # generate a coverage report on terminal output
coverage html # generate a HTML form report in htmlcov/index.html
```
diff --git a/docs/lang/articles/contribution/write_test.md b/docs/lang/articles/contribution/write_test.md
index 475b47e104507..b26a1805d831b 100644
--- a/docs/lang/articles/contribution/write_test.md
+++ b/docs/lang/articles/contribution/write_test.md
@@ -51,7 +51,7 @@ def test_log10():
assert r[None] == 2
```
-Execute `ti test logarithm`, and the functions starting with `test_` in
+Execute `python tests/run_tests.py logarithm`, and the functions starting with `test_` in
`tests/python/test_logarithm.py` will be executed.
## Testing against multiple backends
@@ -229,7 +229,7 @@ exclude them from the test in order to move forward:
```python
# Run this test on all backends except for OpenGL
-@ti.test(excludes=[ti.opengl])
+@ti.test(exclude=[ti.opengl])
def test_sparse_field():
# ... (some tests that requires sparse feature which is not supported by OpenGL)
```
diff --git a/docs/lang/articles/get-started/_category_.json b/docs/lang/articles/get-started/_category_.json
new file mode 100644
index 0000000000000..3562d433d76f3
--- /dev/null
+++ b/docs/lang/articles/get-started/_category_.json
@@ -0,0 +1,4 @@
+{
+ "label": "Getting Started",
+ "position": 1
+}
diff --git a/docs/lang/articles/get-started.md b/docs/lang/articles/get-started/index.md
similarity index 89%
rename from docs/lang/articles/get-started.md
rename to docs/lang/articles/get-started/index.md
index b6d250d157f26..f607ce2f224d6 100644
--- a/docs/lang/articles/get-started.md
+++ b/docs/lang/articles/get-started/index.md
@@ -25,25 +25,15 @@ import TabItem from '@theme/TabItem';
There are a few of extra requirements depend on which operating system you are using:
-
-
- On Ubuntu 19.04+, you need to install `libtinfo5`:
-
- ```sudo apt install libtinfo5```
-
-
- On Arch Linux, you need to install `ncurses5-compat-libs` package from the Arch User Repository:
-
- ```yaourt -S ncurses5-compat-libs```
+ On Arch Linux, you need to install `ncurses5-compat-libs` package from the Arch User Repository: `yaourt -S ncurses5-compat-libs`
@@ -54,13 +44,13 @@ There are a few of extra requirements depend on which operating system you are u
-Please refer to the [Installation Troubleshooting](./misc/install.md) section if you run into any issues when installing Taichi.
+Please refer to the [Installation Troubleshooting](../misc/install.md) section if you run into any issues when installing Taichi.
## Hello, world!
We introduce the Taichi programming language through a very basic _fractal_ example.
-Running the Taichi code below using either `python3 fractal.py` or `ti example fractal` _(you can find more information about the Taichi CLI in the [Command line utilities](./misc/cli_utilities.md) section)_ will give you an animation of [Julia set](https://en.wikipedia.org/wiki/Julia_set):
+Running the Taichi code below using either `python3 fractal.py` or `ti example fractal` _(you can find more information about the Taichi CLI in the [Command line utilities](../misc/cli_utilities.md) section)_ will give you an animation of [Julia set](https://en.wikipedia.org/wiki/Julia_set):
@@ -178,7 +168,7 @@ type-hinted (if any).
Taichi **functions** are defined with the decorator `@ti.func`. They can
**only** be called by Taichi kernels or other Taichi functions.
-See [syntax](./basic/syntax.md) for more details about Taichi
+See [syntax](../basic/syntax.md) for more details about Taichi
kernels and functions.
The language used in Taichi kernels and functions looks exactly like
@@ -273,7 +263,7 @@ over all the pixel coordinates, i.e.,
:::note
-Struct-for is the key to [sparse computation](./advanced/sparse.md) in
+Struct-for is the key to [sparse computation](../advanced/sparse.md) in
Taichi, as it will only loop over active elements in a sparse field. In
dense fields, all elements are active.
:::
@@ -326,6 +316,20 @@ def foo():
:::
+### GUI system
+
+Taichi provides a cpu-based [GUI system](../gui/gui.md) for users to render
+their results on the screen.
+
+```python
+gui = ti.GUI("Julia Set", res=(n * 2, n))
+
+for i in range(1000000):
+ paint(i * 0.03)
+ gui.set_image(pixels)
+ gui.show()
+```
+
### Interacting with other Python packages
#### Python-scope data access
@@ -373,18 +377,18 @@ while gui.running:
gui.show()
```
-See [Interacting with external arrays](./basic/external.md#interacting-with-external-arrays) for more details.
+See [Interacting with external arrays](../basic/external.md#interacting-with-external-arrays) for more details.
## What's next?
Now we have gone through core features of the
Taichi programming language using the fractal example,
feel free to dive into the language concepts in
-the next section, or jump to the advanced topics, such as the [Metaprogramming](./advanced/meta.md) or [Differentiable programming](./advanced/differentiable_programming.md). Remember that you can
+the next section, or jump to the advanced topics, such as the [Metaprogramming](../advanced/meta.md) or [Differentiable programming](../advanced/differentiable_programming.md). Remember that you can
use the search bar at the top right corner to search for topics or keywords
at any time!
If you are interested in joining the Taichi community, we strongly recommend you take some time to
-familiarize yourself with our [contribution guide](./contribution/contributor_guide.md).
+familiarize yourself with our [contribution guide](../contribution/contributor_guide.md).
We hope you enjoy your adventure with Taichi!
diff --git a/docs/lang/articles/gui/_category_.json b/docs/lang/articles/gui/_category_.json
new file mode 100644
index 0000000000000..d05d0602cb17c
--- /dev/null
+++ b/docs/lang/articles/gui/_category_.json
@@ -0,0 +1,4 @@
+{
+ "label": "GUI",
+ "position": 4
+}
diff --git a/docs/lang/articles/misc/ggui.md b/docs/lang/articles/gui/ggui.md
similarity index 83%
rename from docs/lang/articles/misc/ggui.md
rename to docs/lang/articles/gui/ggui.md
index 39a2bb5ed24d7..edf084ba11caf 100644
--- a/docs/lang/articles/misc/ggui.md
+++ b/docs/lang/articles/gui/ggui.md
@@ -1,6 +1,5 @@
---
-sidebar_position: 1
-
+sidebar_position: 2
---
# A New UI system: GGUI
@@ -14,7 +13,7 @@ You also need to install the Vulkan environment: [https://vulkan.lunarg.com/sdk/
A new UI system has been added to Taichi in version `v0.8.0`. The new GUI system uses GPU for rendering, enabling it to be much faster and to render 3d scenes. For these reasons, this new system is sometimes referred to as GGUI. This doc describes the APIs provided.
-Apart from this doc, a good way of getting familiarized with GGUI is to look at the examples. Please checkout the examples provided in [`examples/ggui_examples`](https://github.com/taichi-dev/taichi/tree/master/examples/ggui_examples).
+Apart from this doc, a good way of getting familiarized with GGUI is to look at the examples. Please checkout the examples provided in [`examples/ggui_examples`](https://github.com/taichi-dev/taichi/tree/master/python/taichi/examples/ggui_examples).
## Creating a window
@@ -42,7 +41,7 @@ this retrieves a `Canvas` object that covers the entire window.
### Drawing on the canvas
```python
-canvas.set_back_ground_color(color)
+canvas.set_background_color(color)
canvas.triangles(vertices, color, indices, per_vertex_color)
canvas.circles(vertices, radius, color, per_vertex_color)
canvas.lines(vertices, width, indices, color, per_vertex_color)
@@ -114,7 +113,7 @@ window.GUI.begin(name, x, y, width, height)
window.GUI.text(text)
is_clicked = window.GUI.button(name)
new_value = window.GUI.slider_float(name, old_value, min_value, max_value)
-new_color = window.GUI.slider_float(name, old_color)
+new_color = window.GUI.color_edit_3(name, old_color)
window.GUI.end()
```
@@ -145,7 +144,7 @@ To check if a specific key is currently pressed:
-Here is an input processing example in GGUI version [`mpm128`](https://github.com/taichi-dev/taichi/blob/master/examples/ggui_examples/mpm128_ggui.py):
+Here is an input processing example in GGUI version [`mpm128`](https://github.com/taichi-dev/taichi/blob/master/python/taichi/examples/ggui_examples/mpm128_ggui.py):
```python
while window.running:
@@ -167,3 +166,24 @@ while window.running:
if window.is_pressed(ti.ui.RMB):
attractor_strength[None] = -1
```
+
+
+## Image I/O
+
+To write the current screen content into an image file:
+
+```python
+window.write_image(filename)
+```
+
+Notice that, when the window is showing, you have to call `window.write_image()` before the `window.show()` call.
+
+
+## Off-screen rendering
+
+GGUI supports rendering contents off-screen, that is, writing the results into image files without showing the window at all. This is sometimes referred to as "headless" rendering. To enable this mode, initialize the window with the argument `show_window=False`:
+
+```python
+window = ti.ui.Window('Window Title', (640, 360), show_window = False)
+```
+Then, you can use `window.write_image()` as normal, and remove the `window.show()` call at the end.
diff --git a/docs/lang/articles/gui/gui.md b/docs/lang/articles/gui/gui.md
new file mode 100644
index 0000000000000..0d2495739185c
--- /dev/null
+++ b/docs/lang/articles/gui/gui.md
@@ -0,0 +1,348 @@
+---
+sidebar_position: 1
+---
+
+# GUI system
+
+Taichi has a built-in cpu-based GUI system to help users visualize results.
+
+## Create a window
+
+The following code show how to create a window of resolution `640x360`:
+
+```python
+gui = ti.GUI('Window Title', (640, 360))
+```
+
+:::note
+
+If you are running Taichi on a machine without a GUI environment, consider setting `show_gui` to `False`:
+
+```python
+gui = ti.GUI('Window Title', (640, 360), show_gui=False)
+
+while gui.running:
+ ...
+ gui.show(f'{gui.frame:06d}.png') # save current frame to local file
+```
+
+:::
+
+## Display a window
+
+The following code snippet display frame of the current windows:
+
+```python
+ for frame in range(10000):
+ ...
+ gui.show() # display current frame
+```
+
+:::note
+Current FPS will show besides the title of the window. By default, FPS is limited to 60.
+We can change this number by setting `gui.fps_limit = the_number_we_want`.
+:::
+
+
+## Paint on a window
+Taichi's GUI supports painting simple geometric objects, such as lines, triangles, rectangles, circles, and text.
+
+:::note
+
+The position parameter of every drawing API expects input of 2-element tuples,
+whose values are the relative position of the object range from 0.0 to 1.0.
+(0.0, 0.0) stands for the lower left corner of the window, and (1.0, 1.0) stands for the upper right corner.
+
+Acceptable input for positions are Taichi fields or numpy arrays. Primitive arrays in python are NOT acceptable.
+
+To convert Taichi field to numpy array, use `to_numpy()` on Taichi fields. By doing this, we can also use data
+from Taichi program in other visualization APIs such as matplotlib.
+:::
+
+:::tip
+
+Here we only list the most commonly-used APIs. For a full list of APIs and the detailed API descriptions, please
+see the [API docs](https://api-docs.taichi.graphics/autoapi/taichi/ui/gui/index.html#module-taichi.ui.gui).
+
+:::
+
+```python
+gui.circles(pos, radius=3, palette=[0x068587, 0xED553B, 0xEEEEF0], palette_indices=material)
+```
+draws circles with radius of 1.5 and three different colors differed by `material`, an integer array with the same size as
+`pos`. Each integer in `material` indicates which color the associated circle use (i.e. array [0, 1, 2] indicates these three
+circles are colored separately by the first, second, and third color in `palette`.
+
+
+
+```python
+gui.lines(begin=X, end=Y, radius=2, color=0x068587)
+```
+draws line segments from X positions to Y positions with width of 2 and color in light blue.
+
+
+
+```python
+gui.triangles(a=X, b=Y, c=Z, color=0xED553B)
+```
+draws triangles with color in red and three points positioned at X, Y, and Z.
+
+
+
+## RGB & Hex conversion.
+
+A handy tool for converting colors from RGB to hex and vice versa.
+
+```python
+rgb = (0.4, 0.8, 1.0)
+hex = ti.rgb_to_hex(rgb) # 0x66ccff
+
+rgb = ti.hex_to_rgb(0x007fff) # (0.0, 0.5, 1.0)
+
+rgb = np.array([[0.4, 0.8, 1.0], [0.0, 0.5, 1.0]])
+hex = ti.rgb_to_hex(rgb) # np.array([0x66ccff, 0x007fff])
+```
+
+The return values can be used in GUI drawing APIs.
+
+
+## Event processing
+
+Every event have a key and type.
+
+_Event type_ is the type of event, for now, there are just three type of event:
+
+ ti.GUI.RELEASE # key up or mouse button up
+ ti.GUI.PRESS # key down or mouse button down
+ ti.GUI.MOTION # mouse motion or mouse wheel
+
+_Event key_ is the key that you pressed on keyboard or mouse, can be one of:
+
+ # for ti.GUI.PRESS and ti.GUI.RELEASE event:
+ ti.GUI.ESCAPE # Esc
+ ti.GUI.SHIFT # Shift
+ ti.GUI.LEFT # Left Arrow
+ 'a' # we use lowercase for alphabet
+ 'b'
+ ...
+ ti.GUI.LMB # Left Mouse Button
+ ti.GUI.RMB # Right Mouse Button
+
+ # for ti.GUI.MOTION event:
+ ti.GUI.MOVE # Mouse Moved
+ ti.GUI.WHEEL # Mouse Wheel Scrolling
+
+A _event filter_ is a list combined of _key_, _type_ and _(type, key)_ tuple, e.g.:
+
+```python
+# if ESC pressed or released:
+gui.get_event(ti.GUI.ESCAPE)
+
+# if any key is pressed:
+gui.get_event(ti.GUI.PRESS)
+
+# if ESC pressed or SPACE released:
+gui.get_event((ti.GUI.PRESS, ti.GUI.ESCAPE), (ti.GUI.RELEASE, ti.GUI.SPACE))
+```
+
+`gui.running` checks the state of the window. `ti.GUI.EXIT` occurs when
+you click on the close (X) button of a window. `gui.running` will obtain
+`False` when the GUI is being closed.
+
+For example, loop until the close button is clicked:
+
+ while gui.running:
+ render()
+ gui.set_image(pixels)
+ gui.show()
+
+You can also close the window by manually setting `gui.running` to`False`:
+
+ while gui.running:
+ if gui.get_event(ti.GUI.ESCAPE):
+ gui.running = False
+
+ render()
+ gui.set_image(pixels)
+ gui.show()
+
+`gui.get_event(a, ...)` tries to pop an event from the queue, and stores it into `gui.event`.
+
+For example:
+
+ if gui.get_event():
+ print('Got event, key =', gui.event.key)
+
+For example, loop until ESC is pressed:
+
+ gui = ti.GUI('Title', (640, 480))
+ while not gui.get_event(ti.GUI.ESCAPE):
+ gui.set_image(img)
+ gui.show()
+
+`gui.is_pressed(key, ...)` detects the keys you pressed. You must use it
+together with `gui.get_event`. Otherwise, it is not updated. For example:
+
+ while True:
+ gui.get_event() # must be called before is_pressed
+ if gui.is_pressed('a', ti.GUI.LEFT):
+ print('Go left!')
+ elif gui.is_pressed('d', ti.GUI.RIGHT):
+ print('Go right!')
+
+:::caution
+
+`gui.is_pressed()` must be used together with `gui.get_event()`, or it won't be updated!
+
+:::
+
+For example:
+
+```python
+while True:
+ gui.get_event() # must be called before is_pressed
+ if gui.is_pressed('a', ti.GUI.LEFT):
+ print('Go left!')
+ elif gui.is_pressed('d', ti.GUI.RIGHT):
+ print('Go right!')
+```
+
+`gui.get_cursor_pos()` retrieves the current cursor position on the window. For example:
+
+ mouse_x, mouse_y = gui.get_cursor_pos()
+
+
+## GUI Widgets
+
+Sometimes it's more intuitive to use widgets like slider or button to control the program variables instead of using chaotic keyboard bindings. Taichi GUI provides a set of widgets for that reason:
+
+```python
+import taichi as ti
+
+gui = ti.GUI('GUI widgets')
+
+radius = gui.slider('Radius', 1, 50, step=1)
+xcoor = gui.label('X-coordinate')
+okay = gui.button('OK')
+
+xcoor.value = 0.5
+radius.value = 10
+
+while gui.running:
+ for e in gui.get_events(gui.PRESS):
+ if e.key == gui.ESCAPE:
+ gui.running = False
+ elif e.key == 'a':
+ xcoor.value -= 0.05
+ elif e.key == 'd':
+ xcoor.value += 0.05
+ elif e.key == 's':
+ radius.value -= 1
+ elif e.key == 'w':
+ radius.value += 1
+ elif e.key == okay:
+ print('OK clicked')
+
+ gui.circle((xcoor.value, 0.5), radius=radius.value)
+ gui.show()
+```
+
+
+
+## Image I/O
+
+`ti.imwrite(img, filename)` exports an `np.ndarray` or a Taichi field
+(`ti.Matrix.field`, `ti.Vector.field`, or `ti.field`) to a file with a specified `filename`.
+
+Same as `ti.GUI.show(filename)`, the format of the exported image is determined by **the suffix of** `filename` as well. Now `ti.imwrite` supports exporting images to `png`, `img` and `jpg` and we recommend using `png`.
+
+Please make sure that the input image has **a valid shape**. If you want to export a grayscale image, the input shape of field should be `(height, weight)` or `(height, weight, 1)`. For example:
+
+```python
+import taichi as ti
+
+ti.init()
+
+shape = (512, 512)
+type = ti.u8
+pixels = ti.field(dtype=type, shape=shape)
+
+@ti.kernel
+def draw():
+ for i, j in pixels:
+ pixels[i, j] = ti.random() * 255 # integers between [0, 255] for ti.u8
+
+draw()
+
+ti.imwrite(pixels, f"export_u8.png")
+```
+
+Besides, for RGB or RGBA images, `ti.imwrite` needs to receive a field which has shape `(height, width, 3)` and `(height, width, 4)` individually.
+
+Generally the value of the pixels on each channel of a `png` image is an integer in \[0, 255\]. For this reason, `ti.imwrite` will **cast fields** which has different data types all **into integers between \[0, 255\]**. As a result, `ti.imwrite` has the following requirements for different data types of input fields:
+
+- For float-type (`ti.f16`, `ti.f32`, etc.) input fields, **the value of each pixel should be float between \[0.0, 1.0\]**. Otherwise `ti.imwrite` will first clip them into \[0.0, 1.0\]. Then they are multiplied by 256 and cast to integers ranging from \[0, 255\].
+- For int-type (`ti.u8`, `ti.u16`, etc.) input fields, **the value of each pixel can be any valid integer in its own bounds**. These integers in this field will be scaled to \[0, 255\] by being divided over the upper bound of its basic type accordingly.
+
+Here is another example:
+
+```python
+import taichi as ti
+
+ti.init()
+
+shape = (512, 512)
+channels = 3
+type = ti.f32
+pixels = ti.Matrix.field(channels, dtype=type, shape=shape)
+
+@ti.kernel
+def draw():
+ for i, j in pixels:
+ for k in ti.static(range(channels)):
+ pixels[i, j][k] = ti.random() # floats between [0, 1] for ti.f32
+
+draw()
+
+ti.imwrite(pixels, f"export_f32.png")
+```
+
+## Zero-copying frame buffer
+When the GUI resolution (window size) is large, it sometimes becomes difficult to achieve 60 FPS even without any kernel
+invocations between two frames.
+
+This is mainly due to the copy overhead, where Taichi GUI needs to copy the image buffer from one place to another.
+This process is necessary for the 2D drawing functions, such as `gui.circles`, to work. The larger the image shape is,
+the larger the overhead.
+
+Fortunately, sometimes your program only needs `gui.set_image` alone. In such cases, you can enable the `fast_gui` option
+for better performance. This mode allows Taichi GUI to directly write the image data to the frame buffer without additional
+copying, resulting in a much better FPS.
+
+```python
+gui = ti.GUI(res, title, fast_gui=True)
+```
+
+:::note
+
+Because of the zero-copying mechanism, the image passed into `gui.set_image` must already be in the display-compatible
+format. That is, this field must either be a `ti.Vector(3)` (RGB) or a `ti.Vector(4)` (RGBA). In addition, each channel
+must be of type `ti.f32`, `ti.f64` or `ti.u8`.
+
+:::
+
+:::note
+
+If possible, consider enabling this option, especially when `fullscreen=True`.
+
+:::
+
+:::caution
+
+Despite the performance boost, it has many limitations as trade off:
+
+`gui.set_image` is the only available paint API in this mode.
+
+`gui.set_image` will only take Taichi 3D or 4D vector fields (RGB or RGBA) as input.
+
+:::
diff --git a/docs/lang/articles/misc/_category_.json b/docs/lang/articles/misc/_category_.json
index da3cf059960d1..f32e7be2120fd 100644
--- a/docs/lang/articles/misc/_category_.json
+++ b/docs/lang/articles/misc/_category_.json
@@ -1,4 +1,4 @@
{
"label": "Miscellaneous Topics",
- "position": 4
+ "position": 5
}
diff --git a/docs/lang/articles/misc/debugging.md b/docs/lang/articles/misc/debugging.md
index e87c5e15f1a87..cd8f08da247e5 100644
--- a/docs/lang/articles/misc/debugging.md
+++ b/docs/lang/articles/misc/debugging.md
@@ -221,8 +221,7 @@ def copy(dst: ti.template(), src: ti.template()):
## Pretty Taichi-scope traceback
-Sometimes the Python stack tracebacks resulted from **Taichi-scope** errors
-could be too complicated to read. For example:
+Taichi reports traceback messages when encountered errors in **Taichi-scope**. For example:
```python
import taichi as ti
@@ -247,117 +246,88 @@ def func0():
func0()
```
-The above snippet would result in an `AssertionError`:
+The above snippet would trigger a long and scaring `AssertionError`:
```
Traceback (most recent call last):
- File "misc/demo_excepthook.py", line 20, in
- func0()
- File "/root/taichi/python/taichi/lang/kernel.py", line 559, in wrapped
- return primal(*args, **kwargs)
- File "/root/taichi/python/taichi/lang/kernel.py", line 488, in __call__
- self.materialize(key=key, args=args, arg_features=arg_features)
- File "/root/taichi/python/taichi/lang/kernel.py", line 367, in materialize
- taichi_kernel = taichi_kernel.define(taichi_ast_generator)
- File "/root/taichi/python/taichi/lang/kernel.py", line 364, in taichi_ast_generator
- compiled()
- File "misc/demo_excepthook.py", line 18, in func0
- func1()
- File "/root/taichi/python/taichi/lang/kernel.py", line 39, in decorated
- return fun.__call__(*args)
- File "/root/taichi/python/taichi/lang/kernel.py", line 79, in __call__
- ret = self.compiled(*args)
- File "misc/demo_excepthook.py", line 14, in func1
- func2()
- File "/root/taichi/python/taichi/lang/kernel.py", line 39, in decorated
- return fun.__call__(*args)
- File "/root/taichi/python/taichi/lang/kernel.py", line 79, in __call__
- ret = self.compiled(*args)
- File "misc/demo_excepthook.py", line 10, in func2
- func3()
- File "/root/taichi/python/taichi/lang/kernel.py", line 39, in decorated
- return fun.__call__(*args)
- File "/root/taichi/python/taichi/lang/kernel.py", line 79, in __call__
- ret = self.compiled(*args)
- File "misc/demo_excepthook.py", line 6, in func3
- ti.static_assert(1 + 1 == 3)
- File "/root/taichi/python/taichi/lang/error.py", line 14, in wrapped
- return foo(*args, **kwargs)
- File "/root/taichi/python/taichi/lang/impl.py", line 252, in static_assert
+ File "/Users/lanhaidong/taichi/taichi/python/taichi/lang/ast/ast_transformer_utils.py", line 23, in __call__
+ return method(ctx, node)
+ File "/Users/lanhaidong/taichi/taichi/python/taichi/lang/ast/ast_transformer.py", line 342, in build_Call
+ node.ptr = node.func.ptr(*args, **keywords)
+ File "/Users/lanhaidong/taichi/taichi/python/taichi/lang/impl.py", line 471, in static_assert
assert cond
AssertionError
-```
-Many stack frames are the Taichi compiler implementation details, which
-could be too noisy to read. You could choose to ignore them by using
-`ti.init(excepthook=True)`, which _hooks_ on the exception handler and makes
-the stack traceback from Taichi-scope more intuitive:
+During handling of the above exception, another exception occurred:
+
+Traceback (most recent call last):
+ File "/Users/lanhaidong/taichi/taichi/python/taichi/lang/ast/ast_transformer_utils.py", line 23, in __call__
+ return method(ctx, node)
+ File "/Users/lanhaidong/taichi/taichi/python/taichi/lang/ast/ast_transformer.py", line 360, in build_Call
+ node.ptr = node.func.ptr(*args, **keywords)
+ File "/Users/lanhaidong/taichi/taichi/python/taichi/lang/kernel_impl.py", line 59, in decorated
+ return fun.__call__(*args)
+ File "/Users/lanhaidong/taichi/taichi/python/taichi/lang/kernel_impl.py", line 178, in __call__
+ ret = transform_tree(tree, ctx)
+ File "/Users/lanhaidong/taichi/taichi/python/taichi/lang/ast/transform.py", line 8, in transform_tree
+ ASTTransformer()(ctx, tree)
+ File "/Users/lanhaidong/taichi/taichi/python/taichi/lang/ast/ast_transformer_utils.py", line 26, in __call__
+ raise e
+ File "/Users/lanhaidong/taichi/taichi/python/taichi/lang/ast/ast_transformer_utils.py", line 23, in __call__
+ return method(ctx, node)
+ File "/Users/lanhaidong/taichi/taichi/python/taichi/lang/ast/ast_transformer.py", line 488, in build_Module
+ build_stmt(ctx, stmt)
+ File "/Users/lanhaidong/taichi/taichi/python/taichi/lang/ast/ast_transformer_utils.py", line 26, in __call__
+ raise e
+ File "/Users/lanhaidong/taichi/taichi/python/taichi/lang/ast/ast_transformer_utils.py", line 23, in __call__
+ return method(ctx, node)
+ File "/Users/lanhaidong/taichi/taichi/python/taichi/lang/ast/ast_transformer.py", line 451, in build_FunctionDef
+ build_stmts(ctx, node.body)
+ File "/Users/lanhaidong/taichi/taichi/python/taichi/lang/ast/ast_transformer.py", line 1086, in build_stmts
+ build_stmt(ctx, stmt)
+ File "/Users/lanhaidong/taichi/taichi/python/taichi/lang/ast/ast_transformer_utils.py", line 26, in __call__
+ raise e
+ File "/Users/lanhaidong/taichi/taichi/python/taichi/lang/ast/ast_transformer_utils.py", line 23, in __call__
+ return method(ctx, node)
+ File "/Users/lanhaidong/taichi/taichi/python/taichi/lang/ast/ast_transformer.py", line 964, in build_Expr
+ build_stmt(ctx, node.value)
+ File "/Users/lanhaidong/taichi/taichi/python/taichi/lang/ast/ast_transformer_utils.py", line 32, in __call__
+ raise TaichiCompilationError(msg)
+taichi.lang.exception.TaichiCompilationError: On line 10 of file "misc/demo_traceback.py":
+ ti.static_assert(1 + 1 == 3)
+ ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+AssertionError:
+
+...
+```
+The error message can be verbose and scary. However, many stack frames reveal
+Taichi compiler implementation details, which are too noisy for debugging.
+In current verison, you could choose to supress the level of traceback messages by setting `sys.tracebacklimit`,
+which makes the stack traceback from Taichi-scope more intuitive:
```python {2}
import taichi as ti
-ti.init(excepthook=True)
+import sys
+sys.tracebacklimit=0
...
```
which makes the result look like:
```python
-========== Taichi Stack Traceback ==========
-In () at misc/demo_excepthook.py:21:
---------------------------------------------
-@ti.kernel
-def func0():
- func1()
-
-func0() <--
---------------------------------------------
-In func0() at misc/demo_excepthook.py:19:
---------------------------------------------
- func2()
-
-@ti.kernel
-def func0():
- func1() <--
+AssertionError
-func0()
---------------------------------------------
-In func1() at misc/demo_excepthook.py:15:
---------------------------------------------
- func3()
+During handling of the above exception, another exception occurred:
-@ti.func
-def func1():
- func2() <--
-
-@ti.kernel
---------------------------------------------
-In func2() at misc/demo_excepthook.py:11:
---------------------------------------------
+taichi.lang.exception.TaichiCompilationError: On line 10 of file "misc/demo_traceback.py":
ti.static_assert(1 + 1 == 3)
+ ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+AssertionError:
-@ti.func
-def func2():
- func3() <--
-
-@ti.func
---------------------------------------------
-In func3() at misc/demo_excepthook.py:7:
---------------------------------------------
-ti.enable_excepthook()
-
-@ti.func
-def func3():
- ti.static_assert(1 + 1 == 3) <--
-
-@ti.func
---------------------------------------------
-AssertionError
+...
```
-:::note
-For IPython / Jupyter notebook users, the IPython stack traceback hook
-will be overriden by the Taichi one when `ti.enable_excepthook()` is called.
-:::
+Moreover, when filing an issue, please always unset the `sys.tracebacklimit` value and paste full traceback messages.
## Debugging Tips
diff --git a/docs/lang/articles/misc/export_kernels.md b/docs/lang/articles/misc/export_kernels.md
index 51edb87f7ee01..8847dcb2b569a 100644
--- a/docs/lang/articles/misc/export_kernels.md
+++ b/docs/lang/articles/misc/export_kernels.md
@@ -32,7 +32,7 @@ Linux. In the future, we will support macOS and Windows.
Use `ti.core.start_recording` in the Taichi program you want to export.
Suppose you want to export
-[examples/mpm88.py](https://github.com/taichi-dev/taichi/blob/master/examples/mpm88.py),
+[examples/mpm88.py](https://github.com/taichi-dev/taichi/blob/master/python/taichi/examples/mpm88.py),
here is the workflow:
### Export YAML
diff --git a/docs/lang/articles/misc/export_results.md b/docs/lang/articles/misc/export_results.md
index 9db2c468f7c7e..818928e1cd2da 100644
--- a/docs/lang/articles/misc/export_results.md
+++ b/docs/lang/articles/misc/export_results.md
@@ -94,13 +94,13 @@ print(f'The image has been saved to {filename}')
:::note
All Taichi fields have their own data types, such as `ti.u8` and
`ti.f32`. Different data types can lead to different behaviors of
-`ti.imwrite`. Please check out [GUI system](./gui.md) for
+`ti.imwrite`. Please check out [GUI system](../gui/gui.md) for
more details.
:::
- Taichi offers other helper functions that read and show images in
addition to `ti.imwrite`. They are also demonstrated in
- [GUI system./gui.md).
+ [GUI system](../gui/gui.md).
## Export videos
diff --git a/docs/lang/articles/misc/global_settings.md b/docs/lang/articles/misc/global_settings.md
index 6dd76d93ef0cb..ef87dcdc886b0 100644
--- a/docs/lang/articles/misc/global_settings.md
+++ b/docs/lang/articles/misc/global_settings.md
@@ -20,10 +20,6 @@ sidebar_position: 7
errors: `ti.init(advanced_optimization=False)`.
- Disable fast math to prevent possible undefined math behavior:
`ti.init(fast_math=False)`.
-- To print preprocessed Python code:
- `ti.init(print_preprocessed=True)`.
-- To show pretty Taichi-scope stack traceback:
- `ti.init(excepthook=True)`.
- To print intermediate IR generated: `ti.init(print_ir=True)`.
## Runtime
@@ -46,7 +42,7 @@ sidebar_position: 7
- Cache compiled runtime bitcode in **dev mode** to save start up
time: `export TI_CACHE_RUNTIME_BITCODE=1`.
- To specify how many threads to run test: `export TI_TEST_THREADS=4`
- or `ti test -t4`.
+ or `python tests/run_tests.py -t4`.
## Specifying `ti.init` arguments from environment variables
diff --git a/docs/lang/articles/misc/gui.md b/docs/lang/articles/misc/gui.md
deleted file mode 100644
index fc321deb817dd..0000000000000
--- a/docs/lang/articles/misc/gui.md
+++ /dev/null
@@ -1,529 +0,0 @@
----
-sidebar_position: 1
-
----
-
-# GUI system
-
-Taichi has a built-in GUI system to help users visualize results.
-
-## Create a window
-
-[`ti.GUI(name, res)`](https://api-docs.taichi.graphics/src/taichi.misc.html?highlight=gui%20gui#taichi.misc.gui.GUI)
-creates a window.
-
-The following code show how to create a window of resolution `640x360`:
-
-```python
-gui = ti.GUI('Window Title', (640, 360))
-```
-
-:::note
-
-If you are running Taichi on a machine without a GUI environment, consider setting `show_gui` to `False`:
-
-```python
-gui = ti.GUI('Window Title', (640, 360), show_gui=False)
-
-while gui.running:
- ...
- gui.show(f'{gui.frame:06d}.png') # save a series of screenshot
-```
-
-:::
-
-## Display a window
-
-[`gui.show(filename)`](https://api-docs.taichi.graphics/src/taichi.misc.html?highlight=show#taichi.misc.gui.GUI.show)
-helps display a window. If `filename` is specified, a screenshot will be saved to the path. For example, the following saves frames of the window to `.png`s:
-
- for frame in range(10000):
- render(img)
- gui.set_image(img)
- gui.show(f'{frame:06d}.png')
-
-
-
-## Paint on a window
-Taichi's GUI supports painting simple geometric objects, such as lines, triangles, rectangles, circles, and text.
-
-:::note
-
-The position parameter of every drawing API expects input of 2-element tuples,
-whose values are the relative position of the object range from 0.0 to 1.0.
-(0.0, 0.0) stands for the lower left corner of the window, and (1.0, 1.0) stands for the upper right corner.
-
-Acceptable input for positions are taichi fields or numpy arrays. Primitive arrays in python are NOT acceptable.
-
-For simplicity, we use numpy arrays in the examples below.
-
-:::
-
-:::tip
-
-For detailed API description, please click on the API code. For instance, click on
-`gui.get_image()` to see the description to get a GUI images.
-
-:::
-
-[`gui.set_image(pixels)`](https://api-docs.taichi.graphics/src/taichi.misc.html?highlight=set_image#taichi.misc.gui.GUI.set_image)
-sets an image to display on the window.
-
-The image pixels are set from the values of `img[i, j]`, where `i` indicates the horizontal coordinates (from left to right) and `j` the vertical coordinates (from bottom to top).
-
-If the window size is `(x, y)`, then `img` must be one of:
-
-- `ti.field(shape=(x, y))`, a gray-scale image
-
-- `ti.field(shape=(x, y, 3))`, where `3` is for `(r, g, b)` channels
-
-- `ti.field(shape=(x, y, 2))`, where `2` is for `(r, g)` channels
-
-- `ti.Vector.field(3, shape=(x, y))` `(r, g, b)` channels on each component
-
-- `ti.Vector.field(2, shape=(x, y))` `(r, g)` channels on each component
-
-- `np.ndarray(shape=(x, y))`
-
-- `np.ndarray(shape=(x, y, 3))`
-
-- `np.ndarray(shape=(x, y, 2))`
-
-The data type of `img` must be one of:
-
-- `uint8`, range `[0, 255]`
-
-- `uint16`, range `[0, 65535]`
-
-- `uint32`, range `[0, 4294967295]`
-
-- `float32`, range `[0, 1]`
-
-- `float64`, range `[0, 1]`
-
-:::note
-
-When using `float32` or `float64` as the data type, `img` entries will be clipped into range [0, 1] for display.
-
-:::
-
-[`gui.get_image()`](https://api-docs.taichi.graphics/src/taichi.misc.html?highlight=get_image#taichi.misc.gui.GUI.get_image)
-gets the 4-channel (RGBA) image shown in the current GUI system.
-
-[`gui.circle(pos)`](https://api-docs.taichi.graphics/src/taichi.misc.html?highlight=circle#taichi.misc.gui.GUI.circle)
-draws one solid circle.
-
-The color and radius of circles can be further specified with additional parameters.
-
-[`gui.circles(pos)`](https://api-docs.taichi.graphics/src/taichi.misc.html?highlight=circles#taichi.misc.gui.GUI.circles)
-draws solid circles.
-
-The color and radius of circles can be further specified with additional parameters. For a single color, use the `color` parameter.
-For multiple colors, use `palette` and `palette_indices` instead.
-
-:::note
-
-The unit of raduis in GUI APIs is number of pixels.
-
-:::
-
-For examples:
-```python
-gui.circles(pos, radius=3, color=0x068587)
-```
-draws circles all with radius of 1.5 and blue color positioned at pos array.
-
-
-```python
-gui.circles(pos, radius=3, palette=[0x068587, 0xED553B, 0xEEEEF0], palette_indices=material)
-```
-draws circles with radius of 1.5 and three different colors differed by `material`, an integer array with the same size as
-`pos`. Each integer in `material` indicates which color the associated circle use (i.e. array [0, 1, 2] indicates these three
-circles are colored separately by the first, second, and third color in `palette`.
-
-
-
-[`gui.line(begin, end)`](https://api-docs.taichi.graphics/src/taichi.misc.html?highlight=line#taichi.misc.gui.GUI.line)
-draws one line.
-
-The color and radius of lines can be further specified with additional parameters.
-
-[`gui.lines(begin, end)`](https://api-docs.taichi.graphics/src/taichi.misc.html?highlight=line#taichi.misc.gui.GUI.lines)
-draws lines.
-
-`begin` and `end` both require input of positions.
-
-The color and radius of lines can be further specified with additional parameters.
-
-For example:
-```python
-gui.lines(begin=X, end=Y, radius=2, color=0x068587)
-```
-draws line segments from X positions to Y positions with width of 2 and color in light blue.
-
-
-
-[`gui.triangle(a, b, c)`](https://api-docs.taichi.graphics/src/taichi.misc.html?highlight=triangle#taichi.misc.gui.GUI.triangle)
-draws one solid triangle.
-
-The color of triangles can be further specified with additional parameters.
-
-[`gui.triangles(a, b, c)`](https://api-docs.taichi.graphics/src/taichi.misc.html?highlight=triangles#taichi.misc.gui.GUI.triangles)
-draws solid triangles.
-
-The color of triangles can be further specified with additional parameters.
-
-For example:
-```python
-gui.triangles(a=X, b=Y, c=Z, color=0xED553B)
-```
-draws triangles with color in red and three points positioned at X, Y, and Z.
-
-
-
-[`gui.rect(topleft, bottomright)`](https://api-docs.taichi.graphics/src/taichi.misc.html?highlight=rect#taichi.misc.gui.GUI.rect)
-draws a hollow rectangle.
-
-The color and radius of the stroke of rectangle can be further specified with additional parameters.
-
-For example:
-```python
-gui.rect([0, 0], [0.5, 0.5], radius=1, color=0xED553B)
-```
-draws a rectangle of top left corner at [0, 0] and bottom right corner at [0.5, 0.5], with stroke of radius of 1 and color in red.
-
-
-
-[`gui.arrows(origin, direction)`](https://api-docs.taichi.graphics/src/taichi.misc.html?highlight=arrows#taichi.misc.gui.GUI.arrows)
-draws arrows.
-
-`origin` and `direction` both require input of positions. `origin` refers to the positions of arrows' origins, `direction`
-refers to the directions where the arrows point to relative to their origins.
-
-The color and radius of arrows can be further specified with additional parameters.
-
-For example:
-```python
-x = nunpy.array([[0.1, 0.1], [0.9, 0.1]])
-y = nunpy.array([[0.3, 0.3], [-0.3, 0.3]])
-gui.arrows(x, y, radius=1, color=0xFFFFFF)
-```
-draws two arrow originated at [0.1, 0.1], [0.9, 0.1] and pointing to [0.3, 0.3], [-0.3, 0.3] with radius of 1 and color in white.
-
-
-
-[`gui.arrow_field(direction)`](https://api-docs.taichi.graphics/src/taichi.misc.html?highlight=arrow_field#taichi.misc.gui.GUI.arrow_field)
-draws a field of arrows.
-
-The `direction` requires a field of `shape=(col, row, 2)` where `col` refers to the number of columns of arrow field and `row`
-refers to the number of rows of arrow field.
-
-The color and bound of arrow field can be further specified with additional parameters.
-
-For example:
-```python
-gui.arrow_field(x, bound=0.5, color=0xFFFFFF) # x is a field of shape=(5, 5, 2)
-```
-draws a 5 by 5 arrows pointing to random directions.
-
-
-
-[`gui.point_field(radius)`](https://api-docs.taichi.graphics/src/taichi.misc.html?highlight=point_field#taichi.misc.gui.GUI.point_field)
-draws a field of points.
-
-The `radius` requires a field of `shape=(col, row)` where `col` refers to the number of columns of arrow field and `row`
-refers to the number of rows of arrow field.
-
-The color and bound of point field can be further specified with additional parameters.
-
-For example:
-```python
-x = numpy.array([[3, 5, 7, 9], [9, 7, 5, 3], [6, 6, 6, 6]])
-gui.point_field(radius=x, bound=0.5, color=0xED553B)
-```
-draws a 3 by 4 point field of radius stored in the array.
-
-
-
-[`gui.text(content, pos)`](https://api-docs.taichi.graphics/src/taichi.misc.html?highlight=text#taichi.misc.gui.GUI.text)
-draws a line of text on screen.
-
-The font size and color of text can be further specified with additional parameters.
-
-## RGB & Hex conversion.
-
-[`ti.hex_to_rgb(hex)`](https://api-docs.taichi.graphics/src/taichi.misc.html?highlight=hex_to_rgb#taichi.misc.gui.hex_to_rgb)
-can convert a single integer value to a (R, G, B) tuple of floats.
-
-[`ti.rgb_to_hex(rgb)`](https://api-docs.taichi.graphics/src/taichi.misc.html?highlight=rgb#taichi.misc.gui.rgb_to_hex)
-can convert a (R, G, B) tuple of floats into a single integer value, e.g.,
-
-```python
-rgb = (0.4, 0.8, 1.0)
-hex = ti.rgb_to_hex(rgb) # 0x66ccff
-
-rgb = np.array([[0.4, 0.8, 1.0], [0.0, 0.5, 1.0]])
-hex = ti.rgb_to_hex(rgb) # np.array([0x66ccff, 0x007fff])
-```
-
-The return values can be used in GUI drawing APIs.
-
-
-## Event processing
-
-Every event have a key and type.
-
-_Event type_ is the type of event, for now, there are just three type of event:
-
- ti.GUI.RELEASE # key up or mouse button up
- ti.GUI.PRESS # key down or mouse button down
- ti.GUI.MOTION # mouse motion or mouse wheel
-
-_Event key_ is the key that you pressed on keyboard or mouse, can be one of:
-
- # for ti.GUI.PRESS and ti.GUI.RELEASE event:
- ti.GUI.ESCAPE # Esc
- ti.GUI.SHIFT # Shift
- ti.GUI.LEFT # Left Arrow
- 'a' # we use lowercase for alphabet
- 'b'
- ...
- ti.GUI.LMB # Left Mouse Button
- ti.GUI.RMB # Right Mouse Button
-
- # for ti.GUI.MOTION event:
- ti.GUI.MOVE # Mouse Moved
- ti.GUI.WHEEL # Mouse Wheel Scrolling
-
-A _event filter_ is a list combined of _key_, _type_ and _(type, key)_ tuple, e.g.:
-
-```python
-# if ESC pressed or released:
-gui.get_event(ti.GUI.ESCAPE)
-
-# if any key is pressed:
-gui.get_event(ti.GUI.PRESS)
-
-# if ESC pressed or SPACE released:
-gui.get_event((ti.GUI.PRESS, ti.GUI.ESCAPE), (ti.GUI.RELEASE, ti.GUI.SPACE))
-```
-
-[`gui.running`](https://api-docs.taichi.graphics/src/taichi.misc.html?highlight=running#taichi.misc.gui.GUI.running)
-can help check the state of the window. `ti.GUI.EXIT` occurs when you click on the close (X) button of a window.
- `gui.running` will obtain `False` when the GUI is being closed.
-
-For example, loop until the close button is clicked:
-
- while gui.running:
- render()
- gui.set_image(pixels)
- gui.show()
-
-You can also close the window by manually setting `gui.running` to`False`:
-
- while gui.running:
- if gui.get_event(ti.GUI.ESCAPE):
- gui.running = False
-
- render()
- gui.set_image(pixels)
- gui.show()
-
-[`gui.get_event(a, ...)`](https://api-docs.taichi.graphics/src/taichi.misc.html?highlight=get_event#taichi.misc.gui.GUI.get_event)
-tries to pop an event from the queue, and stores it into `gui.event`.
-
-For example:
-
- if gui.get_event():
- print('Got event, key =', gui.event.key)
-
-For example, loop until ESC is pressed:
-
- gui = ti.GUI('Title', (640, 480))
- while not gui.get_event(ti.GUI.ESCAPE):
- gui.set_image(img)
- gui.show()
-
-[`gui.get_events(a, ...)`](https://api-docs.taichi.graphics/src/taichi.misc.html?highlight=get_event#taichi.misc.gui.GUI.get_events)
-is basically the same as `gui.get_event`, except that it returns a generator of events instead of storing into `gui.event`:
-
- for e in gui.get_events():
- if e.key == ti.GUI.ESCAPE:
- exit()
- elif e.key == ti.GUI.SPACE:
- do_something()
- elif e.key in ['a', ti.GUI.LEFT]:
- ...
-
-[`gui.is_pressed(key, ...)`](https://api-docs.taichi.graphics/src/taichi.misc.html?highlight=is_pressed#taichi.misc.gui.GUI.is_pressed)
-can detect the keys you pressed. It must be used together with `gui.get_event`, or it won't be updated! For
-example:
-
- while True:
- gui.get_event() # must be called before is_pressed
- if gui.is_pressed('a', ti.GUI.LEFT):
- print('Go left!')
- elif gui.is_pressed('d', ti.GUI.RIGHT):
- print('Go right!')
-
-:::caution
-
-`gui.is_pressed()` must be used together with `gui.get_event()`, or it won't be updated!
-
-:::
-
-For example:
-
-```python
-while True:
- gui.get_event() # must be called before is_pressed
- if gui.is_pressed('a', ti.GUI.LEFT):
- print('Go left!')
- elif gui.is_pressed('d', ti.GUI.RIGHT):
- print('Go right!')
-```
-
-[`gui.get_cursor_pos()`](https://api-docs.taichi.graphics/src/taichi.misc.html?highlight=get_cursor#taichi.misc.gui.GUI.get_cursor_pos)
-can return current cursor position within the window. For example:
-
- mouse_x, mouse_y = gui.get_cursor_pos()
-
-[`gui.fps_limit`](https://api-docs.taichi.graphics/src/taichi.misc.html?highlight=fps#taichi.misc.gui.GUI.fps_limit)
-sets the FPS limit for a window. For example, to cap FPS at 24, simply use `gui.fps_limit = 24`. This helps reduce the overload on your hardware especially when you're using OpenGL on your integrated GPU which could make desktop slow to response.
-
-
-
-## GUI Widgets
-
-Sometimes it's more intuitive to use widgets like slider or button to control the program variables instead of using chaotic keyboard bindings. Taichi GUI provides a set of widgets for that reason:
-
-[`gui.slider(text, min, max)`](https://api-docs.taichi.graphics/src/taichi.misc.html?highlight=slider#taichi.misc.gui.GUI.slider)
-creates a slider following the text `{text}: {value:.3f}`.
-
-[`gui.label(text)`](https://api-docs.taichi.graphics/src/taichi.misc.html?highlight=label#taichi.misc.gui.GUI.label)
-displays the label as: `{text}: {value:.3f}`.
-
-[`gui.button(text)`](https://api-docs.taichi.graphics/src/taichi.misc.html?highlight=button#taichi.misc.gui.GUI.button)
-creates a button with text on it.
-
-For example:
-```python
-radius = gui.slider('Radius', 1, 50)
-
-while gui.running:
- print('The radius now is', radius.value)
- ...
- radius.value += 0.01
- ...
- gui.show()
-```
-
-
-
-## Image I/O
-
-[`ti.imwrite(img, filename)`](https://api-docs.taichi.graphics/src/taichi.misc.html?highlight=imwrite#taichi.misc.image.imwrite)
-can export a `np.ndarray` or Taichi field (`ti.Matrix.field`, `ti.Vector.field`, or `ti.field`) to a specified location `filename`.
-
-Same as `ti.GUI.show(filename)`, the format of the exported image is determined by **the suffix of** `filename` as well. Now `ti.imwrite` supports exporting images to `png`, `img` and `jpg` and we recommend using `png`.
-
-Please make sure that the input image has **a valid shape**. If you want to export a grayscale image, the input shape of field should be `(height, weight)` or `(height, weight, 1)`. For example:
-
-```python
-import taichi as ti
-
-ti.init()
-
-shape = (512, 512)
-type = ti.u8
-pixels = ti.field(dtype=type, shape=shape)
-
-@ti.kernel
-def draw():
- for i, j in pixels:
- pixels[i, j] = ti.random() * 255 # integers between [0, 255] for ti.u8
-
-draw()
-
-ti.imwrite(pixels, f"export_u8.png")
-```
-
-Besides, for RGB or RGBA images, `ti.imwrite` needs to receive a field which has shape `(height, width, 3)` and `(height, width, 4)` individually.
-
-Generally the value of the pixels on each channel of a `png` image is an integer in \[0, 255\]. For this reason, `ti.imwrite` will **cast fields** which has different data types all **into integers between \[0, 255\]**. As a result, `ti.imwrite` has the following requirements for different data types of input fields:
-
-- For float-type (`ti.f16`, `ti.f32`, etc.) input fields, **the value of each pixel should be float between \[0.0, 1.0\]**. Otherwise `ti.imwrite` will first clip them into \[0.0, 1.0\]. Then they are multiplied by 256 and cast to integers ranging from \[0, 255\].
-- For int-type (`ti.u8`, `ti.u16`, etc.) input fields, **the value of each pixel can be any valid integer in its own bounds**. These integers in this field will be scaled to \[0, 255\] by being divided over the upper bound of its basic type accordingly.
-
-Here is another example:
-
-```python
-import taichi as ti
-
-ti.init()
-
-shape = (512, 512)
-channels = 3
-type = ti.f32
-pixels = ti.Matrix.field(channels, dtype=type, shape=shape)
-
-@ti.kernel
-def draw():
- for i, j in pixels:
- for k in ti.static(range(channels)):
- pixels[i, j][k] = ti.random() # floats between [0, 1] for ti.f32
-
-draw()
-
-ti.imwrite(pixels, f"export_f32.png")
-```
-
-[`ti.imread(filename)`](https://api-docs.taichi.graphics/src/taichi.misc.html?highlight=imread#taichi.misc.image.imread)
-loads an image from the target filename and returns it as a `np.ndarray(dtype=np.uint8)`.
-Each value in this returned field is an integer in [0, 255].
-
-[`ti.imshow(img, windname)`](https://api-docs.taichi.graphics/src/taichi.misc.html?highlight=imshow#taichi.misc.image.imshow)
-creates an instance of ti.GUI and show the input image on the screen. It has the same logic as `ti.imwrite` for different data types.
-
-[`ti.imresize(img, w)`](https://api-docs.taichi.graphics/src/taichi.misc.html?highlight=imresize#taichi.misc.image.imresize)
-resizes the img specified.
-
-## Zero-copying frame buffer
-When the GUI resolution (window size) is large, it sometimes becomes difficult to achieve 60 FPS even without any kernel
-invocations between two frames.
-
-This is mainly due to the copy overhead, where Taichi GUI needs to copy the image buffer from one place to another.
-This process is necessary for the 2D drawing functions, such as `gui.circles`, to work. The larger the image shape is,
-the larger the overhead.
-
-Fortunately, sometimes your program only needs `gui.set_image` alone. In such cases, you can enable the `fast_gui` option
-for better performance. This mode allows Taichi GUI to directly write the image data to the frame buffer without additional
-copying, resulting in a much better FPS.
-
-```python
-gui = ti.GUI(res, title, fast_gui=True)
-```
-
-:::note
-
-Because of the zero-copying mechanism, the image passed into `gui.set_image` must already be in the display-compatible
-format. That is, this field must either be a `ti.Vector(3)` (RGB) or a `ti.Vector(4)` (RGBA). In addition, each channel
-must be of type `ti.f32`, `ti.f64` or `ti.u8`.
-
-:::
-
-:::note
-
-If possible, consider enabling this option, especially when `fullscreen=True`.
-
-:::
-
-:::caution
-
-Despite the performance boost, it has many limitations as trade off:
-
-`gui.set_image` is the only available paint API in this mode.
-
-`gui.set_image` will only take Taichi 3D or 4D vector fields (RGB or RGBA) as input.
-
-:::
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/docs/lang/articles/tutorials/_category_.json b/docs/lang/articles/tutorials/_category_.json
new file mode 100644
index 0000000000000..6e052e4632618
--- /dev/null
+++ b/docs/lang/articles/tutorials/_category_.json
@@ -0,0 +1,4 @@
+{
+ "label": "Tutorials",
+ "position": 7
+}
diff --git a/docs/lang/articles/tutorials/ndarray_android.md b/docs/lang/articles/tutorials/ndarray_android.md
new file mode 100644
index 0000000000000..0349ed9b2e1a0
--- /dev/null
+++ b/docs/lang/articles/tutorials/ndarray_android.md
@@ -0,0 +1,370 @@
+---
+sidebar_position: 1
+---
+
+# Run a Taichi Program using Ndarray on Android
+
+Taichi's JIT (Just In Time) module compiles a Taichi kernel to the compute shaders according to the specified backend (`arch` in `ti.init()`) and executes these shaders in Taichi's JIT runtime. Taichi's AOT (Ahead of Time) module, however, builds and saves the necessary compute shaders so that you can load and execute these shaders in your own runtime without a Python environment.
+
+Taking a simulation of celestial bodies' orbits as an example, this tutorial walks you through the process of running a Taichi program using Ndarray on Android.
+
+> [Taichi's AOT (Ahead Of Time) module](https://github.com/taichi-dev/taichi/issues/3642) is currently a proof of concept under development and subject to change in the future.
+
+## A definition of Ndarray
+
+Taichi provides a data container called Ndarray. An Ndarray is a multidimensional container of elements of the same type and size; an element in an Ndarray is virtually a scalar or a tensor.
+
+### Ndarray shape
+
+Ndarray shape defines the Ndarray's layout; element shape defines the element's layout. For example:
+
+- An Ndarray with an Ndarray shape of [2, 1024] and an element shape of [] is an array of 2 x 1,024 = 2,048 scalars.
+- An Ndarray with an Ndarray shape of [128, 128] and an element shape of [2, 4] is an array of 128 x 128 = 16,384 2 x 4 matrices.
+
+### Ndarray dimension
+
+The dimension here refers to the number of dimensions of an Ndarray. For example:
+
+- The dimension of an Ndarray with a shape of [1, 2, 3] is three.
+- The dimension of an Ndarray with a shape of [500] is one.
+
+### Benefit of Ndarray
+
+Each Ndarray has a fixed dimension but gives you the flexibility of changing its shape in accordance with its dimension.
+
+Unlike a field's shape, which requires you to rewrite and recompile your Taichi program once it is changed, an Ndarray's shape can be *dynamically* changed without the need to recompile.
+
+Taking the simulation of celestial bodies' orbits as an example, suppose you wish to double the number of your celestial bodies to 2,000:
+
+- With Taichi field, you have to compile twice;
+- With Ndarray, all you need to do is to update your runtime program.
+
+## Run a Taichi program using Ndarray on Android
+
+The following section walks you through the process of running a Taichi program using Ndarray on Android.
+
+1. [Generate necessary compute shaders](#generate-necessary-compute-shaders)
+2. [Parse the generated JSON file](#parse-the-generated-json-file)
+3. [Prepare SSBO and shape information](#prepare-ssbo-and-shape-information)
+4. [Prepare rendering shaders](#prepare-rendering-shaders)
+5. [Execute all shaders](#execute-all-shaders)
+
+:::note
+
+From Step 2, you are required to come up with your own runtime program. We provide an [example Java runtime program for Android](https://github.com/taichi-dev/taichi-aot-demo/blob/master/nbody_ndarray/java_runtime/NbodyNdarray.java) for your reference, but you may need to adapt these codes for your platform and programming language.
+
+:::
+
+### Generate necessary compute shaders
+
+The following Python script defines a Taichi AOT module for generating and saving the necessary compute shaders (GLES shaders in this case) based on the chosen backend (OpenGL).
+
+> Taichi kernels and compute shaders are *not* one-to-one mapping. Each Taichi kernel can generate multiple compute shaders, the number *usually* comparable to that of the loops in the kernel.
+
+
+
+```python
+import taichi as ti
+
+ti.init(arch=ti.opengl, use_gles=True, allow_nv_shader_extension=False)
+
+# Define constants for computation
+G = 1
+PI = 3.141592653
+N = 1000
+m = 5
+galaxy_size = 0.4
+planet_radius = 1
+init_vel = 120
+h = 1e-5
+substepping = 10
+
+# Define Taichi kernels
+@ti.kernel
+def initialize(pos: ti.any_arr(element_dim=1), vel: ti.any_arr(element_dim=1)):
+ center=ti.Vector([0.5, 0.5])
+ for i in pos:
+ theta = ti.random() * 2 * PI
+ r = (ti.sqrt(ti.random()) * 0.7 + 0.3) * galaxy_size
+ offset = r * ti.Vector([ti.cos(theta), ti.sin(theta)])
+ pos[i] = center+offset
+ vel[i] = [-offset.y, offset.x]
+ vel[i] *= init_vel
+
+@ti.kernel
+def compute_force(pos: ti.any_arr(element_dim=1), vel: ti.any_arr(element_dim=1), force: ti.any_arr(element_dim=1)):
+ for i in pos:
+ force[i] = ti.Vector([0.0, 0.0])
+ for i in pos:
+ p = pos[i]
+ for j in pos:
+ if i != j:
+ diff = p-pos[j]
+ r = diff.norm(1e-5)
+ f = -G * m * m * (1.0/r)**3 * diff
+ force[i] += f
+ dt = h/substepping
+ for i in pos:
+ vel[i].atomic_add(dt*force[i]/m)
+ pos[i].atomic_add(dt*vel[i])
+
+# Define Ndarrays
+pos = ti.Vector.ndarray(2, ti.f32, N)
+vel = ti.Vector.ndarray(2, ti.f32, N)
+force = ti.Vector.ndarray(2, ti.f32, N)
+
+# Run the AOT module builder
+def aot():
+ m = ti.aot.Module(ti.opengl)
+ m.add_kernel(initialize, (pos, vel))
+ m.add_kernel(compute_force, (pos, vel, force))
+
+ dir_name = 'nbody_aot'
+ m.save(dir_name, '')
+aot()
+```
+
+**In line 3, you initialize Taichi:**
+
+1. Set `use_gles` to `True` to generate GLES compute shaders for Android.
+2. Set `allow_nv_shader_extension` to `False` to prevent the generated GLES compute shaders from using Nvidia GL extensions on Android.
+
+> This setting is because Android supports GLES APIs but GLES does not support `NV_SHADER_EXTENSION`.
+
+**In line 50-58, you define and build the Taichi AOT module:**
+
+1. Create a Taichi AOT module, specifying its backend as OpenGL:
+
+```python
+ m = ti.aot.Module(ti.opengl)
+```
+
+2. Add the required kernels `initialize` and `compute_force`, each with its own Ndarrays, to the module:
+
+```python
+m.add_kernel(initialize, (pos, vel))
+
+m.add_kernel(compute_force, (pos, vel, force))
+```
+
+3. Specify a folder under your current working directory for holding the files that the module generates:
+
+```python
+dir_name = 'nbody_aot'
+
+m.save(dir_name, '')
+```
+
+*The necessary compute shaders together with a JSON file appear under the specified directory.*
+
+### Parse the generated JSON file
+
+:::note
+
+From this section, you are required to come up with your own runtime program. We provide an [example Java runtime program for Android](https://github.com/taichi-dev/taichi-aot-demo/blob/master/nbody_ndarray/java_runtime/NbodyNdarray.java) for your reference, but you may need to adapt these codes for your platform and programming language.
+
+:::
+
+After generating the necessary GLES compute shaders, you need to write your runtime program to parse the following JSON file to some data structures. The JSON file contains all the necessary information for executing the compute shaders. Organized by Taichi kernel, it provides a clearer image of the compute shaders and Ndarrays in each kernel. Let's take a closer look at the structure.
+
+> Here, the JSON object for the `compute_force` kernel is omitted for brevity. For a complete JSON file, see [metadata.json](https://github.com/taichi-dev/taichi-aot-demo/blob/master/nbody_ndarray/res/metadata.json).
+
+- **Organized by Taichi kernel**
+
+ - `initialize` (line 4)
+ - `compute_force` (line 51)
+
+- **Kernel-specific compute shaders**
+
+ Taking `initialize` as an example, the kernel has generated one compute shader named `initialize_c54_00` (line 7) and the other named `initialize_c54_01` (line 13).
+
+- **Kernel-specific** `args_buff`
+
+ The `initialize` kernel is assigned an `args_buffer` of `128` Bytes (line 21). Note that the size of `args_buffer` is dependent on the number of Ndarrays (`pos` and `vel`) that the kernel takes, (see `arg_count` in line 19). The `initialize` kernel, or each kernel more precisely, has a dedicated `args_buffer` for storing scalar arguments specified in `scalar_args` (line 27) and Ndarray shape information in accordance with what `array_args` (line 28-45) specifies.
+
+ Ndarrays' shape information is organized by their argument index in the `array_args` JSON array: `0` (line 29) corresponds to the `pos` Ndarray, and `1` (line 37) corresponds to the `vel` Ndarray. The argument index is determined by the sequence by which you pass in the Ndarrays when calling `add_kernel()`. See line 53 in the Python script.
+
+ The `pos` Ndarray's shape information in `args_buffer` has an offset of `64` Bytes in `args_buffer` (line 64). According to line 35 and line 43, the `pos` Ndarray's shape information occupies 96 - 64 = 32 Bytes in `args_buffer`.
+
+ :::tip ATTENTION
+ The JSON file only specifies the dimension of the corresponding Ndarray (line 30, 38), allowing you to dynamically update an Ndarray's shape in your runtime program.
+ :::
+
+- **Kernel-specific binding index**
+
+ `used.arr_arg_to_bind_idx` (line 46) maps the SSBO of each Ndarray in the kernel to a "more global" binding index for the compute shaders. For example, `"1": 5` (line 48) binds the `vel` Ndarray to the binding index `5`.
+
+```json
+{
+ "aot_data": {
+ "kernels": {
+ "initialize": {
+ "tasks": [
+ {
+ "name": "initialize_c54_00",
+ "src": "nbody_aot/initialize_c54_00.glsl",
+ "workgroup_size": 1,
+ "num_groups": 1
+ },
+ {
+ "name": "initialize_c54_01",
+ "src": "nbody_aot/initialize_c54_01.glsl",
+ "workgroup_size": 128,
+ "num_groups": 256
+ }
+ ],
+ "arg_count": 2,
+ "ret_count": 0,
+ "args_buf_size": 128,
+ "ret_buf_size": 0,
+ "ext_arr_access": {
+ "0": 2,
+ "1": 3
+ },
+ "scalar_args": {},
+ "arr_args": {
+ "0": {
+ "field_dim": 1,
+ "is_scalar": false,
+ "element_shape": [
+ 2
+ ],
+ "shape_offset_in_bytes_in_args_buf": 64
+ },
+ "1": {
+ "field_dim": 1,
+ "is_scalar": false,
+ "element_shape": [
+ 2
+ ],
+ "shape_offset_in_bytes_in_args_buf": 96
+ }
+ },
+ "used.arr_arg_to_bind_idx": {
+ "0": 4,
+ "1": 5
+ }
+ },
+ "compute_force": {...}
+ },
+ "kernel_tmpls": {},
+ "fields": [],
+ "root_buffer_size": 0
+ }
+}
+```
+The following provides a detailed description of the keys in the generated JSON file:
+
+`aot_data`: The overarching JSON object.
+ - `kernels`: All Taichi kernels.
+ - `$(kernel_name)`: Name of a specific Taichi kernel.
+ - `tasks`: A JSON array of the generated compute shaders.
+ - `name`: Name of a specific compute shader.
+ - `src`: Relative path to the shader file.
+ - `workgroup_size`: N/A
+ - `num_groups`: N/A
+ - `arg_count`: Number of the arguments that the Taichi kernel takes.
+ - `ret_count`: Number of the values that the Taichi kernel returns.
+ - `args_buf_size`: The size of `args_buf` in Bytes.
+ - `ret_buf_size`: The size of `ret_buf` in Bytes.
+ - `scalar_args`: Scalar arguments that the kernel takes.
+ - `arr_args`: Shape information of the Ndarrays in the kernel.
+ - `$(arg_index)`: Argument index of an Ndarray
+ - `field_dim`: The dimension of the Ndarray.
+ - `is_scalar`: Whether the elements in the Ndarray are scalar.
+ - `element_shape`: An `int` array indicating the shape of each element in the Ndarray.
+ - `shape_offset_in_bytes_in_args_buf`: The offset of the Ndarray's shape information in `args_buf`.
+ - `used.arr_arg_to_bind_idx`: A map specifying the SSBO to bind for a given Ndarray. For example, `"1": 5` (line 48) binds the `vel` Ndarray to the binding index `5`.
+
+*Well, we hope you were not overwhelmed with that much information coming in all at once. In the following section, we will revisit the JSON file, as well as provide tables and graphs that help illustrate some of the concepts and notions listed above.*
+
+### Prepare SSBO and shape information
+
+Before executing the GLES compute shaders in your runtime program, you need to get all your resources ready, including:
+
+- Bind SSBO for the corresponding buffer
+- Bind SSBO for each Ndarray
+- Fill `args_buffer` with Ndarray shape information
+
+#### Bind SSBO for the corresponding buffer
+
+The following table lists the buffers commonly used in a Taichi program together with their binding indexes:
+
+| **Buffer** | **Global/kernel-spedific** | **Storing** | **Binding index** |
+| ------------- | -------------------------- | ------------------------------------------------------------ | ----------------- |
+| `root_buffer` | Global | All fields with fixed offsets and of fixed sizes. | `0` |
+| `gtmp_buffer` | Global | Global temporary data | `1` |
+| `args_buffer` | Kernel-specific | Arguments passed to the Taichi kernel - Scalar arguments
- Each Ndarray's shape information:
- Shape of the Ndarray
- Element shape
| `2` |
+
+1. You *only* need to bind an SSBO for `root_buffer` if your Taichi script uses at least one field. Skip this step if your script does not involve field.
+2. Bind a small SSBO, say an SSBO of 1,024 Bytes, to `1`, the binding index of `gtmp_buffer`.
+3. Bind an SSBO of 64 x 5 = 320 Bytes to `2`, the binding index of `args_buffer`.
+
+#### Bind SSBO for each Ndarray
+
+Before running a specific kernel in your runtime program (the `initialize` kernel for example), you must bind SSBO of a proper size for each Ndarray in the kernel in accordance to the value of `used.arr_arg_to_bind_idx`.
+
+The following is a summary of line 29-49 of the above JSON file:
+
+| Ndarray | Taichi kernel | Dimension | Element shape | Argument index | Binding index |
+| ------- | ------------- | --------- | ------------- | -------------- | ------------- |
+| `pos` | `initialize` | `1` | `[2]` | `0` | `4` |
+| `vel` | `initialize` | `1` | `[2]` | `1` | `5` |
+
+If you give each Ndarray a shape [500], and an element shape [2] (meaning that each element is a 2-D vector):
+
+- Each Ndarray has 500 x 2 = 1,000 numbers
+- Because the number type is float (as specified in the above Python script), the size of each Ndarray's SSBO is 1,000 x 4 = 4,000 Bytes.
+
+Therefore you need to:
+
+- Bind an SSBO of 4,000 Bytes to the binding index `4` for the `pos` Ndarray.
+- Bind an SSBO of 4,000 Bytes to the binding index `5` for the `vel` Ndarray.
+
+#### Fill `args_buffer` with Ndarray shape information
+
+When explaining the JSON file, we mention that each kernel has a dedicated `args_buffer` for storing scalar arguments specified in `scalar_args` and Ndarray shape information in accordance with what `array_args` specifies. `array_args` does not specify the Ndarray shape, therefore the final step in your preparation is to fill `args_buffer` with each Ndarray's shape information in your runtime program.
+
+The typical size of an `args_buffer` is 64 + 64 x 4 Bytes. The first 64 Bytes are reserved for scalar arguments; the remaining buffer is then 64 x 4 Bytes. Each Ndarray is allocated 8 x 4 Bytes for storing its shape information (each has *at most* 8 numbers to indicate its shape information), therefore the remaining buffer can store up to 8 Ndarrays' shape information.
+
+- If your Ndarray shape is [100, 200] and element dimension [3, 2], then you fill 100, 200, 3, and 2 in the corresponding location.
+- In this case, both `pos` and `vel` have an Ndarray shape of [500] and an element dimension of [2]. Therefore, you fill 500 and 2 in the corresponding locations.
+
+### Prepare rendering shaders
+
+To perform the rendering (drawing celestial bodies in this case), you are required to write a vertex shader and a fragment shader.
+
+### Execute all shaders
+
+When executing shaders in your runtime program, ensure that you bind SSBOs before executing a Taichi kernel and unbind them when you are done.
+
+ Our [example Android Java runtime](https://github.com/taichi-dev/taichi-aot-demo/blob/master/nbody_ndarray/java_runtime/NbodyNdarray.java) does the following:
+
+1. Run the GLES compute shaders in `initialize` once.
+2. For each frame:
+ 1. Run the GLES compute shaders in `compute_force` 10 times.
+ 2. Run the vertex and fragment shaders once to do the rendering.
+
+
+
+## OpenGL-specific Terms & Definitions
+
+### OpenGL ES (GLES)
+
+OpenGL ES (GLES) is the OpenGL APIs for Embedded Systems. According to [its specifications](https://www.khronos.org/api/opengles), a desktop OpenGL driver supports all GLES APIs.
+
+### OpenGL Shading Language (GLSL)
+
+The OpenGL Shading Language (GLSL) is the primary shading language for OpenGL. GLSL is a C-style language supported directly by OpenGL without extensions.
+
+### Shader
+
+A shader is a user-defined program designed for computing or rendering at a certain stage of a graphics processor.
+
+### SSBO (Shader Storage Buffer Object)
+
+Each Taichi kernel can generate multiple compute shaders, which use SSBO (Shader Storage Buffer Object) as buffer for accessing data.
+
+There are two types of SSBOs: One type corresponds to the buffers maintained by Taichi and includes `root_buffer`, `gtmp_buffer`, and `args_buffer`; the other type corresponds to the Ndarrays maintained by developers and used for sharing data.
+
+> You are required to bind the generated shaders to the corresponding SSBOs in your runtime program. The binding index of an Ndarray's SSBO starts off with `4`.
diff --git a/docs/variable.json b/docs/variable.json
new file mode 100644
index 0000000000000..0967ef424bce6
--- /dev/null
+++ b/docs/variable.json
@@ -0,0 +1 @@
+{}
diff --git a/examples/algorithm/print_offset.py b/examples/algorithm/print_offset.py
deleted file mode 100644
index b905126bd4f9b..0000000000000
--- a/examples/algorithm/print_offset.py
+++ /dev/null
@@ -1,36 +0,0 @@
-import taichi as ti
-
-ti.init(arch=ti.cpu, print_ir=True)
-
-n = 4
-m = 8
-
-a = ti.field(dtype=ti.i32)
-ti.root.dense(ti.ij, (1, 2)).dense(ti.ij, 2).dense(ti.ij, 2).place(a)
-
-
-@ti.kernel
-def fill():
- for i, j in a:
- base = ti.get_addr(a.snode, [0, 0])
- a[i, j] = int(ti.get_addr(a.snode, [i, j]) - base) // 4
-
-
-fill()
-print(a.to_numpy())
-
-ti.get_runtime().prog.visualize_layout('layout.pdf')
-
-gui = ti.GUI('layout', res=(256, 512), background_color=0xFFFFFF)
-
-while True:
- for i in range(1, m):
- gui.line(begin=(0, i / m), end=(1, i / m), radius=2, color=0x000000)
- for i in range(1, n):
- gui.line(begin=(i / n, 0), end=(i / n, 1), radius=2, color=0x000000)
- for i in range(n):
- for j in range(m):
- gui.text(f'{a[i, j]}', ((i + 0.3) / n, (j + 0.75) / m),
- font_size=30,
- color=0x0)
- gui.show()
diff --git a/examples/autodiff/minimization.py b/examples/autodiff/minimization.py
deleted file mode 100644
index c59af83e0cdec..0000000000000
--- a/examples/autodiff/minimization.py
+++ /dev/null
@@ -1,40 +0,0 @@
-import random
-
-import taichi as ti
-
-ti.init(arch=ti.cpu)
-
-n = 8
-x = ti.field(dtype=ti.f32, shape=n, needs_grad=True)
-y = ti.field(dtype=ti.f32, shape=n)
-L = ti.field(dtype=ti.f32, shape=(), needs_grad=True)
-
-
-@ti.kernel
-def reduce():
- for i in range(n):
- L[None] += 0.5 * (x[i] - y[i])**2
-
-
-# Initialize vectors
-for i in range(n):
- x[i] = random.random()
- y[i] = random.random()
-
-
-@ti.kernel
-def gradient_descent():
- for i in x:
- x[i] -= x.grad[i] * 0.1
-
-
-# Optimize with 100 gradient descent iterations
-for k in range(100):
- with ti.Tape(loss=L):
- reduce()
- print('Loss =', L[None])
- gradient_descent()
-
-for i in range(n):
- # Now you should approximately have x[i] == y[i]
- print(x[i], y[i])
diff --git a/examples/autodiff/regression.py b/examples/autodiff/regression.py
deleted file mode 100644
index 2692b8afafa87..0000000000000
--- a/examples/autodiff/regression.py
+++ /dev/null
@@ -1,90 +0,0 @@
-import random
-
-import matplotlib.pyplot as plt
-import numpy as np
-
-import taichi as ti
-import taichi as tc
-
-ti.init(arch=ti.cpu)
-
-tc.set_gdb_trigger(True)
-
-number_coeffs = 4
-learning_rate = 1e-4
-
-N = 32
-x, y = ti.field(ti.f32, shape=N, needs_grad=True), ti.field(ti.f32,
- shape=N,
- needs_grad=True)
-coeffs = ti.field(ti.f32, shape=number_coeffs, needs_grad=True)
-loss = ti.field(ti.f32, shape=(), needs_grad=True)
-
-
-@ti.kernel
-def regress():
- for i in x:
- v = x[i]
- est = 0.0
- for j in ti.static(range(number_coeffs)):
- est += coeffs[j] * (v**j)
- loss[None] += 0.5 * (y[i] - est)**2
-
-
-@ti.kernel
-def update():
- for i in ti.static(range(number_coeffs)):
- coeffs[i] -= learning_rate * coeffs.grad[i]
-
-
-xs = []
-ys = []
-
-for i in range(N):
- v = random.random() * 5 - 2.5
- xs.append(v)
- x[i] = v
- y[i] = (v - 1) * (v - 2) * (v + 2) + random.random() - 0.5
-
-regress()
-
-print('y')
-for i in range(N):
- y.grad[i] = 1
- ys.append(y[i])
-print()
-
-use_tape = True
-
-for i in range(1000):
- if use_tape:
- with ti.Tape(loss=loss):
- regress()
- else:
- ti.clear_all_gradients()
- loss[None] = 0
- loss.grad[None] = 1
- regress()
- regress.grad()
- print('Loss =', loss[None])
- update()
- for i in range(number_coeffs):
- print(coeffs[i], end=', ')
- print()
-
-curve_xs = np.arange(-2.5, 2.5, 0.01)
-curve_ys = curve_xs * 0
-for i in range(number_coeffs):
- curve_ys += coeffs[i] * np.power(curve_xs, i)
-
-plt.title('Nonlinear Regression with Gradient Descent (3rd order polynomial)')
-ax = plt.gca()
-ax.scatter(xs, ys, label='data', color='r')
-ax.plot(curve_xs, curve_ys, label='fitted')
-ax.legend()
-ax.grid(True)
-ax.spines['left'].set_position('zero')
-ax.spines['right'].set_color('none')
-ax.spines['bottom'].set_position('zero')
-ax.spines['top'].set_color('none')
-plt.show()
diff --git a/examples/autodiff/simple_derivative.py b/examples/autodiff/simple_derivative.py
deleted file mode 100644
index 32d3e2bf93eb9..0000000000000
--- a/examples/autodiff/simple_derivative.py
+++ /dev/null
@@ -1,58 +0,0 @@
-import matplotlib.pyplot as plt
-
-import taichi as ti
-
-ti.init(arch=ti.cpu)
-
-N = 2048
-x, y = ti.field(ti.f32), ti.field(ti.f32)
-
-ti.root.dense(ti.i, N).place(x, x.grad, y, y.grad)
-
-
-@ti.kernel
-def poly():
- for i in x:
- v = x[i]
- ret = 0.0
- guard = 0.2
- if v < -guard or v > guard:
- ret = 4 / ti.max(v, 0.1)
- else:
- ret = 0
- y[i] = ret
-
-
-xs = []
-ys = []
-grad_xs = []
-
-for i in range(N):
- v = ((i + 0.5) / N) * 2 - 1
- xs.append(v)
- x[i] = v
-
-poly()
-
-print('y')
-for i in range(N):
- y.grad[i] = 1
- ys.append(y[i])
-print()
-
-poly.grad()
-print('grad_x')
-for i in range(N):
- grad_xs.append(x.grad[i])
-
-plt.title('Auto Diff')
-ax = plt.gca()
-ax.plot(xs, ys, label='f(x)')
-ax.plot(xs, grad_xs, label='f\'(x)')
-ax.legend()
-ax.grid(True)
-ax.spines['left'].set_position('zero')
-ax.spines['right'].set_color('none')
-ax.spines['bottom'].set_position('zero')
-ax.spines['top'].set_color('none')
-plt.show()
diff --git a/examples/chi_examples/CMakeLists.txt b/examples/chi_examples/CMakeLists.txt
deleted file mode 100644
index c83db5cb1895f..0000000000000
--- a/examples/chi_examples/CMakeLists.txt
+++ /dev/null
@@ -1,16 +0,0 @@
-cmake_minimum_required(VERSION 3.12)
-
-set (CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -std=c++17 -fsized-deallocation -Wall -march=nehalem -DTI_ARCH_x64 -DTI_INCLUDED -O3 -DNDEBUG")
-set(CHI_EXAMPLES chi_examples)
-
-project(chi_examples)
-
-add_executable(${CHI_EXAMPLES} main.cpp)
-target_include_directories(${CHI_EXAMPLES}
- PUBLIC $ENV{TAICHI_REPO_DIR}
- PUBLIC $ENV{TAICHI_REPO_DIR}/external/include
- PUBLIC $ENV{TAICHI_REPO_DIR}/external/spdlog/include
- PUBLIC $ENV{TAICHI_REPO_DIR}/external/glad/include
- PUBLIC $ENV{TAICHI_REPO_DIR}/external/glfw/include
-)
-target_link_libraries(${CHI_EXAMPLES} $ENV{TAICHI_REPO_DIR}/build/libtaichi_export_core.so)
diff --git a/examples/chi_examples/README.md b/examples/chi_examples/README.md
deleted file mode 100644
index 280a58c5ad5c2..0000000000000
--- a/examples/chi_examples/README.md
+++ /dev/null
@@ -1,23 +0,0 @@
-# How to use CHI IR Builder
-
-## Build Taichi
-
-Follow the steps in `https://docs.taichi.graphics/lang/articles/contribution/dev_install`
-
-Add option `-DTI_EXPORT_CORE=ON` to your `cmake` command (i.e. use `cmake .. DTI_EXPORT_CORE=ON`).
-
-Make sure taichi is built under `$TAICHI_REPO_DIR/build` directory.
-
-After building, `$TAICHI_REPO_DIR/build/libtaichi_export_core.so` should exist.
-
-## Link with the Taichi Shared Library
-
-`main.cpp` shows how to construct and run Taichi kernels using CHI IR Builder.
-
-```bash
-mkdir build
-cd build
-cmake ..
-make
-./chi_examples
-```
diff --git a/examples/rendering/taichi_logo.py b/examples/rendering/taichi_logo.py
deleted file mode 100644
index a3ad0a5bd67a6..0000000000000
--- a/examples/rendering/taichi_logo.py
+++ /dev/null
@@ -1,22 +0,0 @@
-import taichi as ti
-
-ti.init()
-
-n = 512
-x = ti.field(dtype=ti.f32, shape=(n, n))
-
-
-@ti.kernel
-def paint():
- for i, j in ti.ndrange(n * 4, n * 4):
- # 4x4 super sampling:
- ret = ti.taichi_logo(ti.Vector([i, j]) / (n * 4))
- x[i // 4, j // 4] += ret / 16
-
-
-paint()
-
-gui = ti.GUI('Logo', (n, n))
-while gui.running:
- gui.set_image(x)
- gui.show()
diff --git a/external/FP16 b/external/FP16
new file mode 160000
index 0000000000000..0a92994d729ff
--- /dev/null
+++ b/external/FP16
@@ -0,0 +1 @@
+Subproject commit 0a92994d729ff76a58f692d3028ca1b64b145d91
diff --git a/external/SPIRV-Cross b/external/SPIRV-Cross
index 97a438d214b24..131278458ea8e 160000
--- a/external/SPIRV-Cross
+++ b/external/SPIRV-Cross
@@ -1 +1 @@
-Subproject commit 97a438d214b24e4958ca137a18639670648cedd0
+Subproject commit 131278458ea8eebe6a6e9c476fbcf71278726e1a
diff --git a/external/SPIRV-Headers b/external/SPIRV-Headers
index 5ea2d62e8c0dd..b42ba6d92faf6 160000
--- a/external/SPIRV-Headers
+++ b/external/SPIRV-Headers
@@ -1 +1 @@
-Subproject commit 5ea2d62e8c0ddd9e2a7d0ca5e3f2335e09e5f408
+Subproject commit b42ba6d92faf6b4938e6f22ddd186dbdacc98d78
diff --git a/external/SPIRV-Reflect b/external/SPIRV-Reflect
index 272e050728de8..1aceb6af56e74 160000
--- a/external/SPIRV-Reflect
+++ b/external/SPIRV-Reflect
@@ -1 +1 @@
-Subproject commit 272e050728de8d4a4ce9e7101c1244e6ff56e5b0
+Subproject commit 1aceb6af56e74b92a00378842dda5c5a73f49a4b
diff --git a/external/SPIRV-Tools b/external/SPIRV-Tools
index b46995741b97c..845f3efb8a4eb 160000
--- a/external/SPIRV-Tools
+++ b/external/SPIRV-Tools
@@ -1 +1 @@
-Subproject commit b46995741b97c714e211fe5df8590991ae998475
+Subproject commit 845f3efb8a4eb32e5f484aa6ea9b9e3716d6f7ec
diff --git a/external/Vulkan-Headers b/external/Vulkan-Headers
new file mode 160000
index 0000000000000..5c0fa1d68f48a
--- /dev/null
+++ b/external/Vulkan-Headers
@@ -0,0 +1 @@
+Subproject commit 5c0fa1d68f48ab7fdb331d75f31efc11aa313090
diff --git a/external/VulkanMemoryAllocator b/external/VulkanMemoryAllocator
index b0fce340b6c58..5c710e86a0ce0 160000
--- a/external/VulkanMemoryAllocator
+++ b/external/VulkanMemoryAllocator
@@ -1 +1 @@
-Subproject commit b0fce340b6c581d2bb75ca6c8c6e55235a52d8e2
+Subproject commit 5c710e86a0ce0664bbc2374c39b956b9928b2f29
diff --git a/external/glad b/external/glad
index 7ec5f98f091e5..23685e3caeb24 160000
--- a/external/glad
+++ b/external/glad
@@ -1 +1 @@
-Subproject commit 7ec5f98f091e5cb082e0f1b4f251ee5cb560552b
+Subproject commit 23685e3caeb249de6e082ae513bc1402e68b0643
diff --git a/external/glfw b/external/glfw
index 8bc966bbae496..168d6d24b9ce8 160000
--- a/external/glfw
+++ b/external/glfw
@@ -1 +1 @@
-Subproject commit 8bc966bbae4967a008252f1ac9b625e4dc77ad64
+Subproject commit 168d6d24b9ce85ef3a37114f00b1251b36162c89
diff --git a/external/volk b/external/volk
index b4eb550e25615..96fda88ef5a72 160000
--- a/external/volk
+++ b/external/volk
@@ -1 +1 @@
-Subproject commit b4eb550e2561556db7bf477543d42abd9c4c3217
+Subproject commit 96fda88ef5a7222d7db85e016cbae9343613483b
diff --git a/misc/benchmark_bit_struct_stores.py b/misc/benchmark_bit_struct_stores.py
index 1f80578f1840c..748b41a06bb9b 100644
--- a/misc/benchmark_bit_struct_stores.py
+++ b/misc/benchmark_bit_struct_stores.py
@@ -7,7 +7,7 @@
n = 1024 * 1024 * 256
if quant:
- ci16 = ti.quant.int(16, True)
+ ci16 = ti.types.quantized_types.quant.int(16, True)
x = ti.field(dtype=ci16)
y = ti.field(dtype=ci16)
diff --git a/misc/benchmark_rebuild_graph.py b/misc/benchmark_rebuild_graph.py
index 8d43870e2c83e..416605fa0411e 100644
--- a/misc/benchmark_rebuild_graph.py
+++ b/misc/benchmark_rebuild_graph.py
@@ -1,3 +1,5 @@
+from taichi.lang import impl
+
import taichi as ti
ti.init(arch=ti.cuda, async_mode=True)
@@ -23,4 +25,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/misc/ci_check_pr_title.py b/misc/ci_check_pr_title.py
index b6ac85d445562..90945f8dcc46c 100644
--- a/misc/ci_check_pr_title.py
+++ b/misc/ci_check_pr_title.py
@@ -27,7 +27,8 @@ def get_old_ver():
with open(json_path) as f:
prtags = json.load(f)
-if not title.startswith('['):
+# PR must be properly tagged. The only exception allowed here is the revert pr automatically generated by github.
+if not title.startswith('[') and not title.startswith('Revert '):
exit(f'PR title does not start with any tag: {title}')
if title.endswith(' '):
diff --git a/misc/ci_create_pr_card.py b/misc/ci_create_pr_card.py
new file mode 100644
index 0000000000000..5794e3f88a565
--- /dev/null
+++ b/misc/ci_create_pr_card.py
@@ -0,0 +1,112 @@
+import json
+import os
+from typing import Any, List, Mapping
+
+from github import Github
+from github.Project import Project
+from github.Repository import Repository
+
+
+def load_project_map() -> Mapping[str, str]:
+ with open(os.path.join(os.path.dirname(__file__),
+ 'tag_to_project.json')) as f:
+ return json.load(f)
+
+
+PROJECT_MAP = load_project_map()
+
+
+def extract_tags(title: str) -> List[str]:
+ """
+ Extract tags from PR title like "[ci] [bug] fix a bug"
+ """
+ tags: List[str] = []
+ for x in title.split('] ')[:-1]:
+ if x[0] != '[':
+ raise ValueError(f'No starting [ for tag: {x}]')
+ tags.append(x[1:].lower())
+ return tags
+
+
+def get_project(repo: Repository, name: str) -> Project:
+ """
+ Get project from repository by name
+ """
+ for project in repo.get_projects():
+ if project.name == name:
+ return project
+ raise ValueError(f'No project with name: {name}')
+
+
+def _create_pr_card(pr: dict, project: Project) -> None:
+ to_do_column = next(iter(project.get_columns()))
+ print(f"Creating card for PR #{pr['number']} in project {project.name}")
+ to_do_column.create_card(content_id=pr['id'], content_type="PullRequest")
+
+
+def _remove_pr_card(pr: dict, project: Project) -> None:
+ to_do_column = next(iter(project.get_columns()))
+ for card in to_do_column.get_cards():
+ if not card.content_url:
+ continue
+ if card.content_url.split('/')[-1] == str(pr['number']):
+ print(f"Deleting PR #{pr['number']} from project {project.name}")
+ card.delete()
+ return
+ print(
+ f"PR #{pr['number']} doesn't exist in the To-do column of project {project.name}"
+ )
+
+
+def create_pr_card(event: Mapping[str, Any]) -> None:
+ new_projects = {
+ PROJECT_MAP[tag]
+ for tag in extract_tags(event['pull_request']['title'])
+ if tag in PROJECT_MAP
+ }
+ gh = Github(os.environ['GITHUB_TOKEN'])
+ repo = gh.get_repo(event['repository']['full_name'])
+ pr = event['pull_request']
+ if event['action'] == 'opened':
+ for project_name in new_projects:
+ _create_pr_card(pr, get_project(repo, project_name))
+ else:
+ old_title = event.get("changes", {}).get("title", {}).get("from")
+ if not old_title:
+ print("PR title isn't changed, nothing to do")
+ return
+ old_projects = {
+ PROJECT_MAP[tag]
+ for tag in extract_tags(old_title) if tag in PROJECT_MAP
+ }
+ to_remove = old_projects - new_projects
+ to_add = new_projects - old_projects
+ for project_name in to_remove:
+ _remove_pr_card(pr, get_project(repo, project_name))
+ for project_name in to_add:
+ _create_pr_card(pr, get_project(repo, project_name))
+
+
+def main() -> None:
+ event = json.loads(os.environ['GH_EVENT'])
+ create_pr_card(event)
+
+
+def test():
+ event = {
+ "action": "opened",
+ "repository": {
+ "full_name": "taichi-dev/taichi"
+ },
+ "pull_request": {
+ "id": 841657847,
+ "number": 4224,
+ "title": "[lang] Annotate constants with dtype without casting."
+ }
+ }
+ os.environ["GH_EVENT"] = json.dumps(event)
+ main()
+
+
+if __name__ == '__main__':
+ main()
diff --git a/misc/ci_download.py b/misc/ci_download.py
index 5bbeec6d32b00..3b24ab46bc3ed 100644
--- a/misc/ci_download.py
+++ b/misc/ci_download.py
@@ -1,4 +1,6 @@
import os
+import sys
+import urllib.request
platform = os.environ['CI_PLATFORM']
if platform.startswith('macos'):
@@ -11,5 +13,12 @@
raise Exception(f'Bad CI_PLATFORM={platform}')
llvm_url = f'https://github.com/taichi-dev/taichi_assets/releases/download/llvm10/taichi-llvm-10.0.0-{platform}.zip'
+target_dir = 'taichi-llvm'
print(f'Downloading LLVM from {llvm_url}...')
-os.system(f'wget {llvm_url} --waitretry=3 --tries=5 -O taichi-llvm.zip')
+urllib.request.urlretrieve(llvm_url, "taichi-llvm.zip")
+print(f'Extract zip to local dir {target_dir}...')
+if not os.path.exists(target_dir):
+ os.makedirs(target_dir)
+
+retcode = os.system(f"unzip taichi-llvm.zip -d {target_dir}")
+sys.exit(retcode)
diff --git a/misc/code_format.py b/misc/code_format.py
index 0b08ff8a81935..a7259bd9c297f 100644
--- a/misc/code_format.py
+++ b/misc/code_format.py
@@ -130,6 +130,8 @@ def find_diff_or_empty(s):
continue
if fn.find('docs/build/') != -1:
continue
+ if fn.find(os.path.join('tests', 'python', 'test_exception.py')) != -1:
+ continue
if re.match(r'.*examples\/[a-z_]+\d\d+\.py$', fn):
print(f'Skipping example file "{fn}"...')
continue
diff --git a/misc/demo_external_func.py b/misc/demo_external_func.py
deleted file mode 100644
index 09bef0b1c82d3..0000000000000
--- a/misc/demo_external_func.py
+++ /dev/null
@@ -1,77 +0,0 @@
-import ctypes
-import os
-
-import taichi as ti
-
-ti.init()
-
-N = 1024
-x = ti.field(ti.i32, shape=N)
-y = ti.field(ti.i32, shape=N)
-z = ti.field(ti.i32, shape=N)
-
-source = '''
-extern "C" {
- void add_and_mul(float a, float b, float *c, float *d, int *e) {
- *c = a + b;
- *d = a * b;
- *e = int(a * b + a);
- }
- void pow_int(int a, int b, int *c) {
- int ret = 1;
- for (int i = 0; i < b; i++)
- ret = ret * a;
- *c = ret;
- }
-}
-'''
-
-with open('a.cpp', 'w') as f:
- f.write(source)
-
-os.system("g++ a.cpp -o a.so -fPIC -shared")
-
-so = ctypes.CDLL("./a.so")
-
-
-@ti.kernel
-def call_ext() -> ti.i32:
- a = 2.0
- b = 3.0
- c = 0.0
- d = 0.0
- e = 3
- ti.external_func_call(func=so.add_and_mul, args=(a, b), outputs=(c, d, e))
- p = 0
- ti.external_func_call(func=so.pow_int, args=(int(c + d), e), outputs=(p, ))
- return p
-
-
-# Wrap the external function to make it easier to use
-@ti.func
-def pow_int_wrapper(a, b):
- p = 0
- ti.external_func_call(func=so.pow_int,
- args=(int(a), int(b)),
- outputs=(p, ))
- return p
-
-
-@ti.kernel
-def call_parallel():
- for i in range(N):
- z[i] = pow_int_wrapper(x[i], y[i])
-
-
-assert call_ext() == 11**8
-
-for i in range(N):
- x[i] = i
- y[i] = 3
-
-call_parallel()
-for i in range(N):
- assert z[i] == i**3
-
-os.remove('a.cpp')
-os.remove('a.so')
diff --git a/misc/demo_excepthook.py b/misc/demo_trackback.py
similarity index 90%
rename from misc/demo_excepthook.py
rename to misc/demo_trackback.py
index 7e51a466ccf7e..9f9c3b8d5f735 100644
--- a/misc/demo_excepthook.py
+++ b/misc/demo_trackback.py
@@ -1,7 +1,6 @@
import taichi as ti
ti.init()
-ti.enable_excepthook()
@ti.func
diff --git a/misc/demo_warning.py b/misc/demo_warning.py
deleted file mode 100644
index f41bfb1e99628..0000000000000
--- a/misc/demo_warning.py
+++ /dev/null
@@ -1,14 +0,0 @@
-import taichi as ti
-
-x = ti.Vector([2, 3])
-
-x.transposed(x)
-
-
-@ti.kernel
-def func():
- x = 0
- x = 0.1
-
-
-func()
diff --git a/misc/examples.md b/misc/examples.md
index 909de5a194163..c714afe28a9d1 100644
--- a/misc/examples.md
+++ b/misc/examples.md
@@ -1,11 +1,11 @@
# More examples
-
-
-
-
+
+
+
+
-
-
-
-
+
+
+
+
diff --git a/misc/generate_commit_hash.py b/misc/generate_commit_hash.py
deleted file mode 100644
index 1b23d7968ba54..0000000000000
--- a/misc/generate_commit_hash.py
+++ /dev/null
@@ -1,23 +0,0 @@
-import os
-
-from git import Repo
-
-repo_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), '../')
-repo = Repo(repo_dir)
-commit_hash = str(repo.head.commit)
-print(f"Building commit {commit_hash}")
-
-output_fn = os.path.join(repo_dir, 'taichi/common/commit_hash.h')
-content = f"#define TI_COMMIT_HASH \"{commit_hash}\"\n"
-
-# First read the file to see if an update is needed
-# This reduces unnecessary file changes/linkings
-if os.path.exists(output_fn):
- with open(output_fn, 'r') as f:
- old_content = f.read()
- if old_content == content:
- # No update needed
- exit(0)
-
-with open(output_fn, 'w') as f:
- f.write(content)
diff --git a/misc/generate_example_videos.py b/misc/generate_example_videos.py
new file mode 100644
index 0000000000000..d11520cc3cd45
--- /dev/null
+++ b/misc/generate_example_videos.py
@@ -0,0 +1,24 @@
+import argparse
+import os
+import re
+import subprocess
+
+parser = argparse.ArgumentParser(description='Generate all videos of examples')
+parser.add_argument('output_directory',
+ help='output directory of generated videos')
+output_dir = parser.parse_args().output_directory
+
+example_root = os.path.join('..', 'tests', 'python', 'examples')
+for example_dir in os.listdir(example_root):
+ full_dir = os.path.join(example_root, example_dir)
+ if not os.path.isdir(full_dir):
+ continue
+ for filename in os.listdir(full_dir):
+ match = re.match(r'test_(\w+)\.py', filename)
+ if match:
+ subprocess.run([
+ "python",
+ os.path.join(full_dir, filename),
+ os.path.join(output_dir, match.group(1))
+ ],
+ check=True)
diff --git a/misc/links.md b/misc/links.md
index e8c3ce2c04f78..8b564922781b2 100644
--- a/misc/links.md
+++ b/misc/links.md
@@ -1,14 +1,11 @@
# Links
-- [DiffTaichi](https://github.com/yuanming-hu/difftaichi): 10 differentiable physical simulators built with Taichi differentiable programming, by [Yuanming Hu (yuanming-hu)](https://github.com/yuanming-hu).
+- [DiffTaichi](https://github.com/taichi-dev/difftaichi): 10 differentiable physical simulators built with Taichi differentiable programming, by [Yuanming Hu (yuanming-hu)](https://github.com/yuanming-hu).
- [Taichi elements](https://github.com/taichi-dev/taichi_elements): A high-performance multi-material continuum physics engine based on Taichi (work in progress).
- [taichi-fvm2d-fluid-ns](https://github.com/hejob/taichi-fvm2d-fluid-ns/): 2D FVM Compressible CFD Solver for multiblock structured mesh, by [hejob](https://github.com/hejob).
- [karman_tachi](https://github.com/houkensjtu/karman_taichi): An incompressible fluid solver using FVM that simulates a Karman vortex street, by [houkensjtu](https://github.com/houkensjtu).
- [TaichiMD](https://github.com/victoriacity/taichimd): Interactive, GPU-accelerated Molecular (& Macroscopic) Dynamics using Taichi, by [Andrew Sun (victoriacity)](https://github.com/victoriacity).
- [LBM Taichi](https://github.com/hietwll/LBM_Taichi): A fluid solver based on the Lattice Boltzmann Method (LBM) using Taichi, by [Zhuo Wang (hietwll)](https://github.com/hietwll).
- [Gitee mirror of Taichi](https://gitee.com/mirrors/Taichi): For the convenience of Chinese contributors, clone from the mirror repo hosted on Gitee (码云).
-- [Taichi THREE](https://github.com/taichi-dev/taichi_three): A 3D rendering library based on Taichi.
-- [Taichi GLSL](https://github.com/taichi-dev/taichi_glsl): A Taichi extension library that provides a set of GLSL-style helper functions.
-- [Taichi Blend](https://github.com/taichi-dev/taichi_blend): Taichi Blender intergration for physics-based animations (work in progress)
-- [Taichi.js](https://github.com/taichi-dev/taichi.js): Run compiled Taichi programs in Javascript and WASM (work in progress).
+- [Taichi GLSL](https://github.com/taichi-dev/taichi_glsl): A Taichi extension library, which provides a set of GLSL-style helper functions.
- [Shadertoy in Taichi](https://github.com/Phonicavi/Shadertoy-taichi): Some shadertoy examples implemented in Taichi, by [Qiu Feng (Phonicavi)](https://github.com/Phonicavi).
diff --git a/misc/make_changelog.py b/misc/make_changelog.py
index 6007db5ace5ad..4a2032d10fd34 100644
--- a/misc/make_changelog.py
+++ b/misc/make_changelog.py
@@ -17,7 +17,7 @@ def load_pr_tags():
return details
-def main(ver='master', repo_dir='.'):
+def main(ver=None, repo_dir='.'):
g = Repo(repo_dir)
commits_with_tags = set([tag.commit for tag in g.tags])
commits = list(g.iter_commits(ver, max_count=200))
@@ -33,11 +33,8 @@ def format(c):
for i, c in enumerate(commits):
s = format(c)
- if c in commits_with_tags:
- if i == 0:
- continue
- else:
- break
+ if c in commits_with_tags and i > 0:
+ break
tags = []
while s[0] == '[':
@@ -78,11 +75,11 @@ def format(c):
if __name__ == '__main__':
- ver = sys.argv[1] if len(sys.argv) > 1 else 'master'
+ ver = sys.argv[1] if len(sys.argv) > 1 else None
repo = sys.argv[2] if len(sys.argv) > 2 else '.'
save = sys.argv[3] if len(sys.argv) > 3 else False
res = main(ver, repo)
if save:
- with open('../python/taichi/CHANGELOG.md', 'w') as f:
+ with open('./python/taichi/CHANGELOG.md', 'w') as f:
f.write(res)
print(res)
diff --git a/misc/prtags.json b/misc/prtags.json
index e93b009704ddb..c4e1d3ff73fab 100644
--- a/misc/prtags.json
+++ b/misc/prtags.json
@@ -1,8 +1,10 @@
{
+ "javascript" : "Taichi in javascript",
"ci" : "CI/CD workflow",
"cpu" : "CPU backends",
"cuda" : "CUDA backend",
"doc" : "Documentation",
+ "docs" : "Documentation",
"infra" : "Infrastructure",
"cli" : "Command line interface",
"ir" : "Intermediate representation",
@@ -11,6 +13,8 @@
"metal" : "Metal backend",
"opengl" : "OpenGL backend",
"vulkan" : "Vulkan backend",
+ "dx11" : "DirectX 11 backend",
+ "spirv" : "SPIR-V common codegen",
"wasm" : "WebAssembly backend",
"misc" : "Miscellaneous",
"std" : "Standard library",
@@ -26,8 +30,10 @@
"test" : "Tests",
"benchmark" : "Benchmarking",
"async" : "AsyncEngine",
+ "mesh" : "MeshTaichi",
"workflow" : "GitHub Actions/Workflows",
"linux" : "Linux",
+ "android" : "Android",
"mac" : "Mac OS X",
"windows" : "Windows",
"docker" : "Docker container",
@@ -38,5 +44,6 @@
"blender" : "Blender intergration",
"export" : "Exporting kernels",
"type" : "Type system",
- "release" : "Release"
+ "release" : "Release",
+ "build" : "Build system"
}
diff --git a/misc/save_new_version.py b/misc/save_new_version.py
new file mode 100644
index 0000000000000..1db89916a374e
--- /dev/null
+++ b/misc/save_new_version.py
@@ -0,0 +1,47 @@
+import os
+from datetime import date
+
+import requests
+
+version = os.getenv('RELEASE_VERSION')
+version = version[1:]
+version_num = version.split('.')
+major = int(version_num[0])
+minor = int(version_num[1])
+patch = int(version_num[2])
+release_date = date.today().strftime('%Y-%m-%d')
+
+payload = {
+ 'version': version,
+ 'major': major,
+ 'minor': minor,
+ 'patch': patch,
+ 'date': release_date
+}
+
+username = os.getenv('METADATA_USERNAME')
+password = os.getenv('METADATA_PASSWORD')
+url = os.getenv('METADATA_URL')
+
+try:
+ response = requests.post(f'https://{url}/add_version/main',
+ json=payload,
+ auth=(username, password),
+ timeout=5)
+ response.raise_for_status()
+except requests.exceptions.ConnectionError as err:
+ print('Updating latest version failed: No internet,', err)
+ exit(1)
+except requests.exceptions.HTTPError as err:
+ print('Updating latest version failed: Server error,', err)
+ exit(1)
+except requests.exceptions.Timeout as err:
+ print('Updating latest version failed: Time out when connecting server,',
+ err)
+ exit(1)
+except requests.exceptions.RequestException as err:
+ print('Updating latest version failed:', err)
+ exit(1)
+
+response = response.json()
+print(response['message'])
diff --git a/misc/spMv_linear_solve.py b/misc/spMv_linear_solve.py
index 87bea0728673e..721648f45bb8c 100644
--- a/misc/spMv_linear_solve.py
+++ b/misc/spMv_linear_solve.py
@@ -9,7 +9,7 @@
@ti.kernel
-def fill(A: ti.linalg.sparse_matrix_builder(), b: ti.template(),
+def fill(A: ti.types.sparse_matrix_builder(), b: ti.template(),
interval: ti.i32):
for i in range(n):
A[i, i] += 2.0
diff --git a/misc/sparse_matrix.py b/misc/sparse_matrix.py
index 49cac7af58e87..a56054f7aedcd 100644
--- a/misc/sparse_matrix.py
+++ b/misc/sparse_matrix.py
@@ -9,8 +9,8 @@
@ti.kernel
-def fill(A: ti.linalg.sparse_matrix_builder(),
- b: ti.linalg.sparse_matrix_builder(), interval: ti.i32):
+def fill(A: ti.types.sparse_matrix_builder(),
+ b: ti.types.sparse_matrix_builder(), interval: ti.i32):
for i in range(n):
if i > 0:
A[i - 1, i] += -1.0
diff --git a/misc/tag_to_project.json b/misc/tag_to_project.json
new file mode 100644
index 0000000000000..23bffd761cd9a
--- /dev/null
+++ b/misc/tag_to_project.json
@@ -0,0 +1,12 @@
+{
+ "aot": "AOT",
+ "autodiff": "Autodiff",
+ "benchmark": "Backend Performance",
+ "build": "CI/CD & Build & Tests",
+ "ci": "CI/CD & Build & Tests",
+ "doc": "Docs & Examples & Tutorials",
+ "docs": "Docs & Examples & Tutorials",
+ "gui": "GGUI",
+ "ir": "Compiler Frontend & Middle-end",
+ "lang": "Lang Features & Python"
+}
diff --git a/misc/test_poly_timed.py b/misc/test_poly_timed.py
index 368798185fa04..b104d6c6c8adc 100644
--- a/misc/test_poly_timed.py
+++ b/misc/test_poly_timed.py
@@ -1,7 +1,7 @@
from autograd import grad
+from taichi._testing import approx
import taichi as ti
-from taichi import approx
# Note: test happens at v = 0.2
diff --git a/misc/upload_release.py b/misc/upload_release.py
new file mode 100644
index 0000000000000..d240c7b4bc201
--- /dev/null
+++ b/misc/upload_release.py
@@ -0,0 +1,65 @@
+import os
+import subprocess
+import sys
+
+import requests
+
+
+def upload_taichi_version():
+ username = os.getenv('METADATA_USERNAME')
+ password = os.getenv('METADATA_PASSWORD')
+ url = os.getenv('METADATA_URL')
+ for filename in os.listdir('./dist'):
+ filename = filename[:len(filename) - 4]
+ parts = filename.split('-')
+ payload = {
+ 'version': parts[1],
+ 'platform': parts[4],
+ 'python': parts[2]
+ }
+ try:
+ response = requests.post(f'https://{url}/add_version/detail',
+ json=payload,
+ auth=(username, password),
+ timeout=5)
+ response.raise_for_status()
+ except requests.exceptions.ConnectionError as err:
+ print('Updating latest version failed: No internet,', err)
+ except requests.exceptions.HTTPError as err:
+ print('Updating latest version failed: Server error,', err)
+ except requests.exceptions.Timeout as err:
+ print(
+ 'Updating latest version failed: Time out when connecting server,',
+ err)
+ except requests.exceptions.RequestException as err:
+ print('Updating latest version failed:', err)
+ else:
+ response = response.json()
+ print(response['message'])
+
+
+def upload_artifact(is_taichi):
+ pwd_env = 'PROD_PWD' if is_taichi else 'NIGHT_PWD'
+ twine_password = os.getenv(pwd_env)
+ if not twine_password:
+ sys.exit(f'Missing password env var {pwd_env}')
+ command = [sys.executable, '-m', 'twine', 'upload']
+ if not is_taichi:
+ command.extend(['--repository', 'testpypi'])
+ command.extend(
+ ['--verbose', '-u', '__token__', '-p', twine_password, 'dist/*'])
+ try:
+ subprocess.check_call(command)
+ except subprocess.CalledProcessError as e:
+ sys.exit(f"Twine upload returns error {e.returncode}")
+
+
+if __name__ == '__main__':
+ if os.getenv('GITHUB_REPOSITORY',
+ 'taichi-dev/taichi') != 'taichi-dev/taichi':
+ print('This script should be run from taichi repo')
+ sys.exit(0)
+ is_taichi = os.getenv('PROJECT_NAME', 'taichi') == 'taichi'
+ upload_artifact(is_taichi)
+ if is_taichi:
+ upload_taichi_version()
diff --git a/misc/visualize_quant_types.py b/misc/visualize_quant_types.py
index 24767428cc390..6f51d5ab03ada 100644
--- a/misc/visualize_quant_types.py
+++ b/misc/visualize_quant_types.py
@@ -7,9 +7,9 @@
ti.init()
-f19 = ti.quant.float(exp=6, frac=13, signed=True)
-f16 = ti.quant.float(exp=5, frac=11, signed=True)
-fixed16 = ti.quant.fixed(frac=16, range=2)
+f19 = ti.types.quantized_types.quant.float(exp=6, frac=13, signed=True)
+f16 = ti.types.quantized_types.quant.float(exp=5, frac=11, signed=True)
+fixed16 = ti.types.quantized_types.quant.fixed(frac=16, range=2)
vf19 = ti.Vector.field(2, dtype=f19)
bs_vf19 = ti.root.bit_struct(num_bits=32)
diff --git a/misc/windows_build.py b/misc/windows_build.py
deleted file mode 100644
index b7897054f39cc..0000000000000
--- a/misc/windows_build.py
+++ /dev/null
@@ -1,69 +0,0 @@
-import os
-import shutil
-import sys
-
-
-def execute_cmd(cmd):
- print('Executing', resolve_env(cmd))
- return os.system(cmd)
-
-
-def resolve_env(v):
- # replace `%`
- modified = True
- while modified:
- modified = False
- for i in range(len(v)):
- if v[i] == '%':
- for j in range(i + 1, len(v)):
- if v[j] == '%':
- var = v[i + 1:j]
- v = v[:i] + os.environ[var] + v[j + 1:]
- modified = True
- print(v)
- break
- break
- return v
-
-
-def set_env(**kwargs):
- for k, v in kwargs.items():
- v = resolve_env(v)
- print(f"Setting {k} to '{v}'")
- os.environ[k] = v
-
-
-repo_dir = "E:\\repos\\taichi"
-assert len(
- sys.argv
-) == 3, 'Usage: windows_build.py [python_executable] [cuda_version=10.X]'
-python_executable = sys.argv[1]
-cuda_version = sys.argv[2]
-assert cuda_version in ["10.0", "10.1", "10.2"]
-print("Python =", python_executable)
-print("CUDA =", cuda_version)
-set_env(PYTHON=python_executable)
-set_env(TAICHI_REPO_DIR=repo_dir)
-set_env(PYTHONPATH="%TAICHI_REPO_DIR%\\python")
-set_env(PATH=r"%TAICHI_REPO_DIR%\bin;%PATH%")
-execute_cmd("clang --version")
-
-os.chdir(repo_dir)
-build_dir = os.path.join(repo_dir, 'build')
-if os.path.exists(build_dir):
- shutil.rmtree(build_dir)
-os.mkdir(build_dir)
-os.chdir(build_dir)
-llvm_dir = "E:\\repos\\llvm-8.0.1\\build\\installed\\lib\\cmake\\llvm"
-cuda_dir = f"C:\\Program Files\\NVIDIA GPU Computing Toolkit\\CUDA\\v{cuda_version}"
-execute_cmd(
- f'cmake .. -G"Visual Studio 15 2017 Win64" -DPYTHON_EXECUTABLE="%PYTHON%" -DLLVM_DIR="{llvm_dir}" -DTI_WITH_CUDA:BOOL=True -DCUDA_VERSION={cuda_version} -DCUDA_DIR="f{cuda_dir}"'
-)
-execute_cmd(
- r'msbuild /p:Configuration=RelWithDebInfo /p:Platform=x64 /m taichi.sln')
-os.chdir(repo_dir)
-execute_cmd('%PYTHON% -c "import taichi"')
-execute_cmd('%PYTHON% examples/laplace.py')
-execute_cmd('%PYTHON% bin/taichi test')
-os.chdir(os.path.join(repo_dir, 'python'))
-execute_cmd('%PYTHON% build.py try_upload')
diff --git a/netlify.toml b/netlify.toml
index 461c6e304a23b..d5c51c2b39cee 100644
--- a/netlify.toml
+++ b/netlify.toml
@@ -1,8 +1,7 @@
[build]
- command = "git clone https://github.com/taichi-dev/docs.taichi.graphics.git; rm -rf docs.taichi.graphics/website/docs/lang; cp -rf docs/lang docs.taichi.graphics/website/docs/lang; cd docs.taichi.graphics/website; npm install --global yarn@1.22; yarn install; yarn build"
+ command = "git clone https://github.com/taichi-dev/docs.taichi.graphics.git; rm -rf docs.taichi.graphics/website/docs/lang; cp -rf docs/lang docs.taichi.graphics/website/docs/lang; git clone https://github.com/taichi-dev/docstring-gen docsgen; export DOCSTRING_GEN_PATH=\"$(pwd)/docsgen\"; export TAICHI_PATH=\"$(pwd)/python/taichi\"; export TAICHI_WEBSITE=\"$(pwd)/docs.taichi.graphics\"; pip install sphinx-autoapi==1.8.4 gitpython pydata-sphinx-theme==0.7.2; cd $DOCSTRING_GEN_PATH/experimental; export current_version=master; make clean; make version; make apideploy; cd $TAICHI_WEBSITE/website; npm install --global yarn@1.22; yarn install; yarn build; yarn run apiversion;"
publish = "docs.taichi.graphics/website/build"
-
# Cancel the build if there're no changes detected in docs/ folder.
- ignore = "git remote add upstream https://github.com/taichi-dev/taichi.git; git fetch upstream master; git diff --quiet $COMMIT_REF upstream/master -- docs/"
+ ignore = "git remote add upstream https://github.com/taichi-dev/taichi.git; git fetch upstream master; git diff --quiet $COMMIT_REF upstream/master -- docs/ python/"
diff --git a/pyproject.toml b/pyproject.toml
new file mode 100644
index 0000000000000..facad75d3914b
--- /dev/null
+++ b/pyproject.toml
@@ -0,0 +1,3 @@
+[build-system]
+requires = ["setuptools", "wheel", "numpy", "pybind11", "cmake"]
+build-backend = "setuptools.build_meta"
diff --git a/python/.gitignore b/python/.gitignore
index 12389e8d49a2f..1a30ec8f567f0 100644
--- a/python/.gitignore
+++ b/python/.gitignore
@@ -1,7 +1,9 @@
-lib
+taichi/_lib/runtime
+taichi/_lib/core/*.so
+taichi/_lib/core/*.pyd
taichi.egg-info
taichi/include
-taichi/examples
taichi/assets
taichi/tests
+taichi/tests38
release
diff --git a/python/build.py b/python/build.py
deleted file mode 100644
index e320fb95a1023..0000000000000
--- a/python/build.py
+++ /dev/null
@@ -1,148 +0,0 @@
-import argparse
-import os
-import platform
-import re
-import shutil
-import sys
-
-
-def get_os_name():
- name = platform.platform()
- # in python 3.8, platform.platform() uses mac_ver() on macOS
- # it will return 'macOS-XXXX' instead of 'Darwin-XXXX'
- if name.lower().startswith('darwin') or name.lower().startswith('macos'):
- return 'osx'
- elif name.lower().startswith('windows'):
- return 'win'
- elif name.lower().startswith('linux'):
- return 'linux'
- assert False, "Unknown platform name %s" % name
-
-
-def get_python_executable():
- return '"' + sys.executable.replace('\\', '/') + '"'
-
-
-def build(project_name):
- """Build and package the wheel file in root `dist` dir"""
- if platform.system() == 'Linux':
- if re.search("^clang\+\+-*\d*", str(os.environ.get('CXX'))) is None:
- raise RuntimeError(
- 'Only the wheel with clang will be released to PyPI')
-
- print("Using python executable", get_python_executable())
- os.system(
- '{} -m pip install --user --upgrade twine setuptools wheel'.format(
- get_python_executable()))
-
- os.system(
- f'{get_python_executable()} ../misc/make_changelog.py origin/master ../ True'
- )
-
- # This env var is used in setup.py below.
- os.environ['PROJECT_NAME'] = project_name
- project_tag = ''
- if project_name == 'taichi-nightly':
- project_tag = 'egg_info --tag-date'
- if get_os_name() == 'linux':
- os.system(
- f'cd ..; {get_python_executable()} setup.py {project_tag} bdist_wheel -p manylinux1_x86_64'
- )
- else:
- os.system(
- f'cd .. && {get_python_executable()} setup.py {project_tag} bdist_wheel'
- )
-
- try:
- os.remove('taichi/CHANGELOG.md')
- except FileNotFoundError:
- pass
-
-
-def parse_args():
- parser = argparse.ArgumentParser(description=(
- 'Build and uploads wheels to PyPI. Make sure to run this script '
- 'inside `python/`'))
- parser.add_argument('mode',
- type=str,
- default='',
- help=('Choose one of the modes: '
- '[build, test, try_upload, upload]'))
- parser.add_argument('--skip_build',
- action='store_true',
- help=('Skip the build process if this is enabled'))
- parser.add_argument('--testpypi',
- action='store_true',
- help='Upload to test server if this is enabled')
- parser.add_argument('--project_name',
- action='store',
- dest='project_name',
- default='taichi',
- help='Set the project name')
- return parser.parse_args()
-
-
-def main():
- args = parse_args()
- mode = args.mode
- pypi_user = '__token__'
- pypi_repo = ''
- project_name = args.project_name
-
- env_pypi_pwd = os.environ.get('PYPI_PWD', '')
-
- if not args.skip_build:
- shutil.rmtree('../dist', ignore_errors=True)
-
- if mode == 'try_upload':
- if env_pypi_pwd == '':
- print("Missing environment variable PYPI_PWD")
- print("Giving up and exiting 0 [try_upload mode]")
- exit(0)
- mode = 'upload'
-
- if mode == 'upload' and env_pypi_pwd == '':
- raise RuntimeError("Missing environment variable PYPI_PWD")
-
- os.environ['TWINE_PASSWORD'] = env_pypi_pwd
-
- if mode == 'upload' and args.testpypi:
- pypi_repo = '--repository testpypi'
-
- if not args.skip_build:
- build(project_name)
-
- if mode == 'build':
- return
- elif mode == 'upload':
- os.system('{} -m twine upload {} ../dist/* --verbose -u {}'.format(
- get_python_executable(), pypi_repo, pypi_user))
- elif mode == 'test':
- print('Uninstalling old taichi packages...')
- os.system(
- f'{get_python_executable()} -m pip uninstall -y taichi-nightly')
- os.system(f'{get_python_executable()} -m pip uninstall -y taichi')
- dists = os.listdir('../dist')
- assert len(dists) == 1
- dist = dists[0]
- print('Installing ', dist)
- os.environ['PYTHONPATH'] = ''
- os.makedirs('test_env', exist_ok=True)
- os.system(
- 'cd test_env && {} -m pip install ../../dist/{} --user'.format(
- get_python_executable(), dist))
- print('Entering test environment...')
- if get_os_name() == 'win':
- os.system(
- 'cmd /V /C "set PYTHONPATH=&& set TAICHI_REPO_DIR=&& cd test_env && cmd"'
- )
- else:
- os.system(
- 'cd test_env && PYTHONPATH= TAICHI_REPO_DIR= bash --noprofile --norc '
- )
- else:
- raise ValueError("Unknown mode: %s" % mode)
-
-
-if __name__ == '__main__':
- main()
diff --git a/python/make_release.py b/python/make_release.py
deleted file mode 100644
index 0f61f316e98f2..0000000000000
--- a/python/make_release.py
+++ /dev/null
@@ -1,39 +0,0 @@
-import os
-import shutil
-import zipfile
-
-import requests
-
-projects = ['nightly', 'nightly-cuda-10-0', 'nightly-cuda-10-1']
-
-
-def download(url):
- fn = url.split('/')[-1]
- with requests.get(url, stream=True) as r:
- with open(fn, 'wb') as f:
- shutil.copyfileobj(r.raw, f)
- return fn
-
-
-for p in projects:
- pkg_name_dash = f'taichi-{p}'
- pkg_name_underscore = pkg_name_dash.replace('-', '_')
- package = requests.get(
- f"https://pypi.python.org/pypi/{pkg_name_dash}/json").json()
- version = '0.0.75'
- wheels = package["releases"][version]
- for wheel in wheels:
- py_ver = wheel['python_version']
- print(py_ver, wheel['url'])
- fn = download(wheel['url'])
- folder = wheel['python_version'] + '-' + fn[:-4]
- package_extracted_folder = f"release/{folder}"
- with zipfile.ZipFile(fn, 'r') as zip_ref:
- zip_ref.extractall(package_extracted_folder)
- os.remove(fn)
-
- pkg_ver = f"{pkg_name_underscore}-{version}"
- shutil.make_archive(
- f'release/{folder}', 'zip',
- f'release/{folder}/{pkg_ver}.data/purelib/taichi/lib')
- shutil.rmtree(package_extracted_folder)
diff --git a/python/taichi/__init__.py b/python/taichi/__init__.py
index 8823b50caee62..7d834259bcba4 100644
--- a/python/taichi/__init__.py
+++ b/python/taichi/__init__.py
@@ -1,48 +1,69 @@
import sys
-import taichi.ad as ad
+from taichi._funcs import *
+from taichi._lib import core as _ti_core
from taichi._logging import *
-from taichi.core import get_os_name, package_root, require_version
-from taichi.core import ti_core as core
-from taichi.lang import * # TODO(archibate): It's `taichi.lang.core` overriding `taichi.core`
-from taichi.main import main
-from taichi.misc import *
-from taichi.testing import *
-from taichi.tools import *
-from taichi.torch_io import from_torch, to_torch
-from taichi.type import *
-
-import taichi.ui as ui
+from taichi._snode import *
+from taichi.lang import * # pylint: disable=W0622 # TODO(archibate): It's `taichi.lang.core` overriding `taichi.core`
+from taichi.types.annotations import *
+# Provide a shortcut to types since they're commonly used.
+from taichi.types.primitive_types import *
+
+from taichi import ad, experimental, linalg, tools
+from taichi.ui import GUI, hex_to_rgb, rgb_to_hex, ui
# Issue#2223: Do not reorder, or we're busted with partially initialized module
from taichi import aot # isort:skip
-deprecated_names = {'SOA': 'Layout.SOA', 'AOS': 'Layout.AOS'}
+__deprecated_names__ = {
+ 'SOA': 'Layout.SOA',
+ 'AOS': 'Layout.AOS',
+ 'print_profile_info': 'profiler.print_scoped_profiler_info',
+ 'clear_profile_info': 'profiler.clear_scoped_profiler_info',
+ 'print_memory_profile_info': 'profiler.print_memory_profiler_info',
+ 'CuptiMetric': 'profiler.CuptiMetric',
+ 'get_predefined_cupti_metrics': 'profiler.get_predefined_cupti_metrics',
+ 'print_kernel_profile_info': 'profiler.print_kernel_profiler_info',
+ 'query_kernel_profile_info': 'profiler.query_kernel_profiler_info',
+ 'clear_kernel_profile_info': 'profiler.clear_kernel_profiler_info',
+ 'kernel_profiler_total_time': 'profiler.get_kernel_profiler_total_time',
+ 'set_kernel_profiler_toolkit': 'profiler.set_kernel_profiler_toolkit',
+ 'set_kernel_profile_metrics': 'profiler.set_kernel_profiler_metrics',
+ 'collect_kernel_profile_metrics':
+ 'profiler.collect_kernel_profiler_metrics',
+ 'VideoManager': 'tools.VideoManager',
+ 'PLYWriter': 'tools.PLYWriter',
+ 'imread': 'tools.imread',
+ 'imresize': 'tools.imresize',
+ 'imshow': 'tools.imshow',
+ 'imwrite': 'tools.imwrite',
+ 'quant': 'types.quantized_types.quant',
+ 'type_factory': 'types.quantized_types.type_factory'
+}
+
if sys.version_info.minor < 7:
- for name, alter in deprecated_names.items():
+ for name, alter in __deprecated_names__.items():
exec(f'{name} = {alter}')
else:
def __getattr__(attr):
- if attr in deprecated_names:
- warning('ti.{} is deprecated. Please use ti.{} instead.'.format(
- attr, deprecated_names[attr]),
- DeprecationWarning,
- stacklevel=2)
- exec(f'{attr} = {deprecated_names[attr]}')
+ # There's no easy way to hook accessing attribute with function calls in python3.6.
+ # So let's skip it for now.
+ import warnings # pylint: disable=C0415,W0621
+ if attr == 'cfg':
+ return None if lang.impl.get_runtime(
+ ).prog is None else lang.impl.current_cfg()
+ if attr in __deprecated_names__:
+ warnings.warn(
+ f'ti.{attr} is deprecated. Please use ti.{__deprecated_names__[attr]} instead.',
+ DeprecationWarning)
+ exec(f'{attr} = {__deprecated_names__[attr]}')
return locals()[attr]
raise AttributeError(f"module '{__name__}' has no attribute '{attr}'")
-__all__ = [
- 'ad', 'core', 'misc', 'lang', 'tools', 'main', 'torch_io', 'ui', 'profiler'
-]
-
-complex_kernel = deprecated('ti.complex_kernel',
- 'ti.ad.grad_replaced')(ad.grad_replaced)
-
-complex_kernel_grad = deprecated('ti.complex_kernel_grad',
- 'ti.ad.grad_for')(ad.grad_for)
+__version__ = (_ti_core.get_version_major(), _ti_core.get_version_minor(),
+ _ti_core.get_version_patch())
-__version__ = (core.get_version_major(), core.get_version_minor(),
- core.get_version_patch())
+del sys
+del _ti_core
diff --git a/python/taichi/__main__.py b/python/taichi/__main__.py
index 5d6a8109e6ce6..925fa9116efa4 100644
--- a/python/taichi/__main__.py
+++ b/python/taichi/__main__.py
@@ -1,3 +1,3 @@
-from .main import main
+from ._main import main
main()
diff --git a/python/taichi/_funcs.py b/python/taichi/_funcs.py
new file mode 100644
index 0000000000000..0d9e0392e5223
--- /dev/null
+++ b/python/taichi/_funcs.py
@@ -0,0 +1,437 @@
+import math
+
+from taichi.lang import impl, matrix, ops
+from taichi.lang.impl import expr_init, get_runtime, static
+from taichi.lang.kernel_impl import func, pyfunc
+from taichi.lang.matrix import Matrix, Vector
+from taichi.types import f32, f64
+
+
+@func
+def _randn(dt):
+ """
+ Generate a random float sampled from univariate standard normal
+ (Gaussian) distribution of mean 0 and variance 1, using the
+ Box-Muller transformation.
+ """
+ assert dt == f32 or dt == f64
+ u1 = ops.cast(1.0, dt) - ops.random(dt)
+ u2 = ops.random(dt)
+ r = ops.sqrt(-2 * ops.log(u1))
+ c = ops.cos(math.tau * u2)
+ return r * c
+
+
+def randn(dt=None):
+ """Generate a random float sampled from univariate standard normal
+ (Gaussian) distribution of mean 0 and variance 1, using the
+ Box-Muller transformation. Must be called in Taichi scope.
+
+ Args:
+ dt (DataType): Data type of the required random number. Default to `None`.
+ If set to `None` `dt` will be determined dynamically in runtime.
+
+ Returns:
+ The generated random float.
+
+ Example::
+
+ >>> @ti.kernel
+ >>> def main():
+ >>> print(ti.randn())
+ >>>
+ >>> main()
+ -0.463608
+ """
+ if dt is None:
+ dt = impl.get_runtime().default_fp
+ return _randn(dt)
+
+
+@pyfunc
+def _matrix_transpose(mat):
+ """Permute the first two axes of the matrix.
+
+ Args:
+ mat (:class:`~taichi.lang.matrix.Matrix`): Input matrix.
+
+ Returns:
+ Transpose of the input matrix.
+ """
+ return matrix.Matrix([[mat[i, j] for i in range(mat.n)]
+ for j in range(mat.m)])
+
+
+@pyfunc
+def _matrix_cross3d(self, other):
+ return matrix.Matrix([
+ self[1] * other[2] - self[2] * other[1],
+ self[2] * other[0] - self[0] * other[2],
+ self[0] * other[1] - self[1] * other[0],
+ ])
+
+
+@pyfunc
+def _matrix_cross2d(self, other):
+ return self[0] * other[1] - self[1] * other[0]
+
+
+@pyfunc
+def _matrix_outer_product(self, other):
+ """Perform the outer product with the input Vector (1-D Matrix).
+
+ Args:
+ other (:class:`~taichi.lang.matrix.Matrix`): The input Vector (1-D Matrix) to perform the outer product.
+
+ Returns:
+ :class:`~taichi.lang.matrix.Matrix`: The outer product result (Matrix) of the two Vectors.
+
+ """
+ impl.static(
+ impl.static_assert(self.m == 1,
+ "lhs for outer_product is not a vector"))
+ impl.static(
+ impl.static_assert(other.m == 1,
+ "rhs for outer_product is not a vector"))
+ return matrix.Matrix([[self[i] * other[j] for j in range(other.n)]
+ for i in range(self.n)])
+
+
+@func
+def polar_decompose2d(A, dt):
+ """Perform polar decomposition (A=UP) for 2x2 matrix.
+
+ Mathematical concept refers to https://en.wikipedia.org/wiki/Polar_decomposition.
+
+ Args:
+ A (ti.Matrix(2, 2)): input 2x2 matrix `A`.
+ dt (DataType): date type of elements in matrix `A`, typically accepts ti.f32 or ti.f64.
+
+ Returns:
+ Decomposed 2x2 matrices `U` and `P`.
+ """
+ x, y = A(0, 0) + A(1, 1), A(1, 0) - A(0, 1)
+ scale = (1.0 / ops.sqrt(x * x + y * y))
+ c = x * scale
+ s = y * scale
+ r = Matrix([[c, -s], [s, c]], dt=dt)
+ return r, r.transpose() @ A
+
+
+@func
+def polar_decompose3d(A, dt):
+ """Perform polar decomposition (A=UP) for 3x3 matrix.
+
+ Mathematical concept refers to https://en.wikipedia.org/wiki/Polar_decomposition.
+
+ Args:
+ A (ti.Matrix(3, 3)): input 3x3 matrix `A`.
+ dt (DataType): date type of elements in matrix `A`, typically accepts ti.f32 or ti.f64.
+
+ Returns:
+ Decomposed 3x3 matrices `U` and `P`.
+ """
+ U, sig, V = svd(A, dt)
+ return U @ V.transpose(), V @ sig @ V.transpose()
+
+
+# https://www.seas.upenn.edu/~cffjiang/research/svd/svd.pdf
+@func
+def svd2d(A, dt):
+ """Perform singular value decomposition (A=USV^T) for 2x2 matrix.
+
+ Mathematical concept refers to https://en.wikipedia.org/wiki/Singular_value_decomposition.
+
+ Args:
+ A (ti.Matrix(2, 2)): input 2x2 matrix `A`.
+ dt (DataType): date type of elements in matrix `A`, typically accepts ti.f32 or ti.f64.
+
+ Returns:
+ Decomposed 2x2 matrices `U`, 'S' and `V`.
+ """
+ R, S = polar_decompose2d(A, dt)
+ c, s = ops.cast(0.0, dt), ops.cast(0.0, dt)
+ s1, s2 = ops.cast(0.0, dt), ops.cast(0.0, dt)
+ if abs(S[0, 1]) < 1e-5:
+ c, s = 1, 0
+ s1, s2 = S[0, 0], S[1, 1]
+ else:
+ tao = ops.cast(0.5, dt) * (S[0, 0] - S[1, 1])
+ w = ops.sqrt(tao**2 + S[0, 1]**2)
+ t = ops.cast(0.0, dt)
+ if tao > 0:
+ t = S[0, 1] / (tao + w)
+ else:
+ t = S[0, 1] / (tao - w)
+ c = 1 / ops.sqrt(t**2 + 1)
+ s = -t * c
+ s1 = c**2 * S[0, 0] - 2 * c * s * S[0, 1] + s**2 * S[1, 1]
+ s2 = s**2 * S[0, 0] + 2 * c * s * S[0, 1] + c**2 * S[1, 1]
+ V = Matrix.zero(dt, 2, 2)
+ if s1 < s2:
+ tmp = s1
+ s1 = s2
+ s2 = tmp
+ V = Matrix([[-s, c], [-c, -s]], dt=dt)
+ else:
+ V = Matrix([[c, s], [-s, c]], dt=dt)
+ U = R @ V
+ return U, Matrix([[s1, ops.cast(0, dt)], [ops.cast(0, dt), s2]], dt=dt), V
+
+
+def svd3d(A, dt, iters=None):
+ """Perform singular value decomposition (A=USV^T) for 3x3 matrix.
+
+ Mathematical concept refers to https://en.wikipedia.org/wiki/Singular_value_decomposition.
+
+ Args:
+ A (ti.Matrix(3, 3)): input 3x3 matrix `A`.
+ dt (DataType): date type of elements in matrix `A`, typically accepts ti.f32 or ti.f64.
+ iters (int): iteration number to control algorithm precision.
+
+ Returns:
+ Decomposed 3x3 matrices `U`, 'S' and `V`.
+ """
+ assert A.n == 3 and A.m == 3
+ inputs = tuple([e.ptr for e in A.entries])
+ assert dt in [f32, f64]
+ if iters is None:
+ if dt == f32:
+ iters = 5
+ else:
+ iters = 8
+ if dt == f32:
+ rets = get_runtime().prog.current_ast_builder().sifakis_svd_f32(
+ *inputs, iters)
+ else:
+ rets = get_runtime().prog.current_ast_builder().sifakis_svd_f64(
+ *inputs, iters)
+ assert len(rets) == 21
+ U_entries = rets[:9]
+ V_entries = rets[9:18]
+ sig_entries = rets[18:]
+ U = expr_init(Matrix.zero(dt, 3, 3))
+ V = expr_init(Matrix.zero(dt, 3, 3))
+ sigma = expr_init(Matrix.zero(dt, 3, 3))
+ for i in range(3):
+ for j in range(3):
+ U(i, j)._assign(U_entries[i * 3 + j])
+ V(i, j)._assign(V_entries[i * 3 + j])
+ sigma(i, i)._assign(sig_entries[i])
+ return U, sigma, V
+
+
+@func
+def eig2x2(A, dt):
+ """Compute the eigenvalues and right eigenvectors (Av=lambda v) of a 2x2 real matrix.
+
+ Mathematical concept refers to https://en.wikipedia.org/wiki/Eigendecomposition_of_a_matrix.
+
+ Args:
+ A (ti.Matrix(2, 2)): input 2x2 matrix `A`.
+ dt (DataType): date type of elements in matrix `A`, typically accepts ti.f32 or ti.f64.
+
+ Returns:
+ eigenvalues (ti.Matrix(2, 2)): The eigenvalues in complex form. Each row stores one eigenvalue. The first number of the eigenvalue represents the real part and the second number represents the imaginary part.
+ eigenvectors: (ti.Matrix(4, 2)): The eigenvectors in complex form. Each column stores one eigenvector. Each eigenvector consists of 2 entries, each of which is represented by two numbers for its real part and imaginary part.
+ """
+ tr = A.trace()
+ det = A.determinant()
+ gap = tr**2 - 4 * det
+ lambda1 = Vector.zero(dt, 2)
+ lambda2 = Vector.zero(dt, 2)
+ v1 = Vector.zero(dt, 4)
+ v2 = Vector.zero(dt, 4)
+ if gap > 0:
+ lambda1 = Vector([tr + ops.sqrt(gap), 0.0], dt=dt) * 0.5
+ lambda2 = Vector([tr - ops.sqrt(gap), 0.0], dt=dt) * 0.5
+ A1 = A - lambda1[0] * Matrix.identity(dt, 2)
+ A2 = A - lambda2[0] * Matrix.identity(dt, 2)
+ if all(A1 == Matrix.zero(dt, 2, 2)) and all(
+ A1 == Matrix.zero(dt, 2, 2)):
+ v1 = Vector([0.0, 0.0, 1.0, 0.0]).cast(dt)
+ v2 = Vector([1.0, 0.0, 0.0, 0.0]).cast(dt)
+ else:
+ v1 = Vector([A2[0, 0], 0.0, A2[1, 0], 0.0], dt=dt).normalized()
+ v2 = Vector([A1[0, 0], 0.0, A1[1, 0], 0.0], dt=dt).normalized()
+ else:
+ lambda1 = Vector([tr, ops.sqrt(-gap)], dt=dt) * 0.5
+ lambda2 = Vector([tr, -ops.sqrt(-gap)], dt=dt) * 0.5
+ A1r = A - lambda1[0] * Matrix.identity(dt, 2)
+ A1i = -lambda1[1] * Matrix.identity(dt, 2)
+ A2r = A - lambda2[0] * Matrix.identity(dt, 2)
+ A2i = -lambda2[1] * Matrix.identity(dt, 2)
+ v1 = Vector([A2r[0, 0], A2i[0, 0], A2r[1, 0], A2i[1, 0]],
+ dt=dt).normalized()
+ v2 = Vector([A1r[0, 0], A1i[0, 0], A1r[1, 0], A1i[1, 0]],
+ dt=dt).normalized()
+ eigenvalues = Matrix.rows([lambda1, lambda2])
+ eigenvectors = Matrix.cols([v1, v2])
+
+ return eigenvalues, eigenvectors
+
+
+@func
+def sym_eig2x2(A, dt):
+ """Compute the eigenvalues and right eigenvectors (Av=lambda v) of a 2x2 real symmetric matrix.
+
+ Mathematical concept refers to https://en.wikipedia.org/wiki/Eigendecomposition_of_a_matrix.
+
+ Args:
+ A (ti.Matrix(2, 2)): input 2x2 symmetric matrix `A`.
+ dt (DataType): date type of elements in matrix `A`, typically accepts ti.f32 or ti.f64.
+
+ Returns:
+ eigenvalues (ti.Vector(2)): The eigenvalues. Each entry store one eigen value.
+ eigenvectors (ti.Matrix(2, 2)): The eigenvectors. Each column stores one eigenvector.
+ """
+ tr = A.trace()
+ det = A.determinant()
+ gap = tr**2 - 4 * det
+ lambda1 = (tr + ops.sqrt(gap)) * 0.5
+ lambda2 = (tr - ops.sqrt(gap)) * 0.5
+ eigenvalues = Vector([lambda1, lambda2], dt=dt)
+
+ A1 = A - lambda1 * Matrix.identity(dt, 2)
+ A2 = A - lambda2 * Matrix.identity(dt, 2)
+ v1 = Vector.zero(dt, 2)
+ v2 = Vector.zero(dt, 2)
+ if all(A1 == Matrix.zero(dt, 2, 2)) and all(A1 == Matrix.zero(dt, 2, 2)):
+ v1 = Vector([0.0, 1.0]).cast(dt)
+ v2 = Vector([1.0, 0.0]).cast(dt)
+ else:
+ v1 = Vector([A2[0, 0], A2[1, 0]], dt=dt).normalized()
+ v2 = Vector([A1[0, 0], A1[1, 0]], dt=dt).normalized()
+ eigenvectors = Matrix.cols([v1, v2])
+ return eigenvalues, eigenvectors
+
+
+@func
+def _svd(A, dt):
+ """Perform singular value decomposition (A=USV^T) for arbitrary size matrix.
+
+ Mathematical concept refers to https://en.wikipedia.org/wiki/Singular_value_decomposition.
+ 2D implementation refers to :func:`taichi.svd2d`.
+ 3D implementation refers to :func:`taichi.svd3d`.
+
+ Args:
+ A (ti.Matrix(n, n)): input nxn matrix `A`.
+ dt (DataType): date type of elements in matrix `A`, typically accepts ti.f32 or ti.f64.
+
+ Returns:
+ Decomposed nxn matrices `U`, 'S' and `V`.
+ """
+ if static(A.n == 2): # pylint: disable=R1705
+ ret = svd2d(A, dt)
+ return ret
+ else:
+ return svd3d(A, dt)
+
+
+@func
+def _polar_decompose(A, dt):
+ """Perform polar decomposition (A=UP) for arbitrary size matrix.
+
+ Mathematical concept refers to https://en.wikipedia.org/wiki/Polar_decomposition.
+ 2D implementation refers to :func:`taichi.polar_decompose2d`.
+ 3D implementation refers to :func:`taichi.polar_decompose3d`.
+
+ Args:
+ A (ti.Matrix(n, n)): input nxn matrix `A`.
+ dt (DataType): date type of elements in matrix `A`, typically accepts ti.f32 or ti.f64.
+
+ Returns:
+ Decomposed nxn matrices `U` and `P`.
+ """
+ if static(A.n == 2): # pylint: disable=R1705
+ ret = polar_decompose2d(A, dt)
+ return ret
+ else:
+ return polar_decompose3d(A, dt)
+
+
+def polar_decompose(A, dt=None):
+ """Perform polar decomposition (A=UP) for arbitrary size matrix.
+
+ Mathematical concept refers to https://en.wikipedia.org/wiki/Polar_decomposition.
+ This is only a wrapper for :func:`taichi.polar_decompose`.
+
+ Args:
+ A (ti.Matrix(n, n)): input nxn matrix `A`.
+ dt (DataType): date type of elements in matrix `A`, typically accepts ti.f32 or ti.f64.
+
+ Returns:
+ Decomposed nxn matrices `U` and `P`.
+ """
+ if dt is None:
+ dt = impl.get_runtime().default_fp
+ if A.n != 2 and A.n != 3:
+ raise Exception(
+ "Polar decomposition only supports 2D and 3D matrices.")
+ return _polar_decompose(A, dt)
+
+
+def svd(A, dt=None):
+ """Perform singular value decomposition (A=USV^T) for arbitrary size matrix.
+
+ Mathematical concept refers to https://en.wikipedia.org/wiki/Singular_value_decomposition.
+ This is only a wrappers for :func:`taichi.svd`.
+
+ Args:
+ A (ti.Matrix(n, n)): input nxn matrix `A`.
+ dt (DataType): date type of elements in matrix `A`, typically accepts ti.f32 or ti.f64.
+
+ Returns:
+ Decomposed nxn matrices `U`, 'S' and `V`.
+ """
+ if dt is None:
+ dt = impl.get_runtime().default_fp
+ if A.n != 2 and A.n != 3:
+ raise Exception("SVD only supports 2D and 3D matrices.")
+ return _svd(A, dt)
+
+
+def eig(A, dt=None):
+ """Compute the eigenvalues and right eigenvectors of a real matrix.
+
+ Mathematical concept refers to https://en.wikipedia.org/wiki/Eigendecomposition_of_a_matrix.
+ 2D implementation refers to :func:`taichi.eig2x2`.
+
+ Args:
+ A (ti.Matrix(n, n)): 2D Matrix for which the eigenvalues and right eigenvectors will be computed.
+ dt (DataType): The datatype for the eigenvalues and right eigenvectors.
+
+ Returns:
+ eigenvalues (ti.Matrix(n, 2)): The eigenvalues in complex form. Each row stores one eigenvalue. The first number of the eigenvalue represents the real part and the second number represents the imaginary part.
+ eigenvectors (ti.Matrix(n*2, n)): The eigenvectors in complex form. Each column stores one eigenvector. Each eigenvector consists of n entries, each of which is represented by two numbers for its real part and imaginary part.
+ """
+ if dt is None:
+ dt = impl.get_runtime().default_fp
+ if A.n == 2:
+ return eig2x2(A, dt)
+ raise Exception("Eigen solver only supports 2D matrices.")
+
+
+def sym_eig(A, dt=None):
+ """Compute the eigenvalues and right eigenvectors of a real symmetric matrix.
+
+ Mathematical concept refers to https://en.wikipedia.org/wiki/Eigendecomposition_of_a_matrix.
+ 2D implementation refers to :func:`taichi.sym_eig2x2`.
+
+ Args:
+ A (ti.Matrix(n, n)): Symmetric Matrix for which the eigenvalues and right eigenvectors will be computed.
+ dt (DataType): The datatype for the eigenvalues and right eigenvectors.
+
+ Returns:
+ eigenvalues (ti.Vector(n)): The eigenvalues. Each entry store one eigen value.
+ eigenvectors (ti.Matrix(n, n)): The eigenvectors. Each column stores one eigenvector.
+ """
+ assert all(A == A.transpose()), "A needs to be symmetric"
+ if dt is None:
+ dt = impl.get_runtime().default_fp
+ if A.n == 2:
+ return sym_eig2x2(A, dt)
+ raise Exception("Symmetric eigen solver only supports 2D matrices.")
+
+
+__all__ = ['randn', 'polar_decompose', 'eig', 'sym_eig', 'svd']
diff --git a/python/taichi/_kernels.py b/python/taichi/_kernels.py
new file mode 100644
index 0000000000000..683b703ba14df
--- /dev/null
+++ b/python/taichi/_kernels.py
@@ -0,0 +1,235 @@
+from taichi._lib.utils import get_os_name
+from taichi.lang import ops
+from taichi.lang._ndrange import ndrange
+from taichi.lang.expr import Expr
+from taichi.lang.field import ScalarField
+from taichi.lang.impl import grouped, static, static_assert
+from taichi.lang.kernel_impl import kernel
+from taichi.lang.runtime_ops import sync
+from taichi.lang.snode import deactivate
+from taichi.types.annotations import any_arr, ext_arr, template
+from taichi.types.primitive_types import f16, f32, f64, u8
+
+
+# A set of helper (meta)functions
+@kernel
+def fill_tensor(tensor: template(), val: template()):
+ for I in grouped(tensor):
+ tensor[I] = val
+
+
+@kernel
+def fill_ndarray(ndarray: any_arr(), val: template()):
+ for I in grouped(ndarray):
+ ndarray[I] = val
+
+
+@kernel
+def fill_ndarray_matrix(ndarray: any_arr(), val: template()):
+ for I in grouped(ndarray):
+ ndarray[I].fill(val)
+
+
+@kernel
+def tensor_to_ext_arr(tensor: template(), arr: ext_arr()):
+ for I in grouped(tensor):
+ arr[I] = tensor[I]
+
+
+@kernel
+def ndarray_to_ext_arr(ndarray: any_arr(), arr: ext_arr()):
+ for I in grouped(ndarray):
+ arr[I] = ndarray[I]
+
+
+@kernel
+def ndarray_matrix_to_ext_arr(ndarray: any_arr(), arr: ext_arr(),
+ as_vector: template()):
+ for I in grouped(ndarray):
+ for p in static(range(ndarray[I].n)):
+ for q in static(range(ndarray[I].m)):
+ if static(as_vector):
+ arr[I, p] = ndarray[I][p]
+ else:
+ arr[I, p, q] = ndarray[I][p, q]
+
+
+@kernel
+def vector_to_fast_image(img: template(), out: ext_arr()):
+ # FIXME: Why is ``for i, j in img:`` slower than:
+ for i, j in ndrange(*img.shape):
+ r, g, b = 0, 0, 0
+ color = img[i, img.shape[1] - 1 - j]
+ if static(img.dtype in [f16, f32, f64]):
+ r, g, b = min(255, max(0, int(color * 255)))
+ else:
+ static_assert(img.dtype == u8)
+ r, g, b = color
+ idx = j * img.shape[0] + i
+ # We use i32 for |out| since OpenGL and Metal doesn't support u8 types
+ if static(get_os_name() != 'osx'):
+ out[idx] = (r << 16) + (g << 8) + b
+ else:
+ # What's -16777216?
+ #
+ # On Mac, we need to set the alpha channel to 0xff. Since Mac's GUI
+ # is big-endian, the color is stored in ABGR order, and we need to
+ # add 0xff000000, which is -16777216 in I32's legit range. (Albeit
+ # the clarity, adding 0xff000000 doesn't work.)
+ alpha = -16777216
+ out[idx] = (b << 16) + (g << 8) + r + alpha
+
+
+@kernel
+def tensor_to_image(tensor: template(), arr: ext_arr()):
+ for I in grouped(tensor):
+ t = ops.cast(tensor[I], f32)
+ arr[I, 0] = t
+ arr[I, 1] = t
+ arr[I, 2] = t
+
+
+@kernel
+def vector_to_image(mat: template(), arr: ext_arr()):
+ for I in grouped(mat):
+ for p in static(range(mat.n)):
+ arr[I, p] = ops.cast(mat[I][p], f32)
+ if static(mat.n <= 2):
+ arr[I, 2] = 0
+
+
+@kernel
+def tensor_to_tensor(tensor: template(), other: template()):
+ for I in grouped(tensor):
+ tensor[I] = other[I]
+
+
+@kernel
+def ext_arr_to_tensor(arr: ext_arr(), tensor: template()):
+ for I in grouped(tensor):
+ tensor[I] = arr[I]
+
+
+@kernel
+def ndarray_to_ndarray(ndarray: any_arr(), other: any_arr()):
+ for I in grouped(ndarray):
+ ndarray[I] = other[I]
+
+
+@kernel
+def ext_arr_to_ndarray(arr: ext_arr(), ndarray: any_arr()):
+ for I in grouped(ndarray):
+ ndarray[I] = arr[I]
+
+
+@kernel
+def ext_arr_to_ndarray_matrix(arr: ext_arr(), ndarray: any_arr(),
+ as_vector: template()):
+ for I in grouped(ndarray):
+ for p in static(range(ndarray[I].n)):
+ for q in static(range(ndarray[I].m)):
+ if static(as_vector):
+ ndarray[I][p] = arr[I, p]
+ else:
+ ndarray[I][p, q] = arr[I, p, q]
+
+
+@kernel
+def matrix_to_ext_arr(mat: template(), arr: ext_arr(), as_vector: template()):
+ for I in grouped(mat):
+ for p in static(range(mat.n)):
+ for q in static(range(mat.m)):
+ if static(as_vector):
+ arr[I, p] = mat[I][p]
+ else:
+ arr[I, p, q] = mat[I][p, q]
+
+
+@kernel
+def ext_arr_to_matrix(arr: ext_arr(), mat: template(), as_vector: template()):
+ for I in grouped(mat):
+ for p in static(range(mat.n)):
+ for q in static(range(mat.m)):
+ if static(as_vector):
+ mat[I][p] = arr[I, p]
+ else:
+ mat[I][p, q] = arr[I, p, q]
+
+
+@kernel
+def clear_gradients(_vars: template()):
+ for I in grouped(ScalarField(Expr(_vars[0]))):
+ for s in static(_vars):
+ ScalarField(Expr(s))[I] = 0
+
+
+@kernel
+def clear_loss(l: template()):
+ # Using SNode writers would result in a forced sync, therefore we wrap these
+ # writes into a kernel.
+ l[None] = 0
+ l.grad[None] = 1
+
+
+@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)):
+ mat[I][p, q] = vals[p][q]
+
+
+@kernel
+def snode_deactivate(b: template()):
+ for I in grouped(b):
+ deactivate(b, I)
+
+
+@kernel
+def snode_deactivate_dynamic(b: template()):
+ for I in grouped(b.parent()):
+ deactivate(b, I)
+
+
+# Odd-even merge sort
+# References:
+# https://developer.nvidia.com/gpugems/gpugems2/part-vi-simulation-and-numerical-algorithms/chapter-46-improved-gpu-sorting
+# https://en.wikipedia.org/wiki/Batcher_odd%E2%80%93even_mergesort
+@kernel
+def sort_stage(keys: template(), use_values: int, values: template(), N: int,
+ p: int, k: int, invocations: int):
+ for inv in range(invocations):
+ j = k % p + inv * 2 * k
+ for i in range(0, min(k, N - j - k)):
+ a = i + j
+ b = i + j + k
+ if int(a / (p * 2)) == int(b / (p * 2)):
+ key_a = keys[a]
+ key_b = keys[b]
+ if key_a > key_b:
+ keys[a] = key_b
+ keys[b] = key_a
+ if use_values != 0:
+ temp = values[a]
+ values[a] = values[b]
+ values[b] = temp
+
+
+def parallel_sort(keys, values=None):
+ N = keys.shape[0]
+
+ num_stages = 0
+ p = 1
+ while p < N:
+ k = p
+ while k >= 1:
+ invocations = int((N - k - k % p) / (2 * k)) + 1
+ if values is None:
+ sort_stage(keys, 0, keys, N, p, k, invocations)
+ else:
+ sort_stage(keys, 1, values, N, p, k, invocations)
+ num_stages += 1
+ sync()
+ k = int(k / 2)
+ p = int(p * 2)
+ print(num_stages)
diff --git a/python/taichi/_lib/__init__.py b/python/taichi/_lib/__init__.py
new file mode 100644
index 0000000000000..1334e8635e899
--- /dev/null
+++ b/python/taichi/_lib/__init__.py
@@ -0,0 +1 @@
+from taichi._lib.utils import ti_core as core
diff --git a/python/taichi/_lib/core/__init__.py b/python/taichi/_lib/core/__init__.py
new file mode 100644
index 0000000000000..8b137891791fe
--- /dev/null
+++ b/python/taichi/_lib/core/__init__.py
@@ -0,0 +1 @@
+
diff --git a/python/taichi/core/util.py b/python/taichi/_lib/utils.py
similarity index 57%
rename from python/taichi/core/util.py
rename to python/taichi/_lib/utils.py
index 7341bc6c59d77..8b2db08261479 100644
--- a/python/taichi/core/util.py
+++ b/python/taichi/_lib/utils.py
@@ -1,28 +1,19 @@
-import ctypes
-import datetime
-import multiprocessing
import os
import platform
-import random
-import shutil
import sys
-import time
-from colorama import Back, Fore, Style
+from colorama import Fore, Style
if sys.version_info[0] < 3 or sys.version_info[1] <= 5:
raise RuntimeError(
"\nPlease restart with Python 3.6+\n" + "Current Python version:",
sys.version_info)
-ti_core = None
-
def in_docker():
if os.environ.get("TI_IN_DOCKER", "") == "":
return False
- else:
- return True
+ return True
def get_os_name():
@@ -31,23 +22,26 @@ def get_os_name():
# it will return 'macOS-XXXX' instead of 'Darwin-XXXX'
if name.lower().startswith('darwin') or name.lower().startswith('macos'):
return 'osx'
- elif name.lower().startswith('windows'):
+ if name.lower().startswith('windows'):
return 'win'
- elif name.lower().startswith('linux'):
+ if name.lower().startswith('linux'):
return 'linux'
- assert False, "Unknown platform name %s" % name
+ if 'bsd' in name.lower():
+ return 'unix'
+ assert False, f"Unknown platform name {name}"
def import_ti_core():
- global ti_core
if get_os_name() != 'win':
+ # pylint: disable=E1101
old_flags = sys.getdlopenflags()
sys.setdlopenflags(2 | 8) # RTLD_NOW | RTLD_DEEPBIND
else:
- pyddir = os.path.join(package_root(), 'lib')
- os.environ['PATH'] += ';' + pyddir
+ pyddir = os.path.dirname(os.path.realpath(__file__))
+ os.environ['PATH'] += os.pathsep + pyddir
try:
- import taichi_core as core # pylint: disable=C0415
+ from taichi._lib.core import \
+ taichi_core as core # pylint: disable=C0415
except Exception as e:
if isinstance(e, ImportError):
print(Fore.YELLOW + "Share object taichi_core import failed, "
@@ -55,28 +49,28 @@ def import_ti_core():
"https://docs.taichi.graphics/lang/articles/misc/install" +
Fore.RESET)
if get_os_name() == 'win':
+ # pylint: disable=E1101
e.msg += '\nConsider installing Microsoft Visual C++ Redistributable: https://aka.ms/vs/16/release/vc_redist.x64.exe'
- elif get_os_name() == 'linux':
- e.msg += '\nConsider installing libtinfo5: sudo apt-get install libtinfo5'
raise e from None
- ti_core = core
+
if get_os_name() != 'win':
- sys.setdlopenflags(old_flags)
- lib_dir = os.path.join(package_root(), 'lib')
+ sys.setdlopenflags(old_flags) # pylint: disable=E1101
+ lib_dir = os.path.join(package_root, '_lib', 'runtime')
core.set_lib_dir(locale_encode(lib_dir))
+ return core
def locale_encode(path):
try:
import locale # pylint: disable=C0415
return path.encode(locale.getdefaultlocale()[1])
- except:
+ except (UnicodeEncodeError, TypeError):
try:
return path.encode(sys.getfilesystemencoding())
- except:
+ except UnicodeEncodeError:
try:
return path.encode()
- except:
+ except UnicodeEncodeError:
return path
@@ -84,12 +78,12 @@ def is_ci():
return os.environ.get('TI_CI', '') == '1'
-def package_root():
- return os.path.join(os.path.dirname(os.path.realpath(__file__)), '../')
+package_root = os.path.join(
+ os.path.dirname(os.path.dirname(os.path.realpath(__file__))))
def get_core_shared_object():
- directory = os.path.join(package_root(), 'lib')
+ directory = os.path.join(package_root, '_lib')
return os.path.join(directory, 'libtaichi_core.so')
@@ -106,16 +100,9 @@ def check_exists(src):
)
-def get_unique_task_id():
- return datetime.datetime.now().strftime('task-%Y-%m-%d-%H-%M-%S-r') + (
- '%05d' % random.randint(0, 10000))
+ti_core = import_ti_core()
-
-sys.path.append(os.path.join(package_root(), 'lib'))
-import_ti_core()
-
-ti_core.set_python_package_dir(package_root())
-os.makedirs(ti_core.get_repo_dir(), exist_ok=True)
+ti_core.set_python_package_dir(package_root)
log_level = os.environ.get('TI_LOG_LEVEL', '')
if log_level:
@@ -124,43 +111,42 @@ def get_unique_task_id():
def get_dll_name(name):
if get_os_name() == 'linux':
- return 'libtaichi_%s.so' % name
- elif get_os_name() == 'osx':
- return 'libtaichi_%s.dylib' % name
- elif get_os_name() == 'win':
- return 'taichi_%s.dll' % name
- else:
- raise Exception(f"Unknown OS: {get_os_name()}")
+ return f'libtaichi_{name}.so'
+ if get_os_name() == 'osx':
+ return f'libtaichi_{name}.dylib'
+ if get_os_name() == 'win':
+ return f'taichi_{name}.dll'
+ raise Exception(f"Unknown OS: {get_os_name()}")
def at_startup():
ti_core.set_core_state_python_imported(True)
-def require_version(major, minor=None, patch=None):
- versions = [
- int(ti_core.get_version_major()),
- int(ti_core.get_version_minor()),
- int(ti_core.get_version_patch()),
- ]
- match = major == versions[0] and (
- minor < versions[1] or minor == versions[1] and patch <= versions[2])
- if match:
- return
- else:
- print("Taichi version mismatch. required >= {}.{}.{}".format(
- major, minor, patch))
- print("Installed =", ti_core.get_version_string())
- raise Exception("Taichi version mismatch")
+at_startup()
-at_startup()
+def compare_version(latest, current):
+ latest_num = map(int, latest.split('.'))
+ current_num = map(int, current.split('.'))
+ return tuple(latest_num) > tuple(current_num)
def _print_taichi_header():
header = '[Taichi] '
header += f'version {ti_core.get_version_string()}, '
+ try:
+ timestamp_path = os.path.join(ti_core.get_repo_dir(), 'timestamp')
+ if os.path.exists(timestamp_path):
+ latest_version = ''
+ with open(timestamp_path, 'r') as f:
+ latest_version = f.readlines()[1].rstrip()
+ if compare_version(latest_version, ti_core.get_version_string()):
+ header += f'latest version {latest_version}, '
+ except:
+ pass
+
llvm_version = ti_core.get_llvm_version_string()
header += f'llvm {llvm_version}, '
@@ -177,10 +163,3 @@ def _print_taichi_header():
_print_taichi_header()
-
-__all__ = [
- 'ti_core',
- 'get_os_name',
- 'package_root',
- 'require_version',
-]
diff --git a/python/taichi/_logging.py b/python/taichi/_logging.py
index 927c9b2e59036..29696b9fe062d 100644
--- a/python/taichi/_logging.py
+++ b/python/taichi/_logging.py
@@ -1,10 +1,19 @@
import inspect
import os
-from taichi.core import ti_core
+from taichi._lib import core as ti_core
def _get_logging(name):
+ """Generates a decorator to decorate a specific logger function.
+
+ Args:
+ name (str): The string represents logging level.
+ Effective levels include: 'trace', 'debug', 'info', 'warn', 'error', 'critical'.
+
+ Returns:
+ Callabe: The decorated function.
+ """
def logger(msg, *args, **kwargs):
# Python inspection takes time (~0.1ms) so avoid it as much as possible
if ti_core.logging_effective(name):
@@ -20,19 +29,83 @@ def logger(msg, *args, **kwargs):
def set_logging_level(level):
+ """Setting the logging level to a specified value.
+ Available levels are: 'trace', 'debug', 'info', 'warn', 'error', 'critical'.
+
+ Note that after calling this function, logging levels below the specified one will
+ also be effective. For example if `level` is set to 'warn', then the levels below
+ it, which are 'error' and 'critical' in this case, will also be effective.
+
+ See also https://docs.taichi.graphics/lang/articles/contribution/utilities#logging.
+
+ Args:
+ level (str): Logging level.
+
+ Example::
+
+ >>> set_logging_level('debug')
+ """
ti_core.set_logging_level(level)
def is_logging_effective(level):
+ """Check if the specified logging level is effective.
+ All levels below current level will be effective.
+ The default level is 'info'.
+
+ See also https://docs.taichi.graphics/lang/articles/contribution/utilities#logging.
+
+ Args:
+ level (str): The string represents logging level. \
+ Effective levels include: 'trace', 'debug', 'info', 'warn', 'error', 'critical'.
+
+ Returns:
+ Bool: Indicate whether the logging level is effective.
+
+ Example::
+
+ >>> # assume current level is 'info'
+ >>> print(ti.is_logging_effective("trace")) # False
+ >>> print(ti.is_logging_effective("debug")) # False
+ >>> print(ti.is_logging_effective("info")) # True
+ >>> print(ti.is_logging_effective("warn")) # True
+ >>> print(ti.is_logging_effective("error")) # True
+ >>> print(ti.is_logging_effective("critical")) # True
+ """
return ti_core.logging_effective(level)
+# ------------------------
+
DEBUG = 'debug'
+"""The `str` 'debug', used for the `debug` logging level.
+"""
+# ------------------------
+
TRACE = 'trace'
+"""The `str` 'trace', used for the `debug` logging level.
+"""
+# ------------------------
+
INFO = 'info'
+"""The `str` 'info', used for the `info` logging level.
+"""
+# ------------------------
+
WARN = 'warn'
+"""The `str` 'warn', used for the `warn` logging level.
+"""
+# ------------------------
+
ERROR = 'error'
+"""The `str` 'error', used for the `error` logging level.
+"""
+# ------------------------
+
CRITICAL = 'critical'
+"""The `str` 'critical', used for the `critical` logging level.
+"""
+# ------------------------
supported_log_levels = [DEBUG, TRACE, INFO, WARN, ERROR, CRITICAL]
@@ -44,7 +117,6 @@ def is_logging_effective(level):
critical = _get_logging(CRITICAL)
__all__ = [
- 'DEBUG', 'TRACE', 'INFO', 'WARN', 'ERROR', 'CRITICAL', 'debug', 'trace',
- 'info', 'warn', 'error', 'critical', 'supported_log_levels',
- 'set_logging_level', 'is_logging_effective'
+ 'DEBUG', 'TRACE', 'INFO', 'WARN', 'ERROR', 'CRITICAL', 'set_logging_level',
+ 'is_logging_effective'
]
diff --git a/python/taichi/main.py b/python/taichi/_main.py
similarity index 66%
rename from python/taichi/main.py
rename to python/taichi/_main.py
index f1886b75d9984..254553f3ec841 100644
--- a/python/taichi/main.py
+++ b/python/taichi/_main.py
@@ -1,7 +1,6 @@
import argparse
import math
import os
-import random
import runpy
import shutil
import subprocess
@@ -12,11 +11,10 @@
from pathlib import Path
import numpy as np
-import taichi.cc_compose
-import taichi.diagnose
-from colorama import Back, Fore, Style
-from taichi.core import ti_core as _ti_core
-from taichi.tools import video
+from colorama import Fore
+from taichi._lib import core as _ti_core
+from taichi._lib import utils
+from taichi.tools import cc_compose, diagnose, video
import taichi as ti
@@ -78,7 +76,7 @@ def __call__(self):
# Parse the command
args = self.main_parser.parse_args(sys.argv[1:2])
- if args.command not in self.registered_commands:
+ if args.command not in self.registered_commands: # pylint: disable=E1101
# TODO: do we really need this?
if args.command.endswith(".py"):
TaichiMain._exec_python_file(args.command)
@@ -89,25 +87,19 @@ def __call__(self):
return getattr(self, args.command)(sys.argv[2:])
- def _get_friend_links(self):
- uri = 'en/stable'
- try:
- import locale # pylint: disable=C0415
- if 'zh' in locale.getdefaultlocale()[0]:
- uri = 'zh_CN/latest'
- except:
- pass
+ @staticmethod
+ def _get_friend_links():
return '\n' \
- f'Docs: https://taichi.rtfd.io/{uri}\n' \
- f'GitHub: https://github.com/taichi-dev/taichi\n' \
- f'Forum: https://forum.taichi.graphics\n'
+ 'Docs: https://docs.taichi.graphics/\n' \
+ 'GitHub: https://github.com/taichi-dev/taichi/\n' \
+ 'Forum: https://forum.taichi.graphics/\n'
def _usage(self) -> str:
"""Compose deterministic usage message based on registered_commands."""
# TODO: add some color to commands
msg = "\n"
space = 20
- for command in sorted(self.registered_commands):
+ for command in sorted(self.registered_commands): # pylint: disable=E1101
msg += f" {command}{' ' * (space - len(command))}|-> {getattr(self, command).__doc__}\n"
return msg
@@ -121,7 +113,7 @@ def _exec_python_file(filename: str):
def _get_examples_dir() -> Path:
"""Get the path to the examples directory."""
- root_dir = ti.package_root()
+ root_dir = utils.package_root
examples_dir = Path(root_dir) / 'examples'
return examples_dir
@@ -130,10 +122,7 @@ def _get_available_examples() -> set:
"""Get a set of all available example names."""
examples_dir = TaichiMain._get_examples_dir()
all_examples = examples_dir.rglob('*.py')
- all_example_names = {
- Path(f).stem: Path(f).parent
- for f in all_examples
- }
+ all_example_names = {f.stem: f.parent for f in all_examples}
return all_example_names
@staticmethod
@@ -142,8 +131,7 @@ def support_choice_with_dot_py(choice):
if choice.endswith('.py') and choice.split('.')[0] in choices:
# try to find and remove python file extension
return choice.split('.')[0]
- else:
- return choice
+ return choice
return support_choice_with_dot_py
@@ -156,7 +144,7 @@ def example(self, arguments: list = sys.argv[2:]):
description=f"{self.example.__doc__}")
parser.add_argument(
"name",
- help=f"Name of an example (supports .py extension too)\n",
+ help="Name of an example (supports .py extension too)\n",
type=TaichiMain._example_choices_type(choices.keys()),
choices=sorted(choices.keys()))
parser.add_argument(
@@ -180,6 +168,8 @@ def example(self, arguments: list = sys.argv[2:]):
dest='save',
action='store_true',
help="Save source code to current directory instead of running it")
+
+ # TODO: Pass the arguments to downstream correctly(#3216).
args = parser.parse_args(arguments)
examples_dir = TaichiMain._get_examples_dir()
@@ -189,7 +179,8 @@ def example(self, arguments: list = sys.argv[2:]):
sys.path.append(str((examples_dir / choices[args.name]).resolve()))
# Short circuit for testing
- if self.test_mode: return args
+ if self.test_mode:
+ return args
if args.save:
print(f"Saving example {args.name} to current directory...")
@@ -200,7 +191,7 @@ def example(self, arguments: list = sys.argv[2:]):
try:
import rich.console # pylint: disable=C0415
import rich.syntax # pylint: disable=C0415
- except ImportError as e:
+ except ImportError:
print('To make -P work, please: python3 -m pip install rich')
return 1
# https://rich.readthedocs.io/en/latest/syntax.html
@@ -218,15 +209,19 @@ def example(self, arguments: list = sys.argv[2:]):
runpy.run_path(target, run_name='__main__')
+ return None
+
+ @staticmethod
@register
- def changelog(self, arguments: list = sys.argv[2:]):
+ def changelog(arguments: list = sys.argv[2:]):
"""Display changelog of current version"""
- changelog_md = os.path.join(ti.package_root(), 'CHANGELOG.md')
+ changelog_md = os.path.join(utils.package_root, 'CHANGELOG.md')
with open(changelog_md) as f:
print(f.read())
+ @staticmethod
@register
- def release(self, arguments: list = sys.argv[2:]):
+ def release(arguments: list = sys.argv[2:]):
"""Make source code release"""
raise RuntimeError('TBD')
@@ -257,12 +252,15 @@ def gif(self, arguments: list = sys.argv[2:]):
args = parser.parse_args(arguments)
args.output_file = str(Path(args.input_file).with_suffix('.gif'))
- ti.info(f"Converting {args.input_file} to {args.output_file}")
+ ti._logging.info(f"Converting {args.input_file} to {args.output_file}")
# Short circuit for testing
- if self.test_mode: return args
+ if self.test_mode:
+ return args
video.mp4_to_gif(args.input_file, args.output_file, args.framerate)
+ return None
+
@register
def video_speed(self, arguments: list = sys.argv[2:]):
"""Speed up video in the same directory"""
@@ -290,9 +288,12 @@ def video_speed(self, arguments: list = sys.argv[2:]):
))
# Short circuit for testing
- if self.test_mode: return args
+ if self.test_mode:
+ return args
video.accelerate_video(args.input_file, args.output_file, args.speed)
+ return None
+
@register
def video_crop(self, arguments: list = sys.argv[2:]):
"""Crop video in the same directory"""
@@ -332,10 +333,13 @@ def video_crop(self, arguments: list = sys.argv[2:]):
))
# Short circuit for testing
- if self.test_mode: return args
+ if self.test_mode:
+ return args
video.crop_video(args.input_file, args.output_file, args.x_begin,
args.x_end, args.y_begin, args.y_end)
+ return None
+
@register
def video_scale(self, arguments: list = sys.argv[2:]):
"""Scale video resolution in the same directory"""
@@ -373,10 +377,13 @@ def video_scale(self, arguments: list = sys.argv[2:]):
))
# Short circuit for testing
- if self.test_mode: return args
+ if self.test_mode:
+ return args
video.scale_video(args.input_file, args.output_file, args.ratio_width,
args.ratio_height)
+ return None
+
@register
def video(self, arguments: list = sys.argv[2:]):
"""Make a video using *.png files in the current directory"""
@@ -413,36 +420,42 @@ def video(self, arguments: list = sys.argv[2:]):
assert 1 <= args.crf <= 51, "The range of the CRF scale is 1-51, where 1 is almost lossless, 20 is the default, and 51 is worst quality possible."
- ti.info(f'Making video using {len(args.inputs)} png files...')
- ti.info(f'frame_rate = {args.framerate}')
+ ti._logging.info(f'Making video using {len(args.inputs)} png files...')
+ ti._logging.info(f'frame_rate = {args.framerate}')
# Short circuit for testing
- if self.test_mode: return args
+ if self.test_mode:
+ return args
video.make_video(args.inputs,
output_path=str(args.output_file),
crf=args.crf,
frame_rate=args.framerate)
- ti.info(f'Done! Output video file = {args.output_file}')
+ ti._logging.info(f'Done! Output video file = {args.output_file}')
+ return None
+
+ @staticmethod
@register
- def doc(self, arguments: list = sys.argv[2:]):
+ def doc(arguments: list = sys.argv[2:]):
"""Build documentation"""
raise RuntimeError('TBD')
+ @staticmethod
@register
- def format(self, arguments: list = sys.argv[2:]):
+ def format(arguments: list = sys.argv[2:]):
"""Reformat modified source files"""
raise RuntimeError('Please run python misc/code_format.py instead')
+ @staticmethod
@register
- def format_all(self, arguments: list = sys.argv[2:]):
+ def format_all(arguments: list = sys.argv[2:]):
"""Reformat all source files"""
raise RuntimeError('Please run python misc/code_format.py instead')
@staticmethod
def _display_benchmark_regression(xd, yd, args):
def parse_dat(file):
- dict = {}
+ _dict = {}
with open(file) as f:
for line in f.readlines():
try:
@@ -452,28 +465,27 @@ def parse_dat(file):
b = float(b)
if abs(b % 1.0) < 1e-5: # codegen_*
b = int(b)
- dict[a.strip()] = b
- return dict
+ _dict[a.strip()] = b
+ return _dict
def parse_name(file):
if file[0:5] == 'test_':
return file[5:-4].replace('__test_', '::', 1)
- elif file[0:10] == 'benchmark_':
+ if file[0:10] == 'benchmark_':
return '::'.join(reversed(file[10:-4].split('__arch_')))
- else:
- raise Exception(f'bad benchmark file name {file}')
+ raise Exception(f'bad benchmark file name {file}')
- def get_dats(dir):
- list = []
- for x in os.listdir(dir):
+ def get_dats(directory):
+ _list = []
+ for x in os.listdir(directory):
if x.endswith('.dat'):
- list.append(x)
- dict = {}
- for x in list:
+ _list.append(x)
+ _dict = {}
+ for x in _list:
name = parse_name(x)
- path = os.path.join(dir, x)
- dict[name] = parse_dat(path)
- return dict
+ path = os.path.join(directory, x)
+ _dict[name] = parse_dat(path)
+ return _dict
def plot_in_gui(scatter):
@@ -508,13 +520,16 @@ def plot_in_gui(scatter):
else:
res = b / a
scatter[key].append(res)
- if res == 1: continue
+ if res == 1:
+ continue
if not single_line:
ret += f'{key:<30}'
res -= 1
color = Fore.RESET
- if res > 0: color = Fore.RED
- elif res < 0: color = Fore.GREEN
+ if res > 0:
+ color = Fore.RED
+ elif res < 0:
+ color = Fore.GREEN
if isinstance(a, float):
a = f'{a:>7.2}'
else:
@@ -560,13 +575,16 @@ def regression(self, arguments: list = sys.argv[2:]):
args = parser.parse_args(arguments)
# Short circuit for testing
- if self.test_mode: return args
+ if self.test_mode:
+ return args
baseline_dir = TaichiMain._get_benchmark_baseline_dir()
output_dir = TaichiMain._get_benchmark_output_dir()
TaichiMain._display_benchmark_regression(baseline_dir, output_dir,
args)
+ return None
+
@register
def baseline(self, arguments: list = sys.argv[2:]):
"""Archive current benchmark result as baseline"""
@@ -575,7 +593,8 @@ def baseline(self, arguments: list = sys.argv[2:]):
args = parser.parse_args(arguments)
# Short circuit for testing
- if self.test_mode: return args
+ if self.test_mode:
+ return args
baseline_dir = TaichiMain._get_benchmark_baseline_dir()
output_dir = TaichiMain._get_benchmark_output_dir()
@@ -583,77 +602,7 @@ def baseline(self, arguments: list = sys.argv[2:]):
shutil.copytree(output_dir, baseline_dir)
print(f"[benchmark] baseline data saved to {baseline_dir}")
- @staticmethod
- def _test_python(args):
- print("\nRunning Python tests...\n")
-
- root_dir = ti.package_root()
- test_dir = os.path.join(root_dir, 'tests')
- pytest_args = []
-
- # TODO: use pathlib to deal with suffix and stem name manipulation
- if args.files:
- # run individual tests
- for f in args.files:
- # auto-complete file names
- if not f.startswith('test_'):
- f = 'test_' + f
- if not f.endswith('.py'):
- f = f + '.py'
- pytest_args.append(os.path.join(test_dir, f))
- else:
- # run all the tests
- pytest_args = [test_dir]
- if args.verbose:
- pytest_args += ['-v']
- if args.rerun:
- pytest_args += ['--reruns', args.rerun]
- try:
- if args.coverage:
- pytest_args += ['--cov-branch', '--cov=python/taichi']
- if args.cov_append:
- pytest_args += ['--cov-append']
- if args.keys:
- pytest_args += ['-k', args.keys]
- if args.marks:
- pytest_args += ['-m', args.marks]
- if args.failed_first:
- pytest_args += ['--failed-first']
- if args.fail_fast:
- pytest_args += ['--exitfirst']
- except AttributeError:
- pass
-
- try:
- from multiprocessing import cpu_count # pylint: disable=C0415
- threads = min(8, cpu_count()) # To prevent running out of memory
- except NotImplementedError:
- threads = 2
-
- if not os.environ.get('TI_DEVICE_MEMORY_GB'):
- os.environ['TI_DEVICE_MEMORY_GB'] = '1.0' # Discussion: #769
-
- env_threads = os.environ.get('TI_TEST_THREADS', '')
- threads = args.threads or env_threads or threads
- print(f'Starting {threads} testing thread(s)...')
- if args.show_output:
- pytest_args += ['-s']
- print(
- f'Due to how pytest-xdist is implemented, the -s option does not work with multiple thread...'
- )
- else:
- if int(threads) > 1:
- pytest_args += ['-n', str(threads)]
- import pytest # pylint: disable=C0415
- return int(pytest.main(pytest_args))
-
- @staticmethod
- def _test_cpp(args):
- # Cpp tests use the legacy non LLVM backend
- ti.reset()
- print("Running C++ tests...")
- task = ti.Task('test')
- return int(task.run(*args.files))
+ return None
@register
def benchmark(self, arguments: list = sys.argv[2:]):
@@ -689,7 +638,8 @@ def benchmark(self, arguments: list = sys.argv[2:]):
args = parser.parse_args(arguments)
# Short circuit for testing
- if self.test_mode: return args
+ if self.test_mode:
+ return args
commit_hash = _ti_core.get_commit_hash()
with os.popen('git rev-parse HEAD') as f:
@@ -707,136 +657,18 @@ def benchmark(self, arguments: list = sys.argv[2:]):
os.system('python benchmarks/run.py')
# TODO: benchmark_python(args)
else:
- TaichiMain._test_python(args)
+ # TODO: shall we replace this with the new benchmark tools?
+ os.system('python tests/run_tests.py')
+
+ return None
+ @staticmethod
@register
def test(self, arguments: list = sys.argv[2:]):
- """Run the tests"""
- parser = argparse.ArgumentParser(prog='ti test',
- description=f"{self.test.__doc__}")
- parser.add_argument('files',
- nargs='*',
- help='Test name(s) to be run, e.g. "cli"')
- parser.add_argument('-c',
- '--cpp',
- dest='cpp',
- action='store_true',
- help='Only run the C++ tests')
- parser.add_argument('-s',
- '--show',
- dest='show_output',
- action='store_true',
- help='Show output (do not capture)')
- parser.add_argument('-v',
- '--verbose',
- dest='verbose',
- action='store_true',
- help='Run with verbose outputs')
- parser.add_argument('-r',
- '--rerun',
- required=False,
- default=None,
- dest='rerun',
- type=str,
- help='Rerun failed tests for given times')
- parser.add_argument('-k',
- '--keys',
- required=False,
- default=None,
- dest='keys',
- type=str,
- help='Only run tests that match the keys')
- parser.add_argument('-m',
- '--marks',
- required=False,
- default=None,
- dest='marks',
- type=str,
- help='Only run tests with specific marks')
- parser.add_argument('-f',
- '--failed-first',
- required=False,
- default=None,
- dest='failed_first',
- action='store_true',
- help='Run the previously failed test first')
- parser.add_argument('-x',
- '--fail-fast',
- required=False,
- default=None,
- dest='fail_fast',
- action='store_true',
- help='Exit instantly on the first failed test')
- parser.add_argument('-C',
- '--coverage',
- required=False,
- default=None,
- dest='coverage',
- action='store_true',
- help='Run tests and record the coverage result')
- parser.add_argument(
- '-A',
- '--cov-append',
- required=False,
- default=None,
- dest='cov_append',
- action='store_true',
- help=
- 'Append coverage result to existing one instead of overriding it')
- parser.add_argument(
- '-t',
- '--threads',
- required=False,
- default=None,
- dest='threads',
- type=str,
- help='Custom number of threads for parallel testing')
- parser.add_argument(
- '-a',
- '--arch',
- required=False,
- default=None,
- dest='arch',
- type=str,
- help='Custom the arch(s) (backend) to run tests on')
- parser.add_argument(
- '-n',
- '--exclusive',
- required=False,
- default=False,
- dest='exclusive',
- action='store_true',
- help=
- 'Exclude arch(s) from test instead of include them, together with -a'
+ raise RuntimeError(
+ 'ti test is deprecated. Please run `python tests/run_tests.py` instead.'
)
- args = parser.parse_args(arguments)
-
- if args.arch:
- arch = args.arch
- if args.exclusive:
- arch = '^' + arch
- print(f'Running on Arch={arch}')
- os.environ['TI_WANTED_ARCHS'] = arch
-
- # Short circuit for testing
- if self.test_mode: return args
-
- if args.files:
- if args.cpp:
- return TaichiMain._test_cpp(args)
- else:
- return TaichiMain._test_python(args)
- elif args.cpp:
- # Only run C++ tests
- return TaichiMain._test_cpp(args)
- else:
- # Run both C++ and Python tests
- ret = TaichiMain._test_python(args)
- if ret != 0:
- return ret
- return TaichiMain._test_cpp(args)
-
@register
def run(self, arguments: list = sys.argv[2:]):
"""Run a single script"""
@@ -848,10 +680,13 @@ def run(self, arguments: list = sys.argv[2:]):
args = parser.parse_args(arguments)
# Short circuit for testing
- if self.test_mode: return args
+ if self.test_mode:
+ return args
runpy.run_path(args.filename)
+ return None
+
@register
def debug(self, arguments: list = sys.argv[2:]):
"""Debug a single script"""
@@ -864,56 +699,21 @@ def debug(self, arguments: list = sys.argv[2:]):
args = parser.parse_args(arguments)
# Short circuit for testing
- if self.test_mode: return args
+ if self.test_mode:
+ return args
_ti_core.set_core_trigger_gdb_when_crash(True)
os.environ['TI_DEBUG'] = '1'
runpy.run_path(args.filename, run_name='__main__')
- @register
- def task(self, arguments: list = sys.argv[2:]):
- """Run a specific task"""
- parser = argparse.ArgumentParser(prog='ti task',
- description=f"{self.task.__doc__}")
- parser.add_argument('taskname',
- help='A single task name to run, e.g. test_math')
- parser.add_argument('taskargs',
- nargs='*',
- help='Optional task argument(s) to run with task')
- args = parser.parse_args(arguments)
-
- # Short circuit for testing
- if self.test_mode: return args
-
- task = ti.Task(args.taskname)
- task.run(*args.taskargs)
-
- @register
- def dist(self, arguments: list = sys.argv[2:]):
- """Build package and test in release mode"""
- parser = argparse.ArgumentParser(prog='ti dist',
- description=f"{self.dist.__doc__}")
- parser.add_argument('mode',
- nargs='?',
- default='test',
- choices=['upload', 'try_upload', 'test'],
- help='Which mode shall we run?')
- args = parser.parse_args(arguments)
-
- os.chdir(os.path.join(_ti_core.get_repo_dir(), 'python'))
- sys.argv.pop(0)
- sys.argv.append(args.mode)
- runpy.run_path('build.py')
+ return None
+ @staticmethod
@register
- def diagnose(self, arguments: list = sys.argv[2:]):
+ def diagnose(arguments: list = sys.argv[2:]):
"""System diagnose information"""
- parser = argparse.ArgumentParser(
- prog='ti diagnose', description=f"{self.diagnose.__doc__}")
- args = parser.parse_args(arguments)
-
- taichi.diagnose.main()
+ diagnose.main()
@register
def cc_compose(self, arguments: list = sys.argv[2:]):
@@ -939,16 +739,13 @@ def cc_compose(self, arguments: list = sys.argv[2:]):
help='Generate output C file for Emscripten instead of raw C')
args = parser.parse_args(arguments)
- taichi.cc_compose.main(args.fin_name, args.fout_name, args.hdrout_name,
- args.emscripten)
+ cc_compose.main(args.fin_name, args.fout_name, args.hdrout_name,
+ args.emscripten)
+ @staticmethod
@register
- def repl(self, arguments: list = sys.argv[2:]):
+ def repl(arguments: list = sys.argv[2:]):
"""Start Taichi REPL / Python shell with 'import taichi as ti'"""
- parser = argparse.ArgumentParser(prog='ti repl',
- description=f"{self.repl.__doc__}")
- args = parser.parse_args(arguments)
-
def local_scope():
try:
@@ -956,18 +753,18 @@ def local_scope():
IPython.embed()
except ImportError:
import code # pylint: disable=C0415
- __name__ = '__console__'
+ __name__ = '__console__' # pylint: disable=W0622
code.interact(local=locals())
local_scope()
+ @staticmethod
@register
- def lint(self, arguments: list = sys.argv[2:]):
+ def lint(arguments: list = sys.argv[2:]):
"""Run pylint checker for the Python codebase of Taichi"""
- parser = argparse.ArgumentParser(prog='ti lint',
- description=f"{self.lint.__doc__}")
# TODO: support arguments for lint specific files
- args = parser.parse_args(arguments)
+ # parser = argparse.ArgumentParser(prog='ti lint', description=f"{self.lint.__doc__}")
+ # args = parser.parse_args(arguments)
options = [os.path.dirname(__file__)]
diff --git a/python/taichi/_snode/__init__.py b/python/taichi/_snode/__init__.py
new file mode 100644
index 0000000000000..c323e8e6ab814
--- /dev/null
+++ b/python/taichi/_snode/__init__.py
@@ -0,0 +1,3 @@
+from taichi._snode.fields_builder import FieldsBuilder
+
+__all__ = ['FieldsBuilder']
diff --git a/python/taichi/snode/fields_builder.py b/python/taichi/_snode/fields_builder.py
similarity index 68%
rename from python/taichi/snode/fields_builder.py
rename to python/taichi/_snode/fields_builder.py
index 94fecba268535..4409790e1a298 100644
--- a/python/taichi/snode/fields_builder.py
+++ b/python/taichi/_snode/fields_builder.py
@@ -1,12 +1,10 @@
-import functools
-import types
from typing import Any, Optional, Sequence, Union
-from taichi.core.util import ti_core as _ti_core
+from taichi._lib import core as _ti_core
+from taichi._snode.snode_tree import SNodeTree
from taichi.lang import impl, snode
-from taichi.lang.exception import InvalidOperationError
-from taichi.misc.util import warning
-from taichi.snode.snode_tree import SNodeTree
+from taichi.lang.exception import TaichiRuntimeError
+from taichi.lang.util import warning
_snode_registry = _ti_core.SNodeRegistry()
@@ -36,14 +34,14 @@ class FieldsBuilder:
fb.finalize()
"""
def __init__(self):
- self._ptr = _snode_registry.create_root()
- self._root = snode.SNode(self._ptr)
- self._finalized = False
- self._empty = True
+ self.ptr = _snode_registry.create_root(impl.get_runtime().prog)
+ self.root = snode.SNode(self.ptr)
+ self.finalized = False
+ self.empty = True
# TODO: move this into SNodeTree
@classmethod
- def finalized_roots(cls):
+ def _finalized_roots(cls):
"""Gets all the roots of the finalized SNodeTree.
Returns:
@@ -56,27 +54,11 @@ def finalized_roots(cls):
roots_ptr.append(snode.SNode(res))
return roots_ptr
- @property
- def ptr(self):
- return self._ptr
-
- @property
- def root(self):
- return self._root
-
- @property
- def empty(self):
- return self._empty
-
- @property
- def finalized(self):
- return self._finalized
-
# TODO: move this to SNodeTree class.
def deactivate_all(self):
"""Same as :func:`taichi.lang.snode.SNode.deactivate_all`"""
- if self._finalized:
- self._root.deactivate_all()
+ if self.finalized:
+ self.root.deactivate_all()
else:
warning(
"""'deactivate_all()' would do nothing if FieldsBuilder is not finalized"""
@@ -86,17 +68,17 @@ def dense(self, indices: Union[Sequence[_Axis], _Axis],
dimensions: Union[Sequence[int], int]):
"""Same as :func:`taichi.lang.snode.SNode.dense`"""
self._check_not_finalized()
- self._empty = False
- return self._root.dense(indices, dimensions)
+ self.empty = False
+ return self.root.dense(indices, dimensions)
def pointer(self, indices: Union[Sequence[_Axis], _Axis],
dimensions: Union[Sequence[int], int]):
"""Same as :func:`taichi.lang.snode.SNode.pointer`"""
self._check_not_finalized()
- self._empty = False
- return self._root.pointer(indices, dimensions)
+ self.empty = False
+ return self.root.pointer(indices, dimensions)
- def hash(self, indices, dimensions):
+ def _hash(self, indices, dimensions):
"""Same as :func:`taichi.lang.snode.SNode.hash`"""
raise NotImplementedError()
@@ -106,28 +88,28 @@ def dynamic(self,
chunk_size: Optional[int] = None):
"""Same as :func:`taichi.lang.snode.SNode.dynamic`"""
self._check_not_finalized()
- self._empty = False
- return self._root.dynamic(index, dimension, chunk_size)
+ self.empty = False
+ return self.root.dynamic(index, dimension, chunk_size)
def bitmasked(self, indices: Union[Sequence[_Axis], _Axis],
dimensions: Union[Sequence[int], int]):
"""Same as :func:`taichi.lang.snode.SNode.bitmasked`"""
self._check_not_finalized()
- self._empty = False
- return self._root.bitmasked(indices, dimensions)
+ self.empty = False
+ return self.root.bitmasked(indices, dimensions)
def bit_struct(self, num_bits: int):
"""Same as :func:`taichi.lang.snode.SNode.bit_struct`"""
self._check_not_finalized()
- self._empty = False
- return self._root.bit_struct(num_bits)
+ self.empty = False
+ return self.root.bit_struct(num_bits)
def bit_array(self, indices: Union[Sequence[_Axis], _Axis],
dimensions: Union[Sequence[int], int], num_bits: int):
"""Same as :func:`taichi.lang.snode.SNode.bit_array`"""
self._check_not_finalized()
- self._empty = False
- return self._root.bit_array(indices, dimensions, num_bits)
+ self.empty = False
+ return self.root.bit_array(indices, dimensions, num_bits)
def place(self,
*args: Any,
@@ -135,29 +117,37 @@ def place(self,
shared_exponent: bool = False):
"""Same as :func:`taichi.lang.snode.SNode.place`"""
self._check_not_finalized()
- self._empty = False
- self._root.place(*args, offset=offset, shared_exponent=shared_exponent)
+ self.empty = False
+ self.root.place(*args, offset=offset, shared_exponent=shared_exponent)
def lazy_grad(self):
"""Same as :func:`taichi.lang.snode.SNode.lazy_grad`"""
# TODO: This complicates the implementation. Figure out why we need this
self._check_not_finalized()
- self._empty = False
- self._root.lazy_grad()
+ self.empty = False
+ self.root.lazy_grad()
def finalize(self, raise_warning=True):
"""Constructs the SNodeTree and finalizes this builder.
Args:
raise_warning (bool): Raise warning or not."""
+ return self._finalize(raise_warning, compile_only=False)
+
+ def _finalize_for_aot(self):
+ """Constructs the SNodeTree and compiles the type for AOT purpose."""
+ return self._finalize(raise_warning=False, compile_only=True)
+
+ def _finalize(self, raise_warning, compile_only):
self._check_not_finalized()
- if self._empty and raise_warning:
+ if self.empty and raise_warning:
warning("Finalizing an empty FieldsBuilder!")
- self._finalized = True
+ self.finalized = True
return SNodeTree(
- _ti_core.finalize_snode_tree(_snode_registry, self._ptr,
- impl.get_runtime().prog))
+ _ti_core.finalize_snode_tree(_snode_registry, self.ptr,
+ impl.get_runtime().prog,
+ compile_only))
def _check_not_finalized(self):
- if self._finalized:
- raise InvalidOperationError('FieldsBuilder finalized')
+ if self.finalized:
+ raise TaichiRuntimeError('FieldsBuilder finalized')
diff --git a/python/taichi/snode/snode_tree.py b/python/taichi/_snode/snode_tree.py
similarity index 70%
rename from python/taichi/snode/snode_tree.py
rename to python/taichi/_snode/snode_tree.py
index 72c1dd7ecf400..18887421daf15 100644
--- a/python/taichi/snode/snode_tree.py
+++ b/python/taichi/_snode/snode_tree.py
@@ -3,9 +3,8 @@
# loaded during the import procedure, it's probably still good to delay the
# access to it.
-from taichi.core.util import ti_core as _ti_core
from taichi.lang import impl
-from taichi.lang.exception import InvalidOperationError
+from taichi.lang.exception import TaichiRuntimeError
class SNodeTree:
@@ -15,12 +14,12 @@ def __init__(self, ptr):
def destroy(self):
if self.destroyed:
- raise InvalidOperationError('SNode tree has been destroyed')
+ raise TaichiRuntimeError('SNode tree has been destroyed')
self.ptr.destroy_snode_tree(impl.get_runtime().prog)
self.destroyed = True
@property
def id(self):
if self.destroyed:
- raise InvalidOperationError('SNode tree has been destroyed')
+ raise TaichiRuntimeError('SNode tree has been destroyed')
return self.ptr.id()
diff --git a/python/taichi/_version_check.py b/python/taichi/_version_check.py
new file mode 100644
index 0000000000000..5958aecab2455
--- /dev/null
+++ b/python/taichi/_version_check.py
@@ -0,0 +1,108 @@
+import datetime
+import json
+import os
+import platform
+import threading
+import uuid
+from urllib import request
+
+from taichi._lib import core as _ti_core
+
+
+def check_version(cur_uuid):
+ # Check Taichi version for the user.
+ major = _ti_core.get_version_major()
+ minor = _ti_core.get_version_minor()
+ patch = _ti_core.get_version_patch()
+ version = f'{major}.{minor}.{patch}'
+ payload = {'version': version, 'platform': '', 'python': ''}
+
+ system = platform.system()
+ if system == 'Linux':
+ payload['platform'] = 'manylinux1_x86_64'
+ elif system == 'Windows':
+ payload['platform'] = 'win_amd64'
+ elif system == 'Darwin':
+ if platform.release() < '19.0.0':
+ payload['platform'] = 'macosx_10_14_x86_64'
+ elif platform.machine() == 'x86_64':
+ payload['platform'] = 'macosx_10_15_x86_64'
+ else:
+ payload['platform'] = 'macosx_11_0_arm64'
+
+ python_version = platform.python_version()
+ if python_version.startswith('3.6.'):
+ payload['python'] = 'cp36'
+ elif python_version.startswith('3.7.'):
+ payload['python'] = 'cp37'
+ elif python_version.startswith('3.8.'):
+ payload['python'] = 'cp38'
+ elif python_version.startswith('3.9.'):
+ payload['python'] = 'cp39'
+ elif python_version.startswith('3.10.'):
+ payload['python'] = 'cp310'
+
+ payload['uuid'] = cur_uuid
+ # We do not want request exceptions break users' usage of Taichi.
+ try:
+ payload = json.dumps(payload)
+ payload = payload.encode()
+ req = request.Request('https://metadata.taichi.graphics/check_version',
+ method='POST')
+ req.add_header('Content-Type', 'application/json')
+ with request.urlopen(req, data=payload, timeout=5) as response:
+ response = json.loads(response.read().decode('utf-8'))
+ return response
+ except:
+ return None
+
+
+def write_version_info(response, cur_uuid, version_info_path, cur_date):
+ if response is None:
+ return
+ with open(version_info_path, 'w') as f:
+ f.write((cur_date).strftime('%Y-%m-%d'))
+ f.write('\n')
+ if response['status'] == 1:
+ f.write(response['latest_version'])
+ else:
+ f.write('0.0.0')
+ f.write('\n')
+ f.write(cur_uuid)
+ f.write('\n')
+
+
+def try_check_version():
+ try:
+ os.makedirs(_ti_core.get_repo_dir(), exist_ok=True)
+ version_info_path = os.path.join(_ti_core.get_repo_dir(),
+ 'version_info')
+ cur_date = datetime.date.today()
+ if os.path.exists(version_info_path):
+ with open(version_info_path, 'r') as f:
+ version_info_file = f.readlines()
+ last_time = version_info_file[0].rstrip()
+ cur_uuid = version_info_file[2].rstrip()
+ if cur_date.strftime('%Y-%m-%d') > last_time:
+ response = check_version(cur_uuid)
+ write_version_info(response, cur_uuid, version_info_path,
+ cur_date)
+ else:
+ cur_uuid = str(uuid.uuid4())
+ response = check_version(cur_uuid)
+ write_version_info(response, cur_uuid, version_info_path, cur_date)
+ # Wildcard exception to catch potential file writing errors.
+ except:
+ pass
+
+
+def start_version_check_thread():
+ skip = os.environ.get("TI_SKIP_VERSION_CHECK")
+ if skip != 'ON':
+ # We don't join this thread because we do not wish to block users.
+ check_version_thread = threading.Thread(target=try_check_version,
+ daemon=True)
+ check_version_thread.start()
+
+
+__all__ = []
diff --git a/python/taichi/ad.py b/python/taichi/ad.py
index 04a7ba61e6096..1a896fb3587d8 100644
--- a/python/taichi/ad.py
+++ b/python/taichi/ad.py
@@ -66,7 +66,7 @@ def decorated(*args, **kwargs):
)
if primal.grad is not None:
raise RuntimeError(
- f'Primal function must be a **python** function instead of a taichi kernel. Please wrap the taichi kernel in a @ti.ad.grad_replaced decorated python function instead.'
+ 'Primal function must be a **python** function instead of a taichi kernel. Please wrap the taichi kernel in a @ti.ad.grad_replaced decorated python function instead.'
)
primal.grad = decorated
return decorated
diff --git a/python/taichi/aot/module.py b/python/taichi/aot/module.py
index d96b71a0328ec..696875199033b 100644
--- a/python/taichi/aot/module.py
+++ b/python/taichi/aot/module.py
@@ -1,9 +1,13 @@
from contextlib import contextmanager
+from pathlib import Path, PurePosixPath
from taichi.lang import impl, kernel_impl
+from taichi.lang._ndarray import ScalarNdarray
+from taichi.lang.enums import Layout
from taichi.lang.field import ScalarField
-from taichi.lang.matrix import MatrixField
-from taichi.type.annotations import ArgAnyArray, template
+from taichi.lang.matrix import MatrixField, MatrixNdarray, VectorNdarray
+from taichi.types.annotations import ArgAnyArray, template
+from taichi.types.primitive_types import f32
class KernelTemplate:
@@ -14,11 +18,11 @@ def __init__(self, kernel_fn, aot_module):
@staticmethod
def keygen(v, key_p, fields):
if isinstance(v, (int, float, bool)):
- key_p += '=' + str(v) + '/'
+ key_p += '=' + str(v) + ','
return key_p
for ky, val in fields:
- if (val is v):
- key_p += '=' + ky + '/'
+ if val is v:
+ key_p += '=' + ky + ','
return key_p
raise RuntimeError('Arg type must be of type int/float/boolean' +
'or taichi field. Type ' + str(type(v)) +
@@ -36,8 +40,7 @@ def instantiate(self, **kwargs):
for index, (key, value) in enumerate(kwargs.items()):
template_args[index] = (key, value)
- for i in range(len(kernel.argument_annotations)):
- anno = kernel.argument_annotations[i]
+ for anno in kernel.argument_annotations:
if isinstance(anno, template):
(k, v) = template_args[anno_index]
key_p += k
@@ -74,12 +77,18 @@ class Module:
# for running ``foo`` and ``bar``.
"""
def __init__(self, arch):
+ """Creates a new AOT module instance
+
+ Args:
+ arch: Target backend architecture. This is ignored for now. The AOT
+ backend still uses the one specified in :func:`~taichi.lang.init`.
+ """
self._arch = arch
self._kernels = []
self._fields = {}
- impl.get_runtime().materialize()
- self._aot_builder = impl.get_runtime().prog.make_aot_module_builder(
- arch)
+ rtm = impl.get_runtime()
+ rtm._finalize_root_fb_for_aot()
+ self._aot_builder = rtm.prog.make_aot_module_builder(arch)
def add_field(self, name, field):
"""Add a taichi field to the AOT module.
@@ -88,16 +97,15 @@ def add_field(self, name, field):
name: name of taichi field
field: taichi field
- Example:
- Usage::
-
- a = ti.field(ti.f32, shape=(4,4))
- b = ti.field("something")
+ Example::
- m.add_field(a)
- m.add_field(b)
-
- # Must add in sequence
+ >>> a = ti.field(ti.f32, shape=(4,4))
+ >>> b = ti.field("something")
+ >>>
+ >>> m.add_field(a)
+ >>> m.add_field(b)
+ >>>
+ >>> # Must add in sequence
"""
is_scalar = True
self._fields[name] = field
@@ -109,32 +117,61 @@ def add_field(self, name, field):
column_num = field.n
else:
assert isinstance(field, ScalarField)
- self._aot_builder.add_field(name, is_scalar, field.dtype,
- field.snode.shape, row_num, column_num)
+ self._aot_builder.add_field(name, field.snode.ptr, is_scalar,
+ field.dtype, field.snode.shape, row_num,
+ column_num)
- def add_kernel(self, kernel_fn, name=None):
+ def add_kernel(self, kernel_fn, example_any_arrays=None, name=None):
"""Add a taichi kernel to the AOT module.
Args:
kernel_fn (Function): the function decorated by taichi `kernel`.
+ example_any_arrays (Dict[int, ti.ndarray]): a dict where key is arg_id and key is example any_arr input.
name (str): Name to identify this kernel in the module. If not
provided, uses the built-in ``__name__`` attribute of `kernel_fn`.
- TODO:
- * Support external array
"""
name = name or kernel_fn.__name__
kernel = kernel_fn._primal
assert isinstance(kernel, kernel_impl.Kernel)
injected_args = []
- for i in range(len(kernel.argument_annotations)):
- anno = kernel.argument_annotations[i]
+ num_arr = len([
+ anno for anno in kernel.argument_annotations
+ if isinstance(anno, ArgAnyArray)
+ ])
+ assert example_any_arrays is None or num_arr == len(
+ example_any_arrays
+ ), f'Need {num_arr} example any_arr inputs but got {len(example_any_arrays)}'
+ i = 0
+ for anno in kernel.argument_annotations:
if isinstance(anno, ArgAnyArray):
- raise RuntimeError(
- 'Arg type `ext_arr`/`any_arr` not supported yet')
+ if example_any_arrays:
+ injected_args.append(example_any_arrays[i])
+ else:
+ assert anno.element_shape is not None and anno.field_dim is not None, 'Please either specify element_shape & field_dim in the kernel arg annotation or provide a dict of example ndarrays.'
+ if anno.element_dim == 0:
+ injected_args.append(
+ ScalarNdarray(dtype=f32,
+ shape=(2, ) * anno.field_dim))
+ elif anno.element_dim == 1:
+ injected_args.append(
+ VectorNdarray(anno.element_shape[0],
+ dtype=f32,
+ shape=(2, ) * anno.field_dim,
+ layout=Layout.AOS))
+ elif anno.element_dim == 2:
+ injected_args.append(
+ MatrixNdarray(anno.element_shape[0],
+ anno.element_shape[1],
+ dtype=f32,
+ shape=(2, ) * anno.field_dim,
+ layout=Layout.AOS))
+ else:
+ raise RuntimeError('')
else:
# For primitive types, we can just inject a dummy value.
injected_args.append(0)
+ i = i + 1
kernel.ensure_compiled(*injected_args)
self._aot_builder.add(name, kernel.kernel_cpp)
@@ -148,28 +185,27 @@ def add_kernel_template(self, kernel_fn):
Args:
kernel_fn (Function): the function decorated by taichi `kernel`.
- Example:
- Usage::
-
- @ti.kernel
- def bar_tmpl(a: ti.template()):
- x = a
- # or y = a
- # do something with `x` or `y`
-
- m = ti.aot.Module(arch)
- with m.add_kernel_template(bar_tmpl) as kt:
- kt.instantiate(a=x)
- kt.instantiate(a=y)
-
- @ti.kernel
- def bar_tmpl_multiple_args(a: ti.template(), b: ti.template())
- x = a
- y = b
- # do something with `x` and `y`
-
- with m.add_kernel_template(bar_tmpl) as kt:
- kt.instantiate(a=x, b=y)
+ Example::
+
+ >>> @ti.kernel
+ >>> def bar_tmpl(a: ti.template()):
+ >>> x = a
+ >>> # or y = a
+ >>> # do something with `x` or `y`
+ >>>
+ >>> m = ti.aot.Module(arch)
+ >>> with m.add_kernel_template(bar_tmpl) as kt:
+ >>> kt.instantiate(a=x)
+ >>> kt.instantiate(a=y)
+ >>>
+ >>> @ti.kernel
+ >>> def bar_tmpl_multiple_args(a: ti.template(), b: ti.template())
+ >>> x = a
+ >>> y = b
+ >>> # do something with `x` and `y`
+ >>>
+ >>> with m.add_kernel_template(bar_tmpl) as kt:
+ >>> kt.instantiate(a=x, b=y)
TODO:
* Support external array
@@ -178,4 +214,10 @@ def bar_tmpl_multiple_args(a: ti.template(), b: ti.template())
yield kt
def save(self, filepath, filename):
+ """
+ Args:
+ filepath (str): path to a folder to store aot files.
+ filename (str): filename prefix for stored aot files.
+ """
+ filepath = str(PurePosixPath(Path(filepath)))
self._aot_builder.dump(filepath, filename)
diff --git a/python/taichi/aot/record.py b/python/taichi/aot/record.py
index 0d2882d4efee2..176db9337fc3c 100644
--- a/python/taichi/aot/record.py
+++ b/python/taichi/aot/record.py
@@ -1,6 +1,6 @@
import os
-from taichi.core import ti_core
+from taichi._lib import core as ti_core
def record_action_entry(name, contents):
@@ -52,8 +52,4 @@ def __exit__(self, *args):
__all__ = [
'start_recording',
'stop_recording',
- 'record_action_hint',
- 'record_action_entry',
- 'record_action_config',
- 'RecordKernelGroup',
]
diff --git a/python/taichi/core/__init__.py b/python/taichi/core/__init__.py
deleted file mode 100644
index a19fb81c1e88e..0000000000000
--- a/python/taichi/core/__init__.py
+++ /dev/null
@@ -1,3 +0,0 @@
-from taichi.core.util import *
-
-__all__ = [s for s in dir() if not s.startswith('_')]
diff --git a/examples/algorithm/laplace.py b/python/taichi/examples/algorithm/laplace.py
similarity index 70%
rename from examples/algorithm/laplace.py
rename to python/taichi/examples/algorithm/laplace.py
index fbbb2c16c9c49..694acb30f80b7 100644
--- a/examples/algorithm/laplace.py
+++ b/python/taichi/examples/algorithm/laplace.py
@@ -19,10 +19,15 @@ def laplace():
y[i, j] = 0.0
-for i in range(10):
- x[i, i + 1] = 1.0
+def main():
+ for i in range(10):
+ x[i, i + 1] = 1.0
-laplace()
+ laplace()
-for i in range(10):
- print(y[i, i + 1])
+ for i in range(10):
+ print(y[i, i + 1])
+
+
+if __name__ == '__main__':
+ main()
diff --git a/examples/algorithm/marching_squares.py b/python/taichi/examples/algorithm/marching_squares.py
similarity index 100%
rename from examples/algorithm/marching_squares.py
rename to python/taichi/examples/algorithm/marching_squares.py
diff --git a/examples/algorithm/mciso_advanced.py b/python/taichi/examples/algorithm/mciso_advanced.py
similarity index 99%
rename from examples/algorithm/mciso_advanced.py
rename to python/taichi/examples/algorithm/mciso_advanced.py
index 8cd19a41b9bde..7ea9fa4ea305d 100644
--- a/examples/algorithm/mciso_advanced.py
+++ b/python/taichi/examples/algorithm/mciso_advanced.py
@@ -294,13 +294,6 @@ def __init__(self, N=64, dim=3, blk_size=None):
self.use_sparse = blk_size is not None
self.blk_size = blk_size
- et = [self.et2, self.et3][dim - 2]
- self.et = ti.Vector.field(dim, int, et.shape[:2])
-
- @ti.materialize_callback
- def init_et():
- self.et.from_numpy(et)
-
self.m = ti.field(float) # field to sample
self.g = ti.Vector.field(self.dim, float) # normalized gradient
indices = [ti.ij, ti.ijk][dim - 2]
@@ -316,6 +309,10 @@ def init_et():
dim, float,
(self.N**self.dim, self.dim)) # result buffer, TODO: optimize this
+ et = [self.et2, self.et3][dim - 2]
+ self.et = ti.Vector.field(dim, int, et.shape[:2])
+ self.et.from_numpy(et)
+
@ti.kernel
def compute_grad(self):
for I in ti.grouped(self.g):
diff --git a/examples/algorithm/mgpcg.py b/python/taichi/examples/algorithm/mgpcg.py
similarity index 100%
rename from examples/algorithm/mgpcg.py
rename to python/taichi/examples/algorithm/mgpcg.py
diff --git a/examples/algorithm/mgpcg_advanced.py b/python/taichi/examples/algorithm/mgpcg_advanced.py
similarity index 100%
rename from examples/algorithm/mgpcg_advanced.py
rename to python/taichi/examples/algorithm/mgpcg_advanced.py
diff --git a/python/taichi/examples/algorithm/print_offset.py b/python/taichi/examples/algorithm/print_offset.py
new file mode 100644
index 0000000000000..d7fa00d262194
--- /dev/null
+++ b/python/taichi/examples/algorithm/print_offset.py
@@ -0,0 +1,49 @@
+from taichi.lang import impl
+
+import taichi as ti
+
+ti.init(arch=ti.cpu, print_ir=True)
+
+n = 4
+m = 8
+
+a = ti.field(dtype=ti.i32)
+ti.root.dense(ti.ij, (1, 2)).dense(ti.ij, 2).dense(ti.ij, 2).place(a)
+
+
+@ti.kernel
+def fill():
+ for i, j in a:
+ base = ti.get_addr(a.snode, [0, 0])
+ a[i, j] = int(ti.get_addr(a.snode, [i, j]) - base) // 4
+
+
+def main():
+ fill()
+ print(a.to_numpy())
+
+ impl.get_runtime().prog.visualize_layout('layout.pdf')
+
+ gui = ti.GUI('layout', res=(256, 512), background_color=0xFFFFFF)
+
+ while True:
+ for i in range(1, m):
+ gui.line(begin=(0, i / m),
+ end=(1, i / m),
+ radius=2,
+ color=0x000000)
+ for i in range(1, n):
+ gui.line(begin=(i / n, 0),
+ end=(i / n, 1),
+ radius=2,
+ color=0x000000)
+ for i in range(n):
+ for j in range(m):
+ gui.text(f'{a[i, j]}', ((i + 0.3) / n, (j + 0.75) / m),
+ font_size=30,
+ color=0x0)
+ gui.show()
+
+
+if __name__ == '__main__':
+ main()
diff --git a/python/taichi/examples/autodiff/minimization.py b/python/taichi/examples/autodiff/minimization.py
new file mode 100644
index 0000000000000..1a5f2708e05e5
--- /dev/null
+++ b/python/taichi/examples/autodiff/minimization.py
@@ -0,0 +1,44 @@
+import random
+
+import taichi as ti
+
+ti.init(arch=ti.cpu)
+
+n = 8
+x = ti.field(dtype=ti.f32, shape=n, needs_grad=True)
+y = ti.field(dtype=ti.f32, shape=n)
+L = ti.field(dtype=ti.f32, shape=(), needs_grad=True)
+
+
+@ti.kernel
+def reduce():
+ for i in range(n):
+ L[None] += 0.5 * (x[i] - y[i])**2
+
+
+@ti.kernel
+def gradient_descent():
+ for i in x:
+ x[i] -= x.grad[i] * 0.1
+
+
+def main():
+ # Initialize vectors
+ for i in range(n):
+ x[i] = random.random()
+ y[i] = random.random()
+
+ # Optimize with 100 gradient descent iterations
+ for k in range(100):
+ with ti.Tape(loss=L):
+ reduce()
+ print('Loss =', L[None])
+ gradient_descent()
+
+ for i in range(n):
+ # Now you should approximately have x[i] == y[i]
+ print(x[i], y[i])
+
+
+if __name__ == '__main__':
+ main()
diff --git a/python/taichi/examples/autodiff/regression.py b/python/taichi/examples/autodiff/regression.py
new file mode 100644
index 0000000000000..c70273e973903
--- /dev/null
+++ b/python/taichi/examples/autodiff/regression.py
@@ -0,0 +1,104 @@
+import random
+
+import matplotlib.pyplot as plt
+import numpy as np
+
+import taichi as ti
+
+ti.init(arch=ti.cpu)
+
+number_coeffs = 4
+learning_rate = 1e-4
+
+N = 32
+x, y = ti.field(ti.f32, shape=N, needs_grad=True), ti.field(ti.f32,
+ shape=N,
+ needs_grad=True)
+coeffs = ti.field(ti.f32, shape=number_coeffs, needs_grad=True)
+loss = ti.field(ti.f32, shape=(), needs_grad=True)
+
+
+@ti.kernel
+def regress():
+ for i in x:
+ v = x[i]
+ est = 0.0
+ for j in ti.static(range(number_coeffs)):
+ est += coeffs[j] * (v**j)
+ loss[None] += 0.5 * (y[i] - est)**2
+
+
+@ti.kernel
+def update():
+ for i in ti.static(range(number_coeffs)):
+ coeffs[i] -= learning_rate * coeffs.grad[i]
+
+
+xs = []
+ys = []
+
+
+def initialize():
+ for i in range(N):
+ v = random.random() * 5 - 2.5
+ xs.append(v)
+ x[i] = v
+ y[i] = (v - 1) * (v - 2) * (v + 2) + random.random() - 0.5
+
+ regress()
+
+ print('y')
+ for i in range(N):
+ y.grad[i] = 1
+ ys.append(y[i])
+ print()
+
+
+def regress_raw():
+ use_tape = True
+
+ for i in range(1000):
+ if use_tape:
+ with ti.Tape(loss=loss):
+ regress()
+ else:
+ ti.clear_all_gradients()
+ loss[None] = 0
+ loss.grad[None] = 1
+ regress()
+ regress.grad()
+ print('Loss =', loss[None])
+ update()
+ for i in range(number_coeffs):
+ print(coeffs[i], end=', ')
+ print()
+
+
+def draw():
+ curve_xs = np.arange(-2.5, 2.5, 0.01)
+ curve_ys = curve_xs * 0
+ for i in range(number_coeffs):
+ curve_ys += coeffs[i] * np.power(curve_xs, i)
+
+ plt.title(
+ 'Nonlinear Regression with Gradient Descent (3rd order polynomial)')
+ ax = plt.gca()
+ ax.scatter(xs, ys, label='data', color='r')
+ ax.plot(curve_xs, curve_ys, label='fitted')
+ ax.legend()
+ ax.grid(True)
+ ax.spines['left'].set_position('zero')
+ ax.spines['right'].set_color('none')
+ ax.spines['bottom'].set_position('zero')
+ ax.spines['top'].set_color('none')
+ plt.show()
+
+
+def main():
+ initialize()
+ regress_raw()
+ draw()
+
+
+if __name__ == '__main__':
+ main()
diff --git a/python/taichi/examples/autodiff/simple_derivative.py b/python/taichi/examples/autodiff/simple_derivative.py
new file mode 100644
index 0000000000000..3a049e71e2129
--- /dev/null
+++ b/python/taichi/examples/autodiff/simple_derivative.py
@@ -0,0 +1,69 @@
+import matplotlib.pyplot as plt
+
+import taichi as ti
+
+ti.init(arch=ti.cpu)
+
+N = 2048
+x, y = ti.field(ti.f32), ti.field(ti.f32)
+
+ti.root.dense(ti.i, N).place(x, x.grad, y, y.grad)
+
+
+@ti.kernel
+def poly():
+ for i in x:
+ v = x[i]
+ ret = 0.0
+ guard = 0.2
+ if v < -guard or v > guard:
+ ret = 4 / ti.max(v, 0.1)
+ else:
+ ret = 0
+ y[i] = ret
+
+
+xs = []
+ys = []
+grad_xs = []
+
+
+def initialize():
+ for i in range(N):
+ v = ((i + 0.5) / N) * 2 - 1
+ xs.append(v)
+ x[i] = v
+
+ poly()
+
+ for i in range(N):
+ y.grad[i] = 1
+ ys.append(y[i])
+
+ poly.grad()
+ print('grad_x')
+ for i in range(N):
+ grad_xs.append(x.grad[i])
+
+
+def draw():
+ plt.title('Auto Diff')
+ ax = plt.gca()
+ ax.plot(xs, ys, label='f(x)')
+ ax.plot(xs, grad_xs, label='f\'(x)')
+ ax.legend()
+ ax.grid(True)
+ ax.spines['left'].set_position('zero')
+ ax.spines['right'].set_color('none')
+ ax.spines['bottom'].set_position('zero')
+ ax.spines['top'].set_color('none')
+ plt.show()
+
+
+def main():
+ initialize()
+ draw()
+
+
+if __name__ == '__main__':
+ main()
diff --git a/examples/features/gui/fullscreen.py b/python/taichi/examples/features/gui/fullscreen.py
similarity index 100%
rename from examples/features/gui/fullscreen.py
rename to python/taichi/examples/features/gui/fullscreen.py
diff --git a/examples/features/gui/gui_image_io.py b/python/taichi/examples/features/gui/gui_image_io.py
similarity index 100%
rename from examples/features/gui/gui_image_io.py
rename to python/taichi/examples/features/gui/gui_image_io.py
diff --git a/examples/features/gui/gui_widgets.py b/python/taichi/examples/features/gui/gui_widgets.py
similarity index 100%
rename from examples/features/gui/gui_widgets.py
rename to python/taichi/examples/features/gui/gui_widgets.py
diff --git a/examples/features/gui/keyboard.py b/python/taichi/examples/features/gui/keyboard.py
similarity index 100%
rename from examples/features/gui/keyboard.py
rename to python/taichi/examples/features/gui/keyboard.py
diff --git a/examples/features/io/export_mesh.py b/python/taichi/examples/features/io/export_mesh.py
similarity index 100%
rename from examples/features/io/export_mesh.py
rename to python/taichi/examples/features/io/export_mesh.py
diff --git a/examples/features/io/export_ply.py b/python/taichi/examples/features/io/export_ply.py
similarity index 100%
rename from examples/features/io/export_ply.py
rename to python/taichi/examples/features/io/export_ply.py
diff --git a/examples/features/io/export_videos.py b/python/taichi/examples/features/io/export_videos.py
similarity index 100%
rename from examples/features/io/export_videos.py
rename to python/taichi/examples/features/io/export_videos.py
diff --git a/examples/features/sparse/explicit_activation.py b/python/taichi/examples/features/sparse/explicit_activation.py
similarity index 100%
rename from examples/features/sparse/explicit_activation.py
rename to python/taichi/examples/features/sparse/explicit_activation.py
diff --git a/examples/features/sparse/taichi_bitmasked.py b/python/taichi/examples/features/sparse/taichi_bitmasked.py
similarity index 100%
rename from examples/features/sparse/taichi_bitmasked.py
rename to python/taichi/examples/features/sparse/taichi_bitmasked.py
diff --git a/examples/features/sparse/taichi_dynamic.py b/python/taichi/examples/features/sparse/taichi_dynamic.py
similarity index 100%
rename from examples/features/sparse/taichi_dynamic.py
rename to python/taichi/examples/features/sparse/taichi_dynamic.py
diff --git a/examples/features/sparse/taichi_sparse.py b/python/taichi/examples/features/sparse/taichi_sparse.py
similarity index 93%
rename from examples/features/sparse/taichi_sparse.py
rename to python/taichi/examples/features/sparse/taichi_sparse.py
index 51039a55704b3..ebd8ea9091a6c 100644
--- a/examples/features/sparse/taichi_sparse.py
+++ b/python/taichi/examples/features/sparse/taichi_sparse.py
@@ -1,3 +1,5 @@
+from taichi.examples.patterns import taichi_logo
+
import taichi as ti
ti.init(arch=ti.cuda)
@@ -19,7 +21,7 @@ def activate(t: ti.f32):
p = ti.Vector([i, j]) / n
p = ti.Matrix.rotation2d(ti.sin(t)) @ (p - 0.5) + 0.5
- if ti.taichi_logo(p) == 0:
+ if taichi_logo(p) == 0:
x[i, j] = 1
diff --git a/examples/features/sparse/tutorial.py b/python/taichi/examples/features/sparse/tutorial.py
similarity index 100%
rename from examples/features/sparse/tutorial.py
rename to python/taichi/examples/features/sparse/tutorial.py
diff --git a/examples/ggui_examples/fem128_ggui.py b/python/taichi/examples/ggui_examples/fem128_ggui.py
similarity index 98%
rename from examples/ggui_examples/fem128_ggui.py
rename to python/taichi/examples/ggui_examples/fem128_ggui.py
index 895bd180648c5..187befdcfa205 100644
--- a/examples/ggui_examples/fem128_ggui.py
+++ b/python/taichi/examples/ggui_examples/fem128_ggui.py
@@ -1,6 +1,7 @@
import taichi as ti
-ti.init(arch=ti.gpu)
+arch = ti.vulkan if ti._lib.core.with_vulkan() else ti.cuda
+ti.init(arch=arch)
N = 12
dt = 5e-5
diff --git a/examples/ggui_examples/fractal3d_ggui.py b/python/taichi/examples/ggui_examples/fractal3d_ggui.py
similarity index 98%
rename from examples/ggui_examples/fractal3d_ggui.py
rename to python/taichi/examples/ggui_examples/fractal3d_ggui.py
index 5c283e7e4efb8..deedb922703f4 100644
--- a/examples/ggui_examples/fractal3d_ggui.py
+++ b/python/taichi/examples/ggui_examples/fractal3d_ggui.py
@@ -1,6 +1,7 @@
import taichi as ti
-ti.init(ti.cuda)
+arch = ti.vulkan if ti._lib.core.with_vulkan() else ti.cuda
+ti.init(arch=arch)
@ti.func
diff --git a/examples/ggui_examples/mass_spring_3d_ggui.py b/python/taichi/examples/ggui_examples/mass_spring_3d_ggui.py
similarity index 97%
rename from examples/ggui_examples/mass_spring_3d_ggui.py
rename to python/taichi/examples/ggui_examples/mass_spring_3d_ggui.py
index f9116d2ee0f58..2cfae527a3333 100644
--- a/examples/ggui_examples/mass_spring_3d_ggui.py
+++ b/python/taichi/examples/ggui_examples/mass_spring_3d_ggui.py
@@ -1,6 +1,7 @@
import taichi as ti
-ti.init(arch=ti.cuda) # Alternatively, ti.init(arch=ti.cpu)
+arch = ti.vulkan if ti._lib.core.with_vulkan() else ti.cuda
+ti.init(arch=arch)
N = 128
cell_size = 1.0 / N
diff --git a/examples/ggui_examples/mass_spring_game_ggui.py b/python/taichi/examples/ggui_examples/mass_spring_game_ggui.py
similarity index 98%
rename from examples/ggui_examples/mass_spring_game_ggui.py
rename to python/taichi/examples/ggui_examples/mass_spring_game_ggui.py
index 59314ba062212..d3a392ad73a9e 100644
--- a/examples/ggui_examples/mass_spring_game_ggui.py
+++ b/python/taichi/examples/ggui_examples/mass_spring_game_ggui.py
@@ -1,6 +1,7 @@
import taichi as ti
-ti.init(arch=ti.cuda)
+arch = ti.vulkan if ti._lib.core.with_vulkan() else ti.cuda
+ti.init(arch=arch)
spring_Y = ti.field(dtype=ti.f32, shape=()) # Young's modulus
paused = ti.field(dtype=ti.i32, shape=())
diff --git a/examples/ggui_examples/mpm128_ggui.py b/python/taichi/examples/ggui_examples/mpm128_ggui.py
similarity index 98%
rename from examples/ggui_examples/mpm128_ggui.py
rename to python/taichi/examples/ggui_examples/mpm128_ggui.py
index 279af236b6960..7c1eb5a6345f8 100644
--- a/examples/ggui_examples/mpm128_ggui.py
+++ b/python/taichi/examples/ggui_examples/mpm128_ggui.py
@@ -2,7 +2,8 @@
import taichi as ti
-ti.init(arch=ti.gpu) # Try to run on GPU
+arch = ti.vulkan if ti._lib.core.with_vulkan() else ti.cuda
+ti.init(arch=arch)
quality = 1 # Use a larger value for higher-res simulations
n_particles, n_grid = 9000 * quality**2, 128 * quality
diff --git a/examples/ggui_examples/mpm3d_ggui.py b/python/taichi/examples/ggui_examples/mpm3d_ggui.py
similarity index 95%
rename from examples/ggui_examples/mpm3d_ggui.py
rename to python/taichi/examples/ggui_examples/mpm3d_ggui.py
index 65639580a5366..0e48dfcefba14 100644
--- a/examples/ggui_examples/mpm3d_ggui.py
+++ b/python/taichi/examples/ggui_examples/mpm3d_ggui.py
@@ -2,7 +2,8 @@
import taichi as ti
-ti.init(ti.cuda)
+arch = ti.vulkan if ti._lib.core.with_vulkan() else ti.cuda
+ti.init(arch=arch)
#dim, n_grid, steps, dt = 2, 128, 20, 2e-4
#dim, n_grid, steps, dt = 2, 256, 32, 1e-4
@@ -35,8 +36,8 @@
shape=n_particles) # deformation gradient
Jp = ti.field(float, n_particles)
-colors = ti.Vector.field(3, float, n_particles)
-colors_random = ti.Vector.field(3, float, n_particles)
+colors = ti.Vector.field(4, float, n_particles)
+colors_random = ti.Vector.field(4, float, n_particles)
materials = ti.field(int, n_particles)
grid_v = ti.Vector.field(dim, float, (n_grid, ) * dim)
grid_m = ti.field(float, (n_grid, ) * dim)
@@ -110,8 +111,8 @@ def substep(g_x: float, g_y: float, g_z: float):
if grid_m[I] > 0:
grid_v[I] /= grid_m[I]
grid_v[I] += dt * ti.Vector([g_x, g_y, g_z])
- cond = I < bound and grid_v[I] < 0 or I > n_grid - bound and grid_v[
- I] > 0
+ cond = (I < bound) & (grid_v[I] < 0) | \
+ (I > n_grid - bound) & (grid_v[I] > 0)
grid_v[I] = 0 if cond else grid_v[I]
ti.block_dim(n_grid)
for p in x:
@@ -155,7 +156,9 @@ def init_cube_vol(first_par: int, last_par: int, x_begin: float,
F[i] = ti.Matrix([[1, 0, 0], [0, 1, 0], [0, 0, 1]])
v[i] = ti.Vector([0.0, 0.0, 0.0])
materials[i] = material
- colors_random[i] = ti.Vector([ti.random(), ti.random(), ti.random()])
+ colors_random[i] = ti.Vector(
+ [ti.random(), ti.random(),
+ ti.random(), ti.random()])
used[i] = 1
@@ -199,7 +202,7 @@ def set_color_by_material(material_colors: ti.ext_arr()):
mat = materials[i]
colors[i] = ti.Vector([
material_colors[mat, 0], material_colors[mat, 1],
- material_colors[mat, 2]
+ material_colors[mat, 2], 1.0
])
@@ -292,7 +295,7 @@ def show_options():
"snow color", material_colors[SNOW])
material_colors[JELLY] = window.GUI.color_edit_3(
"jelly color", material_colors[JELLY])
- set_color_by_material(np.array(material_colors))
+ set_color_by_material(np.array(material_colors, dtype=np.float32))
particles_radius = window.GUI.slider_float("particles radius ",
particles_radius, 0, 0.1)
if window.GUI.button("restart"):
diff --git a/examples/ggui_examples/stable_fluid_ggui.py b/python/taichi/examples/ggui_examples/stable_fluid_ggui.py
similarity index 98%
rename from examples/ggui_examples/stable_fluid_ggui.py
rename to python/taichi/examples/ggui_examples/stable_fluid_ggui.py
index caf5fc74a6c6e..16536eb2dc9a6 100644
--- a/examples/ggui_examples/stable_fluid_ggui.py
+++ b/python/taichi/examples/ggui_examples/stable_fluid_ggui.py
@@ -21,7 +21,8 @@
debug = False
paused = False
-ti.init(arch=ti.cuda)
+arch = ti.vulkan if ti._lib.core.with_vulkan() else ti.cuda
+ti.init(arch=arch)
_velocities = ti.Vector.field(2, float, shape=(res, res))
_new_velocities = ti.Vector.field(2, float, shape=(res, res))
diff --git a/examples/minimal.py b/python/taichi/examples/minimal.py
similarity index 100%
rename from examples/minimal.py
rename to python/taichi/examples/minimal.py
diff --git a/python/taichi/tools/patterns.py b/python/taichi/examples/patterns.py
similarity index 100%
rename from python/taichi/tools/patterns.py
rename to python/taichi/examples/patterns.py
diff --git a/examples/rendering/cornell_box.py b/python/taichi/examples/rendering/cornell_box.py
similarity index 94%
rename from examples/rendering/cornell_box.py
rename to python/taichi/examples/rendering/cornell_box.py
index 9560b09eebff5..bc3fff73805b3 100644
--- a/examples/rendering/cornell_box.py
+++ b/python/taichi/examples/rendering/cornell_box.py
@@ -485,33 +485,29 @@ def render():
@ti.kernel
-def tonemap(accumulated: ti.f32) -> ti.f32:
- sum = 0.0
- sum_sq = 0.0
- for i, j in color_buffer:
- luma = color_buffer[i, j][0] * 0.2126 + color_buffer[
- i, j][1] * 0.7152 + color_buffer[i, j][2] * 0.0722
- sum += luma
- sum_sq += ti.pow(luma / accumulated, 2.0)
- mean = sum / (res[0] * res[1])
- var = sum_sq / (res[0] * res[1]) - ti.pow(mean / accumulated, 2.0)
+def tonemap(accumulated: ti.f32):
for i, j in tonemapped_buffer:
- tonemapped_buffer[i, j] = ti.sqrt(color_buffer[i, j] / mean * 0.6)
- return var
-
-
-gui = ti.GUI('Cornell Box', res, fast_gui=True)
-gui.fps_limit = 300
-last_t = time.time()
-i = 0
-while gui.running:
- render()
- interval = 10
- if i % interval == 0:
- var = tonemap(i)
- print("{:.2f} samples/s ({} iters, var={})".format(
- interval / (time.time() - last_t), i, var))
- last_t = time.time()
- gui.set_image(tonemapped_buffer)
- gui.show()
- i += 1
+ tonemapped_buffer[i, j] = ti.sqrt(color_buffer[i, j] / accumulated *
+ 100.0)
+
+
+def main():
+ gui = ti.GUI('Cornell Box', res, fast_gui=True)
+ gui.fps_limit = 300
+ last_t = time.time()
+ i = 0
+ while gui.running:
+ render()
+ interval = 10
+ if i % interval == 0:
+ tonemap(i)
+ print("{:.2f} samples/s ({} iters)".format(
+ interval / (time.time() - last_t), i))
+ last_t = time.time()
+ gui.set_image(tonemapped_buffer)
+ gui.show()
+ i += 1
+
+
+if __name__ == '__main__':
+ main()
diff --git a/examples/rendering/rasterizer.py b/python/taichi/examples/rendering/rasterizer.py
similarity index 100%
rename from examples/rendering/rasterizer.py
rename to python/taichi/examples/rendering/rasterizer.py
diff --git a/examples/rendering/sdf_renderer.py b/python/taichi/examples/rendering/sdf_renderer.py
similarity index 100%
rename from examples/rendering/sdf_renderer.py
rename to python/taichi/examples/rendering/sdf_renderer.py
diff --git a/examples/rendering/simple_uv.py b/python/taichi/examples/rendering/simple_uv.py
similarity index 100%
rename from examples/rendering/simple_uv.py
rename to python/taichi/examples/rendering/simple_uv.py
diff --git a/python/taichi/examples/rendering/taichi_logo.py b/python/taichi/examples/rendering/taichi_logo.py
new file mode 100644
index 0000000000000..40643f1f4e657
--- /dev/null
+++ b/python/taichi/examples/rendering/taichi_logo.py
@@ -0,0 +1,29 @@
+from taichi.examples.patterns import taichi_logo
+
+import taichi as ti
+
+ti.init()
+
+n = 512
+x = ti.field(dtype=ti.f32, shape=(n, n))
+
+
+@ti.kernel
+def paint():
+ for i, j in ti.ndrange(n * 4, n * 4):
+ # 4x4 super sampling:
+ ret = taichi_logo(ti.Vector([i, j]) / (n * 4))
+ x[i // 4, j // 4] += ret / 16
+
+
+def main():
+ paint()
+
+ gui = ti.GUI('Logo', (n, n))
+ while gui.running:
+ gui.set_image(x)
+ gui.show()
+
+
+if __name__ == '__main__':
+ main()
diff --git a/examples/simulation/ad_gravity.py b/python/taichi/examples/simulation/ad_gravity.py
similarity index 82%
rename from examples/simulation/ad_gravity.py
rename to python/taichi/examples/simulation/ad_gravity.py
index 3c3d19089b449..811346dd36756 100644
--- a/examples/simulation/ad_gravity.py
+++ b/python/taichi/examples/simulation/ad_gravity.py
@@ -43,10 +43,15 @@ def init():
x[i] = [ti.random(), ti.random()]
-init()
-gui = ti.GUI('Autodiff gravity')
-while gui.running:
- for i in range(50):
- substep()
- gui.circles(x.to_numpy(), radius=3)
- gui.show()
+def main():
+ init()
+ gui = ti.GUI('Autodiff gravity')
+ while gui.running:
+ for i in range(50):
+ substep()
+ gui.circles(x.to_numpy(), radius=3)
+ gui.show()
+
+
+if __name__ == "__main__":
+ main()
diff --git a/examples/simulation/comet.py b/python/taichi/examples/simulation/comet.py
similarity index 100%
rename from examples/simulation/comet.py
rename to python/taichi/examples/simulation/comet.py
diff --git a/examples/simulation/euler.py b/python/taichi/examples/simulation/euler.py
similarity index 99%
rename from examples/simulation/euler.py
rename to python/taichi/examples/simulation/euler.py
index 994ecf74522ae..f59b8da013c2f 100644
--- a/examples/simulation/euler.py
+++ b/python/taichi/examples/simulation/euler.py
@@ -215,12 +215,12 @@ def HLLC_flux(qL, qR, n):
HLLC = ti.Vector([0.0, 0.0, 0.0, 0.0])
if (0 <= sL):
HLLC = fL
- elif (sL <= 0) and (0 <= sM):
+ elif (0 <= sM):
qsL = rL * (sL-vnL)/(sL-sM) \
* ti.Vector([1.0, sM*nx-vtL*ny,sM*ny+vtL*nx, \
qL[3]/rL + (sM-vnL)*(sM+pL/(rL*(sL-vnL)))])
HLLC = fL + sL * (qsL - qL)
- elif (sM <= 0) and (0 <= sR):
+ elif (0 <= sR):
qsR = rR * (sR-vnR)/(sR-sM) \
* ti.Vector([1.0, sM*nx-vtR*ny,sM*ny+vtR*nx, \
qR[3]/rR + (sM-vnR)*(sM+pR/(rR*(sR-vnR)))])
@@ -283,9 +283,7 @@ def sign(a):
return sgn
-ti.func
-
-
+@ti.func
def cosh(a):
return (ti.exp(a) + ti.exp(-a)) / 2.0
diff --git a/examples/simulation/fem128.py b/python/taichi/examples/simulation/fem128.py
similarity index 98%
rename from examples/simulation/fem128.py
rename to python/taichi/examples/simulation/fem128.py
index 07a2d3b3a56a1..09062dff1cbac 100644
--- a/examples/simulation/fem128.py
+++ b/python/taichi/examples/simulation/fem128.py
@@ -60,7 +60,7 @@ def advance():
if disp2 <= ball_radius**2:
NoV = vel[i].dot(disp)
if NoV < 0: vel[i] -= NoV * disp / disp2
- cond = pos[i] < 0 and vel[i] < 0 or pos[i] > 1 and vel[i] > 0
+ cond = (pos[i] < 0) & (vel[i] < 0) | (pos[i] > 1) & (vel[i] > 0)
# rect boundary condition:
for j in ti.static(range(pos.n)):
if cond[j]: vel[i][j] = 0
diff --git a/examples/simulation/fem99.py b/python/taichi/examples/simulation/fem99.py
similarity index 97%
rename from examples/simulation/fem99.py
rename to python/taichi/examples/simulation/fem99.py
index 304d816c606cb..4683096178288 100644
--- a/examples/simulation/fem99.py
+++ b/python/taichi/examples/simulation/fem99.py
@@ -56,7 +56,7 @@ def advance():
NoV = vel[i].dot(disp)
if NoV < 0: vel[i] -= NoV * disp / disp2
# rect boundary condition:
- cond = pos[i] < 0 and vel[i] < 0 or pos[i] > 1 and vel[i] > 0
+ cond = (pos[i] < 0) & (vel[i] < 0) | (pos[i] > 1) & (vel[i] > 0)
for j in ti.static(range(pos.n)):
if cond[j]: vel[i][j] = 0
pos[i] += dt * vel[i]
diff --git a/examples/simulation/fractal.py b/python/taichi/examples/simulation/fractal.py
similarity index 100%
rename from examples/simulation/fractal.py
rename to python/taichi/examples/simulation/fractal.py
diff --git a/examples/simulation/game_of_life.py b/python/taichi/examples/simulation/game_of_life.py
similarity index 56%
rename from examples/simulation/game_of_life.py
rename to python/taichi/examples/simulation/game_of_life.py
index e5481b072e429..4d38ac319bf4a 100644
--- a/examples/simulation/game_of_life.py
+++ b/python/taichi/examples/simulation/game_of_life.py
@@ -58,31 +58,37 @@ def init():
alive[i, j] = 0
-gui = ti.GUI('Game of Life', (img_size, img_size))
-gui.fps_limit = 15
-
-print('[Hint] Press `r` to reset')
-print('[Hint] Press SPACE to pause')
-print('[Hint] Click LMB, RMB and drag to add alive / dead cells')
-
-init()
-paused = False
-while gui.running:
- for e in gui.get_events(gui.PRESS, gui.MOTION):
- if e.key == gui.ESCAPE:
- gui.running = False
- elif e.key == gui.SPACE:
- paused = not paused
- elif e.key == 'r':
- alive.fill(0)
-
- if gui.is_pressed(gui.LMB, gui.RMB):
- mx, my = gui.get_cursor_pos()
- alive[int(mx * n), int(my * n)] = gui.is_pressed(gui.LMB)
- paused = True
-
- if not paused:
- run()
-
- gui.set_image(ti.imresize(alive, img_size).astype(np.uint8) * 255)
- gui.show()
+def main():
+ gui = ti.GUI('Game of Life', (img_size, img_size))
+ gui.fps_limit = 15
+
+ print('[Hint] Press `r` to reset')
+ print('[Hint] Press SPACE to pause')
+ print('[Hint] Click LMB, RMB and drag to add alive / dead cells')
+
+ init()
+ paused = False
+ while gui.running:
+ for e in gui.get_events(gui.PRESS, gui.MOTION):
+ if e.key == gui.ESCAPE:
+ gui.running = False
+ elif e.key == gui.SPACE:
+ paused = not paused
+ elif e.key == 'r':
+ alive.fill(0)
+
+ if gui.is_pressed(gui.LMB, gui.RMB):
+ mx, my = gui.get_cursor_pos()
+ alive[int(mx * n), int(my * n)] = gui.is_pressed(gui.LMB)
+ paused = True
+
+ if not paused:
+ run()
+
+ gui.set_image(
+ ti.tools.imresize(alive, img_size).astype(np.uint8) * 255)
+ gui.show()
+
+
+if __name__ == '__main__':
+ main()
diff --git a/python/taichi/examples/simulation/implicit_fem.py b/python/taichi/examples/simulation/implicit_fem.py
new file mode 100644
index 0000000000000..ad5c5c6c77637
--- /dev/null
+++ b/python/taichi/examples/simulation/implicit_fem.py
@@ -0,0 +1,353 @@
+import argparse
+
+import numpy as np
+from taichi._lib import core as _ti_core
+
+import taichi as ti
+
+parser = argparse.ArgumentParser()
+parser.add_argument('--exp',
+ choices=['implicit', 'explicit'],
+ default='implicit')
+parser.add_argument('--dim', type=int, default=3)
+parser.add_argument('--gui', choices=['auto', 'ggui', 'cpu'], default='auto')
+parser.add_argument('place_holder', nargs='*')
+args = parser.parse_args()
+
+ti.init(arch=ti.gpu, dynamic_index=True)
+
+if args.gui == 'auto':
+ if _ti_core.GGUI_AVAILABLE:
+ args.gui = 'ggui'
+ else:
+ args.gui = 'cpu'
+
+E, nu = 5e4, 0.0
+mu, la = E / (2 * (1 + nu)), E * nu / ((1 + nu) * (1 - 2 * nu)) # lambda = 0
+density = 1000.0
+dt = 2e-4
+
+if args.exp == 'implicit':
+ dt = 1e-2
+
+n_cube = np.array([5] * 3)
+n_verts = np.product(n_cube)
+n_cells = 5 * np.product(n_cube - 1)
+dx = 1 / (n_cube.max() - 1)
+
+vertices = ti.Vector.field(4, dtype=ti.i32, shape=n_cells)
+
+x = ti.Vector.field(args.dim, dtype=ti.f32, shape=n_verts)
+ox = ti.Vector.field(args.dim, dtype=ti.f32, shape=n_verts)
+v = ti.Vector.field(args.dim, dtype=ti.f32, shape=n_verts)
+f = ti.Vector.field(args.dim, dtype=ti.f32, shape=n_verts)
+mul_ans = ti.Vector.field(args.dim, dtype=ti.f32, shape=n_verts)
+m = ti.field(dtype=ti.f32, shape=n_verts)
+
+n_cells = (n_cube - 1).prod() * 5
+B = ti.Matrix.field(args.dim, args.dim, dtype=ti.f32, shape=n_cells)
+W = ti.field(dtype=ti.f32, shape=n_cells)
+
+
+@ti.func
+def i2p(I):
+ return (I.x * n_cube[1] + I.y) * n_cube[2] + I.z
+
+
+@ti.func
+def set_element(e, I, verts):
+ for i in ti.static(range(args.dim + 1)):
+ vertices[e][i] = i2p(I + (([verts[i] >> k for k in range(3)] ^ I) & 1))
+
+
+@ti.kernel
+def get_vertices():
+ '''
+ This kernel partitions the cube into tetrahedrons.
+ Each unit cube is divided into 5 tetrahedrons.
+ '''
+ for I in ti.grouped(ti.ndrange(*(n_cube - 1))):
+ e = ((I.x * (n_cube[1] - 1) + I.y) * (n_cube[2] - 1) + I.z) * 5
+ for i, j in ti.static(enumerate([0, 3, 5, 6])):
+ set_element(e + i, I, (j, j ^ 1, j ^ 2, j ^ 4))
+ set_element(e + 4, I, (1, 2, 4, 7))
+ for I in ti.grouped(ti.ndrange(*(n_cube))):
+ ox[i2p(I)] = I * dx
+
+
+@ti.func
+def Ds(verts):
+ return ti.Matrix.cols([x[verts[i]] - x[verts[3]] for i in range(3)])
+
+
+@ti.func
+def ssvd(F):
+ U, sig, V = ti.svd(F)
+ if U.determinant() < 0:
+ for i in ti.static(range(3)):
+ U[i, 2] *= -1
+ sig[2, 2] = -sig[2, 2]
+ if V.determinant() < 0:
+ for i in ti.static(range(3)):
+ V[i, 2] *= -1
+ sig[2, 2] = -sig[2, 2]
+ return U, sig, V
+
+
+@ti.func
+def get_force_func(c, verts):
+ F = Ds(verts) @ B[c]
+ P = ti.Matrix.zero(ti.f32, 3, 3)
+ U, sig, V = ssvd(F)
+ P = 2 * mu * (F - U @ V.transpose())
+ H = -W[c] * P @ B[c].transpose()
+ for i in ti.static(range(3)):
+ force = ti.Vector([H[j, i] for j in range(3)])
+ f[verts[i]] += force
+ f[verts[3]] -= force
+
+
+@ti.kernel
+def get_force():
+ for c in vertices:
+ get_force_func(c, vertices[c])
+ for u in f:
+ f[u].y -= 9.8 * m[u]
+
+
+@ti.kernel
+def matmul_cell(ret: ti.template(), vel: ti.template()):
+ for i in ret:
+ ret[i] = vel[i] * m[i]
+ for c in vertices:
+ verts = vertices[c]
+ W_c = W[c]
+ B_c = B[c]
+ for u in range(4):
+ for d in range(3):
+ dD = ti.Matrix.zero(ti.f32, 3, 3)
+ if u == 3:
+ for j in range(3):
+ dD[d, j] = -1
+ else:
+ dD[d, u] = 1
+ dF = dD @ B_c
+ dP = 2.0 * mu * dF
+ dH = -W_c * dP @ B_c.transpose()
+ for i in range(3):
+ for j in range(3):
+ tmp = (vel[verts[i]][j] - vel[verts[3]][j])
+ ret[verts[u]][d] += -dt**2 * dH[j, i] * tmp
+
+
+@ti.kernel
+def add(ans: ti.template(), a: ti.template(), k: ti.f32, b: ti.template()):
+ for i in ans:
+ ans[i] = a[i] + k * b[i]
+
+
+@ti.kernel
+def dot(a: ti.template(), b: ti.template()) -> ti.f32:
+ ans = 0.0
+ for i in a:
+ ans += a[i].dot(b[i])
+ return ans
+
+
+b = ti.Vector.field(3, dtype=ti.f32, shape=n_verts)
+r0 = ti.Vector.field(3, dtype=ti.f32, shape=n_verts)
+p0 = ti.Vector.field(3, dtype=ti.f32, shape=n_verts)
+
+
+@ti.kernel
+def get_b():
+ for i in b:
+ b[i] = m[i] * v[i] + dt * f[i]
+
+
+def cg():
+ def mul(x):
+ matmul_cell(mul_ans, x)
+ return mul_ans
+
+ get_force()
+ get_b()
+ mul(v)
+ add(r0, b, -1, mul(v))
+
+ d = p0
+ d.copy_from(r0)
+ r_2 = dot(r0, r0)
+ n_iter = 50
+ epsilon = 1e-6
+ r_2_init = r_2
+ r_2_new = r_2
+ for iter in range(n_iter):
+ q = mul(d)
+ alpha = r_2_new / dot(d, q)
+ add(v, v, alpha, d)
+ add(r0, r0, -alpha, q)
+ r_2 = r_2_new
+ r_2_new = dot(r0, r0)
+ if r_2_new <= r_2_init * epsilon**2: break
+ beta = r_2_new / r_2
+ add(d, r0, beta, d)
+ f.fill(0)
+ add(x, x, dt, v)
+
+
+@ti.kernel
+def advect():
+ for p in x:
+ v[p] += dt * (f[p] / m[p])
+ x[p] += dt * v[p]
+ f[p] = ti.Vector([0, 0, 0])
+
+
+@ti.kernel
+def init():
+ for u in x:
+ x[u] = ox[u]
+ v[u] = [0.0] * 3
+ f[u] = [0.0] * 3
+ m[u] = 0.0
+ for c in vertices:
+ F = Ds(vertices[c])
+ B[c] = F.inverse()
+ W[c] = ti.abs(F.determinant()) / 6
+ for i in range(4):
+ m[vertices[c][i]] += W[c] / 4 * density
+ for u in x:
+ x[u].y += 1.0
+
+
+@ti.kernel
+def floor_bound():
+ for u in x:
+ if x[u].y < 0:
+ x[u].y = 0
+ if v[u].y < 0:
+ v[u].y = 0
+
+
+@ti.func
+def check(u):
+ ans = 0
+ rest = u
+ for i in ti.static(range(3)):
+ k = rest % n_cube[2 - i]
+ rest = rest // n_cube[2 - i]
+ if k == 0: ans |= (1 << (i * 2))
+ if k == n_cube[2 - i] - 1: ans |= (1 << (i * 2 + 1))
+ return ans
+
+
+su = 0
+for i in range(3):
+ su += (n_cube[i] - 1) * (n_cube[(i + 1) % 3] - 1)
+indices = ti.field(ti.i32, shape=2 * su * 2 * 3)
+
+
+@ti.kernel
+def get_indices():
+ # calculate all the meshes on surface
+ cnt = 0
+ for c in vertices:
+ if c % 5 != 4:
+ for i in ti.static([0, 2, 3]):
+ verts = [vertices[c][(i + j) % 4] for j in range(3)]
+ sum = check(verts[0]) & check(verts[1]) & check(verts[2])
+ if sum:
+ m = ti.atomic_add(cnt, 1)
+ det = ti.Matrix.rows([
+ x[verts[i]] - [0.5, 1.5, 0.5] for i in range(3)
+ ]).determinant()
+ if det < 0:
+ tmp = verts[1]
+ verts[1] = verts[2]
+ verts[2] = tmp
+ indices[m * 3] = verts[0]
+ indices[m * 3 + 1] = verts[1]
+ indices[m * 3 + 2] = verts[2]
+
+
+def substep():
+ if args.exp == 'explicit':
+ for i in range(10):
+ get_force()
+ advect()
+ else:
+ for i in range(1):
+ cg()
+ floor_bound()
+
+
+if __name__ == '__main__':
+ get_vertices()
+ init()
+ get_indices()
+
+ if args.gui == 'ggui':
+ res = (800, 600)
+ window = ti.ui.Window("Implicit FEM", res, vsync=True)
+
+ frame_id = 0
+ canvas = window.get_canvas()
+ scene = ti.ui.Scene()
+ camera = ti.ui.make_camera()
+ camera.position(2.0, 2.0, 3.95)
+ camera.lookat(0.5, 0.5, 0.5)
+ camera.fov(55)
+
+ def render():
+ camera.track_user_inputs(window,
+ movement_speed=0.03,
+ hold_key=ti.ui.RMB)
+ scene.set_camera(camera)
+
+ scene.ambient_light((0.1, ) * 3)
+
+ scene.point_light(pos=(0.5, 10.0, 0.5), color=(0.5, 0.5, 0.5))
+ scene.point_light(pos=(10.0, 10.0, 10.0), color=(0.5, 0.5, 0.5))
+
+ scene.mesh(x, indices, color=(0.73, 0.33, 0.23))
+
+ canvas.scene(scene)
+
+ while window.running:
+ frame_id += 1
+ frame_id = frame_id % 256
+ substep()
+ if window.is_pressed('r'):
+ init()
+ if window.is_pressed(ti.GUI.ESCAPE):
+ break
+
+ render()
+
+ window.show()
+
+ else:
+
+ def T(a):
+
+ phi, theta = np.radians(28), np.radians(32)
+
+ a = a - 0.2
+ x, y, z = a[:, 0], a[:, 1], a[:, 2]
+ c, s = np.cos(phi), np.sin(phi)
+ C, S = np.cos(theta), np.sin(theta)
+ x, z = x * c + z * s, z * c - x * s
+ u, v = x, y * C + z * S
+ return np.array([u, v]).swapaxes(0, 1) + 0.5
+
+ gui = ti.GUI('Implicit FEM')
+ while gui.running:
+ substep()
+ if gui.get_event(ti.GUI.PRESS):
+ if gui.event.key in [ti.GUI.ESCAPE, ti.GUI.EXIT]: break
+ if gui.is_pressed('r'):
+ init()
+ gui.clear(0x000000)
+ gui.circles(T(x.to_numpy() / 3), radius=1.5, color=0xba543a)
+ gui.show()
diff --git a/examples/simulation/implicit_mass_spring.py b/python/taichi/examples/simulation/implicit_mass_spring.py
similarity index 90%
rename from examples/simulation/implicit_mass_spring.py
rename to python/taichi/examples/simulation/implicit_mass_spring.py
index 3ab120f253649..698d73df096de 100644
--- a/examples/simulation/implicit_mass_spring.py
+++ b/python/taichi/examples/simulation/implicit_mass_spring.py
@@ -17,7 +17,7 @@ def __init__(self, N):
self.initPos = ti.Vector.field(2, ti.f32, self.NV)
self.vel = ti.Vector.field(2, ti.f32, self.NV)
self.force = ti.Vector.field(2, ti.f32, self.NV)
- self.invMass = ti.field(ti.f32, self.NV)
+ self.mass = ti.field(ti.f32, self.NV)
self.spring = ti.Vector.field(2, ti.i32, self.NE)
self.indices = ti.field(ti.i32, 2 * self.NE)
@@ -55,9 +55,7 @@ def init_pos(self):
[0.25, 0.25])
self.initPos[k] = self.pos[k]
self.vel[k] = ti.Vector([0, 0])
- self.invMass[k] = 1.0
- self.invMass[self.N] = 0.0
- self.invMass[self.NV - 1] = 0.0
+ self.mass[k] = 1.0
@ti.kernel
def init_edges(self):
@@ -87,12 +85,11 @@ def init_edges(self):
rest_len[idx] = (pos[idx1] - pos[idx2]).norm()
@ti.kernel
- def init_mass_sp(self, M: ti.linalg.sparse_matrix_builder()):
+ def init_mass_sp(self, M: ti.types.sparse_matrix_builder()):
for i in range(self.NV):
- if self.invMass[i] != 0.0:
- mass = 1.0 / self.invMass[i]
- M[2 * i + 0, 2 * i + 0] += mass
- M[2 * i + 1, 2 * i + 1] += mass
+ mass = self.mass[i]
+ M[2 * i + 0, 2 * i + 0] += mass
+ M[2 * i + 1, 2 * i + 1] += mass
@ti.func
def clear_force(self):
@@ -103,8 +100,7 @@ def clear_force(self):
def compute_force(self):
self.clear_force()
for i in self.force:
- if self.invMass[i] != 0.0:
- self.force[i] += self.gravity / self.invMass[i]
+ self.force[i] += self.gravity * self.mass[i]
for i in self.spring:
idx1, idx2 = self.spring[i][0], self.spring[i][1]
@@ -137,11 +133,11 @@ def compute_Jacobians(self):
self.Jv[i] = self.kd * I
# fix point constraint hessian
- self.Jf[0] = ti.Matrix([[self.kf, 0], [0, self.kf]])
- self.Jf[1] = ti.Matrix([[self.kf, 0], [0, self.kf]])
+ self.Jf[0] = ti.Matrix([[-self.kf, 0], [0, -self.kf]])
+ self.Jf[1] = ti.Matrix([[-self.kf, 0], [0, -self.kf]])
@ti.kernel
- def assemble_K(self, K: ti.linalg.sparse_matrix_builder()):
+ def assemble_K(self, K: ti.types.sparse_matrix_builder()):
for i in self.spring:
idx1, idx2 = self.spring[i][0], self.spring[i][1]
for m, n in ti.static(ti.ndrange(2, 2)):
@@ -154,7 +150,7 @@ def assemble_K(self, K: ti.linalg.sparse_matrix_builder()):
K[2 * (self.NV - 1) + m, 2 * (self.NV - 1) + n] += self.Jf[1][m, n]
@ti.kernel
- def assemble_D(self, D: ti.linalg.sparse_matrix_builder()):
+ def assemble_D(self, D: ti.types.sparse_matrix_builder()):
for i in self.spring:
idx1, idx2 = self.spring[i][0], self.spring[i][1]
for m, n in ti.static(ti.ndrange(2, 2)):
@@ -166,9 +162,8 @@ def assemble_D(self, D: ti.linalg.sparse_matrix_builder()):
@ti.kernel
def updatePosVel(self, h: ti.f32, dv: ti.ext_arr()):
for i in self.pos:
- if self.invMass[i] != 0.0:
- self.vel[i] += ti.Vector([dv[2 * i], dv[2 * i + 1]])
- self.pos[i] += h * self.vel[i]
+ self.vel[i] += ti.Vector([dv[2 * i], dv[2 * i + 1]])
+ self.pos[i] += h * self.vel[i]
def update(self, h):
self.compute_force()
@@ -233,7 +228,7 @@ def displayGGUI(self, canvas, radius=0.01, color=(1.0, 1.0, 1.0)):
'--use-ggui',
action='store_true',
help='Display with GGUI')
- args = parser.parse_args()
+ args, unknowns = parser.parse_known_args()
use_ggui = False
use_ggui = args.use_ggui
diff --git a/python/taichi/examples/simulation/inital_value_problem.py b/python/taichi/examples/simulation/inital_value_problem.py
new file mode 100644
index 0000000000000..289c09c7a8e50
--- /dev/null
+++ b/python/taichi/examples/simulation/inital_value_problem.py
@@ -0,0 +1,47 @@
+import math
+import time
+
+import numpy as np
+
+import taichi as ti
+
+
+def init():
+ a = []
+ for i in np.linspace(0, 1, n, False):
+ for j in np.linspace(0, 1, n, False):
+ a.append([i, j])
+ return np.array(a)
+
+
+ti.init(arch=ti.gpu)
+n = 50
+dirs = ti.field(dtype=float, shape=(n * n, 2))
+locations_np = init()
+
+locations = ti.field(dtype=float, shape=(n * n, 2))
+locations.from_numpy(locations_np)
+
+
+@ti.kernel
+def paint(t: float):
+ (o, p) = locations_np.shape
+ for i in range(0, o): # Parallelized over all pixels
+ x = locations[i, 0]
+ y = locations[i, 1]
+ dirs[i, 0] = ti.sin((t * x - y))
+ dirs[i, 1] = ti.cos(t * y - x)
+ len = (dirs[i, 0]**2 + dirs[i, 1]**2)**0.5
+ dirs[i, 0] /= len * 40
+ dirs[i, 1] /= len * 40
+
+
+gui = ti.GUI("Vector Field", res=(500, 500))
+
+begining = time.time_ns()
+for k in range(1000000):
+ start_time = time.time_ns()
+ paint((time.time_ns() - begining) * 0.00000001)
+ dirs_np = dirs.to_numpy()
+ gui.arrows(locations_np, dirs_np, radius=1)
+ gui.show()
diff --git a/examples/simulation/mass_spring_game.py b/python/taichi/examples/simulation/mass_spring_game.py
similarity index 100%
rename from examples/simulation/mass_spring_game.py
rename to python/taichi/examples/simulation/mass_spring_game.py
diff --git a/examples/simulation/mpm128.py b/python/taichi/examples/simulation/mpm128.py
similarity index 100%
rename from examples/simulation/mpm128.py
rename to python/taichi/examples/simulation/mpm128.py
diff --git a/examples/simulation/mpm3d.py b/python/taichi/examples/simulation/mpm3d.py
similarity index 97%
rename from examples/simulation/mpm3d.py
rename to python/taichi/examples/simulation/mpm3d.py
index 0e63972e0120b..f7584c7159a6a 100644
--- a/examples/simulation/mpm3d.py
+++ b/python/taichi/examples/simulation/mpm3d.py
@@ -57,8 +57,8 @@ def substep():
if grid_m[I] > 0:
grid_v[I] /= grid_m[I]
grid_v[I][1] -= dt * gravity
- cond = I < bound and grid_v[I] < 0 or I > n_grid - bound and grid_v[
- I] > 0
+ cond = (I < bound) & (grid_v[I] < 0) | \
+ (I > n_grid - bound) & (grid_v[I] > 0)
grid_v[I] = 0 if cond else grid_v[I]
ti.block_dim(n_grid)
for p in x:
diff --git a/examples/simulation/mpm88.py b/python/taichi/examples/simulation/mpm88.py
similarity index 100%
rename from examples/simulation/mpm88.py
rename to python/taichi/examples/simulation/mpm88.py
diff --git a/examples/simulation/mpm99.py b/python/taichi/examples/simulation/mpm99.py
similarity index 90%
rename from examples/simulation/mpm99.py
rename to python/taichi/examples/simulation/mpm99.py
index cf71c3877bd48..a68f91b75205d 100644
--- a/examples/simulation/mpm99.py
+++ b/python/taichi/examples/simulation/mpm99.py
@@ -117,14 +117,19 @@ def initialize():
Jp[i] = 1
-initialize()
-gui = ti.GUI("Taichi MLS-MPM-99", res=512, background_color=0x112F41)
-while not gui.get_event(ti.GUI.ESCAPE, ti.GUI.EXIT):
- for s in range(int(2e-3 // dt)):
- substep()
- gui.circles(x.to_numpy(),
- radius=1.5,
- palette=[0x068587, 0xED553B, 0xEEEEF0],
- palette_indices=material)
- gui.show(
- ) # Change to gui.show(f'{frame:06d}.png') to write images to disk
+def main():
+ initialize()
+ gui = ti.GUI("Taichi MLS-MPM-99", res=512, background_color=0x112F41)
+ while not gui.get_event(ti.GUI.ESCAPE, ti.GUI.EXIT):
+ for s in range(int(2e-3 // dt)):
+ substep()
+ gui.circles(x.to_numpy(),
+ radius=1.5,
+ palette=[0x068587, 0xED553B, 0xEEEEF0],
+ palette_indices=material)
+ gui.show(
+ ) # Change to gui.show(f'{frame:06d}.png') to write images to disk
+
+
+if __name__ == '__main__':
+ main()
diff --git a/examples/simulation/mpm_lagrangian_forces.py b/python/taichi/examples/simulation/mpm_lagrangian_forces.py
similarity index 100%
rename from examples/simulation/mpm_lagrangian_forces.py
rename to python/taichi/examples/simulation/mpm_lagrangian_forces.py
diff --git a/examples/simulation/nbody.py b/python/taichi/examples/simulation/nbody.py
similarity index 73%
rename from examples/simulation/nbody.py
rename to python/taichi/examples/simulation/nbody.py
index f70d0d5e88684..c6de3021b2c11 100644
--- a/examples/simulation/nbody.py
+++ b/python/taichi/examples/simulation/nbody.py
@@ -80,23 +80,28 @@ def update():
pos[i] += dt * vel[i]
-gui = ti.GUI('N-body problem', (800, 800))
-
-initialize()
-while gui.running:
-
- for e in gui.get_events(ti.GUI.PRESS):
- if e.key in [ti.GUI.ESCAPE, ti.GUI.EXIT]:
- exit()
- elif e.key == 'r':
- initialize()
- elif e.key == ti.GUI.SPACE:
- paused[None] = not paused[None]
-
- if not paused[None]:
- for i in range(substepping):
- compute_force()
- update()
-
- gui.circles(pos.to_numpy(), color=0xffffff, radius=planet_radius)
- gui.show()
+def main():
+ gui = ti.GUI('N-body problem', (800, 800))
+
+ initialize()
+ while gui.running:
+
+ for e in gui.get_events(ti.GUI.PRESS):
+ if e.key in [ti.GUI.ESCAPE, ti.GUI.EXIT]:
+ exit()
+ elif e.key == 'r':
+ initialize()
+ elif e.key == ti.GUI.SPACE:
+ paused[None] = not paused[None]
+
+ if not paused[None]:
+ for i in range(substepping):
+ compute_force()
+ update()
+
+ gui.circles(pos.to_numpy(), color=0xffffff, radius=planet_radius)
+ gui.show()
+
+
+if __name__ == '__main__':
+ main()
diff --git a/examples/simulation/odop_solar.py b/python/taichi/examples/simulation/odop_solar.py
similarity index 100%
rename from examples/simulation/odop_solar.py
rename to python/taichi/examples/simulation/odop_solar.py
diff --git a/examples/simulation/pbf2d.py b/python/taichi/examples/simulation/pbf2d.py
similarity index 100%
rename from examples/simulation/pbf2d.py
rename to python/taichi/examples/simulation/pbf2d.py
diff --git a/examples/simulation/physarum.py b/python/taichi/examples/simulation/physarum.py
similarity index 100%
rename from examples/simulation/physarum.py
rename to python/taichi/examples/simulation/physarum.py
diff --git a/examples/simulation/stable_fluid.py b/python/taichi/examples/simulation/stable_fluid.py
similarity index 98%
rename from examples/simulation/stable_fluid.py
rename to python/taichi/examples/simulation/stable_fluid.py
index eb159f8de7120..864fd4204799a 100644
--- a/examples/simulation/stable_fluid.py
+++ b/python/taichi/examples/simulation/stable_fluid.py
@@ -12,13 +12,13 @@
# How to run:
# `python stable_fluid.py`: use the jacobi iteration to solve the linear system.
-# `python stable_fluid.py -s`: use a sparse matrix to do so.
+# `python stable_fluid.py -S`: use a sparse matrix to do so.
parser = argparse.ArgumentParser()
-parser.add_argument('-s',
+parser.add_argument('-S',
'--use-sp-mat',
action='store_true',
help='Solve Poisson\'s equation by using a sparse matrix')
-args = parser.parse_args()
+args, unknowns = parser.parse_known_args()
res = 512
dt = 0.03
@@ -69,7 +69,7 @@ def swap(self):
if use_sparse_matrix:
# use a sparse matrix to solve Poisson's pressure equation.
@ti.kernel
- def fill_laplacian_matrix(A: ti.linalg.sparse_matrix_builder()):
+ def fill_laplacian_matrix(A: ti.types.sparse_matrix_builder()):
for i, j in ti.ndrange(res, res):
row = i * res + j
center = 0.0
diff --git a/examples/simulation/vortex_rings.py b/python/taichi/examples/simulation/vortex_rings.py
similarity index 100%
rename from examples/simulation/vortex_rings.py
rename to python/taichi/examples/simulation/vortex_rings.py
diff --git a/examples/simulation/waterwave.py b/python/taichi/examples/simulation/waterwave.py
similarity index 100%
rename from examples/simulation/waterwave.py
rename to python/taichi/examples/simulation/waterwave.py
diff --git a/python/taichi/experimental.py b/python/taichi/experimental.py
new file mode 100644
index 0000000000000..13df4c4a9baf3
--- /dev/null
+++ b/python/taichi/experimental.py
@@ -0,0 +1,3 @@
+from taichi.lang.kernel_impl import real_func
+
+__all__ = ["real_func"]
diff --git a/python/taichi/lang/README.md b/python/taichi/lang/README.md
index 26cd7af7eb516..0df5e4f66cf7b 100644
--- a/python/taichi/lang/README.md
+++ b/python/taichi/lang/README.md
@@ -1,3 +1,3 @@
Some notes about the current implementation
-There are lots of `from taichi.lang.meta import xx`. Unfortunately, this cannot be moved into the top of the file, and has to be delayed inside the function. Otherwise, it would result in some cyclic import issue that is not trivially resolvable.
+There are lots of `from taichi._kernels import xx`. Unfortunately, this cannot be moved into the top of the file, and has to be delayed inside the function. Otherwise, it would result in some cyclic import issue that is not trivially resolvable.
diff --git a/python/taichi/lang/__init__.py b/python/taichi/lang/__init__.py
index 9c54dd987875e..f1a171b0cb9eb 100644
--- a/python/taichi/lang/__init__.py
+++ b/python/taichi/lang/__init__.py
@@ -1,1249 +1,32 @@
-import atexit
-import functools
-import os
-import shutil
-import tempfile
-import time
-from contextlib import contextmanager
-from copy import deepcopy as _deepcopy
+import platform
-import taichi.lang.linalg_impl
-import taichi.lang.meta
-from taichi.core.util import locale_encode
-from taichi.core.util import ti_core as _ti_core
-from taichi.lang import _random, impl, types
-from taichi.lang.ast.transformer import TaichiSyntaxError
+from taichi._lib import core as _ti_core
+from taichi.lang import impl
+from taichi.lang._ndarray import *
+from taichi.lang._ndrange import ndrange
from taichi.lang.enums import Layout
-from taichi.lang.exception import InvalidOperationError
+from taichi.lang.exception import *
+from taichi.lang.field import *
from taichi.lang.impl import *
-from taichi.lang.kernel_impl import (KernelArgError, KernelDefError,
- data_oriented, func, kernel, pyfunc)
-from taichi.lang.matrix import Matrix, Vector
-from taichi.lang.ndrange import GroupedNDRange, ndrange
-from taichi.lang.ops import *
-from taichi.lang.quant_impl import quant
-from taichi.lang.runtime_ops import async_flush, sync
-from taichi.lang.struct import Struct
-from taichi.lang.type_factory_impl import type_factory
-from taichi.lang.util import (has_pytorch, is_taichi_class, python_scope,
- taichi_scope, to_numpy_type, to_pytorch_type,
- to_taichi_type)
-from taichi.misc.util import deprecated
-from taichi.profiler import KernelProfiler, get_default_kernel_profiler
-from taichi.profiler.kernelmetrics import (CuptiMetric, default_cupti_metrics,
- get_predefined_cupti_metrics)
-from taichi.snode.fields_builder import FieldsBuilder
-from taichi.type.annotations import any_arr, ext_arr, template
-
-import taichi as ti
-
-runtime = impl.get_runtime()
-
-i = axes(0)
-j = axes(1)
-k = axes(2)
-l = axes(3)
-ij = axes(0, 1)
-ik = axes(0, 2)
-il = axes(0, 3)
-jk = axes(1, 2)
-jl = axes(1, 3)
-kl = axes(2, 3)
-ijk = axes(0, 1, 2)
-ijl = axes(0, 1, 3)
-ikl = axes(0, 2, 3)
-jkl = axes(1, 2, 3)
-ijkl = axes(0, 1, 2, 3)
-
-outer_product = deprecated('ti.outer_product(a, b)',
- 'a.outer_product(b)')(Matrix.outer_product)
-cross = deprecated('ti.cross(a, b)', 'a.cross(b)')(Matrix.cross)
-dot = deprecated('ti.dot(a, b)', 'a.dot(b)')(Matrix.dot)
-normalized = deprecated('ti.normalized(a)',
- 'a.normalized()')(Matrix.normalized)
-
-cfg = default_cfg()
-x86_64 = _ti_core.x64
-"""The x64 CPU backend.
-"""
-x64 = _ti_core.x64
-"""The X64 CPU backend.
-"""
-arm64 = _ti_core.arm64
-"""The ARM CPU backend.
-"""
-cuda = _ti_core.cuda
-"""The CUDA backend.
-"""
-metal = _ti_core.metal
-"""The Apple Metal backend.
-"""
-opengl = _ti_core.opengl
-"""The OpenGL backend. OpenGL 4.3 required.
-"""
-# Skip annotating this one because it is barely maintained.
-cc = _ti_core.cc
-wasm = _ti_core.wasm
-"""The WebAssembly backend.
-"""
-vulkan = _ti_core.vulkan
-"""The Vulkan backend.
-"""
-gpu = [cuda, metal, opengl, vulkan]
-"""A list of GPU backends supported on the current system.
-
-When this is used, Taichi automatically picks the matching GPU backend. If no
-GPU is detected, Taichi falls back to the CPU backend.
-"""
-cpu = _ti_core.host_arch()
-"""A list of CPU backends supported on the current system.
-
-When this is used, Taichi automatically picks the matching CPU backend.
-"""
-timeline_clear = lambda: impl.get_runtime().prog.timeline_clear()
-timeline_save = lambda fn: impl.get_runtime().prog.timeline_save(fn)
-
-# Legacy API
-type_factory_ = _ti_core.get_type_factory_instance()
-
-
-@deprecated('kernel_profiler_print()', 'print_kernel_profile_info()')
-def kernel_profiler_print():
- return print_kernel_profile_info()
-
-
-def print_kernel_profile_info(mode='count'):
- """Print the profiling results of Taichi kernels.
-
- To enable this profiler, set ``kernel_profiler=True`` in ``ti.init()``.
- ``'count'`` mode: print the statistics (min,max,avg time) of launched kernels,
- ``'trace'`` mode: print the records of launched kernels with specific profiling metrics (time, memory load/store and core utilization etc.),
- and defaults to ``'count'``.
-
- Args:
- mode (str): the way to print profiling results.
-
- Example::
-
- >>> import taichi as ti
-
- >>> ti.init(ti.cpu, kernel_profiler=True)
- >>> var = ti.field(ti.f32, shape=1)
-
- >>> @ti.kernel
- >>> def compute():
- >>> var[0] = 1.0
-
- >>> compute()
- >>> ti.print_kernel_profile_info()
- >>> # equivalent calls :
- >>> # ti.print_kernel_profile_info('count')
-
- >>> ti.print_kernel_profile_info('trace')
-
- Note:
- Currently the result of `KernelProfiler` could be incorrect on OpenGL
- backend due to its lack of support for `ti.sync()`.
-
- For advanced mode of `KernelProfiler`, please visit https://docs.taichi.graphics/docs/lang/articles/misc/profiler#advanced-mode.
- """
- get_default_kernel_profiler().print_info(mode)
-
-
-def query_kernel_profile_info(name):
- """Query kernel elapsed time(min,avg,max) on devices using the kernel name.
-
- To enable this profiler, set `kernel_profiler=True` in `ti.init`.
-
- Args:
- name (str): kernel name.
-
- Returns:
- KernelProfilerQueryResult (class): with member variables(counter, min, max, avg)
-
- Example::
-
- >>> import taichi as ti
-
- >>> ti.init(ti.cpu, kernel_profiler=True)
- >>> n = 1024*1024
- >>> var = ti.field(ti.f32, shape=n)
-
- >>> @ti.kernel
- >>> def fill():
- >>> for i in range(n):
- >>> var[i] = 0.1
-
- >>> fill()
- >>> ti.clear_kernel_profile_info() #[1]
- >>> for i in range(100):
- >>> fill()
- >>> query_result = ti.query_kernel_profile_info(fill.__name__) #[2]
- >>> print("kernel excuted times =",query_result.counter)
- >>> print("kernel elapsed time(min_in_ms) =",query_result.min)
- >>> print("kernel elapsed time(max_in_ms) =",query_result.max)
- >>> print("kernel elapsed time(avg_in_ms) =",query_result.avg)
-
- Note:
- [1] To get the correct result, query_kernel_profile_info() must be used in conjunction with
- clear_kernel_profile_info().
-
- [2] Currently the result of `KernelProfiler` could be incorrect on OpenGL
- backend due to its lack of support for `ti.sync()`.
- """
- return get_default_kernel_profiler().query_info(name)
-
-
-@deprecated('kernel_profiler_clear()', 'clear_kernel_profile_info()')
-def kernel_profiler_clear():
- return clear_kernel_profile_info()
-
-
-def clear_kernel_profile_info():
- """Clear all KernelProfiler records."""
- get_default_kernel_profiler().clear_info()
-
-
-def kernel_profiler_total_time():
- """Get elapsed time of all kernels recorded in KernelProfiler.
-
- Returns:
- time (float): total time in second.
- """
- return get_default_kernel_profiler().get_total_time()
-
-
-def set_kernel_profile_metrics(metric_list=default_cupti_metrics):
- """Set metrics that will be collected by the CUPTI toolkit.
-
- Args:
- metric_list (list): a list of :class:`~taichi.lang.CuptiMetric()` instances, default value: :data:`~taichi.lang.default_cupti_metrics`.
-
- Example::
-
- >>> import taichi as ti
-
- >>> ti.init(kernel_profiler=True, arch=ti.cuda)
- >>> num_elements = 128*1024*1024
-
- >>> x = ti.field(ti.f32, shape=num_elements)
- >>> y = ti.field(ti.f32, shape=())
- >>> y[None] = 0
-
- >>> @ti.kernel
- >>> def reduction():
- >>> for i in x:
- >>> y[None] += x[i]
-
- >>> # In the case of not pramater, Taichi will print its pre-defined metrics list
- >>> ti.get_predefined_cupti_metrics()
- >>> # get Taichi pre-defined metrics
- >>> profiling_metrics = ti.get_predefined_cupti_metrics('shared_access')
-
- >>> global_op_atom = ti.CuptiMetric(
- >>> name='l1tex__t_set_accesses_pipe_lsu_mem_global_op_atom.sum',
- >>> header=' global.atom ',
- >>> format=' {:8.0f} ')
- >>> # add user defined metrics
- >>> profiling_metrics += [global_op_atom]
-
- >>> # metrics setting will be retained until the next configuration
- >>> ti.set_kernel_profile_metrics(profiling_metrics)
- >>> for i in range(16):
- >>> reduction()
- >>> ti.print_kernel_profile_info('trace')
-
- Note:
- Metrics setting will be retained until the next configuration.
- """
- get_default_kernel_profiler().set_metrics(metric_list)
-
-
-@contextmanager
-def collect_kernel_profile_metrics(metric_list=default_cupti_metrics):
- """Set temporary metrics that will be collected by the CUPTI toolkit within this context.
-
- Args:
- metric_list (list): a list of :class:`~taichi.lang.CuptiMetric()` instances, default value: :data:`~taichi.lang.default_cupti_metrics`.
-
- Example::
-
- >>> import taichi as ti
-
- >>> ti.init(kernel_profiler=True, arch=ti.cuda)
- >>> num_elements = 128*1024*1024
-
- >>> x = ti.field(ti.f32, shape=num_elements)
- >>> y = ti.field(ti.f32, shape=())
- >>> y[None] = 0
-
- >>> @ti.kernel
- >>> def reduction():
- >>> for i in x:
- >>> y[None] += x[i]
-
- >>> # In the case of not pramater, Taichi will print its pre-defined metrics list
- >>> ti.get_predefined_cupti_metrics()
- >>> # get Taichi pre-defined metrics
- >>> profiling_metrics = ti.get_predefined_cupti_metrics('device_utilization')
-
- >>> global_op_atom = ti.CuptiMetric(
- >>> name='l1tex__t_set_accesses_pipe_lsu_mem_global_op_atom.sum',
- >>> header=' global.atom ',
- >>> format=' {:8.0f} ')
- >>> # add user defined metrics
- >>> profiling_metrics += [global_op_atom]
-
- >>> # metrics setting is temporary, and will be clear when exit from this context.
- >>> with ti.collect_kernel_profile_metrics(profiling_metrics):
- >>> for i in range(16):
- >>> reduction()
- >>> ti.print_kernel_profile_info('trace')
-
- Note:
- The configuration of the ``metric_list`` will be clear when exit from this context.
- """
- get_default_kernel_profiler().set_metrics(metric_list)
- yield get_default_kernel_profiler()
- get_default_kernel_profiler().set_metrics()
-
-
-@deprecated('memory_profiler_print()', 'print_memory_profile_info()')
-def memory_profiler_print():
- return print_memory_profile_info()
-
-
-def print_memory_profile_info():
- """Memory profiling tool for LLVM backends with full sparse support.
-
- This profiler is automatically on.
- """
- impl.get_runtime().materialize()
- impl.get_runtime().prog.print_memory_profiler_info()
-
-
-extension = _ti_core.Extension
-
-
-def is_extension_supported(arch, ext):
- """Checks whether an extension is supported on an arch.
-
- Args:
- arch (taichi_core.Arch): Specified arch.
- ext (taichi_core.Extension): Specified extension.
-
- Returns:
- bool: Whether `ext` is supported on `arch`.
- """
- return _ti_core.is_extension_supported(arch, ext)
-
-
-def reset():
- """Resets Taichi to its initial state.
-
- This would destroy all the fields and kernels.
- """
- _ti_core.reset_snode_access_flag()
- impl.reset()
- global runtime
- runtime = impl.get_runtime()
-
-
-class _EnvironmentConfigurator:
- def __init__(self, kwargs, cfg):
- self.cfg = cfg
- self.kwargs = kwargs
- self.keys = []
-
- def add(self, key, cast=None):
- cast = cast or self.bool_int
-
- self.keys.append(key)
-
- # TI_ASYNC= : no effect
- # TI_ASYNC=0 : False
- # TI_ASYNC=1 : True
- name = 'TI_' + key.upper()
- value = os.environ.get(name, '')
- if len(value):
- self[key] = cast(value)
- if key in self.kwargs:
- _ti_core.warn(
- f'ti.init argument "{key}" overridden by environment variable {name}={value}'
- )
- del self.kwargs[key] # mark as recognized
- elif key in self.kwargs:
- self[key] = self.kwargs[key]
- del self.kwargs[key] # mark as recognized
-
- def __getitem__(self, key):
- return getattr(self.cfg, key)
-
- def __setitem__(self, key, value):
- setattr(self.cfg, key, value)
-
- @staticmethod
- def bool_int(x):
- return bool(int(x))
-
-
-class _SpecialConfig:
- # like CompileConfig in C++, this is the configurations that belong to other submodules
- def __init__(self):
- self.print_preprocessed = False
- self.log_level = 'info'
- self.gdb_trigger = False
- self.excepthook = False
- self.experimental_real_function = False
-
-
-def prepare_sandbox():
- '''
- Returns a temporary directory, which will be automatically deleted on exit.
- It may contain the taichi_core shared object or some misc. files.
- '''
- tmp_dir = tempfile.mkdtemp(prefix='taichi-')
- atexit.register(shutil.rmtree, tmp_dir)
- print(f'[Taichi] preparing sandbox at {tmp_dir}')
- os.mkdir(os.path.join(tmp_dir, 'runtime/'))
- return tmp_dir
-
-
-def init(arch=None,
- default_fp=None,
- default_ip=None,
- _test_mode=False,
- **kwargs):
- """Initializes the Taichi runtime.
-
- This should always be the entry point of your Taichi program. Most
- importantly, it sets the backend used throughout the program.
-
- Args:
- arch: Backend to use. This is usually :const:`~taichi.lang.cpu` or :const:`~taichi.lang.gpu`.
- default_fp (Optional[type]): Default floating-point type.
- default_fp (Optional[type]): Default integral type.
- **kwargs: Taichi provides highly customizable compilation through
- ``kwargs``, which allows for fine grained control of Taichi compiler
- behavior. Below we list some of the most frequently used ones. For a
- complete list, please check out
- https://github.com/taichi-dev/taichi/blob/master/taichi/program/compile_config.h.
-
- * ``cpu_max_num_threads`` (int): Sets the number of threads used by the CPU thread pool.
- * ``debug`` (bool): Enables the debug mode, under which Taichi does a few more things like boundary checks.
- * ``print_ir`` (bool): Prints the CHI IR of the Taichi kernels.
- * ``packed`` (bool): Enables the packed memory layout. See https://docs.taichi.graphics/lang/articles/advanced/layout.
- """
- # Make a deepcopy in case these args reference to items from ti.cfg, which are
- # actually references. If no copy is made and the args are indeed references,
- # ti.reset() could override the args to their default values.
- default_fp = _deepcopy(default_fp)
- default_ip = _deepcopy(default_ip)
- kwargs = _deepcopy(kwargs)
- ti.reset()
-
- spec_cfg = _SpecialConfig()
- env_comp = _EnvironmentConfigurator(kwargs, ti.cfg)
- env_spec = _EnvironmentConfigurator(kwargs, spec_cfg)
-
- # configure default_fp/ip:
- # TODO: move these stuff to _SpecialConfig too:
- env_default_fp = os.environ.get("TI_DEFAULT_FP")
- if env_default_fp:
- if default_fp is not None:
- _ti_core.warn(
- f'ti.init argument "default_fp" overridden by environment variable TI_DEFAULT_FP={env_default_fp}'
- )
- if env_default_fp == '32':
- default_fp = ti.f32
- elif env_default_fp == '64':
- default_fp = ti.f64
- elif env_default_fp is not None:
- raise ValueError(
- f'Invalid TI_DEFAULT_FP={env_default_fp}, should be 32 or 64')
-
- env_default_ip = os.environ.get("TI_DEFAULT_IP")
- if env_default_ip:
- if default_ip is not None:
- _ti_core.warn(
- f'ti.init argument "default_ip" overridden by environment variable TI_DEFAULT_IP={env_default_ip}'
- )
- if env_default_ip == '32':
- default_ip = ti.i32
- elif env_default_ip == '64':
- default_ip = ti.i64
- elif env_default_ip is not None:
- raise ValueError(
- f'Invalid TI_DEFAULT_IP={env_default_ip}, should be 32 or 64')
-
- if default_fp is not None:
- impl.get_runtime().set_default_fp(default_fp)
- if default_ip is not None:
- impl.get_runtime().set_default_ip(default_ip)
-
- # submodule configurations (spec_cfg):
- env_spec.add('print_preprocessed')
- env_spec.add('log_level', str)
- env_spec.add('gdb_trigger')
- env_spec.add('excepthook')
- env_spec.add('experimental_real_function')
-
- # compiler configurations (ti.cfg):
- for key in dir(ti.cfg):
- if key in ['arch', 'default_fp', 'default_ip']:
- continue
- cast = type(getattr(ti.cfg, key))
- if cast is bool:
- cast = None
- env_comp.add(key, cast)
-
- unexpected_keys = kwargs.keys()
-
- if 'use_unified_memory' in unexpected_keys:
- _ti_core.warn(
- f'"use_unified_memory" is a deprecated option, as taichi no longer have the option of using unified memory.'
- )
- del kwargs['use_unified_memory']
-
- if len(unexpected_keys):
- raise KeyError(
- f'Unrecognized keyword argument(s) for ti.init: {", ".join(unexpected_keys)}'
- )
-
- # dispatch configurations that are not in ti.cfg:
- if not _test_mode:
- ti.set_gdb_trigger(spec_cfg.gdb_trigger)
- impl.get_runtime().print_preprocessed = spec_cfg.print_preprocessed
- impl.get_runtime().experimental_real_function = \
- spec_cfg.experimental_real_function
- ti.set_logging_level(spec_cfg.log_level.lower())
- if spec_cfg.excepthook:
- # TODO(#1405): add a way to restore old excepthook
- ti.enable_excepthook()
-
- # select arch (backend):
- env_arch = os.environ.get('TI_ARCH')
- if env_arch is not None:
- ti.info(f'Following TI_ARCH setting up for arch={env_arch}')
- arch = _ti_core.arch_from_name(env_arch)
- ti.cfg.arch = adaptive_arch_select(arch)
- if ti.cfg.arch == cc:
- _ti_core.set_tmp_dir(locale_encode(prepare_sandbox()))
- print(f'[Taichi] Starting on arch={_ti_core.arch_name(ti.cfg.arch)}')
-
- if _test_mode:
- return spec_cfg
-
- get_default_kernel_profiler().set_kernel_profiler_mode(
- ti.cfg.kernel_profiler)
-
- # create a new program:
- impl.get_runtime().create_program()
-
- ti.trace('Materializing runtime...')
- impl.get_runtime().prog.materialize_runtime()
-
- impl._root_fb = FieldsBuilder()
-
-
-def no_activate(*args):
- for v in args:
- _ti_core.no_activate(v.snode.ptr)
-
-
-def block_local(*args):
- """Hints Taichi to cache the fields and to enable the BLS optimization.
-
- Please visit https://docs.taichi.graphics/lang/articles/advanced/performance
- for how BLS is used.
-
- Args:
- *args (List[Field]): A list of sparse Taichi fields.
-
- Raises:
- InvalidOperationError: If the ``dynamic_index`` feature (experimental)
- is enabled.
- """
- if ti.current_cfg().dynamic_index:
- raise InvalidOperationError(
- 'dynamic_index is not allowed when block_local is turned on.')
- for a in args:
- for v in a.get_field_members():
- _ti_core.insert_snode_access_flag(
- _ti_core.SNodeAccessFlag.block_local, v.ptr)
-
-
-@deprecated('ti.cache_shared', 'ti.block_local')
-def cache_shared(*args):
- block_local(*args)
-
-
-def cache_read_only(*args):
- for a in args:
- for v in a.get_field_members():
- _ti_core.insert_snode_access_flag(
- _ti_core.SNodeAccessFlag.read_only, v.ptr)
-
-
-def assume_in_range(val, base, low, high):
- return _ti_core.expr_assume_in_range(
- Expr(val).ptr,
- Expr(base).ptr, low, high)
-
-
-def loop_unique(val, covers=None):
- if covers is None:
- covers = []
- if not isinstance(covers, (list, tuple)):
- covers = [covers]
- covers = [x.snode.ptr if isinstance(x, Expr) else x.ptr for x in covers]
- return _ti_core.expr_loop_unique(Expr(val).ptr, covers)
-
-
-parallelize = _ti_core.parallelize
-serialize = lambda: parallelize(1)
-vectorize = _ti_core.vectorize
-bit_vectorize = _ti_core.bit_vectorize
-block_dim = _ti_core.block_dim
-
-inversed = deprecated('ti.inversed(a)', 'a.inverse()')(Matrix.inversed)
-transposed = deprecated('ti.transposed(a)', 'a.transpose()')(Matrix.transposed)
-
-
-def polar_decompose(A, dt=None):
- """Perform polar decomposition (A=UP) for arbitrary size matrix.
-
- Mathematical concept refers to https://en.wikipedia.org/wiki/Polar_decomposition.
- This is only a wrapper for :func:`taichi.lang.linalg_impl.polar_decompose`.
-
- Args:
- A (ti.Matrix(n, n)): input nxn matrix `A`.
- dt (DataType): date type of elements in matrix `A`, typically accepts ti.f32 or ti.f64.
-
- Returns:
- Decomposed nxn matrices `U` and `P`.
- """
- if dt is None:
- dt = impl.get_runtime().default_fp
- return taichi.lang.linalg_impl.polar_decompose(A, dt)
-
-
-def svd(A, dt=None):
- """Perform singular value decomposition (A=USV^T) for arbitrary size matrix.
-
- Mathematical concept refers to https://en.wikipedia.org/wiki/Singular_value_decomposition.
- This is only a wrappers for :func:`taichi.lang.linalg_impl.svd`.
-
- Args:
- A (ti.Matrix(n, n)): input nxn matrix `A`.
- dt (DataType): date type of elements in matrix `A`, typically accepts ti.f32 or ti.f64.
-
- Returns:
- Decomposed nxn matrices `U`, 'S' and `V`.
- """
- if dt is None:
- dt = impl.get_runtime().default_fp
- return taichi.lang.linalg_impl.svd(A, dt)
-
-
-def eig(A, dt=None):
- """Compute the eigenvalues and right eigenvectors of a real matrix.
-
- Mathematical concept refers to https://en.wikipedia.org/wiki/Eigendecomposition_of_a_matrix.
- 2D implementation refers to :func:`taichi.lang.linalg_impl.eig2x2`.
-
- Args:
- A (ti.Matrix(n, n)): 2D Matrix for which the eigenvalues and right eigenvectors will be computed.
- dt (DataType): The datatype for the eigenvalues and right eigenvectors.
-
- Returns:
- eigenvalues (ti.Matrix(n, 2)): The eigenvalues in complex form. Each row stores one eigenvalue. The first number of the eigenvalue represents the real part and the second number represents the imaginary part.
- eigenvectors (ti.Matrix(n*2, n)): The eigenvectors in complex form. Each column stores one eigenvector. Each eigenvector consists of n entries, each of which is represented by two numbers for its real part and imaginary part.
- """
- if dt is None:
- dt = impl.get_runtime().default_fp
- if A.n == 2:
- return taichi.lang.linalg_impl.eig2x2(A, dt)
- raise Exception("Eigen solver only supports 2D matrices.")
-
-
-def sym_eig(A, dt=None):
- """Compute the eigenvalues and right eigenvectors of a real symmetric matrix.
-
- Mathematical concept refers to https://en.wikipedia.org/wiki/Eigendecomposition_of_a_matrix.
- 2D implementation refers to :func:`taichi.lang.linalg_impl.sym_eig2x2`.
-
- Args:
- A (ti.Matrix(n, n)): Symmetric Matrix for which the eigenvalues and right eigenvectors will be computed.
- dt (DataType): The datatype for the eigenvalues and right eigenvectors.
-
- Returns:
- eigenvalues (ti.Vector(n)): The eigenvalues. Each entry store one eigen value.
- eigenvectors (ti.Matrix(n, n)): The eigenvectors. Each column stores one eigenvector.
- """
- assert all(A == A.transpose()), "A needs to be symmetric"
- if dt is None:
- dt = impl.get_runtime().default_fp
- if A.n == 2:
- return taichi.lang.linalg_impl.sym_eig2x2(A, dt)
- raise Exception("Symmetric eigen solver only supports 2D matrices.")
-
-
-def randn(dt=None):
- """Generates a random number from standard normal distribution.
-
- Implementation refers to :func:`taichi.lang.random.randn`.
-
- Args:
- dt (DataType): The datatype for the generated random number.
-
- Returns:
- The generated random number.
- """
- if dt is None:
- dt = impl.get_runtime().default_fp
- return _random.randn(dt)
-
-
-determinant = deprecated('ti.determinant(a)',
- 'a.determinant()')(Matrix.determinant)
-tr = deprecated('ti.tr(a)', 'a.trace()')(Matrix.trace)
-
-
-def Tape(loss, clear_gradients=True):
- """Return a context manager of :class:`~taichi.lang.tape.TapeImpl`. The
- context manager would catching all of the callings of functions that
- decorated by :func:`~taichi.lang.kernel_impl.kernel` or
- :func:`~taichi.ad.grad_replaced` under `with` statement, and calculate
- all the partial gradients of a given loss variable by calling all of the
- gradient function of the callings caught in reverse order while `with`
- statement ended.
-
- See also :func:`~taichi.lang.kernel_impl.kernel` and
- :func:`~taichi.ad.grad_replaced` for gradient functions.
-
- Args:
- loss(:class:`~taichi.lang.expr.Expr`): The loss field, which shape should be ().
- clear_gradients(Bool): Before `with` body start, clear all gradients or not.
-
- Returns:
- :class:`~taichi.lang.tape.TapeImpl`: The context manager.
-
- Example::
-
- >>> @ti.kernel
- >>> def sum(a: ti.float32):
- >>> for I in ti.grouped(x):
- >>> y[None] += x[I] ** a
- >>>
- >>> with ti.Tape(loss = y):
- >>> sum(2)"""
- impl.get_runtime().materialize()
- if len(loss.shape) != 0:
- raise RuntimeError(
- 'The loss of `Tape` must be a 0-D field, i.e. scalar')
- if not loss.snode.ptr.has_grad():
- raise RuntimeError(
- 'Gradients of loss are not allocated, please use ti.field(..., needs_grad=True)'
- ' for all fields that are required by autodiff.')
- if clear_gradients:
- clear_all_gradients()
-
- taichi.lang.meta.clear_loss(loss)
-
- return runtime.get_tape(loss)
-
-
-def clear_all_gradients():
- """Set all fields' gradients to 0."""
- impl.get_runtime().materialize()
-
- def visit(node):
- places = []
- for i in range(node.ptr.get_num_ch()):
- ch = node.ptr.get_ch(i)
- if not ch.is_place():
- visit(SNode(ch))
- else:
- if not ch.is_primal():
- places.append(ch.get_expr())
-
- places = tuple(places)
- if places:
- taichi.lang.meta.clear_gradients(places)
-
- for root_fb in FieldsBuilder.finalized_roots():
- visit(root_fb)
-
-
-def deactivate_all_snodes():
- """Recursively deactivate all SNodes."""
- for root_fb in FieldsBuilder.finalized_roots():
- root_fb.deactivate_all()
-
-
-def benchmark(func, repeat=300, args=()):
- def run_benchmark():
- compile_time = time.time()
- func(*args) # compile the kernel first
- ti.sync()
- compile_time = time.time() - compile_time
- ti.stat_write('compilation_time', compile_time)
- codegen_stat = _ti_core.stat()
- for line in codegen_stat.split('\n'):
- try:
- a, b = line.strip().split(':')
- except:
- continue
- a = a.strip()
- b = int(float(b))
- if a == 'codegen_kernel_statements':
- ti.stat_write('compiled_inst', b)
- if a == 'codegen_offloaded_tasks':
- ti.stat_write('compiled_tasks', b)
- elif a == 'launched_tasks':
- ti.stat_write('launched_tasks', b)
-
- # Use 3 initial iterations to warm up
- # instruction/data caches. Discussion:
- # https://github.com/taichi-dev/taichi/pull/1002#discussion_r426312136
- for i in range(3):
- func(*args)
- ti.sync()
- ti.clear_kernel_profile_info()
- t = time.time()
- for n in range(repeat):
- func(*args)
- ti.sync()
- elapsed = time.time() - t
- avg = elapsed / repeat
- ti.stat_write('wall_clk_t', avg)
- device_time = ti.kernel_profiler_total_time()
- avg_device_time = device_time / repeat
- ti.stat_write('exec_t', avg_device_time)
-
- run_benchmark()
-
-
-def benchmark_plot(fn=None,
- cases=None,
- columns=None,
- column_titles=None,
- archs=None,
- title=None,
- bars='sync_vs_async',
- bar_width=0.4,
- bar_distance=0,
- left_margin=0,
- size=(12, 8)):
- import matplotlib.pyplot as plt # pylint: disable=C0415
- import yaml # pylint: disable=C0415
- if fn is None:
- fn = os.path.join(_ti_core.get_repo_dir(), 'benchmarks', 'output',
- 'benchmark.yml')
-
- with open(fn, 'r') as f:
- data = yaml.load(f, Loader=yaml.SafeLoader)
- if bars != 'sync_vs_async': # need baseline
- baseline_dir = os.path.join(_ti_core.get_repo_dir(), 'benchmarks',
- 'baseline')
- baseline_file = f'{baseline_dir}/benchmark.yml'
- with open(baseline_file, 'r') as f:
- baseline_data = yaml.load(f, Loader=yaml.SafeLoader)
- if cases is None:
- cases = list(data.keys())
-
- assert len(cases) >= 1
- if len(cases) == 1:
- cases = [cases[0], cases[0]]
- ti.warning(
- 'Function benchmark_plot does not support plotting with only one case for now. Duplicating the item to move on.'
- )
-
- if columns is None:
- columns = list(data[cases[0]].keys())
- if column_titles is None:
- column_titles = columns
- normalize_to_lowest = lambda x: True
- figure, subfigures = plt.subplots(len(cases), len(columns))
- if title is None:
- title = 'Taichi Performance Benchmarks (Higher means more)'
- figure.suptitle(title, fontweight="bold")
- for col_id in range(len(columns)):
- subfigures[0][col_id].set_title(column_titles[col_id])
- for case_id in range(len(cases)):
- case = cases[case_id]
- subfigures[case_id][0].annotate(
- case,
- xy=(0, 0.5),
- xytext=(-subfigures[case_id][0].yaxis.labelpad - 5, 0),
- xycoords=subfigures[case_id][0].yaxis.label,
- textcoords='offset points',
- size='large',
- ha='right',
- va='center')
- for col_id in range(len(columns)):
- col = columns[col_id]
- if archs is None:
- current_archs = data[case][col].keys()
- else:
- current_archs = [
- x for x in archs if x in data[case][col].keys()
- ]
- if bars == 'sync_vs_async':
- y_left = [
- data[case][col][arch]['sync'] for arch in current_archs
- ]
- label_left = 'sync'
- y_right = [
- data[case][col][arch]['async'] for arch in current_archs
- ]
- label_right = 'async'
- elif bars == 'sync_regression':
- y_left = [
- baseline_data[case][col][arch]['sync']
- for arch in current_archs
- ]
- label_left = 'before'
- y_right = [
- data[case][col][arch]['sync'] for arch in current_archs
- ]
- label_right = 'after'
- elif bars == 'async_regression':
- y_left = [
- baseline_data[case][col][arch]['async']
- for arch in current_archs
- ]
- label_left = 'before'
- y_right = [
- data[case][col][arch]['async'] for arch in current_archs
- ]
- label_right = 'after'
- else:
- raise RuntimeError('Unknown bars type')
- if normalize_to_lowest(col):
- for i in range(len(current_archs)):
- maximum = max(y_left[i], y_right[i])
- y_left[i] = y_left[i] / maximum if y_left[i] != 0 else 1
- y_right[i] = y_right[i] / maximum if y_right[i] != 0 else 1
- ax = subfigures[case_id][col_id]
- bar_left = ax.bar(x=[
- i - bar_width / 2 - bar_distance / 2
- for i in range(len(current_archs))
- ],
- height=y_left,
- width=bar_width,
- label=label_left,
- color=(0.47, 0.69, 0.89, 1.0))
- bar_right = ax.bar(x=[
- i + bar_width / 2 + bar_distance / 2
- for i in range(len(current_archs))
- ],
- height=y_right,
- width=bar_width,
- label=label_right,
- color=(0.68, 0.26, 0.31, 1.0))
- ax.set_xticks(range(len(current_archs)))
- ax.set_xticklabels(current_archs)
- figure.legend((bar_left, bar_right), (label_left, label_right),
- loc='lower center')
- figure.subplots_adjust(left=left_margin)
-
- fig = plt.gcf()
- fig.set_size_inches(size)
-
- plt.show()
-
-
-def stat_write(key, value):
- import yaml # pylint: disable=C0415
- case_name = os.environ.get('TI_CURRENT_BENCHMARK')
- if case_name is None:
- return
- if case_name.startswith('benchmark_'):
- case_name = case_name[10:]
- arch_name = _ti_core.arch_name(ti.cfg.arch)
- async_mode = 'async' if ti.cfg.async_mode else 'sync'
- output_dir = os.environ.get('TI_BENCHMARK_OUTPUT_DIR', '.')
- filename = f'{output_dir}/benchmark.yml'
- try:
- with open(filename, 'r') as f:
- data = yaml.load(f, Loader=yaml.SafeLoader)
- except FileNotFoundError:
- data = {}
- data.setdefault(case_name, {})
- data[case_name].setdefault(key, {})
- data[case_name][key].setdefault(arch_name, {})
- data[case_name][key][arch_name][async_mode] = value
- with open(filename, 'w') as f:
- yaml.dump(data, f, Dumper=yaml.SafeDumper)
-
-
-def is_arch_supported(arch):
- """Checks whether an arch is supported on the machine.
-
- Args:
- arch (taichi_core.Arch): Specified arch.
-
- Returns:
- bool: Whether `arch` is supported on the machine.
- """
- arch_table = {
- cuda: _ti_core.with_cuda,
- metal: _ti_core.with_metal,
- opengl: _ti_core.with_opengl,
- cc: _ti_core.with_cc,
- vulkan: lambda: _ti_core.with_vulkan(),
- wasm: lambda: True,
- cpu: lambda: True,
- }
- with_arch = arch_table.get(arch, lambda: False)
- try:
- return with_arch()
- except Exception as e:
- arch = _ti_core.arch_name(arch)
- _ti_core.warn(
- f"{e.__class__.__name__}: '{e}' occurred when detecting "
- f"{arch}, consider add `export TI_WITH_{arch.upper()}=0` "
- f" to environment variables to depress this warning message.")
- return False
-
-
-def supported_archs():
- """Gets all supported archs on the machine.
-
- Returns:
- List[taichi_core.Arch]: All supported archs on the machine.
- """
- archs = set([cpu, cuda, metal, vulkan, opengl, cc])
- archs = set(filter(lambda x: is_arch_supported(x), archs))
-
- wanted_archs = os.environ.get('TI_WANTED_ARCHS', '')
- want_exclude = wanted_archs.startswith('^')
- if want_exclude:
- wanted_archs = wanted_archs[1:]
- wanted_archs = wanted_archs.split(',')
- # Note, ''.split(',') gives you [''], which is not an empty array.
- expanded_wanted_archs = set([])
- for arch in wanted_archs:
- if arch == '':
- continue
- if arch == 'cpu':
- expanded_wanted_archs.add(cpu)
- elif arch == 'gpu':
- expanded_wanted_archs.update(gpu)
- else:
- expanded_wanted_archs.add(_ti_core.arch_from_name(arch))
- if len(expanded_wanted_archs) == 0:
- return list(archs)
- if want_exclude:
- supported = archs - expanded_wanted_archs
- else:
- supported = archs & expanded_wanted_archs
- return list(supported)
-
-
-def adaptive_arch_select(arch):
- if arch is None:
- return cpu
- if not isinstance(arch, (list, tuple)):
- arch = [arch]
- for a in arch:
- if is_arch_supported(a):
- return a
- ti.warn(f'Arch={arch} is not supported, falling back to CPU')
- return cpu
-
-
-class _ArchCheckers(object):
- def __init__(self):
- self._checkers = []
-
- def register(self, c):
- self._checkers.append(c)
-
- def __call__(self, arch):
- assert isinstance(arch, _ti_core.Arch)
- return all([c(arch) for c in self._checkers])
-
-
-_tests_arch_checkers_argname = '_tests_arch_checkers'
-
-
-def _get_or_make_arch_checkers(kwargs):
- k = _tests_arch_checkers_argname
- if k not in kwargs:
- kwargs[k] = _ArchCheckers()
- return kwargs[k]
-
-
-# test with all archs
-def all_archs_with(**kwargs):
- kwargs = _deepcopy(kwargs)
-
- def decorator(test):
- # @pytest.mark.parametrize decorator only knows about regular function args,
- # without *args or **kwargs. By decorating with @functools.wraps, the
- # signature of |test| is preserved, so that @ti.all_archs can be used after
- # the parametrization decorator.
- #
- # Full discussion: https://github.com/pytest-dev/pytest/issues/6810
- @functools.wraps(test)
- def wrapped(*test_args, **test_kwargs):
- can_run_on = test_kwargs.pop(_tests_arch_checkers_argname,
- _ArchCheckers())
- # Filter away archs that don't support 64-bit data.
- fp = kwargs.get('default_fp', ti.f32)
- ip = kwargs.get('default_ip', ti.i32)
- if fp == ti.f64 or ip == ti.i64:
- can_run_on.register(lambda arch: is_extension_supported(
- arch, extension.data64))
-
- for arch in ti.supported_archs():
- if can_run_on(arch):
- print('Running test on arch={}'.format(arch))
- ti.init(arch=arch, **kwargs)
- test(*test_args, **test_kwargs)
- else:
- print('Skipped test on arch={}'.format(arch))
-
- return wrapped
-
- return decorator
-
-
-# test with all archs
-def all_archs(test):
- return all_archs_with()(test)
-
-
-# Exclude the given archs when running the tests
-#
-# Example usage:
-#
-# @ti.archs_excluding(ti.cuda, ti.metal)
-# def test_xx():
-# ...
-#
-# @ti.archs_excluding(ti.cuda, default_fp=ti.f64)
-# def test_yy():
-# ...
-def archs_excluding(*excluded_archs, **kwargs):
- # |kwargs| will be passed to all_archs_with(**kwargs)
- assert all([isinstance(a, _ti_core.Arch) for a in excluded_archs])
- excluded_archs = set(excluded_archs)
-
- def decorator(test):
- @functools.wraps(test)
- def wrapped(*test_args, **test_kwargs):
- def checker(arch):
- return arch not in excluded_archs
-
- _get_or_make_arch_checkers(test_kwargs).register(checker)
- return all_archs_with(**kwargs)(test)(*test_args, **test_kwargs)
-
- return wrapped
-
- return decorator
-
-
-# Specifies the extension features the archs are required to support in order
-# to run the test.
-#
-# Example usage:
-#
-# @ti.require(ti.extension.data64)
-# @ti.all_archs_with(default_fp=ti.f64)
-# def test_xx():
-# ...
-def require(*exts):
- # Because this decorator injects an arch checker, its usage must be followed
- # with all_archs_with(), either directly or indirectly.
- assert all([isinstance(e, _ti_core.Extension) for e in exts])
-
- def decorator(test):
- @functools.wraps(test)
- def wrapped(*test_args, **test_kwargs):
- def checker(arch):
- return all([is_extension_supported(arch, e) for e in exts])
-
- _get_or_make_arch_checkers(test_kwargs).register(checker)
- test(*test_args, **test_kwargs)
-
- return wrapped
-
- return decorator
-
-
-def archs_support_sparse(test, **kwargs):
- wrapped = all_archs_with(**kwargs)(test)
- return require(extension.sparse)(wrapped)
-
-
-def torch_test(func):
- if ti.has_pytorch():
- # OpenGL somehow crashes torch test without a reason, unforturnately
- return ti.test(exclude=[opengl])(func)
- else:
- return lambda: None
-
-
-def get_host_arch_list():
- return [_ti_core.host_arch()]
-
-
-# test with host arch only
-def host_arch_only(func):
- @functools.wraps(func)
- def test(*args, **kwargs):
- archs = [_ti_core.host_arch()]
- for arch in archs:
- ti.init(arch=arch)
- func(*args, **kwargs)
-
- return test
-
-
-def archs_with(archs, **init_kwags):
- """
- Run the test on the given archs with the given init args.
-
- Args:
- archs: a list of Taichi archs
- init_kwargs: kwargs passed to ti.init()
- """
- def decorator(test):
- @functools.wraps(test)
- def wrapped(*test_args, **test_kwargs):
- for arch in archs:
- ti.init(arch=arch, **init_kwags)
- test(*test_args, **test_kwargs)
-
- return wrapped
-
- return decorator
-
-
-def must_throw(ex):
- def decorator(func):
- def func__(*args, **kwargs):
- finishes = False
- try:
- func(*args, **kwargs)
- finishes = True
- except ex:
- # throws. test passed
- pass
- except Exception as err_actual:
- assert False, 'Exception {} instead of {} thrown'.format(
- str(type(err_actual)), str(ex))
- if finishes:
- assert False, 'Test successfully finished instead of throwing {}'.format(
- str(ex))
-
- return func__
-
- return decorator
-
-
-__all__ = [s for s in dir() if not s.startswith('_')]
+from taichi.lang.kernel_impl import *
+from taichi.lang.matrix import *
+from taichi.lang.mesh import *
+from taichi.lang.misc import * # pylint: disable=W0622
+from taichi.lang.ops import * # pylint: disable=W0622
+from taichi.lang.runtime_ops import *
+from taichi.lang.snode import *
+from taichi.lang.source_builder import *
+from taichi.lang.struct import *
+from taichi.types.annotations import any_arr, ext_arr, template
+from taichi.types.primitive_types import f16, f32, f64, i32, i64, u32, u64
+
+from taichi import _logging, _snode
+
+__all__ = [
+ s for s in dir() if not s.startswith('_') and s not in [
+ 'any_array', 'ast', 'common_ops', 'enums', 'exception', 'expr', 'impl',
+ 'inspect', 'kernel_arguments', 'kernel_impl', 'matrix', 'mesh', 'misc',
+ 'ops', 'platform', 'runtime_ops', 'shell', 'snode', 'source_builder',
+ 'struct', 'tape', 'util'
+ ]
+]
diff --git a/python/taichi/lang/_ndarray.py b/python/taichi/lang/_ndarray.py
index f0a59a564cdbb..a1bd856dabf39 100644
--- a/python/taichi/lang/_ndarray.py
+++ b/python/taichi/lang/_ndarray.py
@@ -1,48 +1,35 @@
import numpy as np
-from taichi.core.util import ti_core as _ti_core
+from taichi._lib import core as _ti_core
from taichi.lang import impl
from taichi.lang.enums import Layout
-from taichi.lang.util import (cook_dtype, has_pytorch, python_scope,
- to_pytorch_type, to_taichi_type)
-
-if has_pytorch():
- import torch
+from taichi.lang.util import cook_dtype, python_scope, to_numpy_type
+from taichi.types import primitive_types
class Ndarray:
- """Taichi ndarray class implemented with a torch tensor.
+ """Taichi ndarray class.
Args:
dtype (DataType): Data type of each value.
- shape (Tuple[int]): Shape of the torch tensor.
+ shape (Tuple[int]): Shape of the Ndarray.
"""
- def __init__(self, dtype, shape):
- assert has_pytorch(
- ), "PyTorch must be available if you want to create a Taichi ndarray."
- self.arr = torch.zeros(shape, dtype=to_pytorch_type(cook_dtype(dtype)))
- if impl.current_cfg().arch == _ti_core.Arch.cuda:
- self.arr = self.arr.cuda()
+ def __init__(self, dtype, arr_shape):
+ self.host_accessor = None
+ self.dtype = cook_dtype(dtype)
+ self.arr = _ti_core.Ndarray(impl.get_runtime().prog, cook_dtype(dtype),
+ arr_shape)
@property
- def shape(self):
- """Gets ndarray shape.
+ def element_shape(self):
+ """Gets ndarray element shape.
Returns:
- Tuple[Int]: Ndarray shape.
+ Tuple[Int]: Ndarray element shape.
"""
raise NotImplementedError()
@property
- def dtype(self):
- """Gets data type of each individual value.
-
- Returns:
- DataType: Data type of each individual value.
- """
- return to_taichi_type(self.arr.dtype)
-
- @property
- def data_handle(self):
+ def _data_handle(self):
"""Gets the pointer to underlying data.
Returns:
@@ -79,19 +66,44 @@ def fill(self, val):
Args:
val (Union[int, float]): Value to fill.
"""
- self.arr.fill_(val)
+ if impl.current_cfg().arch != _ti_core.Arch.cuda and impl.current_cfg(
+ ).arch != _ti_core.Arch.x64:
+ self._fill_by_kernel(val)
+ elif self.dtype == primitive_types.f32:
+ self.arr.fill_float(val)
+ elif self.dtype == primitive_types.i32:
+ self.arr.fill_int(val)
+ elif self.dtype == primitive_types.u32:
+ self.arr.fill_uint(val)
+ else:
+ self._fill_by_kernel(val)
- @python_scope
- def to_numpy(self):
+ def _ndarray_to_numpy(self):
"""Converts ndarray to a numpy array.
Returns:
numpy.ndarray: The result numpy array.
"""
- return self.arr.cpu().numpy()
+ arr = np.zeros(shape=self.arr.shape, dtype=to_numpy_type(self.dtype))
+ from taichi._kernels import ndarray_to_ext_arr # pylint: disable=C0415
+ ndarray_to_ext_arr(self, arr)
+ impl.get_runtime().sync()
+ return arr
- @python_scope
- def from_numpy(self, arr):
+ def _ndarray_matrix_to_numpy(self, as_vector):
+ """Converts matrix ndarray to a numpy array.
+
+ Returns:
+ numpy.ndarray: The result numpy array.
+ """
+ arr = np.zeros(shape=self.arr.shape, dtype=to_numpy_type(self.dtype))
+ from taichi._kernels import \
+ ndarray_matrix_to_ext_arr # pylint: disable=C0415
+ ndarray_matrix_to_ext_arr(self, arr, as_vector)
+ impl.get_runtime().sync()
+ return arr
+
+ def _ndarray_from_numpy(self, arr):
"""Loads all values from a numpy array.
Args:
@@ -103,35 +115,169 @@ def from_numpy(self, arr):
raise ValueError(
f"Mismatch shape: {tuple(self.arr.shape)} expected, but {tuple(arr.shape)} provided"
)
- self.arr = torch.from_numpy(arr).to(self.arr.dtype)
+ if hasattr(arr, 'contiguous'):
+ arr = arr.contiguous()
+
+ from taichi._kernels import ext_arr_to_ndarray # pylint: disable=C0415
+ ext_arr_to_ndarray(arr, self)
+ impl.get_runtime().sync()
+
+ def _ndarray_matrix_from_numpy(self, arr, as_vector):
+ """Loads all values from a numpy array.
+
+ Args:
+ arr (numpy.ndarray): The source numpy array.
+ """
+ if not isinstance(arr, np.ndarray):
+ raise TypeError(f"{np.ndarray} expected, but {type(arr)} provided")
+ if tuple(self.arr.shape) != tuple(arr.shape):
+ raise ValueError(
+ f"Mismatch shape: {tuple(self.arr.shape)} expected, but {tuple(arr.shape)} provided"
+ )
+ if hasattr(arr, 'contiguous'):
+ arr = arr.contiguous()
+
+ from taichi._kernels import \
+ ext_arr_to_ndarray_matrix # pylint: disable=C0415
+ ext_arr_to_ndarray_matrix(arr, self, as_vector)
+ impl.get_runtime().sync()
+
+ @python_scope
+ def _get_element_size(self):
+ """Returns the size of one element in bytes.
+
+ Returns:
+ Size in bytes.
+ """
+ return self.arr.element_size()
+
+ @python_scope
+ def _get_nelement(self):
+ """Returns the total number of elements.
+
+ Returns:
+ Total number of elements.
+ """
+ return self.arr.nelement()
+
+ @python_scope
+ def copy_from(self, other):
+ """Copies all elements from another ndarray.
+
+ The shape of the other ndarray needs to be the same as `self`.
+
+ Args:
+ other (Ndarray): The source ndarray.
+ """
+ assert isinstance(other, Ndarray)
+ assert tuple(self.arr.shape) == tuple(other.arr.shape)
+ from taichi._kernels import ndarray_to_ndarray # pylint: disable=C0415
+ ndarray_to_ndarray(self, other)
+ impl.get_runtime().sync()
+
+ def __deepcopy__(self, memo=None):
+ """Copies all elements to a new ndarray.
+
+ Returns:
+ Ndarray: The result ndarray.
+ """
+ raise NotImplementedError()
+
+ def _fill_by_kernel(self, val):
+ """Fills ndarray with a specific scalar value using a ti.kernel.
+
+ Args:
+ val (Union[int, float]): Value to fill.
+ """
+ raise NotImplementedError()
+
+ def _pad_key(self, key):
+ if key is None:
+ key = ()
+ if not isinstance(key, (tuple, list)):
+ key = (key, )
+ assert len(key) == len(self.arr.shape)
+ return key
+
+ def _initialize_host_accessor(self):
+ if self.host_accessor:
+ return
+ impl.get_runtime().materialize()
+ self.host_accessor = NdarrayHostAccessor(self.arr)
class ScalarNdarray(Ndarray):
- """Taichi ndarray with scalar elements implemented with a torch tensor.
+ """Taichi ndarray with scalar elements.
Args:
dtype (DataType): Data type of each value.
shape (Tuple[int]): Shape of the ndarray.
"""
- def __init__(self, dtype, shape):
- super().__init__(dtype, shape)
+ def __init__(self, dtype, arr_shape):
+ super().__init__(dtype, arr_shape)
+ self.shape = tuple(self.arr.shape)
@property
- def shape(self):
- return tuple(self.arr.shape)
+ def element_shape(self):
+ return ()
@python_scope
def __setitem__(self, key, value):
- self.arr.__setitem__(key, value)
+ self._initialize_host_accessor()
+ self.host_accessor.setter(value, *self._pad_key(key))
@python_scope
def __getitem__(self, key):
- return self.arr.__getitem__(key)
+ self._initialize_host_accessor()
+ return self.host_accessor.getter(*self._pad_key(key))
+
+ @python_scope
+ def to_numpy(self):
+ return self._ndarray_to_numpy()
+
+ @python_scope
+ def from_numpy(self, arr):
+ self._ndarray_from_numpy(arr)
+
+ def __deepcopy__(self, memo=None):
+ ret_arr = ScalarNdarray(self.dtype, self.shape)
+ ret_arr.copy_from(self)
+ return ret_arr
+
+ def _fill_by_kernel(self, val):
+ from taichi._kernels import fill_ndarray # pylint: disable=C0415
+ fill_ndarray(self, val)
def __repr__(self):
return ''
+class NdarrayHostAccessor:
+ def __init__(self, ndarray):
+ if _ti_core.is_real(ndarray.dtype):
+
+ def getter(*key):
+ return ndarray.read_float(key)
+
+ def setter(value, *key):
+ ndarray.write_float(key, value)
+ else:
+ if _ti_core.is_signed(ndarray.dtype):
+
+ def getter(*key):
+ return ndarray.read_int(key)
+ else:
+
+ def getter(*key):
+ return ndarray.read_uint(key)
+
+ def setter(value, *key):
+ ndarray.write_int(key, value)
+
+ self.getter = getter
+ self.setter = setter
+
+
class NdarrayHostAccess:
"""Class for accessing VectorNdarray/MatrixNdarray in Python scope.
Args:
@@ -140,14 +286,25 @@ class NdarrayHostAccess:
indices_second (Tuple[Int]): Indices of second-level access (indices in the vector/matrix).
"""
def __init__(self, arr, indices_first, indices_second):
+ self.ndarr = arr
self.arr = arr.arr
if arr.layout == Layout.SOA:
self.indices = indices_second + indices_first
else:
self.indices = indices_first + indices_second
- def getter(self):
- return self.arr[self.indices]
+ def getter():
+ self.ndarr._initialize_host_accessor()
+ return self.ndarr.host_accessor.getter(
+ *self.ndarr._pad_key(self.indices))
+
+ def setter(value):
+ self.ndarr._initialize_host_accessor()
+ self.ndarr.host_accessor.setter(value,
+ *self.ndarr._pad_key(self.indices))
+
+ self.getter = getter
+ self.setter = setter
+
- def setter(self, value):
- self.arr[self.indices] = value
+__all__ = ["Ndarray", "ScalarNdarray"]
diff --git a/python/taichi/lang/ndrange.py b/python/taichi/lang/_ndrange.py
similarity index 57%
rename from python/taichi/lang/ndrange.py
rename to python/taichi/lang/_ndrange.py
index 56b9cb445b63b..a4e8ec85d6de7 100644
--- a/python/taichi/lang/ndrange.py
+++ b/python/taichi/lang/_ndrange.py
@@ -1,20 +1,24 @@
-import taichi as ti
+import collections.abc
+from taichi.lang.exception import TaichiSyntaxError
+from taichi.lang.matrix import _IntermediateMatrix
-class ndrange:
+
+class _Ndrange:
def __init__(self, *args):
args = list(args)
- for i in range(len(args)):
- if isinstance(args[i], list):
- args[i] = tuple(args[i])
- if not isinstance(args[i], tuple):
- args[i] = (0, args[i])
- assert len(args[i]) == 2
+ for i, arg in enumerate(args):
+ if not isinstance(arg, collections.abc.Sequence):
+ args[i] = (0, arg)
+ if len(args[i]) != 2:
+ raise TaichiSyntaxError(
+ "Every argument of ndrange should be a scalar or a tuple/list like (begin, end)"
+ )
self.bounds = args
self.dimensions = [None] * len(args)
- for i in range(len(self.bounds)):
- self.dimensions[i] = self.bounds[i][1] - self.bounds[i][0]
+ for i, bound in enumerate(self.bounds):
+ self.dimensions[i] = bound[1] - bound[0]
self.acc_dimensions = self.dimensions.copy()
for i in reversed(range(len(self.bounds) - 1)):
@@ -38,10 +42,17 @@ def grouped(self):
return GroupedNDRange(self)
+def ndrange(*args):
+ return _Ndrange(*args)
+
+
class GroupedNDRange:
def __init__(self, r):
self.r = r
def __iter__(self):
for ind in self.r:
- yield ti.Vector(list(ind), keep_raw=True)
+ yield _IntermediateMatrix(len(ind), 1, list(ind))
+
+
+__all__ = ['ndrange']
diff --git a/python/taichi/lang/_random.py b/python/taichi/lang/_random.py
deleted file mode 100644
index af92b772083d8..0000000000000
--- a/python/taichi/lang/_random.py
+++ /dev/null
@@ -1,19 +0,0 @@
-import math
-
-from taichi.lang.kernel_impl import func
-
-import taichi as ti
-
-
-@func
-def randn(dt):
- '''
- Generates a random number from standard normal distribution
- using the Box-Muller transform.
- '''
- assert dt == ti.f32 or dt == ti.f64
- u1 = ti.random(dt)
- u2 = ti.random(dt)
- r = ti.sqrt(-2 * ti.log(u1))
- c = ti.cos(math.tau * u2)
- return r * c
diff --git a/python/taichi/lang/any_array.py b/python/taichi/lang/any_array.py
index 1ad0ff8abaebb..5e126f25d9023 100644
--- a/python/taichi/lang/any_array.py
+++ b/python/taichi/lang/any_array.py
@@ -1,4 +1,4 @@
-from taichi.core.util import ti_core as _ti_core
+from taichi._lib import core as _ti_core
from taichi.lang.enums import Layout
from taichi.lang.expr import Expr, make_expr_group
from taichi.lang.util import taichi_scope
@@ -34,13 +34,11 @@ def shape(self):
element_dim = len(self.element_shape)
if element_dim == 0:
return ret
- else:
- return ret[
- element_dim:] if self.layout == Layout.SOA else ret[:
- -element_dim]
+ return ret[
+ element_dim:] if self.layout == Layout.SOA else ret[:-element_dim]
@taichi_scope
- def loop_range(self):
+ def _loop_range(self):
"""Gets the corresponding taichi_core.Expr to serve as loop range.
This is not in use now because struct fors on AnyArrays are not supported yet.
@@ -71,3 +69,6 @@ def subscript(self, i, j):
indices = self.indices_first + indices_second
return Expr(_ti_core.subscript(self.arr.ptr,
make_expr_group(*indices)))
+
+
+__all__ = []
diff --git a/python/taichi/lang/ast/__init__.py b/python/taichi/lang/ast/__init__.py
index 8b137891791fe..5787470b3631b 100644
--- a/python/taichi/lang/ast/__init__.py
+++ b/python/taichi/lang/ast/__init__.py
@@ -1 +1,3 @@
-
+from taichi.lang.ast.ast_transformer_utils import ASTTransformerContext
+from taichi.lang.ast.checkers import KernelSimplicityASTChecker
+from taichi.lang.ast.transform import transform_tree
diff --git a/python/taichi/lang/ast/ast_transformer.py b/python/taichi/lang/ast/ast_transformer.py
new file mode 100644
index 0000000000000..002cc2de76780
--- /dev/null
+++ b/python/taichi/lang/ast/ast_transformer.py
@@ -0,0 +1,1238 @@
+import ast
+import collections.abc
+import itertools
+import warnings
+from collections import ChainMap
+from sys import version_info
+
+import astor
+from taichi._lib import core as _ti_core
+from taichi.lang import expr, impl, kernel_arguments, matrix, mesh
+from taichi.lang import ops as ti_ops
+from taichi.lang._ndrange import _Ndrange, ndrange
+from taichi.lang.ast.ast_transformer_utils import Builder, LoopStatus
+from taichi.lang.ast.symbol_resolver import ASTResolver
+from taichi.lang.exception import TaichiSyntaxError
+from taichi.lang.matrix import MatrixType
+from taichi.lang.util import is_taichi_class, to_taichi_type
+from taichi.types import annotations, primitive_types
+
+if version_info < (3, 9):
+ from astunparse import unparse
+else:
+ from ast import unparse
+
+
+class ASTTransformer(Builder):
+ @staticmethod
+ def build_Name(ctx, node):
+ node.ptr = ctx.get_var_by_name(node.id)
+ return node.ptr
+
+ @staticmethod
+ def build_AnnAssign(ctx, node):
+ build_stmt(ctx, node.value)
+ build_stmt(ctx, node.annotation)
+
+ is_static_assign = isinstance(
+ node.value, ast.Call) and node.value.func.ptr is impl.static
+
+ node.ptr = ASTTransformer.build_assign_annotated(
+ ctx, node.target, node.value.ptr, is_static_assign,
+ node.annotation.ptr)
+ return node.ptr
+
+ @staticmethod
+ def build_assign_annotated(ctx, target, value, is_static_assign,
+ annotation):
+ """Build an annotated assginment like this: target: annotation = value.
+
+ Args:
+ ctx (ast_builder_utils.BuilderContext): The builder context.
+ target (ast.Name): A variable name. `target.id` holds the name as
+ a string.
+ annotation: A type we hope to assign to the target
+ value: A node representing the value.
+ is_static_assign: A boolean value indicating whether this is a static assignment
+ """
+ is_local = isinstance(target, ast.Name)
+ anno = impl.expr_init(annotation)
+ if is_static_assign:
+ raise TaichiSyntaxError(
+ "Static assign cannot be used on annotated assignment")
+ if is_local and not ctx.is_var_declared(target.id):
+ var = ti_ops.cast(value, anno)
+ var = impl.expr_init(var)
+ ctx.create_variable(target.id, var)
+ else:
+ var = build_stmt(ctx, target)
+ if var.ptr.get_ret_type() != anno:
+ raise TaichiSyntaxError(
+ "Static assign cannot have type overloading")
+ var._assign(value)
+ return var
+
+ @staticmethod
+ def build_Assign(ctx, node):
+ build_stmt(ctx, node.value)
+
+ is_static_assign = isinstance(
+ node.value, ast.Call) and node.value.func.ptr is impl.static
+
+ # Keep all generated assign statements and compose single one at last.
+ # The variable is introduced to support chained assignments.
+ # Ref https://github.com/taichi-dev/taichi/issues/2659.
+ for node_target in node.targets:
+ ASTTransformer.build_assign_unpack(ctx, node_target,
+ node.value.ptr,
+ is_static_assign)
+ return None
+
+ @staticmethod
+ def build_assign_unpack(ctx, node_target, values, is_static_assign):
+ """Build the unpack assignments like this: (target1, target2) = (value1, value2).
+ The function should be called only if the node target is a tuple.
+
+ Args:
+ ctx (ast_builder_utils.BuilderContext): The builder context.
+ node_target (ast.Tuple): A list or tuple object. `node_target.elts` holds a
+ list of nodes representing the elements.
+ values: A node/list representing the values.
+ is_static_assign: A boolean value indicating whether this is a static assignment
+ """
+ if not isinstance(node_target, ast.Tuple):
+ return ASTTransformer.build_assign_basic(ctx, node_target, values,
+ is_static_assign)
+ targets = node_target.elts
+ tmp_tuple = values if is_static_assign else impl.expr_init_list(
+ values, len(targets))
+
+ for i, target in enumerate(targets):
+ ASTTransformer.build_assign_basic(ctx, target, tmp_tuple[i],
+ is_static_assign)
+
+ return None
+
+ @staticmethod
+ def build_assign_basic(ctx, target, value, is_static_assign):
+ """Build basic assginment like this: target = value.
+
+ Args:
+ ctx (ast_builder_utils.BuilderContext): The builder context.
+ target (ast.Name): A variable name. `target.id` holds the name as
+ a string.
+ value: A node representing the value.
+ is_static_assign: A boolean value indicating whether this is a static assignment
+ """
+ is_local = isinstance(target, ast.Name)
+ if is_static_assign:
+ if not is_local:
+ raise TaichiSyntaxError(
+ "Static assign cannot be used on elements in arrays")
+ ctx.create_variable(target.id, value)
+ var = value
+ elif is_local and not ctx.is_var_declared(target.id):
+ var = impl.expr_init(value)
+ ctx.create_variable(target.id, var)
+ else:
+ var = build_stmt(ctx, target)
+ try:
+ var._assign(value)
+ except AttributeError:
+ raise TaichiSyntaxError(
+ f"Variable '{unparse(target).strip()}' cannot be assigned. Maybe it is not a Taichi object?"
+ )
+ return var
+
+ @staticmethod
+ def build_NamedExpr(ctx, node):
+ build_stmt(ctx, node.value)
+ is_static_assign = isinstance(
+ node.value, ast.Call) and node.value.func.ptr is impl.static
+ node.ptr = ASTTransformer.build_assign_basic(ctx, node.target,
+ node.value.ptr,
+ is_static_assign)
+ return node.ptr
+
+ @staticmethod
+ def is_tuple(node):
+ if isinstance(node, ast.Tuple):
+ return True
+ if isinstance(node, ast.Index) and isinstance(node.value.ptr, tuple):
+ return True
+ if isinstance(node.ptr, tuple):
+ return True
+ return False
+
+ @staticmethod
+ def build_Subscript(ctx, node):
+ build_stmt(ctx, node.value)
+ build_stmt(ctx, node.slice)
+ if not ASTTransformer.is_tuple(node.slice):
+ node.slice.ptr = [node.slice.ptr]
+ node.ptr = impl.subscript(node.value.ptr, *node.slice.ptr)
+ return node.ptr
+
+ @staticmethod
+ def build_Slice(ctx, node):
+ if node.lower is not None:
+ build_stmt(ctx, node.lower)
+ if node.upper is not None:
+ build_stmt(ctx, node.upper)
+ if node.step is not None:
+ build_stmt(ctx, node.step)
+
+ node.ptr = slice(node.lower.ptr if node.lower else None,
+ node.upper.ptr if node.upper else None,
+ node.step.ptr if node.step else None)
+ return node.ptr
+
+ @staticmethod
+ def build_ExtSlice(ctx, node):
+ build_stmts(ctx, node.dims)
+ node.ptr = tuple(dim.ptr for dim in node.dims)
+ return node.ptr
+
+ @staticmethod
+ def build_Tuple(ctx, node):
+ build_stmts(ctx, node.elts)
+ node.ptr = tuple(elt.ptr for elt in node.elts)
+ return node.ptr
+
+ @staticmethod
+ def build_List(ctx, node):
+ build_stmts(ctx, node.elts)
+ node.ptr = [elt.ptr for elt in node.elts]
+ return node.ptr
+
+ @staticmethod
+ def build_Dict(ctx, node):
+ dic = {}
+ for key, value in zip(node.keys, node.values):
+ if key is None:
+ dic.update(build_stmt(ctx, value))
+ else:
+ dic[build_stmt(ctx, key)] = build_stmt(ctx, value)
+ node.ptr = dic
+ return node.ptr
+
+ @staticmethod
+ def process_listcomp(ctx, node, result):
+ result.append(build_stmt(ctx, node.elt))
+
+ @staticmethod
+ def process_dictcomp(ctx, node, result):
+ key = build_stmt(ctx, node.key)
+ value = build_stmt(ctx, node.value)
+ result[key] = value
+
+ @staticmethod
+ def process_generators(ctx, node, now_comp, func, result):
+ if now_comp >= len(node.generators):
+ return func(ctx, node, result)
+ with ctx.static_scope_guard():
+ _iter = build_stmt(ctx, node.generators[now_comp].iter)
+ for value in _iter:
+ with ctx.variable_scope_guard():
+ ASTTransformer.build_assign_unpack(
+ ctx, node.generators[now_comp].target, value, True)
+ with ctx.static_scope_guard():
+ build_stmts(ctx, node.generators[now_comp].ifs)
+ ASTTransformer.process_ifs(ctx, node, now_comp, 0, func,
+ result)
+ return None
+
+ @staticmethod
+ def process_ifs(ctx, node, now_comp, now_if, func, result):
+ if now_if >= len(node.generators[now_comp].ifs):
+ return ASTTransformer.process_generators(ctx, node, now_comp + 1,
+ func, result)
+ cond = node.generators[now_comp].ifs[now_if].ptr
+ if cond:
+ ASTTransformer.process_ifs(ctx, node, now_comp, now_if + 1, func,
+ result)
+
+ return None
+
+ @staticmethod
+ def build_ListComp(ctx, node):
+ result = []
+ ASTTransformer.process_generators(ctx, node, 0,
+ ASTTransformer.process_listcomp,
+ result)
+ node.ptr = result
+ return node.ptr
+
+ @staticmethod
+ def build_DictComp(ctx, node):
+ result = {}
+ ASTTransformer.process_generators(ctx, node, 0,
+ ASTTransformer.process_dictcomp,
+ result)
+ node.ptr = result
+ return node.ptr
+
+ @staticmethod
+ def build_Index(ctx, node):
+
+ node.ptr = build_stmt(ctx, node.value)
+ return node.ptr
+
+ @staticmethod
+ def build_Constant(ctx, node):
+ node.ptr = node.value
+ return node.ptr
+
+ @staticmethod
+ def build_Num(ctx, node):
+ node.ptr = node.n
+ return node.ptr
+
+ @staticmethod
+ def build_Str(ctx, node):
+ node.ptr = node.s
+ return node.ptr
+
+ @staticmethod
+ def build_Bytes(ctx, node):
+ node.ptr = node.s
+ return node.ptr
+
+ @staticmethod
+ def build_NameConstant(ctx, node):
+ node.ptr = node.value
+ return node.ptr
+
+ @staticmethod
+ def build_keyword(ctx, node):
+ build_stmt(ctx, node.value)
+ if node.arg is None:
+ node.ptr = node.value.ptr
+ else:
+ node.ptr = {node.arg: node.value.ptr}
+ return node.ptr
+
+ @staticmethod
+ def build_Starred(ctx, node):
+ node.ptr = build_stmt(ctx, node.value)
+ return node.ptr
+
+ @staticmethod
+ def build_JoinedStr(ctx, node):
+ str_spec = ''
+ args = []
+ for sub_node in node.values:
+ if isinstance(sub_node, ast.FormattedValue):
+ str_spec += '{}'
+ args.append(build_stmt(ctx, sub_node.value))
+ elif isinstance(sub_node, ast.Constant):
+ str_spec += sub_node.value
+ elif isinstance(sub_node, ast.Str):
+ str_spec += sub_node.s
+ else:
+ raise TaichiSyntaxError("Invalid value for fstring.")
+
+ args.insert(0, str_spec)
+ node.ptr = impl.ti_format(*args)
+ return node.ptr
+
+ @staticmethod
+ def build_call_if_is_builtin(ctx, node, args, keywords):
+ func = node.func.ptr
+ replace_func = {
+ id(print): impl.ti_print,
+ id(min): ti_ops.min,
+ id(max): ti_ops.max,
+ id(int): impl.ti_int,
+ id(float): impl.ti_float,
+ id(any): ti_ops.ti_any,
+ id(all): ti_ops.ti_all,
+ id(abs): abs,
+ id(pow): pow,
+ }
+ if id(func) in replace_func:
+ node.ptr = replace_func[id(func)](*args, **keywords)
+ if func is min or func is max:
+ name = "min" if func is min else "max"
+ warnings.warn_explicit(
+ f'Calling builtin function "{name}" in Taichi scope is deprecated. '
+ f'Please use "ti.{name}" instead.', DeprecationWarning,
+ ctx.file, node.lineno + ctx.lineno_offset)
+ return True
+ return False
+
+ @staticmethod
+ def build_call_if_is_type(ctx, node, args, keywords):
+ func = node.func.ptr
+ if id(func) in primitive_types.type_ids:
+ if len(args) != 1 or keywords or isinstance(args[0], expr.Expr):
+ raise TaichiSyntaxError(
+ "Type annotation can only be given to a single literal.")
+ node.ptr = expr.Expr(args[0], dtype=func)
+ return True
+ return False
+
+ @staticmethod
+ def warn_if_is_external_func(ctx, node):
+ func = node.func.ptr
+ if ctx.is_in_static_scope(): # allow external function in static scope
+ return
+ if hasattr(func, "_is_taichi_function") or hasattr(
+ func, "_is_wrapped_kernel"): # taichi func/kernel
+ return
+ if hasattr(
+ func, "__module__"
+ ) and func.__module__ and func.__module__.startswith("taichi."):
+ return
+ name = unparse(node.func).strip()
+ warnings.warn_explicit(
+ f'Calling non-taichi function "{name}". '
+ f'Scope inside the function is not processed by the Taichi AST transformer. '
+ f'The function may not work as expected. Proceed with caution! '
+ f'Maybe you can consider turning it into a @ti.func?', UserWarning,
+ ctx.file, node.lineno + ctx.lineno_offset)
+
+ @staticmethod
+ def build_Call(ctx, node):
+ if ASTTransformer.get_decorator(ctx, node) == 'static':
+ with ctx.static_scope_guard():
+ build_stmt(ctx, node.func)
+ build_stmts(ctx, node.args)
+ build_stmts(ctx, node.keywords)
+ else:
+ build_stmt(ctx, node.func)
+ build_stmts(ctx, node.args)
+ build_stmts(ctx, node.keywords)
+
+ args = []
+ for arg in node.args:
+ if isinstance(arg, ast.Starred):
+ for i in arg.ptr:
+ args.append(i)
+ else:
+ args.append(arg.ptr)
+ keywords = dict(ChainMap(*[keyword.ptr for keyword in node.keywords]))
+ func = node.func.ptr
+
+ 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)
+ node.ptr = impl.ti_format(*args, **keywords)
+ return node.ptr
+
+ if ASTTransformer.build_call_if_is_builtin(ctx, node, args, keywords):
+ return node.ptr
+
+ if ASTTransformer.build_call_if_is_type(ctx, node, args, keywords):
+ return node.ptr
+
+ node.ptr = func(*args, **keywords)
+ ASTTransformer.warn_if_is_external_func(ctx, node)
+
+ return node.ptr
+
+ @staticmethod
+ def build_FunctionDef(ctx, node):
+ if ctx.visited_funcdef:
+ raise TaichiSyntaxError(
+ f"Function definition is not allowed in 'ti.{'kernel' if ctx.is_kernel else 'func'}'."
+ )
+ ctx.visited_funcdef = True
+
+ args = node.args
+ assert args.vararg is None
+ assert args.kwonlyargs == []
+ assert args.kw_defaults == []
+ assert args.kwarg is None
+
+ def transform_as_kernel():
+ # Treat return type
+ if node.returns is not None:
+ kernel_arguments.decl_ret(ctx.func.return_type)
+
+ for i, arg in enumerate(args.args):
+ if isinstance(ctx.func.argument_annotations[i],
+ annotations.template):
+ ctx.create_variable(arg.arg, ctx.global_vars[arg.arg])
+ elif isinstance(ctx.func.argument_annotations[i],
+ annotations.sparse_matrix_builder):
+ ctx.create_variable(
+ arg.arg,
+ kernel_arguments.decl_sparse_matrix(
+ to_taichi_type(ctx.arg_features[i])))
+ elif isinstance(ctx.func.argument_annotations[i],
+ annotations.any_arr):
+ ctx.create_variable(
+ arg.arg,
+ kernel_arguments.decl_any_arr_arg(
+ to_taichi_type(ctx.arg_features[i][0]),
+ ctx.arg_features[i][1], ctx.arg_features[i][2],
+ ctx.arg_features[i][3]))
+ elif isinstance(ctx.func.argument_annotations[i], MatrixType):
+ ctx.create_variable(
+ arg.arg,
+ kernel_arguments.decl_matrix_arg(
+ ctx.func.argument_annotations[i]))
+ else:
+ ctx.global_vars[
+ arg.arg] = kernel_arguments.decl_scalar_arg(
+ ctx.func.argument_annotations[i])
+ # remove original args
+ node.args.args = []
+
+ if ctx.is_kernel: # ti.kernel
+ transform_as_kernel()
+
+ else: # ti.func
+ if ctx.is_real_function:
+ transform_as_kernel()
+ else:
+ len_args = len(args.args)
+ len_default = len(args.defaults)
+ len_provided = len(ctx.argument_data)
+ len_minimum = len_args - len_default
+ if len_args < len_provided or len_args - len_default > len_provided:
+ if len(args.defaults):
+ raise TaichiSyntaxError(
+ f"Function receives {len_minimum} to {len_args} argument(s) and {len_provided} provided."
+ )
+ else:
+ raise TaichiSyntaxError(
+ f"Function receives {len_args} argument(s) and {len_provided} provided."
+ )
+ # Transform as force-inlined func
+ default_start = len_provided - len_minimum
+ ctx.argument_data = list(ctx.argument_data)
+ for arg in args.defaults[default_start:]:
+ ctx.argument_data.append(build_stmt(ctx, arg))
+ assert len(args.args) == len(ctx.argument_data)
+ for i, (arg,
+ data) in enumerate(zip(args.args, ctx.argument_data)):
+ # Remove annotations because they are not used.
+ args.args[i].annotation = None
+ # Template arguments are passed by reference.
+ if isinstance(ctx.func.argument_annotations[i],
+ annotations.template):
+ ctx.create_variable(ctx.func.argument_names[i], data)
+ continue
+ # Create a copy for non-template arguments,
+ # so that they are passed by value.
+ ctx.create_variable(arg.arg, impl.expr_init_func(data))
+
+ with ctx.variable_scope_guard():
+ build_stmts(ctx, node.body)
+
+ return None
+
+ @staticmethod
+ def build_Return(ctx, node):
+ if not ctx.is_real_function:
+ if ctx.is_in_non_static_control_flow():
+ raise TaichiSyntaxError(
+ "Return inside non-static if/for is not supported")
+ build_stmt(ctx, node.value)
+ if ctx.is_kernel or ctx.is_real_function:
+ # TODO: check if it's at the end of a kernel, throw TaichiSyntaxError if not
+ if node.value is not None:
+ if ctx.func.return_type is None:
+ raise TaichiSyntaxError(
+ f'A {"kernel" if ctx.is_kernel else "function"} '
+ 'with a return value must be annotated '
+ 'with a return type, e.g. def func() -> ti.f32')
+ if id(ctx.func.return_type) in primitive_types.type_ids:
+ ctx.ast_builder.create_kernel_exprgroup_return(
+ expr.make_expr_group(
+ ti_ops.cast(expr.Expr(node.value.ptr),
+ ctx.func.return_type).ptr))
+ elif isinstance(ctx.func.return_type, MatrixType):
+ ctx.ast_builder.create_kernel_exprgroup_return(
+ expr.make_expr_group([
+ ti_ops.cast(exp, ctx.func.return_type.dtype)
+ for exp in itertools.chain.from_iterable(
+ node.value.ptr.to_list())
+ ]))
+ else:
+ raise TaichiSyntaxError(
+ "The return type is not supported now!")
+ # For args[0], it is an ast.Attribute, because it loads the
+ # attribute, |ptr|, of the expression |ret_expr|. Therefore we
+ # only need to replace the object part, i.e. args[0].value
+ else:
+ ctx.return_data = node.value.ptr
+ if not ctx.is_real_function:
+ ctx.returned = True
+ return None
+
+ @staticmethod
+ def build_Module(ctx, node):
+ with ctx.variable_scope_guard():
+ # Do NOT use |build_stmts| which inserts 'del' statements to the
+ # end and deletes parameters passed into the module
+ for stmt in node.body:
+ build_stmt(ctx, stmt)
+ return None
+
+ @staticmethod
+ def build_Attribute(ctx, node):
+ build_stmt(ctx, node.value)
+ node.ptr = getattr(node.value.ptr, node.attr)
+ return node.ptr
+
+ @staticmethod
+ def build_BinOp(ctx, node):
+ build_stmt(ctx, node.left)
+ build_stmt(ctx, node.right)
+ op = {
+ ast.Add: lambda l, r: l + r,
+ ast.Sub: lambda l, r: l - r,
+ ast.Mult: lambda l, r: l * r,
+ ast.Div: lambda l, r: l / r,
+ ast.FloorDiv: lambda l, r: l // r,
+ ast.Mod: lambda l, r: l % r,
+ ast.Pow: lambda l, r: l**r,
+ ast.LShift: lambda l, r: l << r,
+ ast.RShift: lambda l, r: l >> r,
+ ast.BitOr: lambda l, r: l | r,
+ ast.BitXor: lambda l, r: l ^ r,
+ ast.BitAnd: lambda l, r: l & r,
+ ast.MatMult: lambda l, r: l @ r,
+ }.get(type(node.op))
+ node.ptr = op(node.left.ptr, node.right.ptr)
+ return node.ptr
+
+ @staticmethod
+ def build_AugAssign(ctx, node):
+ build_stmt(ctx, node.target)
+ build_stmt(ctx, node.value)
+ node.ptr = node.target.ptr._augassign(node.value.ptr,
+ type(node.op).__name__)
+ return node.ptr
+
+ @staticmethod
+ def build_UnaryOp(ctx, node):
+ build_stmt(ctx, node.operand)
+ op = {
+ ast.UAdd: lambda l: l,
+ ast.USub: lambda l: -l,
+ ast.Not: ti_ops.logical_not,
+ ast.Invert: lambda l: ~l,
+ }.get(type(node.op))
+ node.ptr = op(node.operand.ptr)
+ return node.ptr
+
+ @staticmethod
+ def build_short_circuit_and(ast_builder, operands):
+ if len(operands) == 1:
+ return operands[0].ptr
+
+ val = impl.expr_init(None)
+ lhs = operands[0].ptr
+ impl.begin_frontend_if(ast_builder, lhs)
+
+ ast_builder.begin_frontend_if_true()
+ rhs = ASTTransformer.build_short_circuit_and(ast_builder, operands[1:])
+ val._assign(rhs)
+ ast_builder.pop_scope()
+
+ ast_builder.begin_frontend_if_false()
+ val._assign(0)
+ ast_builder.pop_scope()
+
+ return val
+
+ @staticmethod
+ def build_short_circuit_or(ast_builder, operands):
+ if len(operands) == 1:
+ return operands[0].ptr
+
+ val = impl.expr_init(None)
+ lhs = operands[0].ptr
+ impl.begin_frontend_if(ast_builder, lhs)
+
+ ast_builder.begin_frontend_if_true()
+ val._assign(1)
+ ast_builder.pop_scope()
+
+ ast_builder.begin_frontend_if_false()
+ rhs = ASTTransformer.build_short_circuit_or(ast_builder, operands[1:])
+ val._assign(rhs)
+ ast_builder.pop_scope()
+
+ return val
+
+ @staticmethod
+ def build_normal_bool_op(op):
+ def inner(ast_builder, operands):
+ result = op(operands[0].ptr, operands[1].ptr)
+ for i in range(2, len(operands)):
+ result = op(result, operands[i].ptr)
+ return result
+
+ return inner
+
+ @staticmethod
+ def build_static_short_circuit_and(ast_builder, operands):
+ for operand in operands:
+ if not operand.ptr:
+ return operand.ptr
+ return operands[-1].ptr
+
+ @staticmethod
+ def build_static_short_circuit_or(ast_builder, operands):
+ for operand in operands:
+ if operand.ptr:
+ return operand.ptr
+ return operands[-1].ptr
+
+ @staticmethod
+ def build_BoolOp(ctx, node):
+ build_stmts(ctx, node.values)
+ if ctx.is_in_static_scope():
+ ops = {
+ ast.And: ASTTransformer.build_static_short_circuit_and,
+ ast.Or: ASTTransformer.build_static_short_circuit_or,
+ }
+ elif impl.get_runtime().short_circuit_operators:
+ ops = {
+ ast.And: ASTTransformer.build_short_circuit_and,
+ ast.Or: ASTTransformer.build_short_circuit_or,
+ }
+ else:
+ ops = {
+ ast.And:
+ ASTTransformer.build_normal_bool_op(ti_ops.logical_and),
+ ast.Or: ASTTransformer.build_normal_bool_op(ti_ops.logical_or),
+ }
+ op = ops.get(type(node.op))
+ node.ptr = op(ctx.ast_builder, node.values)
+ return node.ptr
+
+ @staticmethod
+ def build_Compare(ctx, node):
+ build_stmt(ctx, node.left)
+ build_stmts(ctx, node.comparators)
+ ops = {
+ ast.Eq: lambda l, r: l == r,
+ ast.NotEq: lambda l, r: l != r,
+ ast.Lt: lambda l, r: l < r,
+ ast.LtE: lambda l, r: l <= r,
+ ast.Gt: lambda l, r: l > r,
+ ast.GtE: lambda l, r: l >= r,
+ }
+ ops_static = {
+ ast.In: lambda l, r: l in r,
+ ast.NotIn: lambda l, r: l not in r,
+ ast.Is: lambda l, r: l is r,
+ ast.IsNot: lambda l, r: l is not r,
+ }
+ if ctx.is_in_static_scope():
+ ops = {**ops, **ops_static}
+ operands = [node.left.ptr
+ ] + [comparator.ptr for comparator in node.comparators]
+ val = True
+ for i, node_op in enumerate(node.ops):
+ l = operands[i]
+ r = operands[i + 1]
+ op = ops.get(type(node_op))
+ if isinstance(node_op, (ast.Is, ast.IsNot)):
+ name = "is" if isinstance(node_op, ast.Is) else "is not"
+ warnings.warn_explicit(
+ f'Operator "{name}" in Taichi scope is deprecated. Please avoid using it.',
+ DeprecationWarning, ctx.file,
+ node.lineno + ctx.lineno_offset)
+ if op is None:
+ if type(node_op) in ops_static:
+ raise TaichiSyntaxError(
+ f'"{type(node_op).__name__}" is only supported inside `ti.static`.'
+ )
+ else:
+ raise TaichiSyntaxError(
+ f'"{type(node_op).__name__}" is not supported in Taichi kernels.'
+ )
+ val = ti_ops.logical_and(val, op(l, r))
+ node.ptr = val
+ return node.ptr
+
+ @staticmethod
+ def get_decorator(ctx, node):
+ if not isinstance(node, ast.Call):
+ return ''
+ for wanted, name in [
+ (impl.static, 'static'),
+ (impl.grouped, 'grouped'),
+ (ndrange, 'ndrange'),
+ ]:
+ if ASTResolver.resolve_to(node.func, wanted, ctx.global_vars):
+ return name
+ return ''
+
+ @staticmethod
+ def get_for_loop_targets(node):
+ """
+ Returns the list of indices of the for loop |node|.
+ See also: https://docs.python.org/3/library/ast.html#ast.For
+ """
+ if isinstance(node.target, ast.Name):
+ return [node.target.id]
+ assert isinstance(node.target, ast.Tuple)
+ return [name.id for name in node.target.elts]
+
+ @staticmethod
+ def build_static_for(ctx, node, is_grouped):
+ if is_grouped:
+ assert len(node.iter.args[0].args) == 1
+ ndrange_arg = build_stmt(ctx, node.iter.args[0].args[0])
+ if not isinstance(ndrange_arg, _Ndrange):
+ raise TaichiSyntaxError(
+ "Only 'ti.ndrange' is allowed in 'ti.static(ti.grouped(...))'."
+ )
+ targets = ASTTransformer.get_for_loop_targets(node)
+ if len(targets) != 1:
+ raise TaichiSyntaxError(
+ f"Group for should have 1 loop target, found {len(targets)}"
+ )
+ target = targets[0]
+ for value in impl.grouped(ndrange_arg):
+ with ctx.variable_scope_guard():
+ ctx.create_variable(target, value)
+ build_stmts(ctx, node.body)
+ status = ctx.loop_status()
+ if status == LoopStatus.Break:
+ break
+ elif status == LoopStatus.Continue:
+ ctx.set_loop_status(LoopStatus.Normal)
+ else:
+ build_stmt(ctx, node.iter)
+ targets = ASTTransformer.get_for_loop_targets(node)
+ for target_values in node.iter.ptr:
+ if not isinstance(
+ target_values,
+ collections.abc.Sequence) or len(targets) == 1:
+ target_values = [target_values]
+ with ctx.variable_scope_guard():
+ for target, target_value in zip(targets, target_values):
+ ctx.create_variable(target, target_value)
+ build_stmts(ctx, node.body)
+ status = ctx.loop_status()
+ if status == LoopStatus.Break:
+ break
+ elif status == LoopStatus.Continue:
+ ctx.set_loop_status(LoopStatus.Normal)
+ return None
+
+ @staticmethod
+ def build_range_for(ctx, node):
+ with ctx.variable_scope_guard():
+ loop_name = node.target.id
+ ctx.check_loop_var(loop_name)
+ loop_var = expr.Expr(_ti_core.make_id_expr(''))
+ ctx.create_variable(loop_name, loop_var)
+ if len(node.iter.args) not in [1, 2]:
+ raise TaichiSyntaxError(
+ f"Range should have 1 or 2 arguments, found {len(node.iter.args)}"
+ )
+ if len(node.iter.args) == 2:
+ begin = ti_ops.cast(
+ expr.Expr(build_stmt(ctx, node.iter.args[0])),
+ primitive_types.i32)
+ end = ti_ops.cast(
+ expr.Expr(build_stmt(ctx, node.iter.args[1])),
+ primitive_types.i32)
+ else:
+ begin = ti_ops.cast(expr.Expr(0), primitive_types.i32)
+ end = ti_ops.cast(
+ expr.Expr(build_stmt(ctx, node.iter.args[0])),
+ primitive_types.i32)
+ ctx.ast_builder.begin_frontend_range_for(loop_var.ptr, begin.ptr,
+ end.ptr)
+ build_stmts(ctx, node.body)
+ ctx.ast_builder.end_frontend_range_for()
+ return None
+
+ @staticmethod
+ def build_ndrange_for(ctx, node):
+ with ctx.variable_scope_guard():
+ ndrange_var = impl.expr_init(build_stmt(ctx, node.iter))
+ ndrange_begin = ti_ops.cast(expr.Expr(0), primitive_types.i32)
+ ndrange_end = ti_ops.cast(
+ expr.Expr(impl.subscript(ndrange_var.acc_dimensions, 0)),
+ primitive_types.i32)
+ ndrange_loop_var = expr.Expr(_ti_core.make_id_expr(''))
+ ctx.ast_builder.begin_frontend_range_for(ndrange_loop_var.ptr,
+ ndrange_begin.ptr,
+ ndrange_end.ptr)
+ I = impl.expr_init(ndrange_loop_var)
+ targets = ASTTransformer.get_for_loop_targets(node)
+ for i, target in enumerate(targets):
+ if i + 1 < len(targets):
+ target_tmp = impl.expr_init(
+ I // ndrange_var.acc_dimensions[i + 1])
+ else:
+ target_tmp = impl.expr_init(I)
+ ctx.create_variable(
+ target,
+ impl.expr_init(target_tmp + impl.subscript(
+ impl.subscript(ndrange_var.bounds, i), 0)))
+ if i + 1 < len(targets):
+ I._assign(I -
+ target_tmp * ndrange_var.acc_dimensions[i + 1])
+ build_stmts(ctx, node.body)
+ ctx.ast_builder.end_frontend_range_for()
+ return None
+
+ @staticmethod
+ def build_grouped_ndrange_for(ctx, node):
+ with ctx.variable_scope_guard():
+ ndrange_var = impl.expr_init(build_stmt(ctx, node.iter.args[0]))
+ ndrange_begin = ti_ops.cast(expr.Expr(0), primitive_types.i32)
+ ndrange_end = ti_ops.cast(
+ expr.Expr(impl.subscript(ndrange_var.acc_dimensions, 0)),
+ primitive_types.i32)
+ ndrange_loop_var = expr.Expr(_ti_core.make_id_expr(''))
+ ctx.ast_builder.begin_frontend_range_for(ndrange_loop_var.ptr,
+ ndrange_begin.ptr,
+ ndrange_end.ptr)
+
+ targets = ASTTransformer.get_for_loop_targets(node)
+ if len(targets) != 1:
+ raise TaichiSyntaxError(
+ f"Group for should have 1 loop target, found {len(targets)}"
+ )
+ target = targets[0]
+ target_var = impl.expr_init(
+ matrix.Vector([0] * len(ndrange_var.dimensions),
+ dt=primitive_types.i32))
+ ctx.create_variable(target, target_var)
+ I = impl.expr_init(ndrange_loop_var)
+ for i in range(len(ndrange_var.dimensions)):
+ if i + 1 < len(ndrange_var.dimensions):
+ target_tmp = I // ndrange_var.acc_dimensions[i + 1]
+ else:
+ target_tmp = I
+ impl.subscript(target_var, i)._assign(target_tmp +
+ ndrange_var.bounds[i][0])
+ if i + 1 < len(ndrange_var.dimensions):
+ I._assign(I -
+ target_tmp * ndrange_var.acc_dimensions[i + 1])
+ build_stmts(ctx, node.body)
+ ctx.ast_builder.end_frontend_range_for()
+ return None
+
+ @staticmethod
+ def build_struct_for(ctx, node, is_grouped):
+ # for i, j in x
+ # for I in ti.grouped(x)
+ targets = ASTTransformer.get_for_loop_targets(node)
+
+ for target in targets:
+ ctx.check_loop_var(target)
+
+ with ctx.variable_scope_guard():
+ if is_grouped:
+ if len(targets) != 1:
+ raise TaichiSyntaxError(
+ f"Group for should have 1 loop target, found {len(targets)}"
+ )
+ target = targets[0]
+ loop_var = build_stmt(ctx, node.iter)
+ loop_indices = expr.make_var_list(size=len(loop_var.shape))
+ expr_group = expr.make_expr_group(loop_indices)
+ impl.begin_frontend_struct_for(ctx.ast_builder, expr_group,
+ loop_var)
+ ctx.create_variable(
+ target, matrix.Vector(loop_indices,
+ dt=primitive_types.i32))
+ build_stmts(ctx, node.body)
+ ctx.ast_builder.end_frontend_struct_for()
+ else:
+ _vars = []
+ for name in targets:
+ var = expr.Expr(_ti_core.make_id_expr(""))
+ _vars.append(var)
+ ctx.create_variable(name, var)
+ loop_var = node.iter.ptr
+ expr_group = expr.make_expr_group(*_vars)
+ impl.begin_frontend_struct_for(ctx.ast_builder, expr_group,
+ loop_var)
+ build_stmts(ctx, node.body)
+ ctx.ast_builder.end_frontend_struct_for()
+ return None
+
+ @staticmethod
+ def build_mesh_for(ctx, node):
+ targets = ASTTransformer.get_for_loop_targets(node)
+ if len(targets) != 1:
+ raise TaichiSyntaxError(
+ "Mesh for should have 1 loop target, found {len(targets)}")
+ target = targets[0]
+
+ with ctx.variable_scope_guard():
+ var = expr.Expr(_ti_core.make_id_expr(""))
+ ctx.mesh = node.iter.ptr.mesh
+ assert isinstance(ctx.mesh, impl.MeshInstance)
+ mesh_idx = mesh.MeshElementFieldProxy(ctx.mesh,
+ node.iter.ptr._type, var.ptr)
+ ctx.create_variable(target, mesh_idx)
+ ctx.ast_builder.begin_frontend_mesh_for(mesh_idx.ptr,
+ ctx.mesh.mesh_ptr,
+ node.iter.ptr._type)
+ build_stmts(ctx, node.body)
+ ctx.mesh = None
+ ctx.ast_builder.end_frontend_mesh_for()
+ return None
+
+ @staticmethod
+ def build_nested_mesh_for(ctx, node):
+ targets = ASTTransformer.get_for_loop_targets(node)
+ if len(targets) != 1:
+ raise TaichiSyntaxError(
+ "Nested-mesh for should have 1 loop target, found {len(targets)}"
+ )
+ target = targets[0]
+
+ with ctx.variable_scope_guard():
+ ctx.mesh = node.iter.ptr.mesh
+ assert isinstance(ctx.mesh, impl.MeshInstance)
+ loop_name = node.target.id + '_index__'
+ loop_var = expr.Expr(_ti_core.make_id_expr(''))
+ ctx.create_variable(loop_name, loop_var)
+ begin = expr.Expr(0)
+ end = node.iter.ptr.size
+ ctx.ast_builder.begin_frontend_range_for(loop_var.ptr, begin.ptr,
+ end.ptr)
+ entry_expr = _ti_core.get_relation_access(
+ ctx.mesh.mesh_ptr, node.iter.ptr.from_index.ptr,
+ node.iter.ptr.to_element_type, loop_var.ptr)
+ entry_expr.type_check(impl.get_runtime().prog.config)
+ mesh_idx = mesh.MeshElementFieldProxy(
+ ctx.mesh, node.iter.ptr.to_element_type, entry_expr)
+ ctx.create_variable(target, mesh_idx)
+ build_stmts(ctx, node.body)
+ ctx.ast_builder.end_frontend_range_for()
+
+ return None
+
+ @staticmethod
+ def build_For(ctx, node):
+ if node.orelse:
+ raise TaichiSyntaxError(
+ "'else' clause for 'for' not supported in Taichi kernels")
+ decorator = ASTTransformer.get_decorator(ctx, node.iter)
+ double_decorator = ''
+ if decorator != '' and len(node.iter.args) == 1:
+ double_decorator = ASTTransformer.get_decorator(
+ ctx, node.iter.args[0])
+
+ if decorator == 'static':
+ if double_decorator == 'static':
+ raise TaichiSyntaxError("'ti.static' cannot be nested")
+ with ctx.loop_scope_guard(is_static=True):
+ return ASTTransformer.build_static_for(
+ ctx, node, double_decorator == 'grouped')
+ with ctx.loop_scope_guard():
+ if decorator == 'ndrange':
+ if double_decorator != '':
+ raise TaichiSyntaxError(
+ "No decorator is allowed inside 'ti.ndrange")
+ return ASTTransformer.build_ndrange_for(ctx, node)
+ if decorator == 'grouped':
+ if double_decorator == 'static':
+ raise TaichiSyntaxError(
+ "'ti.static' is not allowed inside 'ti.grouped'")
+ elif double_decorator == 'ndrange':
+ return ASTTransformer.build_grouped_ndrange_for(ctx, node)
+ elif double_decorator == 'grouped':
+ raise TaichiSyntaxError("'ti.grouped' cannot be nested")
+ else:
+ return ASTTransformer.build_struct_for(ctx,
+ node,
+ is_grouped=True)
+ elif isinstance(node.iter, ast.Call) and isinstance(
+ node.iter.func, ast.Name) and node.iter.func.id == 'range':
+ return ASTTransformer.build_range_for(ctx, node)
+ else:
+ build_stmt(ctx, node.iter)
+ if isinstance(node.iter.ptr, mesh.MeshElementField):
+ if not _ti_core.is_extension_supported(
+ impl.default_cfg().arch, _ti_core.Extension.mesh):
+ raise Exception(
+ 'Backend ' + str(impl.default_cfg().arch) +
+ ' doesn\'t support MeshTaichi extension')
+ return ASTTransformer.build_mesh_for(ctx, node)
+ if isinstance(node.iter.ptr, mesh.MeshRelationAccessProxy):
+ return ASTTransformer.build_nested_mesh_for(ctx, node)
+ # Struct for
+ return ASTTransformer.build_struct_for(ctx,
+ node,
+ is_grouped=False)
+
+ @staticmethod
+ def build_While(ctx, node):
+ if node.orelse:
+ raise TaichiSyntaxError(
+ "'else' clause for 'while' not supported in Taichi kernels")
+
+ with ctx.loop_scope_guard():
+ ctx.ast_builder.begin_frontend_while(expr.Expr(1).ptr)
+ while_cond = build_stmt(ctx, node.test)
+ impl.begin_frontend_if(ctx.ast_builder, while_cond)
+ ctx.ast_builder.begin_frontend_if_true()
+ ctx.ast_builder.pop_scope()
+ ctx.ast_builder.begin_frontend_if_false()
+ ctx.ast_builder.insert_break_stmt()
+ ctx.ast_builder.pop_scope()
+ build_stmts(ctx, node.body)
+ ctx.ast_builder.pop_scope()
+ return None
+
+ @staticmethod
+ def build_If(ctx, node):
+ build_stmt(ctx, node.test)
+ is_static_if = (ASTTransformer.get_decorator(ctx,
+ node.test) == "static")
+
+ if is_static_if:
+ if node.test.ptr:
+ build_stmts(ctx, node.body)
+ else:
+ build_stmts(ctx, node.orelse)
+ return node
+
+ with ctx.non_static_control_flow_guard():
+ impl.begin_frontend_if(ctx.ast_builder, node.test.ptr)
+ ctx.ast_builder.begin_frontend_if_true()
+ build_stmts(ctx, node.body)
+ ctx.ast_builder.pop_scope()
+ ctx.ast_builder.begin_frontend_if_false()
+ build_stmts(ctx, node.orelse)
+ ctx.ast_builder.pop_scope()
+ return None
+
+ @staticmethod
+ def build_Expr(ctx, node):
+ build_stmt(ctx, node.value)
+ if not isinstance(node.value, ast.Call):
+ return None
+ is_taichi_function = getattr(node.value.func.ptr,
+ '_is_taichi_function', False)
+ if is_taichi_function and node.value.func.ptr._is_real_function:
+ func_call_result = node.value.ptr
+ ctx.ast_builder.insert_expr_stmt(func_call_result.ptr)
+ return None
+
+ @staticmethod
+ def build_IfExp(ctx, node):
+ build_stmt(ctx, node.test)
+ build_stmt(ctx, node.body)
+ build_stmt(ctx, node.orelse)
+
+ if is_taichi_class(node.test.ptr) or is_taichi_class(
+ node.body.ptr) or is_taichi_class(node.orelse.ptr):
+ node.ptr = ti_ops.select(node.test.ptr, node.body.ptr,
+ node.orelse.ptr)
+ warnings.warn_explicit(
+ 'Using conditional expression for element-wise select operation on '
+ 'Taichi vectors/matrices is deprecated. '
+ 'Please use "ti.select" instead.', DeprecationWarning,
+ ctx.file, node.lineno + ctx.lineno_offset)
+ return node.ptr
+
+ is_static_if = (ASTTransformer.get_decorator(ctx,
+ node.test) == "static")
+
+ if is_static_if:
+ if node.test.ptr:
+ node.ptr = build_stmt(ctx, node.body)
+ else:
+ node.ptr = build_stmt(ctx, node.orelse)
+ return node.ptr
+
+ val = impl.expr_init(None)
+
+ impl.begin_frontend_if(ctx.ast_builder, node.test.ptr)
+ ctx.ast_builder.begin_frontend_if_true()
+ val._assign(node.body.ptr)
+ ctx.ast_builder.pop_scope()
+ ctx.ast_builder.begin_frontend_if_false()
+ val._assign(node.orelse.ptr)
+ ctx.ast_builder.pop_scope()
+
+ node.ptr = val
+ return node.ptr
+
+ @staticmethod
+ def _is_string_mod_args(msg):
+ # 1. str % (a, b, c, ...)
+ # 2. str % single_item
+ # Note that |msg.right| may not be a tuple.
+ if not isinstance(msg, ast.BinOp):
+ return False
+ if not isinstance(msg.op, ast.Mod):
+ return False
+ if isinstance(msg.left, ast.Str):
+ return True
+ if isinstance(msg.left, ast.Constant) and isinstance(
+ msg.left.value, str):
+ return True
+ return False
+
+ @staticmethod
+ def _handle_string_mod_args(ctx, node):
+ msg = build_stmt(ctx, node.left)
+ args = build_stmt(ctx, node.right)
+ if not isinstance(args, collections.abc.Sequence):
+ args = (args, )
+ return msg, args
+
+ @staticmethod
+ def build_Assert(ctx, node):
+ extra_args = []
+ if node.msg is not None:
+ if isinstance(node.msg, ast.Constant):
+ msg = node.msg.value
+ elif isinstance(node.msg, ast.Str):
+ msg = node.msg.s
+ elif ASTTransformer._is_string_mod_args(node.msg):
+ msg, extra_args = ASTTransformer._handle_string_mod_args(
+ ctx, node.msg)
+ else:
+ raise ValueError(
+ f"assert info must be constant, not {ast.dump(node.msg)}")
+ else:
+ msg = astor.to_source(node.test)
+ test = build_stmt(ctx, node.test)
+ impl.ti_assert(test, msg.strip(), extra_args)
+ return None
+
+ @staticmethod
+ def build_Break(ctx, node):
+ if ctx.is_in_static_for():
+ ctx.set_loop_status(LoopStatus.Break)
+ else:
+ ctx.ast_builder.insert_break_stmt()
+ return None
+
+ @staticmethod
+ def build_Continue(ctx, node):
+ if ctx.is_in_static_for():
+ ctx.set_loop_status(LoopStatus.Continue)
+ else:
+ ctx.ast_builder.insert_continue_stmt()
+ return None
+
+ @staticmethod
+ def build_Pass(ctx, node):
+ return None
+
+
+build_stmt = ASTTransformer()
+
+
+def build_stmts(ctx, stmts):
+ with ctx.variable_scope_guard():
+ for stmt in stmts:
+ if ctx.returned or ctx.loop_status() != LoopStatus.Normal:
+ break
+ else:
+ build_stmt(ctx, stmt)
+ return stmts
diff --git a/python/taichi/lang/ast/ast_transformer_utils.py b/python/taichi/lang/ast/ast_transformer_utils.py
new file mode 100644
index 0000000000000..900351cc75555
--- /dev/null
+++ b/python/taichi/lang/ast/ast_transformer_utils.py
@@ -0,0 +1,264 @@
+import ast
+import builtins
+import traceback
+from enum import Enum
+from sys import version_info
+from textwrap import TextWrapper
+
+from taichi.lang.exception import (TaichiCompilationError, TaichiNameError,
+ TaichiSyntaxError,
+ handle_exception_from_cpp)
+
+
+class Builder:
+ def __call__(self, ctx, node):
+ method = getattr(self, 'build_' + node.__class__.__name__, None)
+ try:
+ if method is None:
+ error_msg = f'Unsupported node "{node.__class__.__name__}"'
+ raise TaichiSyntaxError(error_msg)
+ return method(ctx, node)
+ except Exception as e:
+ if ctx.raised or not isinstance(node, (ast.stmt, ast.expr)):
+ raise e.with_traceback(None)
+ ctx.raised = True
+ e = handle_exception_from_cpp(e)
+ if not isinstance(e, TaichiCompilationError):
+ msg = ctx.get_pos_info(node) + traceback.format_exc()
+ raise TaichiCompilationError(msg) from None
+ msg = ctx.get_pos_info(node) + str(e)
+ raise type(e)(msg) from None
+
+
+class VariableScopeGuard:
+ def __init__(self, scopes):
+ self.scopes = scopes
+
+ def __enter__(self):
+ self.scopes.append({})
+
+ def __exit__(self, exc_type, exc_val, exc_tb):
+ self.scopes.pop()
+
+
+class StaticScopeStatus:
+ def __init__(self):
+ self.is_in_static_scope = False
+
+
+class StaticScopeGuard:
+ def __init__(self, status):
+ self.status = status
+
+ def __enter__(self):
+ self.prev = self.status.is_in_static_scope
+ self.status.is_in_static_scope = True
+
+ def __exit__(self, exc_type, exc_val, exc_tb):
+ self.status.is_in_static_scope = self.prev
+
+
+class NonStaticControlFlowStatus:
+ def __init__(self):
+ self.is_in_non_static_control_flow = False
+
+
+class NonStaticControlFlowGuard:
+ def __init__(self, status):
+ self.status = status
+
+ def __enter__(self):
+ self.prev = self.status.is_in_non_static_control_flow
+ self.status.is_in_non_static_control_flow = True
+
+ def __exit__(self, exc_type, exc_val, exc_tb):
+ self.status.is_in_non_static_control_flow = self.prev
+
+
+class LoopStatus(Enum):
+ Normal = 0
+ Break = 1
+ Continue = 2
+
+
+class LoopScopeAttribute:
+ def __init__(self, is_static):
+ self.is_static = is_static
+ self.status = LoopStatus.Normal
+
+
+class LoopScopeGuard:
+ def __init__(self, scopes, non_static_guard=None):
+ self.scopes = scopes
+ self.non_static_guard = non_static_guard
+
+ def __enter__(self):
+ self.scopes.append(LoopScopeAttribute(self.non_static_guard is None))
+ if self.non_static_guard:
+ self.non_static_guard.__enter__()
+
+ def __exit__(self, exc_type, exc_val, exc_tb):
+ self.scopes.pop()
+ if self.non_static_guard:
+ self.non_static_guard.__exit__(exc_type, exc_val, exc_tb)
+
+
+class ASTTransformerContext:
+ def __init__(self,
+ excluded_parameters=(),
+ is_kernel=True,
+ func=None,
+ arg_features=None,
+ global_vars=None,
+ argument_data=None,
+ file=None,
+ src=None,
+ start_lineno=None,
+ ast_builder=None,
+ is_real_function=False):
+ self.func = func
+ self.local_scopes = []
+ self.loop_scopes = []
+ self.excluded_parameters = excluded_parameters
+ self.is_kernel = is_kernel
+ self.arg_features = arg_features
+ self.returns = None
+ self.global_vars = global_vars
+ self.argument_data = argument_data
+ self.return_data = None
+ self.file = file
+ self.src = src
+ self.indent = 0
+ for c in self.src[0]:
+ if c == ' ':
+ self.indent += 1
+ else:
+ break
+ self.lineno_offset = start_lineno - 1
+ self.raised = False
+ self.non_static_control_flow_status = NonStaticControlFlowStatus()
+ self.static_scope_status = StaticScopeStatus()
+ self.returned = False
+ self.ast_builder = ast_builder
+ self.visited_funcdef = False
+ self.is_real_function = is_real_function
+
+ # e.g.: FunctionDef, Module, Global
+ def variable_scope_guard(self):
+ return VariableScopeGuard(self.local_scopes)
+
+ # e.g.: For, While
+ def loop_scope_guard(self, is_static=False):
+ if is_static:
+ return LoopScopeGuard(self.loop_scopes)
+ return LoopScopeGuard(self.loop_scopes,
+ self.non_static_control_flow_guard())
+
+ def non_static_control_flow_guard(self):
+ return NonStaticControlFlowGuard(self.non_static_control_flow_status)
+
+ def static_scope_guard(self):
+ return StaticScopeGuard(self.static_scope_status)
+
+ def current_scope(self):
+ return self.local_scopes[-1]
+
+ def current_loop_scope(self):
+ return self.loop_scopes[-1]
+
+ def loop_status(self):
+ if self.loop_scopes:
+ return self.loop_scopes[-1].status
+ return LoopStatus.Normal
+
+ def set_loop_status(self, status):
+ self.loop_scopes[-1].status = status
+
+ def is_in_static_for(self):
+ if self.loop_scopes:
+ return self.loop_scopes[-1].is_static
+ return False
+
+ def is_in_non_static_control_flow(self):
+ return self.non_static_control_flow_status.is_in_non_static_control_flow
+
+ def is_in_static_scope(self):
+ return self.static_scope_status.is_in_static_scope
+
+ def is_var_declared(self, name):
+ for s in self.local_scopes:
+ if name in s:
+ return True
+ return False
+
+ def create_variable(self, name, var):
+ if name in self.current_scope():
+ raise TaichiSyntaxError("Recreating variables is not allowed")
+ self.current_scope()[name] = var
+
+ def check_loop_var(self, loop_var):
+ if self.is_var_declared(loop_var):
+ raise TaichiSyntaxError(
+ f"Variable '{loop_var}' is already declared in the outer scope and cannot be used as loop variable"
+ )
+
+ def get_var_by_name(self, name):
+ for s in reversed(self.local_scopes):
+ if name in s:
+ return s[name]
+ if name in self.global_vars:
+ return self.global_vars[name]
+ try:
+ return getattr(builtins, name)
+ except AttributeError:
+ raise TaichiNameError(f'Name "{name}" is not defined')
+
+ def get_pos_info(self, node):
+ msg = f'On line {node.lineno + self.lineno_offset} of file "{self.file}", in {self.func.func.__name__}:\n'
+ if version_info < (3, 8):
+ msg += self.src[node.lineno - 1] + "\n"
+ return msg
+ col_offset = self.indent + node.col_offset
+ end_col_offset = self.indent + node.end_col_offset
+
+ wrapper = TextWrapper(width=80)
+
+ def gen_line(code, hint):
+ hint += ' ' * (len(code) - len(hint))
+ code = wrapper.wrap(code)
+ hint = wrapper.wrap(hint)
+ if not len(code):
+ return "\n\n"
+ return "".join([c + '\n' + h + '\n' for c, h in zip(code, hint)])
+
+ if node.lineno == node.end_lineno:
+ hint = ' ' * col_offset + '^' * (end_col_offset - col_offset)
+ msg += gen_line(self.src[node.lineno - 1], hint)
+ else:
+ node_type = node.__class__.__name__
+
+ if node_type in ["For", "While", "FunctionDef", "If"]:
+ end_lineno = max(node.body[0].lineno - 1, node.lineno)
+ else:
+ end_lineno = node.end_lineno
+
+ for i in range(node.lineno - 1, end_lineno):
+ last = len(self.src[i])
+ while last > 0 and (self.src[i][last - 1].isspace() or
+ not self.src[i][last - 1].isprintable()):
+ last -= 1
+ first = 0
+ while first < len(self.src[i]) and (
+ self.src[i][first].isspace()
+ or not self.src[i][first].isprintable()):
+ first += 1
+ if i == node.lineno - 1:
+ hint = ' ' * col_offset + '^' * (last - col_offset)
+ elif i == node.end_lineno - 1:
+ hint = ' ' * first + '^' * (end_col_offset - first)
+ elif first < last:
+ hint = ' ' * first + '^' * (last - first)
+ else:
+ hint = ''
+ msg += gen_line(self.src[i], hint)
+ return msg
diff --git a/python/taichi/lang/ast/checkers.py b/python/taichi/lang/ast/checkers.py
index fb8a3d62ba418..825a13ae6ab9e 100644
--- a/python/taichi/lang/ast/checkers.py
+++ b/python/taichi/lang/ast/checkers.py
@@ -1,6 +1,6 @@
import ast
-import taichi.lang.kernel_impl
+from taichi.lang.exception import TaichiSyntaxError
from taichi.lang.shell import oinspect
@@ -55,7 +55,8 @@ def get_error_location(self, node):
lineno = self._func_lineno + node.lineno - 1
return f'file={self._func_file} kernel={self._func_name} line={lineno}'
- def should_check(self, node):
+ @staticmethod
+ def should_check(node):
if not isinstance(node, ast.stmt):
return False
# TODO(#536): Frontend pass should help make sure |func| is a valid AST for
@@ -69,7 +70,7 @@ def generic_visit(self, node):
return
if not (self.top_level or self.current_scope.allows_more_stmt):
- raise taichi.lang.kernel_impl.KernelDefError(
+ raise TaichiSyntaxError(
f'No more statements allowed, at {self.get_error_location(node)}'
)
old_top_level = self.top_level
@@ -83,25 +84,23 @@ def generic_visit(self, node):
if old_top_level:
self._scope_guards.pop()
- def visit_For(self, node):
+ @staticmethod
+ def visit_for(node):
# TODO: since autodiff is enhanced, AST checker rules should be relaxed. This part should be updated.
+ # original code is #def visit_For(self, node) without #@staticmethod before fix pylint R0201
return
- if (isinstance(node.iter, ast.Call)
- and isinstance(node.iter.func, ast.Attribute)
- and isinstance(node.iter.func.value, ast.Name)
- and node.iter.func.value.id == 'ti'
- and node.iter.func.attr == 'static'):
- is_static = True
- else:
- is_static = False
- if not (self.top_level or self.current_scope.allows_for_loop
- or is_static):
- raise taichi.lang.kernel_impl.KernelDefError(
- f'No more for loops allowed, at {self.get_error_location(node)}'
- )
-
- with self.new_scope():
- super().generic_visit(node)
-
- if not (self.top_level or is_static):
- self.current_scope.mark_no_more_stmt()
+ # is_static = (isinstance(node.iter, ast.Call)
+ # and isinstance(node.iter.func, ast.Attribute)
+ # and isinstance(node.iter.func.value, ast.Name)
+ # and node.iter.func.value.id == 'ti'
+ # and node.iter.func.attr == 'static')
+ # if not (self.top_level or self.current_scope.allows_for_loop
+ # or is_static):
+ # raise TaichiSyntaxError(
+ # f'No more for loops allowed, at {self.get_error_location(node)}'
+ # )
+ # with self.new_scope():
+ # super().generic_visit(node)
+ #
+ # if not (self.top_level or is_static):
+ # self.current_scope.mark_no_more_stmt()
diff --git a/python/taichi/lang/ast/transform.py b/python/taichi/lang/ast/transform.py
new file mode 100644
index 0000000000000..f9ca13b49f43e
--- /dev/null
+++ b/python/taichi/lang/ast/transform.py
@@ -0,0 +1,7 @@
+from taichi.lang.ast.ast_transformer import ASTTransformer
+from taichi.lang.ast.ast_transformer_utils import ASTTransformerContext
+
+
+def transform_tree(tree, ctx: ASTTransformerContext):
+ ASTTransformer()(ctx, tree)
+ return ctx.return_data
diff --git a/python/taichi/lang/ast/transformer.py b/python/taichi/lang/ast/transformer.py
deleted file mode 100644
index 476a3cd9c7855..0000000000000
--- a/python/taichi/lang/ast/transformer.py
+++ /dev/null
@@ -1,98 +0,0 @@
-import ast
-
-import astor
-from taichi.lang import impl
-from taichi.lang.ast.symbol_resolver import ASTResolver
-from taichi.lang.ast_builder_utils import BuilderContext
-from taichi.lang.exception import TaichiSyntaxError
-from taichi.lang.stmt_builder import build_stmt
-
-import taichi as ti
-
-
-# Total transform
-class ASTTransformerTotal(object):
- def __init__(self,
- func=None,
- excluded_parameters=(),
- is_kernel=True,
- arg_features=None):
- self.func = func
- self.excluded_parameters = excluded_parameters
- self.is_kernel = is_kernel
- self.arg_features = arg_features
- self.pass_Checks = ASTTransformerChecks(func=func)
-
- @staticmethod
- def print_ast(tree, title=None):
- if not impl.get_runtime().print_preprocessed:
- return
- if title is not None:
- ti.info(f'{title}:')
- print(astor.to_source(tree.body[0], indent_with=' '), flush=True)
-
- def visit(self, tree):
- self.print_ast(tree, 'Initial AST')
- ctx = BuilderContext(func=self.func,
- excluded_parameters=self.excluded_parameters,
- is_kernel=self.is_kernel,
- arg_features=self.arg_features)
- # Convert Python AST to Python code that generates Taichi C++ AST.
- tree = build_stmt(ctx, tree)
- ast.fix_missing_locations(tree)
- self.print_ast(tree, 'Preprocessed')
- self.pass_Checks.visit(tree) # does not modify the AST
-
-
-class ASTTransformerBase(ast.NodeTransformer):
- def __init__(self, func):
- super().__init__()
- self.func = func
-
- @staticmethod
- def get_decorator(node):
- if not isinstance(node, ast.Call):
- return ''
- for wanted, name in [
- (ti.static, 'static'),
- (ti.grouped, 'grouped'),
- (ti.ndrange, 'ndrange'),
- ]:
- if ASTResolver.resolve_to(node.func, wanted, globals()):
- return name
- return ''
-
-
-# Performs checks at the Python AST level. Does not modify the AST.
-class ASTTransformerChecks(ASTTransformerBase):
- def __init__(self, func):
- super().__init__(func)
- self.has_return = False
- self.in_static_if = False
-
- def visit_If(self, node):
- node.test = self.visit(node.test)
-
- old_in_static_if = self.in_static_if
- self.in_static_if = self.get_decorator(node.test) == 'static'
-
- node.body = list(map(self.visit, node.body))
- if node.orelse is not None:
- node.orelse = list(map(self.visit, node.orelse))
-
- self.in_static_if = old_in_static_if
-
- return node
-
- def visit_Return(self, node):
- if self.in_static_if: # we can have multiple return in static-if branches
- return node
-
- if not self.has_return:
- self.has_return = True
- else:
- raise TaichiSyntaxError(
- 'Taichi functions/kernels cannot have multiple returns!'
- ' Consider using a local variable to walk around.')
-
- return node
diff --git a/python/taichi/lang/ast_builder_utils.py b/python/taichi/lang/ast_builder_utils.py
deleted file mode 100644
index 103bc3a4fd368..0000000000000
--- a/python/taichi/lang/ast_builder_utils.py
+++ /dev/null
@@ -1,91 +0,0 @@
-import ast
-
-from taichi.lang.exception import TaichiSyntaxError
-
-
-class Builder(object):
- def __call__(self, ctx, node):
- method = getattr(self, 'build_' + node.__class__.__name__, None)
- if method is None:
- try:
- import astpretty # pylint: disable=C0415
- error_msg = f'Unsupported node {node}:\n{astpretty.pformat(node)}'
- except:
- error_msg = f'Unsupported node {node}'
- raise TaichiSyntaxError(error_msg)
- return method(ctx, node)
-
-
-def parse_stmt(stmt):
- return ast.parse(stmt).body[0]
-
-
-def parse_expr(expr):
- return ast.parse(expr).body[0].value
-
-
-class ScopeGuard:
- def __init__(self, scopes, stmt_block=None):
- self.scopes = scopes
- self.stmt_block = stmt_block
-
- def __enter__(self):
- self.scopes.append([])
-
- def __exit__(self, exc_type, exc_val, exc_tb):
- local = self.scopes[-1]
- if self.stmt_block is not None:
- for var in reversed(local):
- stmt = parse_stmt('del var')
- stmt.targets[0].id = var
- self.stmt_block.append(stmt)
- self.scopes.pop()
-
-
-class BuilderContext:
- def __init__(self,
- excluded_parameters=(),
- is_kernel=True,
- func=None,
- arg_features=None):
- self.func = func
- self.local_scopes = []
- self.control_scopes = []
- self.excluded_parameters = excluded_parameters
- self.is_kernel = is_kernel
- self.arg_features = arg_features
- self.returns = None
-
- # e.g.: FunctionDef, Module, Global
- def variable_scope(self, *args):
- return ScopeGuard(self.local_scopes, *args)
-
- # e.g.: For, While
- def control_scope(self):
- return ScopeGuard(self.control_scopes)
-
- def current_scope(self):
- return self.local_scopes[-1]
-
- def current_control_scope(self):
- return self.control_scopes[-1]
-
- def var_declared(self, name):
- for s in self.local_scopes:
- if name in s:
- return True
- return False
-
- def is_creation(self, name):
- return not self.var_declared(name)
-
- def create_variable(self, name):
- assert name not in self.current_scope(
- ), "Recreating variables is not allowed"
- self.current_scope().append(name)
-
- def check_loop_var(self, loop_var):
- if self.var_declared(loop_var):
- raise TaichiSyntaxError(
- "Variable '{}' is already declared in the outer scope and cannot be used as loop variable"
- .format(loop_var))
diff --git a/python/taichi/lang/common_ops.py b/python/taichi/lang/common_ops.py
index 28b36964b2440..38b13a29573ce 100644
--- a/python/taichi/lang/common_ops.py
+++ b/python/taichi/lang/common_ops.py
@@ -1,167 +1,132 @@
-import taichi as ti
+import warnings
+
+from taichi.lang import ops
class TaichiOperations:
"""The base class of taichi operations of expressions. Subclasses: :class:`~taichi.lang.expr.Expr`, :class:`~taichi.lang.matrix.Matrix`"""
+
+ __deprecated_atomic_ops__ = {
+ "atomic_add": "_atomic_add",
+ "atomic_and": "_atomic_and",
+ "atomic_or": "_atomic_or",
+ "atomic_sub": "_atomic_sub",
+ "atomic_xor": "_atomic_xor",
+ }
+
+ def __getattr__(self, item):
+ if item in TaichiOperations.__deprecated_atomic_ops__:
+ warnings.warn(
+ f"a.{item}(b) is deprecated. Please use ti.{item}(a, b) instead.",
+ DeprecationWarning)
+ return getattr(self,
+ TaichiOperations.__deprecated_atomic_ops__[item])
+ raise AttributeError(
+ f"'{type(self).__name__}' object has no attribute '{item}'")
+
def __neg__(self):
- _taichi_skip_traceback = 1
- return ti.neg(self)
+ return ops.neg(self)
def __abs__(self):
- _taichi_skip_traceback = 1
- return ti.abs(self)
+ return ops.abs(self)
def __add__(self, other):
- _taichi_skip_traceback = 1
- return ti.add(self, other)
+ return ops.add(self, other)
def __radd__(self, other):
- _taichi_skip_traceback = 1
- return ti.add(other, self)
+ return ops.add(other, self)
def __sub__(self, other):
- _taichi_skip_traceback = 1
- return ti.sub(self, other)
+ return ops.sub(self, other)
def __rsub__(self, other):
- _taichi_skip_traceback = 1
- return ti.sub(other, self)
+ return ops.sub(other, self)
def __mul__(self, other):
- _taichi_skip_traceback = 1
- return ti.mul(self, other)
+ return ops.mul(self, other)
def __rmul__(self, other):
- _taichi_skip_traceback = 1
- return ti.mul(other, self)
+ return ops.mul(other, self)
def __truediv__(self, other):
- _taichi_skip_traceback = 1
- return ti.truediv(self, other)
+ return ops.truediv(self, other)
def __rtruediv__(self, other):
- _taichi_skip_traceback = 1
- return ti.truediv(other, self)
+ return ops.truediv(other, self)
def __floordiv__(self, other):
- _taichi_skip_traceback = 1
- return ti.floordiv(self, other)
+ return ops.floordiv(self, other)
def __rfloordiv__(self, other):
- _taichi_skip_traceback = 1
- return ti.floordiv(other, self)
+ return ops.floordiv(other, self)
def __mod__(self, other):
- _taichi_skip_traceback = 1
- return ti.mod(self, other)
+ return ops.mod(self, other)
def __rmod__(self, other):
- _taichi_skip_traceback = 1
- return ti.mod(other, self)
+ return ops.mod(other, self)
def __pow__(self, other, modulo=None):
- _taichi_skip_traceback = 1
- return ti.pow(self, other)
+ return ops.pow(self, other)
def __rpow__(self, other, modulo=None):
- _taichi_skip_traceback = 1
- return ti.pow(other, self)
+ return ops.pow(other, self)
def __le__(self, other):
- _taichi_skip_traceback = 1
- return ti.cmp_le(self, other)
+ return ops.cmp_le(self, other)
def __lt__(self, other):
- _taichi_skip_traceback = 1
- return ti.cmp_lt(self, other)
+ return ops.cmp_lt(self, other)
def __ge__(self, other):
- _taichi_skip_traceback = 1
- return ti.cmp_ge(self, other)
+ return ops.cmp_ge(self, other)
def __gt__(self, other):
- _taichi_skip_traceback = 1
- return ti.cmp_gt(self, other)
+ return ops.cmp_gt(self, other)
def __eq__(self, other):
- _taichi_skip_traceback = 1
- return ti.cmp_eq(self, other)
+ return ops.cmp_eq(self, other)
def __ne__(self, other):
- _taichi_skip_traceback = 1
- return ti.cmp_ne(self, other)
+ return ops.cmp_ne(self, other)
def __and__(self, other):
- _taichi_skip_traceback = 1
- return ti.bit_and(self, other)
+ return ops.bit_and(self, other)
def __rand__(self, other):
- _taichi_skip_traceback = 1
- return ti.bit_and(other, self)
+ return ops.bit_and(other, self)
def __or__(self, other):
- _taichi_skip_traceback = 1
- return ti.bit_or(self, other)
+ return ops.bit_or(self, other)
def __ror__(self, other):
- _taichi_skip_traceback = 1
- return ti.bit_or(other, self)
+ return ops.bit_or(other, self)
def __xor__(self, other):
- _taichi_skip_traceback = 1
- return ti.bit_xor(self, other)
+ return ops.bit_xor(self, other)
def __rxor__(self, other):
- _taichi_skip_traceback = 1
- return ti.bit_xor(other, self)
+ return ops.bit_xor(other, self)
def __lshift__(self, other):
- _taichi_skip_traceback = 1
- return ti.bit_shl(self, other)
+ return ops.bit_shl(self, other)
def __rlshift__(self, other):
- _taichi_skip_traceback = 1
- return ti.bit_shl(other, self)
+ return ops.bit_shl(other, self)
def __rshift__(self, other):
- _taichi_skip_traceback = 1
- return ti.bit_sar(self, other)
+ return ops.bit_sar(self, other)
def __rrshift__(self, other):
- _taichi_skip_traceback = 1
- return ti.bit_sar(other, self)
-
- def logical_and(self, other):
- """Return the new expression of computing logical and between self and a given operand.
-
- Args:
- other (Any): Given operand.
-
- Returns:
- :class:`~taichi.lang.expr.Expr`: The computing expression of logical and."""
- _taichi_skip_traceback = 1
- return ti.logical_and(self, other)
-
- def logical_or(self, other):
- """Return the new expression of computing logical or between self and a given operand.
-
- Args:
- other (Any): Given operand.
-
- Returns:
- :class:`~taichi.lang.expr.Expr`: The computing expression of logical or."""
- _taichi_skip_traceback = 1
- return ti.logical_or(self, other)
+ return ops.bit_sar(other, self)
def __invert__(self): # ~a => a.__invert__()
- _taichi_skip_traceback = 1
- return ti.bit_not(self)
+ return ops.bit_not(self)
def __not__(self): # not a => a.__not__()
- _taichi_skip_traceback = 1
- return ti.logical_not(self)
+ return ops.logical_not(self)
- def atomic_add(self, other):
+ def _atomic_add(self, other):
"""Return the new expression of computing atomic add between self and a given operand.
Args:
@@ -169,10 +134,9 @@ def atomic_add(self, other):
Returns:
:class:`~taichi.lang.expr.Expr`: The computing expression of atomic add."""
- _taichi_skip_traceback = 1
- return ti.atomic_add(self, other)
+ return ops.atomic_add(self, other)
- def atomic_sub(self, other):
+ def _atomic_sub(self, other):
"""Return the new expression of computing atomic sub between self and a given operand.
Args:
@@ -180,10 +144,9 @@ def atomic_sub(self, other):
Returns:
:class:`~taichi.lang.expr.Expr`: The computing expression of atomic sub."""
- _taichi_skip_traceback = 1
- return ti.atomic_sub(self, other)
+ return ops.atomic_sub(self, other)
- def atomic_and(self, other):
+ def _atomic_and(self, other):
"""Return the new expression of computing atomic and between self and a given operand.
Args:
@@ -191,10 +154,9 @@ def atomic_and(self, other):
Returns:
:class:`~taichi.lang.expr.Expr`: The computing expression of atomic and."""
- _taichi_skip_traceback = 1
- return ti.atomic_and(self, other)
+ return ops.atomic_and(self, other)
- def atomic_xor(self, other):
+ def _atomic_xor(self, other):
"""Return the new expression of computing atomic xor between self and a given operand.
Args:
@@ -202,10 +164,9 @@ def atomic_xor(self, other):
Returns:
:class:`~taichi.lang.expr.Expr`: The computing expression of atomic xor."""
- _taichi_skip_traceback = 1
- return ti.atomic_xor(self, other)
+ return ops.atomic_xor(self, other)
- def atomic_or(self, other):
+ def _atomic_or(self, other):
"""Return the new expression of computing atomic or between self and a given operand.
Args:
@@ -213,66 +174,58 @@ def atomic_or(self, other):
Returns:
:class:`~taichi.lang.expr.Expr`: The computing expression of atomic or."""
- _taichi_skip_traceback = 1
- return ti.atomic_or(self, other)
+ return ops.atomic_or(self, other)
def __iadd__(self, other):
- _taichi_skip_traceback = 1
- self.atomic_add(other)
+ self._atomic_add(other)
return self
def __isub__(self, other):
- _taichi_skip_traceback = 1
- self.atomic_sub(other)
+ self._atomic_sub(other)
return self
def __iand__(self, other):
- _taichi_skip_traceback = 1
- self.atomic_and(other)
+ self._atomic_and(other)
return self
def __ixor__(self, other):
- _taichi_skip_traceback = 1
- self.atomic_xor(other)
+ self._atomic_xor(other)
return self
def __ior__(self, other):
- _taichi_skip_traceback = 1
- self.atomic_or(other)
+ self._atomic_or(other)
return self
# we don't support atomic_mul/truediv/floordiv/mod yet:
def __imul__(self, other):
- _taichi_skip_traceback = 1
- self.assign(ti.mul(self, other))
+ self._assign(ops.mul(self, other))
return self
def __itruediv__(self, other):
- _taichi_skip_traceback = 1
- self.assign(ti.truediv(self, other))
+ self._assign(ops.truediv(self, other))
return self
def __ifloordiv__(self, other):
- _taichi_skip_traceback = 1
- self.assign(ti.floordiv(self, other))
+ self._assign(ops.floordiv(self, other))
return self
def __imod__(self, other):
- _taichi_skip_traceback = 1
- self.assign(ti.mod(self, other))
+ self._assign(ops.mod(self, other))
return self
def __ilshift__(self, other):
- _taichi_skip_traceback = 1
- self.assign(ti.bit_shl(self, other))
+ self._assign(ops.bit_shl(self, other))
return self
def __irshift__(self, other):
- _taichi_skip_traceback = 1
- self.assign(ti.bit_shr(self, other))
+ self._assign(ops.bit_shr(self, other))
+ return self
+
+ def __ipow__(self, other):
+ self._assign(ops.pow(self, other))
return self
- def assign(self, other):
+ def _assign(self, other):
"""Assign the expression of the given operand to self.
Args:
@@ -280,16 +233,15 @@ def assign(self, other):
Returns:
:class:`~taichi.lang.expr.Expr`: The expression after assigning."""
- _taichi_skip_traceback = 1
- return ti.assign(self, other)
+ return ops.assign(self, other)
- def augassign(self, x, op):
+ # pylint: disable=R0201
+ def _augassign(self, x, op):
"""Generate the computing expression between self and the given operand of given operator and assigned to self.
Args:
x (Any): Given operand.
op (str): The name of operator."""
- _taichi_skip_traceback = 1
if op == 'Add':
self += x
elif op == 'Sub':
@@ -312,13 +264,13 @@ def augassign(self, x, op):
self >>= x
elif op == 'LShift':
self <<= x
+ elif op == 'Pow':
+ self **= x
else:
assert False, op
def __ti_int__(self):
- _taichi_skip_traceback = 1
- return ti.cast(self, int)
+ return ops.cast(self, int)
def __ti_float__(self):
- _taichi_skip_traceback = 1
- return ti.cast(self, float)
+ return ops.cast(self, float)
diff --git a/python/taichi/lang/enums.py b/python/taichi/lang/enums.py
index e7649322c9d9c..43f14d50dcccb 100644
--- a/python/taichi/lang/enums.py
+++ b/python/taichi/lang/enums.py
@@ -9,3 +9,6 @@ class Layout(Enum):
"""
AOS = 1
SOA = 2
+
+
+__all__ = ['Layout']
diff --git a/python/taichi/lang/exception.py b/python/taichi/lang/exception.py
index 149279f6eea91..d76290638e632 100644
--- a/python/taichi/lang/exception.py
+++ b/python/taichi/lang/exception.py
@@ -1,7 +1,51 @@
-class TaichiSyntaxError(Exception):
- def __init__(self, *args):
- super().__init__(*args)
+from taichi._lib import core
-class InvalidOperationError(Exception):
+class TaichiCompilationError(Exception):
+ """Base class for all compilation exceptions.
+ """
pass
+
+
+class TaichiSyntaxError(TaichiCompilationError, SyntaxError):
+ """Thrown when a syntax error is found during compilation.
+ """
+ pass
+
+
+class TaichiNameError(TaichiCompilationError, NameError):
+ """Thrown when an undefine name is found during compilation.
+ """
+ pass
+
+
+class TaichiTypeError(TaichiCompilationError, TypeError):
+ """Thrown when a type mismatch is found during compilation.
+ """
+ pass
+
+
+class TaichiRuntimeError(RuntimeError):
+ """Thrown when the compiled program cannot be executed due to unspecified reasons.
+ """
+ pass
+
+
+class TaichiRuntimeTypeError(TaichiRuntimeError, TypeError):
+ def __init__(self, pos, needed, provided):
+ message = f'Argument {pos} (type={provided}) cannot be converted into required type {needed}'
+ super().__init__(message)
+
+
+def handle_exception_from_cpp(exc):
+ if isinstance(exc, core.TaichiTypeError):
+ return TaichiTypeError(str(exc))
+ if isinstance(exc, core.TaichiSyntaxError):
+ return TaichiSyntaxError(str(exc))
+ return exc
+
+
+__all__ = [
+ 'TaichiSyntaxError', 'TaichiTypeError', 'TaichiCompilationError',
+ 'TaichiNameError', 'TaichiRuntimeError', 'TaichiRuntimeTypeError'
+]
diff --git a/python/taichi/lang/expr.py b/python/taichi/lang/expr.py
index 0f9b14957cc5b..05582cd6be9cb 100644
--- a/python/taichi/lang/expr.py
+++ b/python/taichi/lang/expr.py
@@ -1,17 +1,16 @@
import numpy as np
-from taichi.core.util import ti_core as _ti_core
+from taichi._lib import core as _ti_core
from taichi.lang import impl
from taichi.lang.common_ops import TaichiOperations
-from taichi.lang.util import is_taichi_class, python_scope
-
-import taichi as ti
+from taichi.lang.exception import TaichiTypeError
+from taichi.lang.util import is_taichi_class, to_numpy_type, to_taichi_type
+from taichi.types.primitive_types import integer_types, real_types
# Scalar, basic data type
class Expr(TaichiOperations):
"""A Python-side Expr wrapper, whose member variable `ptr` is an instance of C++ Expr class. A C++ Expr object contains member variable `expr` which holds an instance of C++ Expression class."""
- def __init__(self, *args, tb=None):
- _taichi_skip_traceback = 1
+ def __init__(self, *args, tb=None, dtype=None):
self.tb = tb
if len(args) == 1:
if isinstance(args[0], _ti_core.Expr):
@@ -20,21 +19,24 @@ def __init__(self, *args, tb=None):
self.ptr = args[0].ptr
self.tb = args[0].tb
elif is_taichi_class(args[0]):
- raise ValueError('cannot initialize scalar expression from '
- f'taichi class: {type(args[0])}')
+ raise TaichiTypeError(
+ 'Cannot initialize scalar expression from '
+ f'taichi class: {type(args[0])}')
else:
# assume to be constant
arg = args[0]
- try:
- if isinstance(arg, np.ndarray):
- arg = arg.dtype(arg)
- except:
- pass
- self.ptr = impl.make_constant_expr(arg).ptr
+ if isinstance(arg, np.ndarray):
+ if arg.shape:
+ raise TaichiTypeError(
+ "Only 0-dimensional numpy array can be used to initialize a scalar expression"
+ )
+ arg = arg.dtype.type(arg)
+ self.ptr = make_constant_expr(arg, dtype).ptr
else:
assert False
if self.tb:
self.ptr.set_tb(self.tb)
+ self.ptr.type_check(impl.get_runtime().prog.config)
def __hash__(self):
return self.ptr.get_raw_address()
@@ -46,22 +48,74 @@ def __repr__(self):
return ''
-def make_var_vector(size):
+def _check_in_range(npty, val):
+ iif = np.iinfo(npty)
+ if not iif.min <= val <= iif.max:
+ # This isn't the case we want to deal with: |val| does't fall into the valid range of either
+ # the signed or the unsigned type.
+ raise TaichiTypeError(
+ f'Constant {val} has exceeded the range of {to_taichi_type(npty)}: [{iif.min}, {iif.max}]'
+ )
+
+
+def _clamp_unsigned_to_range(npty, val):
+ # npty: np.int32 or np.int64
+ iif = np.iinfo(npty)
+ if iif.min <= val <= iif.max:
+ return val
+ cap = (1 << iif.bits)
+ assert 0 <= val < cap
+ new_val = val - cap
+ return new_val
+
+
+def make_constant_expr(val, dtype):
+ if isinstance(val, (int, np.integer)):
+ constant_dtype = impl.get_runtime(
+ ).default_ip if dtype is None else dtype
+ if constant_dtype not in integer_types:
+ raise TaichiTypeError(
+ 'Integer literals must be annotated with a integer type. For type casting, use `ti.cast`.'
+ )
+ _check_in_range(to_numpy_type(constant_dtype), val)
+ return Expr(
+ _ti_core.make_const_expr_int(
+ constant_dtype, _clamp_unsigned_to_range(np.int64, val)))
+ if isinstance(val, (float, np.floating)):
+ constant_dtype = impl.get_runtime(
+ ).default_fp if dtype is None else dtype
+ if constant_dtype not in real_types:
+ raise TaichiTypeError(
+ 'Floating-point literals must be annotated with a floating-point type. For type casting, use `ti.cast`.'
+ )
+ return Expr(_ti_core.make_const_expr_fp(constant_dtype, val))
+ raise TaichiTypeError(f'Invalid constant scalar data type: {type(val)}')
+
+
+def make_var_list(size):
exprs = []
for _ in range(size):
exprs.append(_ti_core.make_id_expr(''))
- return ti.Vector(exprs, disable_local_tensor=True)
+ return exprs
def make_expr_group(*exprs):
+ from taichi.lang.matrix import Matrix # pylint: disable=C0415
if len(exprs) == 1:
if isinstance(exprs[0], (list, tuple)):
exprs = exprs[0]
- elif isinstance(exprs[0], ti.Matrix):
+ elif isinstance(exprs[0], Matrix):
mat = exprs[0]
assert mat.m == 1
exprs = mat.entries
expr_group = _ti_core.ExprGroup()
for i in exprs:
- expr_group.push_back(Expr(i).ptr)
+ if isinstance(i, Matrix):
+ assert i.local_tensor_proxy is not None
+ expr_group.push_back(i.local_tensor_proxy)
+ else:
+ expr_group.push_back(Expr(i).ptr)
return expr_group
+
+
+__all__ = []
diff --git a/python/taichi/lang/expr_builder.py b/python/taichi/lang/expr_builder.py
deleted file mode 100644
index 8cf2082fbcaf0..0000000000000
--- a/python/taichi/lang/expr_builder.py
+++ /dev/null
@@ -1,277 +0,0 @@
-import ast
-import warnings
-
-from taichi.lang.ast.symbol_resolver import ASTResolver
-from taichi.lang.ast_builder_utils import *
-from taichi.lang.exception import TaichiSyntaxError
-
-import taichi as ti
-
-
-class ExprBuilder(Builder):
- @staticmethod
- def build_Subscript(ctx, node):
- def get_subscript_index(node):
- assert isinstance(node, ast.Subscript), type(node)
- # ast.Index has been deprecated in Python 3.9,
- # use the index value directly instead :)
- if isinstance(node.slice, ast.Index):
- return build_expr(ctx, node.slice.value)
- return build_expr(ctx, node.slice)
-
- value = build_expr(ctx, node.value)
- indices = get_subscript_index(node)
- if isinstance(indices, ast.Tuple):
- indices = indices.elts
- else:
- indices = [indices]
-
- call = ast.Call(func=parse_expr('ti.subscript'),
- args=[value] + indices,
- keywords=[])
- return ast.copy_location(call, node)
-
- @staticmethod
- def build_Compare(ctx, node):
- operands = build_exprs(ctx, [node.left] + list(node.comparators))
- operators = []
- for i in range(len(node.ops)):
- if isinstance(node.ops[i], ast.Lt):
- op_str = 'Lt'
- elif isinstance(node.ops[i], ast.LtE):
- op_str = 'LtE'
- elif isinstance(node.ops[i], ast.Gt):
- op_str = 'Gt'
- elif isinstance(node.ops[i], ast.GtE):
- op_str = 'GtE'
- elif isinstance(node.ops[i], ast.Eq):
- op_str = 'Eq'
- elif isinstance(node.ops[i], ast.NotEq):
- op_str = 'NotEq'
- elif isinstance(node.ops[i], ast.In):
- raise TaichiSyntaxError(
- '"in" is not supported in Taichi kernels.')
- elif isinstance(node.ops[i], ast.NotIn):
- raise TaichiSyntaxError(
- '"not in" is not supported in Taichi kernels.')
- elif isinstance(node.ops[i], ast.Is):
- raise TaichiSyntaxError(
- '"is" is not supported in Taichi kernels.')
- elif isinstance(node.ops[i], ast.IsNot):
- raise TaichiSyntaxError(
- '"is not" is not supported in Taichi kernels.')
- else:
- raise Exception(f'Unknown operator {node.ops[i]}')
- operators += [
- ast.copy_location(ast.Str(s=op_str, kind=None), node)
- ]
-
- call = ast.Call(
- func=parse_expr('ti.chain_compare'),
- args=[
- ast.copy_location(ast.List(elts=operands, ctx=ast.Load()),
- node),
- ast.copy_location(ast.List(elts=operators, ctx=ast.Load()),
- node)
- ],
- keywords=[])
- call = ast.copy_location(call, node)
- return call
-
- @staticmethod
- def build_Call(ctx, node):
- if ASTResolver.resolve_to(node.func, ti.static, globals()):
- # Do not modify the expression if the function called is ti.static
- return node
- node.func = build_expr(ctx, node.func)
- node.args = build_exprs(ctx, node.args)
- for i in range(len(node.keywords)):
- node.keywords[i].value = build_expr(ctx, node.keywords[i].value)
- if isinstance(node.func, ast.Attribute):
- attr_name = node.func.attr
- if attr_name == 'format':
- node.args.insert(0, node.func.value)
- node.func = parse_expr('ti.ti_format')
- if isinstance(node.func, ast.Name):
- func_name = node.func.id
- if func_name == 'print':
- node.func = parse_expr('ti.ti_print')
- elif func_name == 'min':
- node.func = parse_expr('ti.ti_min')
- elif func_name == 'max':
- node.func = parse_expr('ti.ti_max')
- elif func_name == 'int':
- node.func = parse_expr('ti.ti_int')
- elif func_name == 'float':
- node.func = parse_expr('ti.ti_float')
- elif func_name == 'any':
- node.func = parse_expr('ti.ti_any')
- elif func_name == 'all':
- node.func = parse_expr('ti.ti_all')
- else:
- pass
-
- _taichi_skip_traceback = 1
- ti_func = node.func
- if '_sitebuiltins' == getattr(ti_func, '__module__', '') and getattr(
- getattr(ti_func, '__class__', ''), '__name__',
- '') == 'Quitter':
- raise TaichiSyntaxError(
- f'exit or quit not supported in Taichi-scope')
- if getattr(ti_func, '__module__', '') == '__main__' and not getattr(
- ti_func, '__wrapped__', ''):
- warnings.warn(
- f'Calling into non-Taichi function {ti_func.__name__}.'
- ' This means that scope inside that function will not be processed'
- ' by the Taichi transformer. Proceed with caution! '
- ' Maybe you want to decorate it with @ti.func?',
- UserWarning,
- stacklevel=2)
-
- return node
-
- @staticmethod
- def build_IfExp(ctx, node):
- node.test = build_expr(ctx, node.test)
- node.body = build_expr(ctx, node.body)
- node.orelse = build_expr(ctx, node.orelse)
-
- call = ast.Call(func=parse_expr('ti.select'),
- args=[node.test, node.body, node.orelse],
- keywords=[])
- return ast.copy_location(call, node)
-
- @staticmethod
- def build_UnaryOp(ctx, node):
- node.operand = build_expr(ctx, node.operand)
- if isinstance(node.op, ast.Not):
- # Python does not support overloading logical and & or
- new_node = parse_expr('ti.logical_not(0)')
- new_node.args[0] = node.operand
- node = new_node
- return node
-
- @staticmethod
- def build_BoolOp(ctx, node):
- node.values = build_exprs(ctx, node.values)
-
- def make_node(a, b, token):
- new_node = parse_expr('ti.logical_{}(0, 0)'.format(token))
- new_node.args[0] = a
- new_node.args[1] = b
- return new_node
-
- token = ''
- if isinstance(node.op, ast.And):
- token = 'and'
- elif isinstance(node.op, ast.Or):
- token = 'or'
- else:
- print(node.op)
- print("BoolOp above not implemented")
- exit(0)
-
- new_node = node.values[0]
- for i in range(1, len(node.values)):
- new_node = make_node(new_node, node.values[i], token)
-
- return new_node
-
- @staticmethod
- def build_BinOp(ctx, node):
- node.left = build_expr(ctx, node.left)
- node.right = build_expr(ctx, node.right)
- return node
-
- @staticmethod
- def build_Attribute(ctx, node):
- node.value = build_expr(ctx, node.value)
- return node
-
- @staticmethod
- def build_List(ctx, node):
- node.elts = build_exprs(ctx, node.elts)
- return node
-
- @staticmethod
- def build_Tuple(ctx, node):
- node.elts = build_exprs(ctx, node.elts)
- return node
-
- @staticmethod
- def build_Dict(ctx, node):
- node.keys = build_exprs(ctx, node.keys)
- node.values = build_exprs(ctx, node.values)
- return node
-
- @staticmethod
- def build_ListComp(ctx, node):
- node.elt = build_expr(ctx, node.elt)
- node.generators = build_exprs(ctx, node.generators)
- return node
-
- @staticmethod
- def build_DictComp(ctx, node):
- node.key = build_expr(ctx, node.value)
- node.value = build_expr(ctx, node.value)
- node.generators = build_exprs(ctx, node.generators)
- return node
-
- @staticmethod
- def build_comprehension(ctx, node):
- node.target = build_expr(ctx, node.target)
- node.iter = build_expr(ctx, node.iter)
- node.ifs = build_exprs(ctx, node.ifs)
- return node
-
- @staticmethod
- def build_Starred(ctx, node):
- node.value = build_expr(ctx, node.value)
- return node
-
- @staticmethod
- def build_Set(ctx, node):
- raise TaichiSyntaxError(
- 'Python set is not supported in Taichi kernels.')
-
- @staticmethod
- def build_Name(ctx, node):
- return node
-
- @staticmethod
- def build_NamedExpr(ctx, node):
- node.value = build_expr(ctx, node.value)
- return node
-
- @staticmethod
- def build_Constant(ctx, node):
- return node
-
- # Methods for Python 3.7 or lower
- @staticmethod
- def build_Num(ctx, node):
- return node
-
- @staticmethod
- def build_Str(ctx, node):
- return node
-
- @staticmethod
- def build_Bytes(ctx, node):
- return node
-
- @staticmethod
- def build_NameConstant(ctx, node):
- return node
-
-
-build_expr = ExprBuilder()
-
-
-def build_exprs(ctx, exprs):
- result = []
- # TODO(#2495): check if we really need this variable scope
- with ctx.variable_scope(result):
- for expr in list(exprs):
- result.append(build_expr(ctx, expr))
- return result
diff --git a/python/taichi/lang/field.py b/python/taichi/lang/field.py
index 19f6e0ab0b939..7afe75bdc996d 100644
--- a/python/taichi/lang/field.py
+++ b/python/taichi/lang/field.py
@@ -1,9 +1,7 @@
import taichi.lang
-from taichi.core.util import ti_core as _ti_core
+from taichi._lib import core as _ti_core
from taichi.lang.util import python_scope, to_numpy_type, to_pytorch_type
-import taichi as ti
-
class Field:
"""Taichi field with SNode implementation.
@@ -16,8 +14,8 @@ class Field:
Args:
vars (List[Expr]): Field members.
"""
- def __init__(self, vars):
- self.vars = vars
+ def __init__(self, _vars):
+ self.vars = _vars
self.host_accessors = None
self.grad = None
@@ -25,6 +23,15 @@ def __init__(self, vars):
def snode(self):
"""Gets representative SNode for info purposes.
+ Returns:
+ SNode: Representative SNode (SNode of first field member).
+ """
+ return self._snode
+
+ @property
+ def _snode(self):
+ """Gets representative SNode for info purposes.
+
Returns:
SNode: Representative SNode (SNode of first field member).
"""
@@ -37,7 +44,7 @@ def shape(self):
Returns:
Tuple[Int]: Field shape.
"""
- return self.snode.shape
+ return self._snode.shape
@property
def dtype(self):
@@ -46,16 +53,16 @@ def dtype(self):
Returns:
DataType: Data type of each individual value.
"""
- return self.snode.dtype
+ return self._snode._dtype
@property
- def name(self):
+ def _name(self):
"""Gets field name.
Returns:
str: Field name.
"""
- return self.snode.name
+ return self._snode._name
def parent(self, n=1):
"""Gets an ancestor of the representative SNode in the SNode tree.
@@ -68,7 +75,7 @@ def parent(self, n=1):
"""
return self.snode.parent(n)
- def get_field_members(self):
+ def _get_field_members(self):
"""Gets field members.
Returns:
@@ -76,7 +83,7 @@ def get_field_members(self):
"""
return self.vars
- def loop_range(self):
+ def _loop_range(self):
"""Gets representative field member for loop range info.
Returns:
@@ -84,7 +91,7 @@ def loop_range(self):
"""
return self.vars[0].ptr
- def set_grad(self, grad):
+ def _set_grad(self, grad):
"""Sets corresponding gradient field.
Args:
@@ -156,9 +163,13 @@ def copy_from(self, other):
Args:
other (Field): The source field.
"""
- assert isinstance(other, Field)
- assert len(self.shape) == len(other.shape)
- taichi.lang.meta.tensor_to_tensor(self, other)
+ if not isinstance(other, Field):
+ raise TypeError('Cannot copy from a non-field object')
+ if self.shape != other.shape:
+ raise ValueError(f"ti.field shape {self.shape} does not match"
+ f" the source field shape {other.shape}")
+ from taichi._kernels import tensor_to_tensor # pylint: disable=C0415
+ tensor_to_tensor(self, other)
@python_scope
def __setitem__(self, key, value):
@@ -185,12 +196,11 @@ def __getitem__(self, key):
def __str__(self):
if taichi.lang.impl.inside_kernel():
return self.__repr__() # make pybind11 happy, see Matrix.__str__
- if self.snode.ptr is None:
+ if self._snode.ptr is None:
return ''
- else:
- return str(self.to_numpy())
+ return str(self.to_numpy())
- def pad_key(self, key):
+ def _pad_key(self, key):
if key is None:
key = ()
if not isinstance(key, (tuple, list)):
@@ -198,7 +208,7 @@ def pad_key(self, key):
assert len(key) == len(self.shape)
return key + ((0, ) * (_ti_core.get_max_num_indices() - len(key)))
- def initialize_host_accessors(self):
+ def _initialize_host_accessors(self):
if self.host_accessors:
return
taichi.lang.impl.get_runtime().materialize()
@@ -206,7 +216,7 @@ def initialize_host_accessors(self):
SNodeHostAccessor(e.ptr.snode()) for e in self.vars
]
- def host_access(self, key):
+ def _host_access(self, key):
return [SNodeHostAccess(e, key) for e in self.host_accessors]
@@ -221,7 +231,8 @@ def __init__(self, var):
@python_scope
def fill(self, val):
- taichi.lang.meta.fill_tensor(self, val)
+ from taichi._kernels import fill_tensor # pylint: disable=C0415
+ fill_tensor(self, val)
@python_scope
def to_numpy(self, dtype=None):
@@ -229,18 +240,22 @@ def to_numpy(self, dtype=None):
dtype = to_numpy_type(self.dtype)
import numpy as np # pylint: disable=C0415
arr = np.zeros(shape=self.shape, dtype=dtype)
- taichi.lang.meta.tensor_to_ext_arr(self, arr)
- ti.sync()
+ from taichi._kernels import tensor_to_ext_arr # pylint: disable=C0415
+ tensor_to_ext_arr(self, arr)
+ taichi.lang.runtime_ops.sync()
return arr
@python_scope
def to_torch(self, device=None):
import torch # pylint: disable=C0415
+
+ # pylint: disable=E1101
arr = torch.zeros(size=self.shape,
dtype=to_pytorch_type(self.dtype),
device=device)
- taichi.lang.meta.tensor_to_ext_arr(self, arr)
- ti.sync()
+ from taichi._kernels import tensor_to_ext_arr # pylint: disable=C0415
+ tensor_to_ext_arr(self, arr)
+ taichi.lang.runtime_ops.sync()
return arr
@python_scope
@@ -254,18 +269,19 @@ def from_numpy(self, arr):
f" the numpy array shape {arr.shape}")
if hasattr(arr, 'contiguous'):
arr = arr.contiguous()
- taichi.lang.meta.ext_arr_to_tensor(arr, self)
- ti.sync()
+ from taichi._kernels import ext_arr_to_tensor # pylint: disable=C0415
+ ext_arr_to_tensor(arr, self)
+ taichi.lang.runtime_ops.sync()
@python_scope
def __setitem__(self, key, value):
- self.initialize_host_accessors()
- self.host_accessors[0].setter(value, *self.pad_key(key))
+ self._initialize_host_accessors()
+ self.host_accessors[0].setter(value, *self._pad_key(key))
@python_scope
def __getitem__(self, key):
- self.initialize_host_accessors()
- return self.host_accessors[0].getter(*self.pad_key(key))
+ self._initialize_host_accessors()
+ return self.host_accessors[0].getter(*self._pad_key(key))
def __repr__(self):
# make interactive shell happy, prevent materialization
@@ -307,3 +323,6 @@ class SNodeHostAccess:
def __init__(self, accessor, key):
self.accessor = accessor
self.key = key
+
+
+__all__ = ["Field", "ScalarField"]
diff --git a/python/taichi/lang/impl.py b/python/taichi/lang/impl.py
index 4811a3b399059..742cccf31792c 100644
--- a/python/taichi/lang/impl.py
+++ b/python/taichi/lang/impl.py
@@ -1,80 +1,84 @@
import numbers
-import warnings
from types import FunctionType, MethodType
from typing import Iterable
import numpy as np
-from taichi.core.util import ti_core as _ti_core
+from taichi._lib import core as _ti_core
+from taichi._snode.fields_builder import FieldsBuilder
from taichi.lang._ndarray import ScalarNdarray
+from taichi.lang._ndrange import GroupedNDRange, _Ndrange
from taichi.lang.any_array import AnyArray, AnyArrayAccess
-from taichi.lang.exception import InvalidOperationError, TaichiSyntaxError
+from taichi.lang.exception import TaichiRuntimeError
from taichi.lang.expr import Expr, make_expr_group
from taichi.lang.field import Field, ScalarField
from taichi.lang.kernel_arguments import SparseMatrixProxy
-from taichi.lang.matrix import MatrixField
+from taichi.lang.matrix import (Matrix, MatrixField, _IntermediateMatrix,
+ _MatrixFieldElement)
+from taichi.lang.mesh import (ConvType, MeshElementFieldProxy, MeshInstance,
+ MeshRelationAccessProxy,
+ MeshReorderedMatrixFieldProxy,
+ MeshReorderedScalarFieldProxy, element_type_name)
from taichi.lang.snode import SNode
-from taichi.lang.struct import StructField
+from taichi.lang.struct import Struct, StructField, _IntermediateStruct
from taichi.lang.tape import TapeImpl
-from taichi.lang.util import (cook_dtype, has_pytorch, is_taichi_class,
- python_scope, taichi_scope, to_pytorch_type)
-from taichi.misc.util import deprecated, get_traceback, warning
-from taichi.snode.fields_builder import FieldsBuilder
-from taichi.type.primitive_types import f32, f64, i32, i64, u32, u64
-
-import taichi as ti
+from taichi.lang.util import (cook_dtype, get_traceback, is_taichi_class,
+ python_scope, taichi_scope, warning)
+from taichi.types.primitive_types import f16, f32, f64, i32, i64
@taichi_scope
def expr_init_local_tensor(shape, element_type, elements):
- return _ti_core.expr_alloca_local_tensor(shape, element_type, elements)
+ return get_runtime().prog.current_ast_builder().expr_alloca_local_tensor(
+ shape, element_type, elements)
@taichi_scope
def expr_init(rhs):
if rhs is None:
- return Expr(_ti_core.expr_alloca())
- if is_taichi_class(rhs):
- if rhs.local_tensor_proxy is not None:
- return rhs
- else:
- return rhs.variable()
- else:
- if isinstance(rhs, list):
- return [expr_init(e) for e in rhs]
- elif isinstance(rhs, tuple):
- return tuple(expr_init(e) for e in rhs)
- elif isinstance(rhs, dict):
- return dict((key, expr_init(val)) for key, val in rhs.items())
- elif isinstance(rhs, _ti_core.DataType):
- return rhs
- elif isinstance(rhs, _ti_core.Arch):
- return rhs
- elif isinstance(rhs, ti.ndrange):
- return rhs
- elif hasattr(rhs, '_data_oriented'):
- return rhs
- else:
- return Expr(_ti_core.expr_var(Expr(rhs).ptr))
+ return Expr(get_runtime().prog.current_ast_builder().expr_alloca())
+ if isinstance(rhs, Matrix):
+ return Matrix(rhs.to_list())
+ if isinstance(rhs, Struct):
+ return Struct(rhs.to_dict())
+ if isinstance(rhs, list):
+ return [expr_init(e) for e in rhs]
+ if isinstance(rhs, tuple):
+ return tuple(expr_init(e) for e in rhs)
+ if isinstance(rhs, dict):
+ return dict((key, expr_init(val)) for key, val in rhs.items())
+ if isinstance(rhs, _ti_core.DataType):
+ return rhs
+ if isinstance(rhs, _ti_core.Arch):
+ return rhs
+ if isinstance(rhs, _Ndrange):
+ return rhs
+ if isinstance(rhs, MeshElementFieldProxy):
+ return rhs
+ if isinstance(rhs, MeshRelationAccessProxy):
+ return rhs
+ if hasattr(rhs, '_data_oriented'):
+ return rhs
+ return Expr(get_runtime().prog.current_ast_builder().expr_var(
+ Expr(rhs).ptr))
@taichi_scope
def expr_init_list(xs, expected):
- if not isinstance(xs, (list, tuple, ti.Matrix)):
+ if not isinstance(xs, (list, tuple, Matrix)):
raise TypeError(f'Cannot unpack type: {type(xs)}')
- if isinstance(xs, ti.Matrix):
+ if isinstance(xs, Matrix):
if not xs.m == 1:
raise ValueError(
- f'Matrices with more than one columns cannot be unpacked')
+ 'Matrices with more than one columns cannot be unpacked')
xs = xs.entries
if expected != len(xs):
raise ValueError(
f'Tuple assignment size mismatch: {expected} != {len(xs)}')
if isinstance(xs, list):
return [expr_init(e) for e in xs]
- elif isinstance(xs, tuple):
+ if isinstance(xs, tuple):
return tuple(expr_init(e) for e in xs)
- else:
- raise ValueError(f'Cannot unpack from {type(xs)}')
+ raise ValueError(f'Cannot unpack from {type(xs)}')
@taichi_scope
@@ -85,7 +89,7 @@ def expr_init_func(
return expr_init(rhs)
-def begin_frontend_struct_for(group, loop_range):
+def begin_frontend_struct_for(ast_builder, group, loop_range):
if not isinstance(loop_range, (AnyArray, Field, SNode, _Root)):
raise TypeError(
'Can only iterate through Taichi fields/snodes (via template) or dense arrays (via any_arr)'
@@ -96,10 +100,11 @@ def begin_frontend_struct_for(group, loop_range):
f'({group.size()} != {len(loop_range.shape)}). Maybe you wanted to '
'use "for I in ti.grouped(x)" to group all indices into a single vector I?'
)
- _ti_core.begin_frontend_struct_for(group, loop_range.loop_range())
+ ast_builder.begin_frontend_struct_for(group, loop_range._loop_range())
-def begin_frontend_if(cond):
+def begin_frontend_if(ast_builder, cond):
+ assert ast_builder is not None
if is_taichi_class(cond):
raise ValueError(
'The truth value of vectors/matrices is ambiguous.\n'
@@ -107,70 +112,85 @@ def begin_frontend_if(cond):
' if all(x == y):\n'
'or\n'
' if any(x != y):\n')
- _ti_core.begin_frontend_if(Expr(cond).ptr)
-
-
-def wrap_scalar(x):
- if type(x) in [int, float]:
- return Expr(x)
- else:
- return x
+ ast_builder.begin_frontend_if(Expr(cond).ptr)
@taichi_scope
-def subscript(value, *indices):
- _taichi_skip_traceback = 1
+def subscript(value, *_indices, skip_reordered=False):
if isinstance(value, np.ndarray):
- return value.__getitem__(*indices)
+ return value.__getitem__(_indices)
if isinstance(value, (tuple, list, dict)):
- assert len(indices) == 1
- return value[indices[0]]
+ assert len(_indices) == 1
+ return value[_indices[0]]
+ has_slice = False
flattened_indices = []
- for i in range(len(indices)):
- if is_taichi_class(indices[i]):
- ind = indices[i].entries
+ for _index in _indices:
+ if is_taichi_class(_index):
+ ind = _index.entries
+ elif isinstance(_index, slice):
+ ind = [_index]
+ has_slice = True
else:
- ind = [indices[i]]
+ ind = [_index]
flattened_indices += ind
- indices = tuple(flattened_indices)
- if isinstance(indices, tuple) and len(indices) == 1 and indices[0] is None:
- indices = ()
- indices_expr_group = make_expr_group(*indices)
- index_dim = indices_expr_group.size()
+ _indices = tuple(flattened_indices)
+ if isinstance(_indices,
+ tuple) and len(_indices) == 1 and _indices[0] is None:
+ _indices = ()
+
+ if has_slice:
+ if not isinstance(value, Matrix):
+ raise SyntaxError(
+ f"The type {type(value)} do not support index of slice type")
+ else:
+ indices_expr_group = make_expr_group(*_indices)
+ index_dim = indices_expr_group.size()
if is_taichi_class(value):
- return value.subscript(*indices)
- elif isinstance(value, SparseMatrixProxy):
- return value.subscript(*indices)
- elif isinstance(value, Field):
- var = value.get_field_members()[0].ptr
- if var.snode() is None:
- if var.is_primal():
+ return value._subscript(*_indices)
+ if isinstance(value, MeshElementFieldProxy):
+ return value.subscript(*_indices)
+ if isinstance(value, MeshRelationAccessProxy):
+ return value.subscript(*_indices)
+ if isinstance(value,
+ (MeshReorderedScalarFieldProxy,
+ MeshReorderedMatrixFieldProxy)) and not skip_reordered:
+ assert index_dim == 1
+ reordered_index = tuple([
+ Expr(
+ _ti_core.get_index_conversion(value.mesh_ptr,
+ value.element_type,
+ Expr(_indices[0]).ptr,
+ ConvType.g2r))
+ ])
+ return subscript(value, *reordered_index, skip_reordered=True)
+ if isinstance(value, SparseMatrixProxy):
+ return value.subscript(*_indices)
+ if isinstance(value, Field):
+ _var = value._get_field_members()[0].ptr
+ if _var.snode() is None:
+ if _var.is_primal():
raise RuntimeError(
- f"{var.get_expr_name()} has not been placed.")
+ f"{_var.get_expr_name()} has not been placed.")
else:
raise RuntimeError(
- f"Gradient {var.get_expr_name()} has not been placed, check whether `needs_grad=True`"
+ f"Gradient {_var.get_expr_name()} has not been placed, check whether `needs_grad=True`"
)
- field_dim = int(var.get_attribute("dim"))
+ field_dim = int(_var.get_attribute("dim"))
if field_dim != index_dim:
raise IndexError(
f'Field with dim {field_dim} accessed with indices of dim {index_dim}'
)
if isinstance(value, MatrixField):
- return ti.Matrix.with_entries(value.n, value.m, [
- Expr(_ti_core.subscript(e.ptr, indices_expr_group))
- for e in value.get_field_members()
- ])
- elif isinstance(value, StructField):
- return ti.Struct(
- {k: subscript(v, *indices)
- for k, v in value.items})
- else:
- return Expr(_ti_core.subscript(var, indices_expr_group))
- elif isinstance(value, AnyArray):
+ return _MatrixFieldElement(value, indices_expr_group)
+ if isinstance(value, StructField):
+ return _IntermediateStruct(
+ {k: subscript(v, *_indices)
+ for k, v in value._items})
+ return Expr(_ti_core.subscript(_var, indices_expr_group))
+ if isinstance(value, AnyArray):
# TODO: deprecate using get_attribute to get dim
field_dim = int(value.ptr.get_attribute("dim"))
element_dim = len(value.element_shape)
@@ -182,14 +202,14 @@ def subscript(value, *indices):
return Expr(_ti_core.subscript(value.ptr, indices_expr_group))
n = value.element_shape[0]
m = 1 if element_dim == 1 else value.element_shape[1]
- any_array_access = AnyArrayAccess(value, indices)
- ret = ti.Matrix.with_entries(n, m, [
+ any_array_access = AnyArrayAccess(value, _indices)
+ ret = _IntermediateMatrix(n, m, [
any_array_access.subscript(i, j) for i in range(n)
for j in range(m)
])
ret.any_array_access = any_array_access
return ret
- elif isinstance(value, SNode):
+ if isinstance(value, SNode):
# When reading bit structure we only support the 0-D case for now.
field_dim = 0
if field_dim != index_dim:
@@ -197,105 +217,40 @@ def subscript(value, *indices):
f'Field with dim {field_dim} accessed with indices of dim {index_dim}'
)
return Expr(_ti_core.subscript(value.ptr, indices_expr_group))
- else: # Directly evaluate in Python for non-Taichi types
- return value.__getitem__(*indices)
-
-
-@taichi_scope
-def local_subscript_with_offset(var, indices, shape):
- return Expr(
- _ti_core.local_subscript_with_offset(var, make_expr_group(*indices),
- shape))
+ # Directly evaluate in Python for non-Taichi types
+ return value.__getitem__(*_indices)
@taichi_scope
-def global_subscript_with_offset(var, indices, shape, is_aos):
+def make_tensor_element_expr(_var, _indices, shape, stride):
return Expr(
- _ti_core.global_subscript_with_offset(var.ptr,
- make_expr_group(*indices), shape,
- is_aos))
-
-
-@taichi_scope
-def chain_compare(comparators, ops):
- _taichi_skip_traceback = 1
- assert len(comparators) == len(ops) + 1, \
- f'Chain comparison invoked with {len(comparators)} comparators but {len(ops)} operators'
- ret = True
- for i in range(len(ops)):
- lhs = comparators[i]
- rhs = comparators[i + 1]
- if ops[i] == 'Lt':
- now = lhs < rhs
- elif ops[i] == 'LtE':
- now = lhs <= rhs
- elif ops[i] == 'Gt':
- now = lhs > rhs
- elif ops[i] == 'GtE':
- now = lhs >= rhs
- elif ops[i] == 'Eq':
- now = lhs == rhs
- elif ops[i] == 'NotEq':
- now = lhs != rhs
- else:
- assert False, f'Unknown operator {ops[i]}'
- ret = ti.logical_and(ret, now)
- return ret
-
-
-@taichi_scope
-def insert_expr_stmt_if_ti_func(func, *args, **kwargs):
- """This method is used only for real functions. It inserts a
- FrontendExprStmt to the C++ AST to hold the function call if `func` is a
- Taichi function.
-
- Args:
- func: The function to be called.
- args: The arguments of the function call.
- kwargs: The keyword arguments of the function call.
-
- Returns:
- The return value of the function call if it's a non-Taichi function.
- Returns None if it's a Taichi function."""
- is_taichi_function = getattr(func, '_is_taichi_function', False)
- # If is_taichi_function is true: call a decorated Taichi function
- # in a Taichi kernel/function.
-
- if is_taichi_function:
- # Compiles the function here.
- # Invokes Func.__call__.
- func_call_result = func(*args, **kwargs)
- # Insert FrontendExprStmt here.
- return _ti_core.insert_expr_stmt(func_call_result.ptr)
- else:
- # Call the non-Taichi function directly.
- return func(*args, **kwargs)
+ _ti_core.make_tensor_element_expr(_var, make_expr_group(*_indices),
+ shape, stride))
class PyTaichi:
def __init__(self, kernels=None):
self.materialized = False
self.prog = None
- self.materialize_callbacks = []
self.compiled_functions = {}
self.compiled_grad_functions = {}
self.scope_stack = []
self.inside_kernel = False
self.current_kernel = None
self.global_vars = []
- self.print_preprocessed = False
- self.experimental_real_function = False
+ self.matrix_fields = []
self.default_fp = f32
self.default_ip = i32
self.target_tape = None
self.grad_replaced = False
self.kernels = kernels or []
+ self._signal_handler_registry = None
def get_num_compiled_functions(self):
return len(self.compiled_functions) + len(self.compiled_grad_functions)
def set_default_fp(self, fp):
- assert fp in [f32, f64]
+ assert fp in [f16, f32, f64]
self.default_fp = fp
default_cfg().default_fp = self.default_fp
@@ -308,28 +263,37 @@ def create_program(self):
if self.prog is None:
self.prog = _ti_core.Program()
- def materialize_root_fb(self, first):
- if not root.finalized and not root.empty:
- root.finalize()
- elif first:
- root.finalize(raise_warning=False)
-
+ @staticmethod
+ def materialize_root_fb(is_first_call):
if root.finalized:
- global _root_fb
- _root_fb = FieldsBuilder()
+ return
+ if not is_first_call and root.empty:
+ # We have to forcefully finalize when `is_first_call` is True (even
+ # if the root itself is empty), so that there is a valid struct
+ # llvm::Module, if no field has been declared before the first kernel
+ # invocation. Example case:
+ # https://github.com/taichi-dev/taichi/blob/27bb1dc3227d9273a79fcb318fdb06fd053068f5/tests/python/test_ad_basics.py#L260-L266
+ return
+ root.finalize(raise_warning=not is_first_call)
+ global _root_fb
+ _root_fb = FieldsBuilder()
- def materialize(self):
- self.materialize_root_fb(not self.materialized)
+ @staticmethod
+ def _finalize_root_fb_for_aot():
+ if _root_fb.finalized:
+ raise RuntimeError(
+ 'AOT: can only finalize the root FieldsBuilder once')
+ _root_fb._finalize_for_aot()
- if self.materialized:
- return
+ @staticmethod
+ def _get_tb(_var):
+ return getattr(_var, 'declaration_tb', str(_var.ptr))
- self.materialized = True
+ def _check_field_not_placed(self):
not_placed = []
- for var in self.global_vars:
- if var.ptr.snode() is None:
- tb = getattr(var, 'declaration_tb', str(var.ptr))
- not_placed.append(tb)
+ for _var in self.global_vars:
+ if _var.ptr.snode() is None:
+ not_placed.append(self._get_tb(_var))
if len(not_placed):
bar = '=' * 44 + '\n'
@@ -339,14 +303,41 @@ def materialize(self):
f'{bar}Please consider specifying a shape for them. E.g.,' +
'\n\n x = ti.field(float, shape=(2, 3))')
- for callback in self.materialize_callbacks:
- callback()
- self.materialize_callbacks = []
+ def _check_matrix_field_member_shape(self):
+ for _field in self.matrix_fields:
+ shapes = [
+ _field.get_scalar_field(i, j).shape for i in range(_field.n)
+ for j in range(_field.m)
+ ]
+ if any(shape != shapes[0] for shape in shapes):
+ raise RuntimeError(
+ 'Members of the following field have different shapes ' +
+ f'{shapes}:\n{self._get_tb(_field._get_field_members()[0])}'
+ )
+
+ def _calc_matrix_field_dynamic_index_stride(self):
+ for _field in self.matrix_fields:
+ _field._calc_dynamic_index_stride()
+
+ def materialize(self):
+ self.materialize_root_fb(not self.materialized)
+ self.materialized = True
+
+ self._check_field_not_placed()
+ self._check_matrix_field_member_shape()
+ self._calc_matrix_field_dynamic_index_stride()
+ self.global_vars = []
+ self.matrix_fields = []
+
+ def _register_signal_handlers(self):
+ if self._signal_handler_registry is None:
+ self._signal_handler_registry = _ti_core.HackedSignalRegister()
def clear(self):
if self.prog:
self.prog.finalize()
self.prog = None
+ self._signal_handler_registry = None
self.materialized = False
def get_tape(self, loss=None):
@@ -364,56 +355,6 @@ def get_runtime():
return pytaichi
-def materialize_callback(foo):
- get_runtime().materialize_callbacks.append(foo)
-
-
-def _clamp_unsigned_to_range(npty, val):
- # npty: np.int32 or np.int64
- iif = np.iinfo(npty)
- if iif.min <= val <= iif.max:
- return val
- cap = (1 << iif.bits)
- if not (0 <= val < cap):
- # We let pybind11 fail intentionally, because this isn't the case we want
- # to deal with: |val| does't fall into the valid range of either
- # the signed or the unsigned type.
- return val
- new_val = val - cap
- ti.warn(
- f'Constant {val} has exceeded the range of {iif.bits} int, clamped to {new_val}'
- )
- return new_val
-
-
-@taichi_scope
-def make_constant_expr(val):
- _taichi_skip_traceback = 1
- if isinstance(val, (int, np.integer)):
- if pytaichi.default_ip in {i32, u32}:
- # It is not always correct to do such clamp without the type info on
- # the LHS, but at least this makes assigning constant to unsigned
- # int work. See https://github.com/taichi-dev/taichi/issues/2060
- return Expr(
- _ti_core.make_const_expr_i32(
- _clamp_unsigned_to_range(np.int32, val)))
- elif pytaichi.default_ip in {i64, u64}:
- return Expr(
- _ti_core.make_const_expr_i64(
- _clamp_unsigned_to_range(np.int64, val)))
- else:
- assert False
- elif isinstance(val, (float, np.floating, np.ndarray)):
- if pytaichi.default_fp == f32:
- return Expr(_ti_core.make_const_expr_f32(val))
- elif pytaichi.default_fp == f64:
- return Expr(_ti_core.make_const_expr_f64(val))
- else:
- assert False
- else:
- raise ValueError(f'Invalid constant scalar expression: {type(val)}')
-
-
def reset():
global pytaichi
old_kernels = pytaichi.kernels
@@ -431,7 +372,6 @@ def static_print(*args, __p=print, **kwargs):
# we don't add @taichi_scope decorator for @ti.pyfunc to work
def static_assert(cond, msg=None):
- _taichi_skip_traceback = 1
if msg is not None:
assert cond, msg
else:
@@ -451,7 +391,7 @@ def __getattr__(self, item):
if item == '__qualname__':
# For sphinx docstring extraction.
return '_UninitializedRootFieldsBuilder'
- raise InvalidOperationError('Please call init() first')
+ raise TaichiRuntimeError('Please call init() first')
# `root` initialization must be delayed until after the program is
@@ -469,26 +409,36 @@ def __getattr__(self, item):
_root_fb = _UninitializedRootFieldsBuilder()
+def deactivate_all_snodes():
+ """Recursively deactivate all SNodes."""
+ for root_fb in FieldsBuilder._finalized_roots():
+ root_fb.deactivate_all()
+
+
class _Root:
"""Wrapper around the default root FieldsBuilder instance."""
- def parent(self, n=1):
+ @staticmethod
+ def parent(n=1):
"""Same as :func:`taichi.SNode.parent`"""
return _root_fb.root.parent(n)
- def loop_range(self):
+ @staticmethod
+ def _loop_range():
"""Same as :func:`taichi.SNode.loop_range`"""
- return _root_fb.root.loop_range()
+ return _root_fb.root._loop_range()
- def get_children(self):
+ @staticmethod
+ def _get_children():
"""Same as :func:`taichi.SNode.get_children`"""
- return _root_fb.root.get_children()
+ return _root_fb.root._get_children()
# TODO: Record all of the SNodeTrees that finalized under 'ti.root'
- def deactivate_all(self):
+ @staticmethod
+ def deactivate_all():
warning(
"""'ti.root.deactivate_all()' would deactivate all finalized snodes."""
)
- ti.deactivate_all_snodes()
+ deactivate_all_snodes()
@property
def shape(self):
@@ -496,8 +446,8 @@ def shape(self):
return _root_fb.root.shape
@property
- def id(self):
- return _root_fb.root.id
+ def _id(self):
+ return _root_fb.root._id
def __getattr__(self, item):
return getattr(_root_fb, item)
@@ -524,7 +474,7 @@ def create_field_member(dtype, name):
# primal
x = Expr(_ti_core.make_id_expr(""))
- x.declaration_tb = get_traceback(stacklevel=2)
+ x.declaration_tb = get_traceback(stacklevel=4)
x.ptr = _ti_core.global_new(x.ptr, dtype)
x.ptr.set_name(name)
x.ptr.set_is_primal(True)
@@ -542,12 +492,6 @@ def create_field_member(dtype, name):
return x, x_grad
-@deprecated('ti.var', 'ti.field')
-def var(dt, shape=None, offset=None, needs_grad=False):
- _taichi_skip_traceback = 1
- return field(dt, shape, offset, needs_grad)
-
-
@python_scope
def field(dtype, shape=None, name="", offset=None, needs_grad=False):
"""Defines a Taichi field
@@ -576,7 +520,6 @@ def field(dtype, shape=None, name="", offset=None, needs_grad=False):
>>> x2 = ti.field(ti.f32)
>>> ti.root.dense(ti.ij, shape=(16, 8)).place(x2)
"""
- _taichi_skip_traceback = 1
if isinstance(shape, numbers.Number):
shape = (shape, )
@@ -590,13 +533,11 @@ def field(dtype, shape=None, name="", offset=None, needs_grad=False):
), f'The dimensionality of shape and offset must be the same ({len(shape)} != {len(offset)})'
assert (offset is None or shape
- is not None), f'The shape cannot be None when offset is being set'
-
- del _taichi_skip_traceback
+ is not None), 'The shape cannot be None when offset is being set'
x, x_grad = create_field_member(dtype, name)
x, x_grad = ScalarField(x), ScalarField(x_grad)
- x.set_grad(x_grad)
+ x._set_grad(x_grad)
if shape is not None:
dim = len(shape)
@@ -625,45 +566,43 @@ def ndarray(dtype, shape):
@taichi_scope
-def ti_print(*vars, sep=' ', end='\n'):
- def entry2content(var):
- if isinstance(var, str):
- return var
- else:
- return Expr(var).ptr
+def ti_print(*_vars, sep=' ', end='\n'):
+ def entry2content(_var):
+ if isinstance(_var, str):
+ return _var
+ return Expr(_var).ptr
- def list_ti_repr(var):
+ def list_ti_repr(_var):
yield '[' # distinguishing tuple & list will increase maintainance cost
- for i, v in enumerate(var):
+ for i, v in enumerate(_var):
if i:
yield ', '
yield v
yield ']'
- def vars2entries(vars):
- for var in vars:
- if hasattr(var, '__ti_repr__'):
- res = var.__ti_repr__()
- elif isinstance(var, (list, tuple)):
- res = var
+ def vars2entries(_vars):
+ for _var in _vars:
+ if hasattr(_var, '__ti_repr__'):
+ res = _var.__ti_repr__()
+ elif isinstance(_var, (list, tuple)):
# If the first element is '__ti_format__', this list is the result of ti_format.
- if len(var) > 0 and isinstance(
- var[0], str) and var[0] == '__ti_format__':
- res = var[1:]
+ if len(_var) > 0 and isinstance(
+ _var[0], str) and _var[0] == '__ti_format__':
+ res = _var[1:]
else:
- res = list_ti_repr(var)
+ res = list_ti_repr(_var)
else:
- yield var
+ yield _var
continue
for v in vars2entries(res):
yield v
- def add_separators(vars):
- for i, var in enumerate(vars):
+ def add_separators(_vars):
+ for i, _var in enumerate(_vars):
if i:
yield sep
- yield var
+ yield _var
yield end
def fused_string(entries):
@@ -679,28 +618,34 @@ def fused_string(entries):
if accumated:
yield accumated
- vars = add_separators(vars)
- entries = vars2entries(vars)
+ _vars = add_separators(_vars)
+ entries = vars2entries(_vars)
entries = fused_string(entries)
contentries = [entry2content(entry) for entry in entries]
- _ti_core.create_print(contentries)
+ get_runtime().prog.current_ast_builder().create_print(contentries)
@taichi_scope
-def ti_format(*args):
+def ti_format(*args, **kwargs):
content = args[0]
mixed = args[1:]
new_mixed = []
+ new_mixed_kwargs = {}
args = []
for x in mixed:
- if isinstance(x, ti.Expr):
+ if isinstance(x, Expr):
new_mixed.append('{}')
args.append(x)
else:
new_mixed.append(x)
-
+ for k, v in kwargs.items():
+ if isinstance(v, Expr):
+ new_mixed_kwargs[k] = '{}'
+ args.append(v)
+ else:
+ new_mixed_kwargs[k] = v
try:
- content = content.format(*new_mixed)
+ content = content.format(*new_mixed, **new_mixed_kwargs)
except ValueError:
print('Number formatting is not supported with Taichi fields')
exit(1)
@@ -709,8 +654,8 @@ def ti_format(*args):
args
) + 1, 'Number of args is different from number of positions provided in string'
- for i in range(len(args)):
- res.insert(i * 2 + 1, args[i])
+ for i, arg in enumerate(args):
+ res.insert(i * 2 + 1, arg)
res.insert(0, '__ti_format__')
return res
@@ -719,26 +664,22 @@ def ti_format(*args):
def ti_assert(cond, msg, extra_args):
# Mostly a wrapper to help us convert from Expr (defined in Python) to
# _ti_core.Expr (defined in C++)
- _ti_core.create_assert_stmt(
+ get_runtime().prog.current_ast_builder().create_assert_stmt(
Expr(cond).ptr, msg, [Expr(x).ptr for x in extra_args])
@taichi_scope
-def ti_int(var):
- _taichi_skip_traceback = 1
- if hasattr(var, '__ti_int__'):
- return var.__ti_int__()
- else:
- return int(var)
+def ti_int(_var):
+ if hasattr(_var, '__ti_int__'):
+ return _var.__ti_int__()
+ return int(_var)
@taichi_scope
-def ti_float(var):
- _taichi_skip_traceback = 1
- if hasattr(var, '__ti_float__'):
- return var.__ti_float__()
- else:
- return float(var)
+def ti_float(_var):
+ if hasattr(_var, '__ti_float__'):
+ return _var.__ti_float__()
+ return float(_var)
@taichi_scope
@@ -782,14 +723,6 @@ def axes(*x: Iterable[int]):
return [_ti_core.Axis(i) for i in x]
-@deprecated("ti.indices", "ti.axes")
-def indices(*x):
- """Same as :func:`~taichi.lang.impl.axes`."""
- return [_ti_core.Axis(i) for i in x]
-
-
-index = indices
-
Axis = _ti_core.Axis
@@ -833,24 +766,22 @@ def static(x, *xs):
>>> print(1)
>>> print(2)
"""
- _taichi_skip_traceback = 1
if len(xs): # for python-ish pointer assign: x, y = ti.static(y, x)
return [static(x)] + [static(x) for x in xs]
if isinstance(x,
- (bool, int, float, range, list, tuple, enumerate, ti.ndrange,
- ti.GroupedNDRange, zip, filter, map)) or x is None:
+ (bool, int, float, range, list, tuple, enumerate, _Ndrange,
+ GroupedNDRange, zip, filter, map)) or x is None:
return x
- elif isinstance(x, AnyArray):
+ if isinstance(x, AnyArray):
return x
- elif isinstance(x, Field):
+ if isinstance(x, Field):
return x
- elif isinstance(x, (FunctionType, MethodType)):
+ if isinstance(x, (FunctionType, MethodType)):
return x
- else:
- raise ValueError(
- f'Input to ti.static must be compile-time constants or global pointers, instead of {type(x)}'
- )
+ raise ValueError(
+ f'Input to ti.static must be compile-time constants or global pointers, instead of {type(x)}'
+ )
@taichi_scope
@@ -862,21 +793,20 @@ def grouped(x):
Example::
- >>> for I in ti.grouped(ti.ndrange(8, 16)):
+ >>> for I in ti.grouped(ndrange(8, 16)):
>>> print(I[0] + I[1])
"""
- if isinstance(x, ti.ndrange):
+ if isinstance(x, _Ndrange):
return x.grouped()
- else:
- return x
+ return x
def stop_grad(x):
- _ti_core.stop_grad(x.snode.ptr)
+ get_runtime().prog.current_ast_builder().stop_grad(x.snode.ptr)
def current_cfg():
- return _ti_core.current_compile_config()
+ return get_runtime().prog.config
def default_cfg():
@@ -886,3 +816,19 @@ def default_cfg():
def call_internal(name, *args):
return expr_init(
_ti_core.insert_internal_func_call(name, make_expr_group(args)))
+
+
+@taichi_scope
+def mesh_relation_access(mesh, from_index, to_element_type):
+ # to support ti.mesh_local and access mesh attribute as field
+ if isinstance(from_index, MeshInstance):
+ return getattr(from_index, element_type_name(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/python/taichi/lang/kernel_arguments.py b/python/taichi/lang/kernel_arguments.py
index e603a87d63c6d..aa1bf8a3c7324 100644
--- a/python/taichi/lang/kernel_arguments.py
+++ b/python/taichi/lang/kernel_arguments.py
@@ -1,64 +1,78 @@
import taichi.lang
-from taichi.core.util import ti_core as _ti_core
+from taichi._lib import core as _ti_core
+from taichi.lang import impl, ops
from taichi.lang.any_array import AnyArray
from taichi.lang.enums import Layout
from taichi.lang.expr import Expr
+from taichi.lang.matrix import Matrix, MatrixType
from taichi.lang.util import cook_dtype
-from taichi.linalg import SparseMatrixBuilder
-from taichi.type.primitive_types import u64
+from taichi.types.primitive_types import u64
class SparseMatrixEntry:
- def __init__(self, ptr, i, j):
+ def __init__(self, ptr, i, j, dtype):
self.ptr = ptr
self.i = i
self.j = j
+ self.dtype = dtype
- def augassign(self, value, op):
+ def _augassign(self, value, op):
+ call_func = f"insert_triplet_{self.dtype}"
if op == 'Add':
- taichi.lang.impl.call_internal("insert_triplet", self.ptr, self.i,
- self.j,
- taichi.lang.impl.ti_float(value))
+ taichi.lang.impl.call_internal(call_func, self.ptr, self.i, self.j,
+ ops.cast(value, self.dtype))
elif op == 'Sub':
- taichi.lang.impl.call_internal("insert_triplet", self.ptr, self.i,
- self.j,
- -taichi.lang.impl.ti_float(value))
+ taichi.lang.impl.call_internal(call_func, self.ptr, self.i, self.j,
+ -ops.cast(value, self.dtype))
else:
- assert False, f"Only operations '+=' and '-=' are supported on sparse matrices."
+ assert False, "Only operations '+=' and '-=' are supported on sparse matrices."
class SparseMatrixProxy:
- def __init__(self, ptr):
+ def __init__(self, ptr, dtype):
self.ptr = ptr
+ self.dtype = dtype
def subscript(self, i, j):
- return SparseMatrixEntry(self.ptr, i, j)
+ return SparseMatrixEntry(self.ptr, i, j, self.dtype)
def decl_scalar_arg(dtype):
dtype = cook_dtype(dtype)
- arg_id = _ti_core.decl_arg(dtype, False)
+ arg_id = impl.get_runtime().prog.decl_arg(dtype, False)
return Expr(_ti_core.make_arg_load_expr(arg_id, dtype))
-def decl_sparse_matrix():
+def decl_matrix_arg(matrixtype):
+ return Matrix(
+ [[decl_scalar_arg(matrixtype.dtype) for _ in range(matrixtype.m)]
+ for _ in range(matrixtype.n)])
+
+
+def decl_sparse_matrix(dtype):
+ value_type = cook_dtype(dtype)
ptr_type = cook_dtype(u64)
# Treat the sparse matrix argument as a scalar since we only need to pass in the base pointer
- arg_id = _ti_core.decl_arg(ptr_type, False)
- return SparseMatrixProxy(_ti_core.make_arg_load_expr(arg_id, ptr_type))
+ arg_id = impl.get_runtime().prog.decl_arg(ptr_type, False)
+ return SparseMatrixProxy(_ti_core.make_arg_load_expr(arg_id, ptr_type),
+ value_type)
def decl_any_arr_arg(dtype, dim, element_shape, layout):
dtype = cook_dtype(dtype)
- arg_id = _ti_core.decl_arg(dtype, True)
element_dim = len(element_shape)
+ arg_id = impl.get_runtime().prog.decl_arr_arg(dtype, dim, element_shape)
if layout == Layout.AOS:
element_dim = -element_dim
return AnyArray(
- _ti_core.make_external_tensor_expr(dtype, dim, arg_id, element_dim),
- element_shape, layout)
+ _ti_core.make_external_tensor_expr(dtype, dim, arg_id, element_dim,
+ element_shape), element_shape,
+ layout)
-def decl_scalar_ret(dtype):
- dtype = cook_dtype(dtype)
- return _ti_core.decl_ret(dtype)
+def decl_ret(dtype):
+ if isinstance(dtype, MatrixType):
+ dtype = _ti_core.decl_tensor_type([dtype.n, dtype.m], dtype.dtype)
+ else:
+ dtype = cook_dtype(dtype)
+ return impl.get_runtime().prog.decl_ret(dtype)
diff --git a/python/taichi/lang/kernel_impl.py b/python/taichi/lang/kernel_impl.py
index 1e5913dc2afc5..7ccad655421ff 100644
--- a/python/taichi/lang/kernel_impl.py
+++ b/python/taichi/lang/kernel_impl.py
@@ -1,49 +1,33 @@
import ast
-import copy
import functools
import inspect
import re
+import sys
+import textwrap
import numpy as np
import taichi.lang
-from taichi.core.util import ti_core as _ti_core
-from taichi.lang import impl, util
-from taichi.lang.ast.checkers import KernelSimplicityASTChecker
-from taichi.lang.ast.transformer import ASTTransformerTotal
+from taichi._lib import core as _ti_core
+from taichi.lang import impl, runtime_ops
+from taichi.lang.ast import (ASTTransformerContext, KernelSimplicityASTChecker,
+ transform_tree)
from taichi.lang.enums import Layout
-from taichi.lang.exception import TaichiSyntaxError
+from taichi.lang.exception import (TaichiCompilationError, TaichiRuntimeError,
+ TaichiRuntimeTypeError, TaichiSyntaxError)
+from taichi.lang.expr import Expr
+from taichi.lang.matrix import Matrix, MatrixType
from taichi.lang.shell import _shell_pop_print, oinspect
-from taichi.lang.util import to_taichi_type
-from taichi.linalg.sparse_matrix import sparse_matrix_builder
-from taichi.misc.util import obsolete
-from taichi.type import any_arr, primitive_types, template
+from taichi.lang.util import has_pytorch, to_taichi_type
+from taichi.types import (any_arr, primitive_types, sparse_matrix_builder,
+ template)
-import taichi as ti
+from taichi import _logging
-if util.has_pytorch():
+if has_pytorch():
import torch
-def _remove_indent(lines):
- lines = lines.split('\n')
- to_remove = 0
- for i in range(len(lines[0])):
- if lines[0][i] == ' ':
- to_remove = i + 1
- else:
- break
-
- cleaned = []
- for l in lines:
- cleaned.append(l[to_remove:])
- if len(l) >= to_remove:
- for i in range(to_remove):
- assert l[i] == ' '
-
- return '\n'.join(cleaned)
-
-
-def func(fn):
+def func(fn, is_real_function=False):
"""Marks a function as callable in Taichi-scope.
This decorator transforms a Python function into a Taichi one. Taichi
@@ -51,6 +35,7 @@ def func(fn):
Args:
fn (Callable): The Python function to be decorated
+ is_real_function (bool): Whether the function is a real function
Returns:
Callable: The decorated function
@@ -67,18 +52,21 @@ def func(fn):
"""
is_classfunc = _inside_class(level_of_class_stackframe=3)
- _taichi_skip_traceback = 1
- fun = Func(fn, classfunc=is_classfunc)
+ fun = Func(fn, _classfunc=is_classfunc, is_real_function=is_real_function)
@functools.wraps(fn)
def decorated(*args):
- _taichi_skip_traceback = 1
return fun.__call__(*args)
decorated._is_taichi_function = True
+ decorated._is_real_function = is_real_function
return decorated
+def real_func(fn):
+ return func(fn, is_real_function=True)
+
+
def pyfunc(fn):
"""Marks a function as callable in both Taichi and Python scopes.
@@ -95,54 +83,99 @@ def pyfunc(fn):
Callable: The decorated function
"""
is_classfunc = _inside_class(level_of_class_stackframe=3)
- fun = Func(fn, classfunc=is_classfunc, pyfunc=True)
+ fun = Func(fn, _classfunc=is_classfunc, _pyfunc=True)
@functools.wraps(fn)
def decorated(*args):
- _taichi_skip_traceback = 1
return fun.__call__(*args)
decorated._is_taichi_function = True
return decorated
+def _get_tree_and_ctx(self,
+ excluded_parameters=(),
+ is_kernel=True,
+ arg_features=None,
+ args=None,
+ ast_builder=None,
+ is_real_function=False):
+ file = oinspect.getsourcefile(self.func)
+ src, start_lineno = oinspect.getsourcelines(self.func)
+ src = [textwrap.fill(line, tabsize=4, width=9999) for line in src]
+ tree = ast.parse(textwrap.dedent("\n".join(src)))
+
+ func_body = tree.body[0]
+ func_body.decorator_list = []
+
+ global_vars = _get_global_vars(self.func)
+
+ for i, arg in enumerate(func_body.args.args):
+ anno = arg.annotation
+ if isinstance(anno, ast.Name):
+ global_vars[anno.id] = self.argument_annotations[i]
+
+ if isinstance(func_body.returns, ast.Name):
+ global_vars[func_body.returns.id] = self.return_type
+
+ if is_kernel or is_real_function:
+ # inject template parameters into globals
+ for i in self.template_slot_locations:
+ template_var_name = self.argument_names[i]
+ global_vars[template_var_name] = args[i]
+
+ return tree, ASTTransformerContext(excluded_parameters=excluded_parameters,
+ is_kernel=is_kernel,
+ func=self,
+ arg_features=arg_features,
+ global_vars=global_vars,
+ argument_data=args,
+ src=src,
+ start_lineno=start_lineno,
+ file=file,
+ ast_builder=ast_builder,
+ is_real_function=is_real_function)
+
+
class Func:
function_counter = 0
- def __init__(self, func, classfunc=False, pyfunc=False):
- self.func = func
+ def __init__(self,
+ _func,
+ _classfunc=False,
+ _pyfunc=False,
+ is_real_function=False):
+ self.func = _func
self.func_id = Func.function_counter
Func.function_counter += 1
self.compiled = None
- self.classfunc = classfunc
- self.pyfunc = pyfunc
+ self.classfunc = _classfunc
+ self.pyfunc = _pyfunc
+ self.is_real_function = is_real_function
self.argument_annotations = []
self.argument_names = []
- _taichi_skip_traceback = 1
+ self.return_type = None
self.extract_arguments()
self.template_slot_locations = []
- for i in range(len(self.argument_annotations)):
- if isinstance(self.argument_annotations[i], template):
+ for i, anno in enumerate(self.argument_annotations):
+ if isinstance(anno, template):
self.template_slot_locations.append(i)
self.mapper = TaichiCallableTemplateMapper(
self.argument_annotations, self.template_slot_locations)
self.taichi_functions = {} # The |Function| class in C++
def __call__(self, *args):
- _taichi_skip_traceback = 1
if not impl.inside_kernel():
if not self.pyfunc:
raise TaichiSyntaxError(
- "Taichi functions cannot be called from Python-scope."
- " Use @ti.pyfunc if you wish to call Taichi functions "
- "from both Python-scope and Taichi-scope.")
+ "Taichi functions cannot be called from Python-scope.")
return self.func(*args)
- if impl.get_runtime().experimental_real_function:
+ if self.is_real_function:
if impl.get_runtime().current_kernel.is_grad:
raise TaichiSyntaxError(
"Real function in gradient kernels unsupported.")
- instance_id, arg_features = self.mapper.lookup(args)
+ instance_id, _ = self.mapper.lookup(args)
key = _ti_core.FunctionKey(self.func.__name__, self.func_id,
instance_id)
if self.compiled is None:
@@ -150,58 +183,46 @@ def __call__(self, *args):
if key.instance_id not in self.compiled:
self.do_compile(key=key, args=args)
return self.func_call_rvalue(key=key, args=args)
- else:
- if self.compiled is None:
- self.do_compile(key=None, args=args)
- ret = self.compiled(*args)
- return ret
+ tree, ctx = _get_tree_and_ctx(
+ self,
+ is_kernel=False,
+ args=args,
+ ast_builder=impl.get_runtime().prog.current_ast_builder(),
+ is_real_function=self.is_real_function)
+ ret = transform_tree(tree, ctx)
+ if not self.is_real_function:
+ if self.return_type and not ctx.returned:
+ raise TaichiSyntaxError(
+ "Function has a return type but does not have a return statement"
+ )
+ return ret
def func_call_rvalue(self, key, args):
# Skip the template args, e.g., |self|
- assert impl.get_runtime().experimental_real_function
+ assert self.is_real_function
non_template_args = []
- for i in range(len(self.argument_annotations)):
- if not isinstance(self.argument_annotations[i], template):
+ for i, anno in enumerate(self.argument_annotations):
+ if not isinstance(anno, template):
non_template_args.append(args[i])
non_template_args = impl.make_expr_group(non_template_args)
- return ti.Expr(
+ return Expr(
_ti_core.make_func_call_expr(
self.taichi_functions[key.instance_id], non_template_args))
def do_compile(self, key, args):
- src = _remove_indent(oinspect.getsource(self.func))
- tree = ast.parse(src)
-
- func_body = tree.body[0]
- func_body.decorator_list = []
-
- visitor = ASTTransformerTotal(is_kernel=False, func=self)
- visitor.visit(tree)
-
- ast.increment_lineno(tree, oinspect.getsourcelines(self.func)[1] - 1)
+ tree, ctx = _get_tree_and_ctx(self,
+ is_kernel=False,
+ args=args,
+ is_real_function=self.is_real_function)
+ fn = impl.get_runtime().prog.create_function(key)
- local_vars = {}
- global_vars = _get_global_vars(self.func)
+ def func_body():
+ ctx.ast_builder = fn.ast_builder()
+ transform_tree(tree, ctx)
- if impl.get_runtime().experimental_real_function:
- # inject template parameters into globals
- for i in self.template_slot_locations:
- template_var_name = self.argument_names[i]
- global_vars[template_var_name] = args[i]
-
- exec(
- compile(tree,
- filename=oinspect.getsourcefile(self.func),
- mode='exec'), global_vars, local_vars)
-
- if impl.get_runtime().experimental_real_function:
- self.compiled[key.instance_id] = local_vars[self.func.__name__]
- self.taichi_functions[key.instance_id] = _ti_core.create_function(
- key)
- self.taichi_functions[key.instance_id].set_function_body(
- self.compiled[key.instance_id])
- else:
- self.compiled = local_vars[self.func.__name__]
+ self.taichi_functions[key.instance_id] = fn
+ self.compiled[key.instance_id] = func_body
+ self.taichi_functions[key.instance_id].set_function_body(func_body)
def extract_arguments(self):
sig = inspect.signature(self.func)
@@ -212,18 +233,18 @@ def extract_arguments(self):
for i, arg_name in enumerate(arg_names):
param = params[arg_name]
if param.kind == inspect.Parameter.VAR_KEYWORD:
- raise KernelDefError(
+ raise TaichiSyntaxError(
'Taichi functions do not support variable keyword parameters (i.e., **kwargs)'
)
if param.kind == inspect.Parameter.VAR_POSITIONAL:
- raise KernelDefError(
+ raise TaichiSyntaxError(
'Taichi functions do not support variable positional parameters (i.e., *args)'
)
if param.kind == inspect.Parameter.KEYWORD_ONLY:
- raise KernelDefError(
+ raise TaichiSyntaxError(
'Taichi functions do not support keyword parameters')
if param.kind != inspect.Parameter.POSITIONAL_OR_KEYWORD:
- raise KernelDefError(
+ raise TaichiSyntaxError(
'Taichi functions only support "positional or keyword" parameters'
)
annotation = param.annotation
@@ -232,16 +253,15 @@ def extract_arguments(self):
annotation = template()
# TODO: pyfunc also need type annotation check when real function is enabled,
# but that has to happen at runtime when we know which scope it's called from.
- elif not self.pyfunc and impl.get_runtime(
- ).experimental_real_function:
- raise KernelDefError(
+ elif not self.pyfunc and self.is_real_function:
+ raise TaichiSyntaxError(
f'Taichi function `{self.func.__name__}` parameter `{arg_name}` must be type annotated'
)
else:
if not id(annotation
) in primitive_types.type_ids and not isinstance(
annotation, template):
- raise KernelDefError(
+ raise TaichiSyntaxError(
f'Invalid type annotation (argument {i}) of Taichi function: {annotation}'
)
self.argument_annotations.append(annotation)
@@ -269,17 +289,23 @@ def extract_arg(arg, anno):
TaichiCallableTemplateMapper.extract_arg(item, anno)
for item in arg)
return arg
- elif isinstance(anno, any_arr):
+ if isinstance(anno, any_arr):
if isinstance(arg, taichi.lang._ndarray.ScalarNdarray):
- anno.check_element_dim(arg, 0)
+ anno._check_element_dim(arg, 0)
+ anno._check_element_shape(())
+ anno._check_field_dim(len(arg.shape))
return arg.dtype, len(arg.shape), (), Layout.AOS
if isinstance(arg, taichi.lang.matrix.VectorNdarray):
- anno.check_element_dim(arg, 1)
- anno.check_layout(arg)
+ anno._check_element_dim(arg, 1)
+ anno._check_element_shape((arg.n, ))
+ anno._check_field_dim(len(arg.shape))
+ anno._check_layout(arg)
return arg.dtype, len(arg.shape) + 1, (arg.n, ), arg.layout
if isinstance(arg, taichi.lang.matrix.MatrixNdarray):
- anno.check_element_dim(arg, 2)
- anno.check_layout(arg)
+ anno._check_element_dim(arg, 2)
+ anno._check_element_shape((arg.n, arg.m))
+ anno._check_field_dim(len(arg.shape))
+ anno._check_layout(arg)
return arg.dtype, len(arg.shape) + 2, (arg.n,
arg.m), arg.layout
# external arrays
@@ -288,14 +314,17 @@ def extract_arg(arg, anno):
shape = tuple(arg.shape)
if len(shape) < element_dim:
raise ValueError(
- f"Invalid argument into ti.any_arr() - required element_dim={element_dim}, but the argument has only {len(shape)} dimensions"
- )
+ f"Invalid argument into ti.any_arr() - required element_dim={element_dim}, "
+ f"but the argument has only {len(shape)} dimensions")
element_shape = (
) if element_dim == 0 else shape[:
element_dim] if layout == Layout.SOA else shape[
-element_dim:]
return to_taichi_type(arg.dtype), len(shape), element_shape, layout
- return (type(arg).__name__, )
+ if isinstance(anno, sparse_matrix_builder):
+ return arg.dtype
+ # Use '#' as a placeholder because other kinds of arguments are not involved in template instantiation
+ return '#'
def extract(self, args):
extracted = []
@@ -305,7 +334,6 @@ def extract(self, args):
def lookup(self, args):
if len(args) != self.num_args:
- _taichi_skip_traceback = 1
raise TypeError(
f'{self.num_args} argument(s) needed but {len(args)} provided.'
)
@@ -317,30 +345,25 @@ def lookup(self, args):
return self.mapping[key], key
-class KernelDefError(Exception):
- def __init__(self, msg):
- super().__init__(msg)
-
-
-class KernelArgError(Exception):
- def __init__(self, pos, needed, provided):
- message = f'Argument {pos} (type={provided}) cannot be converted into required type {needed}'
- super().__init__(message)
- self.pos = pos
- self.needed = needed
- self.provided = provided
+def _get_global_vars(_func):
+ # Discussions: https://github.com/taichi-dev/taichi/issues/282
+ global_vars = _func.__globals__.copy()
+ freevar_names = _func.__code__.co_freevars
+ closure = _func.__closure__
+ if closure:
+ freevar_values = list(map(lambda x: x.cell_contents, closure))
+ for name, value in zip(freevar_names, freevar_values):
+ global_vars[name] = value
-def _get_global_vars(func):
- closure_vars = inspect.getclosurevars(func)
- return {**closure_vars.globals, **closure_vars.nonlocals}
+ return global_vars
class Kernel:
counter = 0
- def __init__(self, func, is_grad, classkernel=False):
- self.func = func
+ def __init__(self, _func, is_grad, _classkernel=False):
+ self.func = _func
self.kernel_counter = Kernel.counter
Kernel.counter += 1
self.is_grad = is_grad
@@ -348,13 +371,11 @@ def __init__(self, func, is_grad, classkernel=False):
self.argument_annotations = []
self.argument_names = []
self.return_type = None
- self.classkernel = classkernel
- _taichi_skip_traceback = 1
+ self.classkernel = _classkernel
self.extract_arguments()
- del _taichi_skip_traceback
self.template_slot_locations = []
- for i in range(len(self.argument_annotations)):
- if isinstance(self.argument_annotations[i], template):
+ for i, anno in enumerate(self.argument_annotations):
+ if isinstance(anno, template):
self.template_slot_locations.append(i)
self.mapper = TaichiCallableTemplateMapper(
self.argument_annotations, self.template_slot_locations)
@@ -378,22 +399,22 @@ def extract_arguments(self):
for i, arg_name in enumerate(arg_names):
param = params[arg_name]
if param.kind == inspect.Parameter.VAR_KEYWORD:
- raise KernelDefError(
+ raise TaichiSyntaxError(
'Taichi kernels do not support variable keyword parameters (i.e., **kwargs)'
)
if param.kind == inspect.Parameter.VAR_POSITIONAL:
- raise KernelDefError(
+ raise TaichiSyntaxError(
'Taichi kernels do not support variable positional parameters (i.e., *args)'
)
if param.default is not inspect.Parameter.empty:
- raise KernelDefError(
+ raise TaichiSyntaxError(
'Taichi kernels do not support default values for arguments'
)
if param.kind == inspect.Parameter.KEYWORD_ONLY:
- raise KernelDefError(
+ raise TaichiSyntaxError(
'Taichi kernels do not support keyword parameters')
if param.kind != inspect.Parameter.POSITIONAL_OR_KEYWORD:
- raise KernelDefError(
+ raise TaichiSyntaxError(
'Taichi kernels only support "positional or keyword" parameters'
)
annotation = param.annotation
@@ -401,8 +422,7 @@ def extract_arguments(self):
if i == 0 and self.classkernel: # The |self| parameter
annotation = template()
else:
- _taichi_skip_traceback = 1
- raise KernelDefError(
+ raise TaichiSyntaxError(
'Taichi kernels parameters must be type annotated')
else:
if isinstance(annotation, (template, any_arr)):
@@ -411,16 +431,16 @@ def extract_arguments(self):
pass
elif isinstance(annotation, sparse_matrix_builder):
pass
+ elif isinstance(annotation, MatrixType):
+ pass
else:
- _taichi_skip_traceback = 1
- raise KernelDefError(
+ raise TaichiSyntaxError(
f'Invalid type annotation (argument {i}) of Taichi kernel: {annotation}'
)
self.argument_annotations.append(annotation)
self.argument_names.append(param.name)
def materialize(self, key=None, args=None, arg_features=None):
- _taichi_skip_traceback = 1
if key is None:
key = (self.func, 0)
self.runtime.materialize()
@@ -429,86 +449,100 @@ def materialize(self, key=None, args=None, arg_features=None):
grad_suffix = ""
if self.is_grad:
grad_suffix = "_grad"
- kernel_name = "{}_c{}_{}{}".format(self.func.__name__,
- self.kernel_counter, key[1],
- grad_suffix)
- ti.trace("Compiling kernel {}...".format(kernel_name))
-
- src = _remove_indent(oinspect.getsource(self.func))
- tree = ast.parse(src)
-
- func_body = tree.body[0]
- func_body.decorator_list = []
-
- local_vars = {}
- global_vars = _get_global_vars(self.func)
-
- for i, arg in enumerate(func_body.args.args):
- anno = arg.annotation
- if isinstance(anno, ast.Name):
- global_vars[anno.id] = self.argument_annotations[i]
-
- if isinstance(func_body.returns, ast.Name):
- global_vars[func_body.returns.id] = self.return_type
+ kernel_name = f"{self.func.__name__}_c{self.kernel_counter}_{key[1]}{grad_suffix}"
+ _logging.trace(f"Compiling kernel {kernel_name}...")
- if self.is_grad:
- KernelSimplicityASTChecker(self.func).visit(tree)
-
- visitor = ASTTransformerTotal(
+ tree, ctx = _get_tree_and_ctx(
+ self,
+ args=args,
excluded_parameters=self.template_slot_locations,
- func=self,
arg_features=arg_features)
- visitor.visit(tree)
-
- ast.increment_lineno(tree, oinspect.getsourcelines(self.func)[1] - 1)
-
- # inject template parameters into globals
- for i in self.template_slot_locations:
- template_var_name = self.argument_names[i]
- global_vars[template_var_name] = args[i]
-
- exec(
- compile(tree,
- filename=oinspect.getsourcefile(self.func),
- mode='exec'), global_vars, local_vars)
- compiled = local_vars[self.func.__name__]
+ if self.is_grad:
+ KernelSimplicityASTChecker(self.func).visit(tree)
# Do not change the name of 'taichi_ast_generator'
# The warning system needs this identifier to remove unnecessary messages
- def taichi_ast_generator():
- _taichi_skip_traceback = 1
+ def taichi_ast_generator(kernel_cxx):
if self.runtime.inside_kernel:
raise TaichiSyntaxError(
- "Kernels cannot call other kernels. I.e., nested kernels are not allowed. Please check if you have direct/indirect invocation of kernels within kernels. Note that some methods provided by the Taichi standard library may invoke kernels, and please move their invocations to Python-scope."
- )
+ "Kernels cannot call other kernels. I.e., nested kernels are not allowed. "
+ "Please check if you have direct/indirect invocation of kernels within kernels. "
+ "Note that some methods provided by the Taichi standard library may invoke kernels, "
+ "and please move their invocations to Python-scope.")
self.runtime.inside_kernel = True
self.runtime.current_kernel = self
try:
- compiled()
+ ctx.ast_builder = kernel_cxx.ast_builder()
+ transform_tree(tree, ctx)
+ if not ctx.is_real_function:
+ if self.return_type and not ctx.returned:
+ raise TaichiSyntaxError(
+ "Kernel has a return type but does not have a return statement"
+ )
finally:
self.runtime.inside_kernel = False
self.runtime.current_kernel = None
- taichi_kernel = _ti_core.create_kernel(taichi_ast_generator,
- kernel_name, self.is_grad)
+ taichi_kernel = impl.get_runtime().prog.create_kernel(
+ taichi_ast_generator, kernel_name, self.is_grad)
self.kernel_cpp = taichi_kernel
assert key not in self.compiled_functions
self.compiled_functions[key] = self.get_function_body(taichi_kernel)
+ def get_torch_callbacks(self, v, has_torch, is_ndarray=True):
+ callbacks = []
+
+ def get_call_back(u, v):
+ def call_back():
+ u.copy_(v)
+
+ return call_back
+
+ assert has_torch
+ assert isinstance(v, torch.Tensor)
+ if not v.is_contiguous():
+ raise ValueError(
+ "Non contiguous tensors are not supported, please call tensor.contiguous() before passing it into taichi kernel."
+ )
+ tmp = v
+ taichi_arch = self.runtime.prog.config.arch
+ # Ndarray means its memory is allocated on the specified taichi arch.
+ # Since torch only supports CPU & CUDA, torch-base ndarray only supports
+ # taichi cpu/cuda backend as well.
+ # Note I put x64/arm64/cuda here to be more specific.
+ assert not is_ndarray or taichi_arch in (
+ _ti_core.Arch.cuda, _ti_core.Arch.x64, _ti_core.Arch.arm64
+ ), "Torch-based ndarray is only supported on taichi x64/arm64/cuda backend."
+
+ if str(v.device).startswith('cuda'):
+ # External tensor on cuda
+ if taichi_arch != _ti_core.Arch.cuda:
+ # copy data back to cpu
+ host_v = v.to(device='cpu', copy=True)
+ tmp = host_v
+ callbacks.append(get_call_back(v, host_v))
+ else:
+ # External tensor on cpu
+ if taichi_arch == _ti_core.Arch.cuda:
+ gpu_v = v.cuda()
+ tmp = gpu_v
+ callbacks.append(get_call_back(v, gpu_v))
+ return tmp, callbacks
+
def get_function_body(self, t_kernel):
# The actual function body
def func__(*args):
assert len(args) == len(
self.argument_annotations
- ), '{} arguments needed but {} provided'.format(
- len(self.argument_annotations), len(args))
+ ), f'{len(self.argument_annotations)} arguments needed but {len(args)} provided'
tmps = []
callbacks = []
has_external_arrays = False
+ has_torch = has_pytorch()
actual_argument_slot = 0
launch_ctx = t_kernel.make_launch_context()
@@ -520,67 +554,67 @@ def func__(*args):
# Note: do not use sth like "needed == f32". That would be slow.
if id(needed) in primitive_types.real_type_ids:
if not isinstance(v, (float, int)):
- raise KernelArgError(i, needed.to_string(), provided)
+ raise TaichiRuntimeTypeError(i, needed.to_string(),
+ provided)
launch_ctx.set_arg_float(actual_argument_slot, float(v))
elif id(needed) in primitive_types.integer_type_ids:
if not isinstance(v, int):
- raise KernelArgError(i, needed.to_string(), provided)
+ raise TaichiRuntimeTypeError(i, needed.to_string(),
+ provided)
launch_ctx.set_arg_int(actual_argument_slot, int(v))
elif isinstance(needed, sparse_matrix_builder):
- # Pass only the base pointer of the ti.linalg.sparse_matrix_builder() argument
- launch_ctx.set_arg_int(actual_argument_slot, v.get_addr())
- elif isinstance(needed, any_arr) and (
- self.match_ext_arr(v)
- or isinstance(v, taichi.lang._ndarray.Ndarray)):
- if isinstance(v, taichi.lang._ndarray.Ndarray):
- v = v.arr
+ # Pass only the base pointer of the ti.types.sparse_matrix_builder() argument
+ launch_ctx.set_arg_int(actual_argument_slot, v._get_addr())
+ elif isinstance(needed, any_arr) and isinstance(
+ v, taichi.lang._ndarray.Ndarray):
+ has_external_arrays = True
+ v = v.arr
+ launch_ctx.set_arg_ndarray(actual_argument_slot, v)
+ elif isinstance(needed, any_arr) and (self.match_ext_arr(v)):
has_external_arrays = True
is_numpy = isinstance(v, np.ndarray)
if is_numpy:
tmp = np.ascontiguousarray(v)
# Purpose: DO NOT GC |tmp|!
tmps.append(tmp)
- launch_ctx.set_arg_external_array(
+ launch_ctx.set_arg_external_array_with_shape(
actual_argument_slot, int(tmp.ctypes.data),
- tmp.nbytes)
+ tmp.nbytes, v.shape)
else:
-
- def get_call_back(u, v):
- def call_back():
- u.copy_(v)
-
- return call_back
-
- assert util.has_pytorch()
- assert isinstance(v, torch.Tensor)
- tmp = v
- taichi_arch = self.runtime.prog.config.arch
-
- if str(v.device).startswith('cuda'):
- # External tensor on cuda
- if taichi_arch != _ti_core.Arch.cuda:
- # copy data back to cpu
- host_v = v.to(device='cpu', copy=True)
- tmp = host_v
- callbacks.append(get_call_back(v, host_v))
- else:
- # External tensor on cpu
- if taichi_arch == _ti_core.Arch.cuda:
- gpu_v = v.cuda()
- tmp = gpu_v
- callbacks.append(get_call_back(v, gpu_v))
- launch_ctx.set_arg_external_array(
+ is_ndarray = False
+ tmp, torch_callbacks = self.get_torch_callbacks(
+ v, has_torch, is_ndarray)
+ callbacks += torch_callbacks
+ launch_ctx.set_arg_external_array_with_shape(
actual_argument_slot, int(tmp.data_ptr()),
- tmp.element_size() * tmp.nelement())
- shape = v.shape
- max_num_indices = _ti_core.get_max_num_indices()
- assert len(
- shape
- ) <= max_num_indices, "External array cannot have > {} indices".format(
- max_num_indices)
- for ii, s in enumerate(shape):
- launch_ctx.set_extra_arg_int(actual_argument_slot, ii,
- s)
+ tmp.element_size() * tmp.nelement(), v.shape)
+
+ elif isinstance(needed, MatrixType):
+ if id(needed.dtype) in primitive_types.real_type_ids:
+ for a in range(needed.n):
+ for b in range(needed.m):
+ if not isinstance(v[a, b], (int, float)):
+ raise TaichiRuntimeTypeError(
+ i, needed.dtype.to_string(),
+ type(v[a, b]))
+ launch_ctx.set_arg_float(
+ actual_argument_slot, float(v[a, b]))
+ actual_argument_slot += 1
+ elif id(needed.dtype) in primitive_types.integer_type_ids:
+ for a in range(needed.n):
+ for b in range(needed.m):
+ if not isinstance(v[a, b], int):
+ raise TaichiRuntimeTypeError(
+ i, needed.dtype.to_string(),
+ type(v[a, b]))
+ launch_ctx.set_arg_int(actual_argument_slot,
+ int(v[a, b]))
+ actual_argument_slot += 1
+ else:
+ raise ValueError(
+ f'Matrix dtype {needed.dtype} is not integer type or real type.'
+ )
+ continue
else:
raise ValueError(
f'Argument type mismatch. Expecting {needed}, got {type(v)}.'
@@ -592,21 +626,43 @@ def call_back():
if not self.is_grad and self.runtime.target_tape and not self.runtime.grad_replaced:
self.runtime.target_tape.insert(self, args)
+ if actual_argument_slot > 8 and (
+ impl.current_cfg().arch == _ti_core.opengl
+ or impl.current_cfg().arch == _ti_core.cc):
+ raise TaichiRuntimeError(
+ f"The number of elements in kernel arguments is too big! Do not exceed 8 on {_ti_core.arch_name(impl.current_cfg().arch)} backend."
+ )
+
+ if actual_argument_slot > 64 and (
+ (impl.current_cfg().arch != _ti_core.opengl
+ and impl.current_cfg().arch != _ti_core.cc)):
+ raise TaichiRuntimeError(
+ f"The number of elements in kernel arguments is too big! Do not exceed 64 on {_ti_core.arch_name(impl.current_cfg().arch)} backend."
+ )
+
t_kernel(launch_ctx)
ret = None
ret_dt = self.return_type
has_ret = ret_dt is not None
- if has_external_arrays or has_ret:
- ti.sync()
+ if has_ret or (impl.current_cfg().async_mode
+ and has_external_arrays):
+ runtime_ops.sync()
if has_ret:
if id(ret_dt) in primitive_types.integer_type_ids:
ret = t_kernel.get_ret_int(0)
- else:
+ elif id(ret_dt) in primitive_types.real_type_ids:
ret = t_kernel.get_ret_float(0)
-
+ elif id(ret_dt.dtype) in primitive_types.integer_type_ids:
+ it = iter(t_kernel.get_ret_int_tensor(0))
+ ret = Matrix([[next(it) for _ in range(ret_dt.m)]
+ for _ in range(ret_dt.n)])
+ else:
+ it = iter(t_kernel.get_ret_float_tensor(0))
+ ret = Matrix([[next(it) for _ in range(ret_dt.m)]
+ for _ in range(ret_dt.n)])
if callbacks:
for c in callbacks:
c()
@@ -615,9 +671,10 @@ def call_back():
return func__
- def match_ext_arr(self, v):
+ @staticmethod
+ def match_ext_arr(v):
has_array = isinstance(v, np.ndarray)
- if not has_array and util.has_pytorch():
+ if not has_array and has_pytorch():
has_array = isinstance(v, torch.Tensor)
return has_array
@@ -631,7 +688,11 @@ def ensure_compiled(self, *args):
# Thus this part needs to be fast. (i.e. < 3us on a 4 GHz x64 CPU)
@_shell_pop_print
def __call__(self, *args, **kwargs):
- _taichi_skip_traceback = 1
+ if self.is_grad and impl.current_cfg().opt_level == 0:
+ _logging.warn(
+ """opt_level = 1 is enforced to enable gradient computation."""
+ )
+ impl.current_cfg().opt_level = 1
assert len(kwargs) == 0, 'kwargs not supported for Taichi kernels'
key = self.ensure_compiled(*args)
return self.compiled_functions[key](*args)
@@ -657,10 +718,9 @@ def __call__(self, *args, **kwargs):
def _inside_class(level_of_class_stackframe):
- frames = oinspect.stack()
try:
- maybe_class_frame = frames[level_of_class_stackframe]
- statement_list = maybe_class_frame[4]
+ maybe_class_frame = sys._getframe(level_of_class_stackframe)
+ statement_list = inspect.getframeinfo(maybe_class_frame)[3]
first_statment = statement_list[0].strip()
for pat in _KERNEL_CLASS_STACKFRAME_STMT_RES:
if pat.match(first_statment):
@@ -670,16 +730,15 @@ def _inside_class(level_of_class_stackframe):
return False
-def _kernel_impl(func, level_of_class_stackframe, verbose=False):
+def _kernel_impl(_func, level_of_class_stackframe, verbose=False):
# Can decorators determine if a function is being defined inside a class?
# https://stackoverflow.com/a/8793684/12003165
is_classkernel = _inside_class(level_of_class_stackframe + 1)
- _taichi_skip_traceback = 1
if verbose:
- print(f'kernel={func.__name__} is_classkernel={is_classkernel}')
- primal = Kernel(func, is_grad=False, classkernel=is_classkernel)
- adjoint = Kernel(func, is_grad=True, classkernel=is_classkernel)
+ print(f'kernel={_func.__name__} is_classkernel={is_classkernel}')
+ primal = Kernel(_func, is_grad=False, _classkernel=is_classkernel)
+ adjoint = Kernel(_func, is_grad=True, _classkernel=is_classkernel)
# Having |primal| contains |grad| makes the tape work.
primal.grad = adjoint
@@ -691,22 +750,23 @@ def _kernel_impl(func, level_of_class_stackframe, verbose=False):
# owning the kernel, which is not known until the kernel is accessed.
#
# See also: _BoundedDifferentiableMethod, data_oriented.
- @functools.wraps(func)
+ @functools.wraps(_func)
def wrapped(*args, **kwargs):
- _taichi_skip_traceback = 1
# If we reach here (we should never), it means the class is not decorated
# with @ti.data_oriented, otherwise getattr would have intercepted the call.
clsobj = type(args[0])
assert not hasattr(clsobj, '_data_oriented')
- raise KernelDefError(
+ raise TaichiSyntaxError(
f'Please decorate class {clsobj.__name__} with @ti.data_oriented'
)
else:
- @functools.wraps(func)
+ @functools.wraps(_func)
def wrapped(*args, **kwargs):
- _taichi_skip_traceback = 1
- return primal(*args, **kwargs)
+ try:
+ return primal(*args, **kwargs)
+ except TaichiCompilationError as e:
+ raise type(e)('\n' + str(e)) from None
wrapped.grad = adjoint
@@ -745,34 +805,28 @@ def kernel(fn):
>>> for i in x:
>>> x[i] = i
"""
- _taichi_skip_traceback = 1
return _kernel_impl(fn, level_of_class_stackframe=3)
-classfunc = obsolete('@ti.classfunc', '@ti.func directly')
-classkernel = obsolete('@ti.classkernel', '@ti.kernel directly')
-
-
class _BoundedDifferentiableMethod:
def __init__(self, kernel_owner, wrapped_kernel_func):
clsobj = type(kernel_owner)
if not getattr(clsobj, '_data_oriented', False):
- raise KernelDefError(
+ raise TaichiSyntaxError(
f'Please decorate class {clsobj.__name__} with @ti.data_oriented'
)
self._kernel_owner = kernel_owner
self._primal = wrapped_kernel_func._primal
self._adjoint = wrapped_kernel_func._adjoint
self._is_staticmethod = wrapped_kernel_func._is_staticmethod
+ self.__name__ = None
def __call__(self, *args, **kwargs):
- _taichi_skip_traceback = 1
if self._is_staticmethod:
return self._primal(*args, **kwargs)
return self._primal(self._kernel_owner, *args, **kwargs)
def grad(self, *args, **kwargs):
- _taichi_skip_traceback = 1
return self._adjoint(self._kernel_owner, *args, **kwargs)
@@ -806,7 +860,6 @@ def data_oriented(cls):
The decorated class.
"""
def _getattr(self, item):
- _taichi_skip_traceback = 1
method = cls.__dict__.get(item, None)
is_property = method.__class__ == property
is_staticmethod = method.__class__ == staticmethod
@@ -835,3 +888,6 @@ def _getattr(self, item):
cls._data_oriented = True
return cls
+
+
+__all__ = ["data_oriented", "func", "kernel"]
diff --git a/python/taichi/lang/linalg_impl.py b/python/taichi/lang/linalg_impl.py
deleted file mode 100644
index 8e2f8c2a96fa7..0000000000000
--- a/python/taichi/lang/linalg_impl.py
+++ /dev/null
@@ -1,262 +0,0 @@
-from taichi.core.util import ti_core as _ti_core
-from taichi.lang.impl import expr_init
-from taichi.lang.kernel_impl import func
-
-import taichi as ti
-
-
-@func
-def polar_decompose2d(A, dt):
- """Perform polar decomposition (A=UP) for 2x2 matrix.
-
- Mathematical concept refers to https://en.wikipedia.org/wiki/Polar_decomposition.
-
- Args:
- A (ti.Matrix(2, 2)): input 2x2 matrix `A`.
- dt (DataType): date type of elements in matrix `A`, typically accepts ti.f32 or ti.f64.
-
- Returns:
- Decomposed 2x2 matrices `U` and `P`.
- """
- x, y = A(0, 0) + A(1, 1), A(1, 0) - A(0, 1)
- scale = (1.0 / ti.sqrt(x * x + y * y))
- c = x * scale
- s = y * scale
- r = ti.Matrix([[c, -s], [s, c]], dt=dt)
- return r, r.transpose() @ A
-
-
-@func
-def polar_decompose3d(A, dt):
- """Perform polar decomposition (A=UP) for 3x3 matrix.
-
- Mathematical concept refers to https://en.wikipedia.org/wiki/Polar_decomposition.
-
- Args:
- A (ti.Matrix(3, 3)): input 3x3 matrix `A`.
- dt (DataType): date type of elements in matrix `A`, typically accepts ti.f32 or ti.f64.
-
- Returns:
- Decomposed 3x3 matrices `U` and `P`.
- """
- U, sig, V = ti.svd(A, dt)
- return U @ V.transpose(), V @ sig @ V.transpose()
-
-
-# https://www.seas.upenn.edu/~cffjiang/research/svd/svd.pdf
-@func
-def svd2d(A, dt):
- """Perform singular value decomposition (A=USV^T) for 2x2 matrix.
-
- Mathematical concept refers to https://en.wikipedia.org/wiki/Singular_value_decomposition.
-
- Args:
- A (ti.Matrix(2, 2)): input 2x2 matrix `A`.
- dt (DataType): date type of elements in matrix `A`, typically accepts ti.f32 or ti.f64.
-
- Returns:
- Decomposed 2x2 matrices `U`, 'S' and `V`.
- """
- R, S = polar_decompose2d(A, dt)
- c, s = ti.cast(0.0, dt), ti.cast(0.0, dt)
- s1, s2 = ti.cast(0.0, dt), ti.cast(0.0, dt)
- if abs(S[0, 1]) < 1e-5:
- c, s = 1, 0
- s1, s2 = S[0, 0], S[1, 1]
- else:
- tao = ti.cast(0.5, dt) * (S[0, 0] - S[1, 1])
- w = ti.sqrt(tao**2 + S[0, 1]**2)
- t = ti.cast(0.0, dt)
- if tao > 0:
- t = S[0, 1] / (tao + w)
- else:
- t = S[0, 1] / (tao - w)
- c = 1 / ti.sqrt(t**2 + 1)
- s = -t * c
- s1 = c**2 * S[0, 0] - 2 * c * s * S[0, 1] + s**2 * S[1, 1]
- s2 = s**2 * S[0, 0] + 2 * c * s * S[0, 1] + c**2 * S[1, 1]
- V = ti.Matrix.zero(dt, 2, 2)
- if s1 < s2:
- tmp = s1
- s1 = s2
- s2 = tmp
- V = ti.Matrix([[-s, c], [-c, -s]], dt=dt)
- else:
- V = ti.Matrix([[c, s], [-s, c]], dt=dt)
- U = R @ V
- return U, ti.Matrix([[s1, ti.cast(0, dt)], [ti.cast(0, dt), s2]], dt=dt), V
-
-
-def svd3d(A, dt, iters=None):
- """Perform singular value decomposition (A=USV^T) for 3x3 matrix.
-
- Mathematical concept refers to https://en.wikipedia.org/wiki/Singular_value_decomposition.
-
- Args:
- A (ti.Matrix(3, 3)): input 3x3 matrix `A`.
- dt (DataType): date type of elements in matrix `A`, typically accepts ti.f32 or ti.f64.
- iters (int): iteration number to control algorithm precision.
-
- Returns:
- Decomposed 3x3 matrices `U`, 'S' and `V`.
- """
- assert A.n == 3 and A.m == 3
- inputs = tuple([e.ptr for e in A.entries])
- assert dt in [ti.f32, ti.f64]
- if iters is None:
- if dt == ti.f32:
- iters = 5
- else:
- iters = 8
- if dt == ti.f32:
- rets = _ti_core.sifakis_svd_f32(*inputs, iters)
- else:
- rets = _ti_core.sifakis_svd_f64(*inputs, iters)
- assert len(rets) == 21
- U_entries = rets[:9]
- V_entries = rets[9:18]
- sig_entries = rets[18:]
- U = expr_init(ti.Matrix.zero(dt, 3, 3))
- V = expr_init(ti.Matrix.zero(dt, 3, 3))
- sigma = expr_init(ti.Matrix.zero(dt, 3, 3))
- for i in range(3):
- for j in range(3):
- U(i, j).assign(U_entries[i * 3 + j])
- V(i, j).assign(V_entries[i * 3 + j])
- sigma(i, i).assign(sig_entries[i])
- return U, sigma, V
-
-
-@func
-def eig2x2(A, dt):
- """Compute the eigenvalues and right eigenvectors (Av=lambda v) of a 2x2 real matrix.
-
- Mathematical concept refers to https://en.wikipedia.org/wiki/Eigendecomposition_of_a_matrix.
-
- Args:
- A (ti.Matrix(2, 2)): input 2x2 matrix `A`.
- dt (DataType): date type of elements in matrix `A`, typically accepts ti.f32 or ti.f64.
-
- Returns:
- eigenvalues (ti.Matrix(2, 2)): The eigenvalues in complex form. Each row stores one eigenvalue. The first number of the eigenvalue represents the real part and the second number represents the imaginary part.
- eigenvectors: (ti.Matrix(4, 2)): The eigenvectors in complex form. Each column stores one eigenvector. Each eigenvector consists of 2 entries, each of which is represented by two numbers for its real part and imaginary part.
- """
- tr = A.trace()
- det = A.determinant()
- gap = tr**2 - 4 * det
- lambda1 = ti.Vector.zero(dt, 2)
- lambda2 = ti.Vector.zero(dt, 2)
- v1 = ti.Vector.zero(dt, 4)
- v2 = ti.Vector.zero(dt, 4)
- if gap > 0:
- lambda1 = ti.Vector([tr + ti.sqrt(gap), 0.0], dt=dt) * 0.5
- lambda2 = ti.Vector([tr - ti.sqrt(gap), 0.0], dt=dt) * 0.5
- A1 = A - lambda1[0] * ti.Matrix.identity(dt, 2)
- A2 = A - lambda2[0] * ti.Matrix.identity(dt, 2)
- if all(A1 == ti.Matrix.zero(dt, 2, 2)) and all(
- A1 == ti.Matrix.zero(dt, 2, 2)):
- v1 = ti.Vector([0.0, 0.0, 1.0, 0.0]).cast(dt)
- v2 = ti.Vector([1.0, 0.0, 0.0, 0.0]).cast(dt)
- else:
- v1 = ti.Vector([A2[0, 0], 0.0, A2[1, 0], 0.0], dt=dt).normalized()
- v2 = ti.Vector([A1[0, 0], 0.0, A1[1, 0], 0.0], dt=dt).normalized()
- else:
- lambda1 = ti.Vector([tr, ti.sqrt(-gap)], dt=dt) * 0.5
- lambda2 = ti.Vector([tr, -ti.sqrt(-gap)], dt=dt) * 0.5
- A1r = A - lambda1[0] * ti.Matrix.identity(dt, 2)
- A1i = -lambda1[1] * ti.Matrix.identity(dt, 2)
- A2r = A - lambda2[0] * ti.Matrix.identity(dt, 2)
- A2i = -lambda2[1] * ti.Matrix.identity(dt, 2)
- v1 = ti.Vector([A2r[0, 0], A2i[0, 0], A2r[1, 0], A2i[1, 0]],
- dt=dt).normalized()
- v2 = ti.Vector([A1r[0, 0], A1i[0, 0], A1r[1, 0], A1i[1, 0]],
- dt=dt).normalized()
- eigenvalues = ti.Matrix.rows([lambda1, lambda2])
- eigenvectors = ti.Matrix.cols([v1, v2])
-
- return eigenvalues, eigenvectors
-
-
-@func
-def sym_eig2x2(A, dt):
- """Compute the eigenvalues and right eigenvectors (Av=lambda v) of a 2x2 real symmetric matrix.
-
- Mathematical concept refers to https://en.wikipedia.org/wiki/Eigendecomposition_of_a_matrix.
-
- Args:
- A (ti.Matrix(2, 2)): input 2x2 symmetric matrix `A`.
- dt (DataType): date type of elements in matrix `A`, typically accepts ti.f32 or ti.f64.
-
- Returns:
- eigenvalues (ti.Vector(2)): The eigenvalues. Each entry store one eigen value.
- eigenvectors (ti.Matrix(2, 2)): The eigenvectors. Each column stores one eigenvector.
- """
- tr = A.trace()
- det = A.determinant()
- gap = tr**2 - 4 * det
- lambda1 = (tr + ti.sqrt(gap)) * 0.5
- lambda2 = (tr - ti.sqrt(gap)) * 0.5
- eigenvalues = ti.Vector([lambda1, lambda2], dt=dt)
-
- A1 = A - lambda1 * ti.Matrix.identity(dt, 2)
- A2 = A - lambda2 * ti.Matrix.identity(dt, 2)
- v1 = ti.Vector.zero(dt, 2)
- v2 = ti.Vector.zero(dt, 2)
- if all(A1 == ti.Matrix.zero(dt, 2, 2)) and all(
- A1 == ti.Matrix.zero(dt, 2, 2)):
- v1 = ti.Vector([0.0, 1.0]).cast(dt)
- v2 = ti.Vector([1.0, 0.0]).cast(dt)
- else:
- v1 = ti.Vector([A2[0, 0], A2[1, 0]], dt=dt).normalized()
- v2 = ti.Vector([A1[0, 0], A1[1, 0]], dt=dt).normalized()
- eigenvectors = ti.Matrix.cols([v1, v2])
- return eigenvalues, eigenvectors
-
-
-@func
-def svd(A, dt):
- """Perform singular value decomposition (A=USV^T) for arbitrary size matrix.
-
- Mathematical concept refers to https://en.wikipedia.org/wiki/Singular_value_decomposition.
- 2D implementation refers to :func:`taichi.lang.linalg_impl.svd2d`.
- 3D implementation refers to :func:`taichi.lang.linalg_impl.svd3d`.
-
- Args:
- A (ti.Matrix(n, n)): input nxn matrix `A`.
- dt (DataType): date type of elements in matrix `A`, typically accepts ti.f32 or ti.f64.
-
- Returns:
- Decomposed nxn matrices `U`, 'S' and `V`.
- """
- if ti.static(A.n == 2):
- ret = svd2d(A, dt)
- return ret
- elif ti.static(A.n == 3):
- return svd3d(A, dt)
- else:
- raise Exception("SVD only supports 2D and 3D matrices.")
-
-
-@func
-def polar_decompose(A, dt):
- """Perform polar decomposition (A=UP) for arbitrary size matrix.
-
- Mathematical concept refers to https://en.wikipedia.org/wiki/Polar_decomposition.
- 2D implementation refers to :func:`taichi.lang.linalg_impl.polar_decompose2d`.
- 3D implementation refers to :func:`taichi.lang.linalg_impl.polar_decompose3d`.
-
- Args:
- A (ti.Matrix(n, n)): input nxn matrix `A`.
- dt (DataType): date type of elements in matrix `A`, typically accepts ti.f32 or ti.f64.
-
- Returns:
- Decomposed nxn matrices `U` and `P`.
- """
- if ti.static(A.n == 2):
- ret = polar_decompose2d(A, dt)
- return ret
- elif ti.static(A.n == 3):
- return polar_decompose3d(A, dt)
- else:
- raise Exception(
- "Polar decomposition only supports 2D and 3D matrices.")
diff --git a/python/taichi/lang/matrix.py b/python/taichi/lang/matrix.py
index ec27c56bce229..02efa2e6aedbd 100644
--- a/python/taichi/lang/matrix.py
+++ b/python/taichi/lang/matrix.py
@@ -1,47 +1,38 @@
-import copy
import numbers
from collections.abc import Iterable
import numpy as np
-import taichi.lang
+from taichi._lib import core as ti_core
from taichi.lang import expr, impl
-from taichi.lang import kernel_impl as kern_mod
from taichi.lang import ops as ops_mod
+from taichi.lang import runtime_ops
from taichi.lang._ndarray import Ndarray, NdarrayHostAccess
from taichi.lang.common_ops import TaichiOperations
from taichi.lang.enums import Layout
from taichi.lang.exception import TaichiSyntaxError
from taichi.lang.field import Field, ScalarField, SNodeHostAccess
-from taichi.lang.ops import cast
-from taichi.lang.types import CompoundType
-from taichi.lang.util import (cook_dtype, in_python_scope, is_taichi_class,
- python_scope, taichi_scope, to_numpy_type,
- to_pytorch_type)
-from taichi.misc.util import deprecated, warning
-
-import taichi as ti
+from taichi.lang.util import (cook_dtype, in_python_scope, python_scope,
+ taichi_scope, to_numpy_type, to_pytorch_type,
+ warning)
+from taichi.types import primitive_types
+from taichi.types.compound_types import CompoundType
class Matrix(TaichiOperations):
"""The matrix class.
Args:
- n (int): the first dimension of a matrix.
+ n (Union[int, list, tuple, np.ndarray]): the first dimension of a matrix.
m (int): the second dimension of a matrix.
- dt (DataType): the elmement data type.
- keep_raw (Bool, optional): Keep the contents in `n` as is.
+ dt (DataType): the element data type.
"""
- is_taichi_class = True
-
- def __init__(self,
- n=1,
- m=1,
- dt=None,
- keep_raw=False,
- disable_local_tensor=False):
+ _is_taichi_class = True
+
+ def __init__(self, n=1, m=1, dt=None, suppress_warning=False):
self.local_tensor_proxy = None
self.any_array_access = None
self.grad = None
+ self.dynamic_index_stride = None
if isinstance(n, (list, tuple, np.ndarray)):
if len(n) == 0:
@@ -49,60 +40,69 @@ def __init__(self,
elif isinstance(n[0], Matrix):
raise Exception(
'cols/rows required when using list of vectors')
- elif not isinstance(n[0], Iterable):
- if impl.inside_kernel():
- # wrap potential constants with Expr
- if keep_raw:
- mat = [list([x]) for x in n]
- else:
- if in_python_scope(
- ) or disable_local_tensor or not ti.current_cfg(
- ).dynamic_index:
- mat = [list([expr.Expr(x)]) for x in n]
- else:
- if not ti.is_extension_supported(
- ti.cfg.arch, ti.extension.dynamic_index):
- raise Exception(
- 'Backend ' + str(ti.cfg.arch) +
- ' doesn\'t support dynamic index')
- if dt is None:
- if isinstance(n[0], int):
- dt = impl.get_runtime().default_ip
- elif isinstance(n[0], float):
- dt = impl.get_runtime().default_fp
- else:
- raise Exception(
- 'dt required when using dynamic_index for local tensor'
- )
- self.local_tensor_proxy = impl.expr_init_local_tensor(
- [len(n)], dt,
- expr.make_expr_group([expr.Expr(x)
- for x in n]))
- mat = []
- for i in range(len(n)):
- mat.append(
- list([
- ti.local_subscript_with_offset(
- self.local_tensor_proxy, (i, ),
- (len(n), ))
- ]))
- else:
+ elif not isinstance(n[0], Iterable): # now init a Vector
+ if in_python_scope():
mat = [[x] for x in n]
- else:
- if in_python_scope(
- ) or disable_local_tensor or not ti.current_cfg(
- ).dynamic_index:
- mat = [list(r) for r in n]
+ elif not impl.current_cfg().dynamic_index:
+ mat = [[impl.expr_init(x)] for x in n]
else:
- if not ti.is_extension_supported(
- ti.cfg.arch, ti.extension.dynamic_index):
- raise Exception('Backend ' + str(ti.cfg.arch) +
- ' doesn\'t support dynamic index')
+ if not ti_core.is_extension_supported(
+ impl.current_cfg().arch,
+ ti_core.Extension.dynamic_index):
+ raise Exception(
+ f"Backend {impl.current_cfg().arch} doesn't support dynamic index"
+ )
if dt is None:
- if isinstance(n[0][0], int):
+ if isinstance(n[0], (int, np.integer)):
+ dt = impl.get_runtime().default_ip
+ elif isinstance(n[0], float):
+ dt = impl.get_runtime().default_fp
+ elif isinstance(n[0], expr.Expr):
+ dt = n[0].ptr.get_ret_type()
+ if dt == ti_core.DataType_unknown:
+ raise TypeError(
+ 'Element type of the matrix cannot be inferred. Please set dt instead for now.'
+ )
+ else:
+ raise Exception(
+ 'dt required when using dynamic_index for local tensor'
+ )
+ self.local_tensor_proxy = impl.expr_init_local_tensor(
+ [len(n)], dt,
+ expr.make_expr_group([expr.Expr(x) for x in n]))
+ self.dynamic_index_stride = 1
+ mat = []
+ for i in range(len(n)):
+ mat.append(
+ list([
+ impl.make_tensor_element_expr(
+ self.local_tensor_proxy, (expr.Expr(
+ i, dtype=primitive_types.i32), ),
+ (len(n), ), self.dynamic_index_stride)
+ ]))
+ else: # now init a Matrix
+ if in_python_scope():
+ mat = [list(row) for row in n]
+ elif not impl.current_cfg().dynamic_index:
+ mat = [[impl.expr_init(x) for x in row] for row in n]
+ else:
+ if not ti_core.is_extension_supported(
+ impl.current_cfg().arch,
+ ti_core.Extension.dynamic_index):
+ raise Exception(
+ f"Backend {impl.current_cfg().arch} doesn't support dynamic index"
+ )
+ if dt is None:
+ if isinstance(n[0][0], (int, np.integer)):
dt = impl.get_runtime().default_ip
elif isinstance(n[0][0], float):
dt = impl.get_runtime().default_fp
+ elif isinstance(n[0][0], expr.Expr):
+ dt = n[0][0].ptr.get_ret_type()
+ if dt == ti_core.DataType_unknown:
+ raise TypeError(
+ 'Element type of the matrix cannot be inferred. Please set dt instead for now.'
+ )
else:
raise Exception(
'dt required when using dynamic_index for local tensor'
@@ -111,14 +111,18 @@ def __init__(self,
[len(n), len(n[0])], dt,
expr.make_expr_group(
[expr.Expr(x) for row in n for x in row]))
+ self.dynamic_index_stride = 1
mat = []
for i in range(len(n)):
mat.append([])
for j in range(len(n[0])):
mat[i].append(
- ti.local_subscript_with_offset(
- self.local_tensor_proxy, (i, j),
- (len(n), len(n[0]))))
+ impl.make_tensor_element_expr(
+ self.local_tensor_proxy,
+ (expr.Expr(i, dtype=primitive_types.i32),
+ expr.Expr(j, dtype=primitive_types.i32)),
+ (len(n), len(n[0])),
+ self.dynamic_index_stride))
self.n = len(mat)
if len(mat) > 0:
self.m = len(mat[0])
@@ -134,10 +138,10 @@ def __init__(self,
self.m = m
else:
raise ValueError(
- "Declaring matrix fields using `ti.Matrix(n, m, dt, shape)` is no longer supported. Use `ti.Matrix.field(n, m, dtype, shape)` instead."
- )
+ "Declaring matrix fields using `ti.Matrix(n, m, dt, shape)` is no longer supported. "
+ "Use `ti.Matrix.field(n, m, dtype, shape)` instead.")
- if self.n * self.m > 32:
+ if self.n * self.m > 32 and not suppress_warning:
warning(
f'Taichi matrices/vectors with {self.n}x{self.m} > 32 entries are not suggested.'
' Matrices/vectors will be automatically unrolled at compile-time for performance.'
@@ -149,65 +153,42 @@ def __init__(self,
UserWarning,
stacklevel=2)
- def element_wise_binary(self, foo, other):
- _taichi_skip_traceback = 1
- ret = self.empty_copy()
- if isinstance(other, (list, tuple)):
- other = Matrix(other)
- if isinstance(other, Matrix):
- assert self.m == other.m and self.n == other.n, f"Dimension mismatch between shapes ({self.n}, {self.m}), ({other.n}, {other.m})"
- for i in range(self.n * self.m):
- ret.entries[i] = foo(self.entries[i], other.entries[i])
- else: # assumed to be scalar
- for i in range(self.n * self.m):
- ret.entries[i] = foo(self.entries[i], other)
- return ret
+ def _element_wise_binary(self, foo, other):
+ other = self._broadcast_copy(other)
+ return Matrix([[foo(self(i, j), other(i, j)) for j in range(self.m)]
+ for i in range(self.n)])
- def broadcast_copy(self, other):
+ def _broadcast_copy(self, other):
if isinstance(other, (list, tuple)):
other = Matrix(other)
if not isinstance(other, Matrix):
- ret = self.empty_copy()
- ret.entries = [other for _ in ret.entries]
- other = ret
+ other = Matrix([[other for _ in range(self.m)]
+ for _ in range(self.n)])
assert self.m == other.m and self.n == other.n, f"Dimension mismatch between shapes ({self.n}, {self.m}), ({other.n}, {other.m})"
return other
- def element_wise_ternary(self, foo, other, extra):
- ret = self.empty_copy()
- other = self.broadcast_copy(other)
- extra = self.broadcast_copy(extra)
- for i in range(self.n * self.m):
- ret.entries[i] = foo(self.entries[i], other.entries[i],
- extra.entries[i])
- return ret
+ def _element_wise_ternary(self, foo, other, extra):
+ other = self._broadcast_copy(other)
+ extra = self._broadcast_copy(extra)
+ return Matrix([[
+ foo(self(i, j), other(i, j), extra(i, j)) for j in range(self.m)
+ ] for i in range(self.n)])
- def element_wise_writeback_binary(self, foo, other):
- ret = self.empty_copy()
- if isinstance(other, (list, tuple)):
- other = Matrix(other)
- if is_taichi_class(other):
- other = other.variable()
- if foo.__name__ == 'assign' and not isinstance(other, Matrix):
+ def _element_wise_writeback_binary(self, foo, other):
+ if foo.__name__ == 'assign' and not isinstance(other,
+ (list, tuple, Matrix)):
raise TaichiSyntaxError(
'cannot assign scalar expr to '
f'taichi class {type(self)}, maybe you want to use `a.fill(b)` instead?'
)
- if isinstance(other, Matrix):
- assert self.m == other.m and self.n == other.n, f"Dimension mismatch between shapes ({self.n}, {self.m}), ({other.n}, {other.m})"
- for i in range(self.n * self.m):
- ret.entries[i] = foo(self.entries[i], other.entries[i])
- else: # assumed to be scalar
- for i in range(self.n * self.m):
- ret.entries[i] = foo(self.entries[i], other)
- return ret
+ other = self._broadcast_copy(other)
+ entries = [[foo(self(i, j), other(i, j)) for j in range(self.m)]
+ for i in range(self.n)]
+ return self if foo.__name__ == 'assign' else Matrix(entries)
- def element_wise_unary(self, foo):
- _taichi_skip_traceback = 1
- ret = self.empty_copy()
- for i in range(self.n * self.m):
- ret.entries[i] = foo(self.entries[i])
- return ret
+ def _element_wise_unary(self, foo):
+ return Matrix([[foo(self(i, j)) for j in range(self.m)]
+ for i in range(self.n)])
def __matmul__(self, other):
"""Matrix-matrix or matrix-vector multiply.
@@ -219,26 +200,24 @@ def __matmul__(self, other):
The matrix-matrix product or matrix-vector product.
"""
- _taichi_skip_traceback = 1
assert isinstance(other, Matrix), "rhs of `@` is not a matrix / vector"
assert self.m == other.n, f"Dimension mismatch between shapes ({self.n}, {self.m}), ({other.n}, {other.m})"
- del _taichi_skip_traceback
- ret = Matrix.new(self.n, other.m)
+ entries = []
for i in range(self.n):
+ entries.append([])
for j in range(other.m):
acc = self(i, 0) * other(0, j)
for k in range(1, other.n):
acc = acc + self(i, k) * other(k, j)
- ret.set_entry(i, j, acc)
- return ret
+ entries[i].append(acc)
+ return Matrix(entries)
- def linearize_entry_id(self, *args):
+ def _linearize_entry_id(self, *args):
assert 1 <= len(args) <= 2
if len(args) == 1 and isinstance(args[0], (list, tuple)):
args = args[0]
if len(args) == 1:
args = args + (0, )
- _taichi_skip_traceback = 1
# TODO(#1004): See if it's possible to support indexing at runtime
for i, a in enumerate(args):
if not isinstance(a, int):
@@ -253,25 +232,24 @@ def linearize_entry_id(self, *args):
'See https://docs.taichi.graphics/lang/articles/advanced/meta#when-to-use-for-loops-with-tistatic for more details.'
)
assert 0 <= args[0] < self.n, \
- f"The 0-th matrix index is out of range: 0 <= {args[0]} < {self.n}"
+ f"The 0-th matrix index is out of range: 0 <= {args[0]} < {self.n}"
assert 0 <= args[1] < self.m, \
- f"The 1-th matrix index is out of range: 0 <= {args[1]} < {self.m}"
+ f"The 1-th matrix index is out of range: 0 <= {args[1]} < {self.m}"
return args[0] * self.m + args[1]
def __call__(self, *args, **kwargs):
- _taichi_skip_traceback = 1
assert kwargs == {}
- ret = self.entries[self.linearize_entry_id(*args)]
+ ret = self.entries[self._linearize_entry_id(*args)]
if isinstance(ret, SNodeHostAccess):
ret = ret.accessor.getter(*ret.key)
elif isinstance(ret, NdarrayHostAccess):
ret = ret.getter()
return ret
- def set_entry(self, i, j, e):
- idx = self.linearize_entry_id(i, j)
+ def _set_entry(self, i, j, e):
+ idx = self._linearize_entry_id(i, j)
if impl.inside_kernel():
- self.entries[idx].assign(e)
+ self.entries[idx]._assign(e)
else:
if isinstance(self.entries[idx], SNodeHostAccess):
self.entries[idx].accessor.setter(e, *self.entries[idx].key)
@@ -280,106 +258,103 @@ def set_entry(self, i, j, e):
else:
self.entries[idx] = e
+ def _get_slice(self, a, b):
+ if not isinstance(a, slice):
+ a = [a]
+ else:
+ a = range(a.start or 0, a.stop or self.n, a.step or 1)
+ if not isinstance(b, slice):
+ b = [b]
+ else:
+ b = range(b.start or 0, b.stop or self.m, b.step or 1)
+ return Matrix([[self(i, j) for j in b] for i in a])
+
@taichi_scope
- def subscript(self, *indices):
- _taichi_skip_traceback = 1
+ def _subscript(self, *indices):
assert len(indices) in [1, 2]
i = indices[0]
j = 0 if len(indices) == 1 else indices[1]
+ if isinstance(i, slice) or isinstance(j, slice):
+ for a in (i, j):
+ if isinstance(a, slice):
+ if isinstance(a.start, expr.Expr) or isinstance(
+ a.step, expr.Expr) or isinstance(
+ a.stop, expr.Expr):
+ raise TaichiSyntaxError(
+ "The element type of slice of Matrix/Vector index must be a compile-time constant integer!"
+ )
+ return self._get_slice(i, j)
if self.any_array_access:
return self.any_array_access.subscript(i, j)
- elif self.local_tensor_proxy is not None:
+ if self.local_tensor_proxy is not None:
+ assert self.dynamic_index_stride is not None
if len(indices) == 1:
- return ti.local_subscript_with_offset(self.local_tensor_proxy,
- (i, ), (self.n, ))
- else:
- return ti.local_subscript_with_offset(self.local_tensor_proxy,
- (i, j), (self.n, self.m))
- # ptr.is_global_ptr() will check whether it's an element in the field (which is different from ptr.is_global_var()).
- elif isinstance(self.entries[0],
- ti.Expr) and self.entries[0].ptr.is_global_ptr(
- ) and ti.current_cfg().dynamic_index:
- # TODO: Add API to query whether AOS or SOA
- return ti.global_subscript_with_offset(self.entries[0], (i, j),
- (self.n, self.m), True)
- else:
- return self(i, j)
+ return impl.make_tensor_element_expr(self.local_tensor_proxy,
+ (i, ), (self.n, ),
+ self.dynamic_index_stride)
+ return impl.make_tensor_element_expr(self.local_tensor_proxy,
+ (i, j), (self.n, self.m),
+ self.dynamic_index_stride)
+ if impl.current_cfg().dynamic_index and isinstance(
+ self,
+ _MatrixFieldElement) and self.dynamic_index_stride is not None:
+ return impl.make_tensor_element_expr(self.entries[0].ptr, (i, j),
+ (self.n, self.m),
+ self.dynamic_index_stride)
+ return self(i, j)
@property
def x(self):
"""Get the first element of a matrix."""
- _taichi_skip_traceback = 1
if impl.inside_kernel():
- return self.subscript(0)
- else:
- return self[0]
+ return self._subscript(0)
+ return self[0]
@property
def y(self):
"""Get the second element of a matrix."""
- _taichi_skip_traceback = 1
if impl.inside_kernel():
- return self.subscript(1)
- else:
- return self[1]
+ return self._subscript(1)
+ return self[1]
@property
def z(self):
"""Get the third element of a matrix."""
- _taichi_skip_traceback = 1
if impl.inside_kernel():
- return self.subscript(2)
- else:
- return self[2]
+ return self._subscript(2)
+ return self[2]
@property
def w(self):
"""Get the fourth element of a matrix."""
- _taichi_skip_traceback = 1
if impl.inside_kernel():
- return self.subscript(3)
- else:
- return self[3]
+ return self._subscript(3)
+ return self[3]
# since Taichi-scope use v.x.assign() instead
@x.setter
@python_scope
def x(self, value):
- _taichi_skip_traceback = 1
self[0] = value
@y.setter
@python_scope
def y(self, value):
- _taichi_skip_traceback = 1
self[1] = value
@z.setter
@python_scope
def z(self, value):
- _taichi_skip_traceback = 1
self[2] = value
@w.setter
@python_scope
def w(self, value):
- _taichi_skip_traceback = 1
self[3] = value
- @property
- @python_scope
- def value(self):
- if isinstance(self.entries[0], SNodeHostAccess):
- # fetch values from SNodeHostAccessor
- ret = self.empty_copy()
- for i in range(self.n):
- for j in range(self.m):
- ret.entries[i * self.m + j] = self(i, j)
- else:
- # is local python-scope matrix
- ret = self.entries
- return ret
+ def to_list(self):
+ return [[self(i, j) for j in range(self.m)] for i in range(self.n)]
# host access & python scope operation
@python_scope
@@ -398,6 +373,8 @@ def __getitem__(self, indices):
assert len(indices) in [1, 2]
i = indices[0]
j = 0 if len(indices) == 1 else indices[1]
+ if isinstance(i, slice) or isinstance(j, slice):
+ return self._get_slice(i, j)
return self(i, j)
@python_scope
@@ -413,7 +390,7 @@ def __setitem__(self, indices, item):
assert len(indices) in [1, 2]
i = indices[0]
j = 0 if len(indices) == 1 else indices[1]
- self.set_entry(i, j, item)
+ self._set_entry(i, j, item)
def __len__(self):
"""Get the length of each row of a matrix"""
@@ -422,11 +399,10 @@ def __len__(self):
def __iter__(self):
if self.m == 1:
return (self(i) for i in range(self.n))
- else:
- return ([self(i, j) for j in range(self.m)] for i in range(self.n))
+ return ([self(i, j) for j in range(self.m)] for i in range(self.n))
@python_scope
- def set_entries(self, value):
+ def _set_entries(self, value):
if not isinstance(value, (list, tuple)):
value = list(value)
if not isinstance(value[0], (list, tuple)):
@@ -435,20 +411,6 @@ def set_entries(self, value):
for j in range(self.m):
self[i, j] = value[i][j]
- def empty_copy(self):
- return Matrix.empty(self.n, self.m)
-
- def copy(self):
- ret = self.empty_copy()
- ret.entries = copy.copy(self.entries)
- return ret
-
- @taichi_scope
- def variable(self):
- ret = self.copy()
- ret.entries = [impl.expr_init(e) for e in ret.entries]
- return ret
-
@taichi_scope
def cast(self, dtype):
"""Cast the matrix element data type.
@@ -460,11 +422,9 @@ def cast(self, dtype):
A new matrix with each element's type is dtype.
"""
- _taichi_skip_traceback = 1
- ret = self.copy()
- for i in range(len(self.entries)):
- ret.entries[i] = ops_mod.cast(ret.entries[i], dtype)
- return ret
+ return Matrix(
+ [[ops_mod.cast(self(i, j), dtype) for j in range(self.m)]
+ for i in range(self.n)])
def trace(self):
"""The sum of a matrix diagonal elements.
@@ -474,10 +434,10 @@ def trace(self):
"""
assert self.n == self.m
- sum = self(0, 0)
+ _sum = self(0, 0)
for i in range(1, self.n):
- sum = sum + self(i, i)
- return sum
+ _sum = _sum + self(i, i)
+ return _sum
@taichi_scope
def inverse(self):
@@ -495,14 +455,12 @@ def inverse(self):
"""
assert self.n == self.m, 'Only square matrices are invertible'
if self.n == 1:
- return Matrix([1 / self(0, 0)], disable_local_tensor=True)
- elif self.n == 2:
- inv_det = impl.expr_init(1.0 / self.determinant())
- # Discussion: https://github.com/taichi-dev/taichi/pull/943#issuecomment-626344323
- return inv_det * Matrix([[self(1, 1), -self(0, 1)],
- [-self(1, 0), self(0, 0)]],
- disable_local_tensor=True).variable()
- elif self.n == 3:
+ return Matrix([1 / self(0, 0)])
+ if self.n == 2:
+ inv_determinant = impl.expr_init(1.0 / self.determinant())
+ return inv_determinant * Matrix([[self(
+ 1, 1), -self(0, 1)], [-self(1, 0), self(0, 0)]])
+ if self.n == 3:
n = 3
inv_determinant = impl.expr_init(1.0 / self.determinant())
entries = [[0] * n for _ in range(n)]
@@ -512,11 +470,11 @@ def E(x, y):
for i in range(n):
for j in range(n):
- entries[j][i] = impl.expr_init(
- inv_determinant * (E(i + 1, j + 1) * E(i + 2, j + 2) -
- E(i + 2, j + 1) * E(i + 1, j + 2)))
- return Matrix(entries, disable_local_tensor=True)
- elif self.n == 4:
+ entries[j][i] = inv_determinant * (
+ E(i + 1, j + 1) * E(i + 2, j + 2) -
+ E(i + 2, j + 1) * E(i + 1, j + 2))
+ return Matrix(entries)
+ if self.n == 4:
n = 4
inv_determinant = impl.expr_init(1.0 / self.determinant())
entries = [[0] * n for _ in range(n)]
@@ -526,25 +484,18 @@ def E(x, y):
for i in range(n):
for j in range(n):
- entries[j][i] = impl.expr_init(
- inv_determinant * (-1)**(i + j) *
- ((E(i + 1, j + 1) *
- (E(i + 2, j + 2) * E(i + 3, j + 3) -
- E(i + 3, j + 2) * E(i + 2, j + 3)) -
- E(i + 2, j + 1) *
- (E(i + 1, j + 2) * E(i + 3, j + 3) -
- E(i + 3, j + 2) * E(i + 1, j + 3)) +
- E(i + 3, j + 1) *
- (E(i + 1, j + 2) * E(i + 2, j + 3) -
- E(i + 2, j + 2) * E(i + 1, j + 3)))))
- return Matrix(entries, disable_local_tensor=True)
- else:
- raise Exception(
- "Inversions of matrices with sizes >= 5 are not supported")
-
- inversed = deprecated('a.inversed()', 'a.inverse()')(inverse)
+ entries[j][i] = inv_determinant * (-1)**(i + j) * ((
+ E(i + 1, j + 1) *
+ (E(i + 2, j + 2) * E(i + 3, j + 3) -
+ E(i + 3, j + 2) * E(i + 2, j + 3)) - E(i + 2, j + 1) *
+ (E(i + 1, j + 2) * E(i + 3, j + 3) -
+ E(i + 3, j + 2) * E(i + 1, j + 3)) + E(i + 3, j + 1) *
+ (E(i + 1, j + 2) * E(i + 2, j + 3) -
+ E(i + 2, j + 2) * E(i + 1, j + 3))))
+ return Matrix(entries)
+ raise Exception(
+ "Inversions of matrices with sizes >= 5 are not supported")
- @kern_mod.pyfunc
def normalized(self, eps=0):
"""Normalize a vector.
@@ -567,16 +518,6 @@ def normalized(self, eps=0):
invlen = 1 / (self.norm() + eps)
return invlen * self
- @staticmethod
- @deprecated('ti.Matrix.transposed(a)', 'a.transpose()')
- def transposed(a):
- return a.transpose()
-
- @deprecated('a.T()', 'a.transpose()')
- def T(self):
- return self.transpose()
-
- @kern_mod.pyfunc
def transpose(self):
"""Get the transpose of a matrix.
@@ -584,10 +525,8 @@ def transpose(self):
Get the transpose of a matrix.
"""
- ret = Matrix([[self[i, j] for i in range(self.n)]
- for j in range(self.m)],
- disable_local_tensor=True)
- return ret
+ from taichi._funcs import _matrix_transpose # pylint: disable=C0415
+ return _matrix_transpose(self)
@taichi_scope
def determinant(a):
@@ -605,11 +544,11 @@ def determinant(a):
"""
if a.n == 2 and a.m == 2:
return a(0, 0) * a(1, 1) - a(0, 1) * a(1, 0)
- elif a.n == 3 and a.m == 3:
+ if a.n == 3 and a.m == 3:
return a(0, 0) * (a(1, 1) * a(2, 2) - a(2, 1) * a(1, 2)) - a(
1, 0) * (a(0, 1) * a(2, 2) - a(2, 1) * a(0, 2)) + a(
2, 0) * (a(0, 1) * a(1, 2) - a(1, 1) * a(0, 2))
- elif a.n == 4 and a.m == 4:
+ if a.n == 4 and a.m == 4:
n = 4
def E(x, y):
@@ -626,9 +565,8 @@ def E(x, y):
E(i + 3, 1) *
(E(i + 1, 2) * E(i + 2, 3) - E(i + 2, 2) * E(i + 1, 3))))
return det
- else:
- raise Exception(
- "Determinants of matrices with sizes >= 5 are not supported")
+ raise Exception(
+ "Determinants of matrices with sizes >= 5 are not supported")
@staticmethod
def diag(dim, val):
@@ -636,7 +574,7 @@ def diag(dim, val):
Args:
dim (int): the dimension of a square matrix.
- val (TypeVar): the diagonal elment value.
+ val (TypeVar): the diagonal element value.
Returns:
The constructed diagonal square matrix.
@@ -646,9 +584,9 @@ def diag(dim, val):
for i in range(dim):
for j in range(dim):
if i == j:
- ret.set_entry(i, j, val)
+ ret._set_entry(i, j, val)
else:
- ret.set_entry(i, j, 0 * val)
+ ret._set_entry(i, j, 0 * val)
# TODO: need a more systematic way to create a "0" with the right type
return ret
@@ -659,7 +597,6 @@ def sum(self):
ret = ret + self.entries[i]
return ret
- @kern_mod.pyfunc
def norm(self, eps=0):
"""Return the square root of the sum of the absolute squares of its elements.
@@ -678,7 +615,6 @@ def norm(self, eps=0):
"""
return ops_mod.sqrt(self.norm_sqr() + eps)
- @kern_mod.pyfunc
def norm_inv(self, eps=0):
"""Return the inverse of the matrix/vector `norm`. For `norm`: please see :func:`~taichi.lang.matrix.Matrix.norm`.
@@ -691,20 +627,17 @@ def norm_inv(self, eps=0):
"""
return ops_mod.rsqrt(self.norm_sqr() + eps)
- @kern_mod.pyfunc
def norm_sqr(self):
"""Return the sum of the absolute squares of its elements."""
- return (self**2).sum()
+ return (self * self).sum()
- @kern_mod.pyfunc
def max(self):
"""Return the maximum element value."""
- return ops_mod.ti_max(*self.entries)
+ return ops_mod.max(*self.entries)
- @kern_mod.pyfunc
def min(self):
- """Return the minumum element value."""
- return ops_mod.ti_min(*self.entries)
+ """Return the minimum element value."""
+ return ops_mod.min(*self.entries)
def any(self):
"""Test whether any element not equal zero.
@@ -713,10 +646,10 @@ def any(self):
bool: True if any element is not equal zero, False otherwise.
"""
- ret = ti.cmp_ne(self.entries[0], 0)
+ ret = ops_mod.cmp_ne(self.entries[0], 0)
for i in range(1, len(self.entries)):
- ret = ret + ti.cmp_ne(self.entries[i], 0)
- return -ti.cmp_lt(ret, 0)
+ ret = ret + ops_mod.cmp_ne(self.entries[i], 0)
+ return -ops_mod.cmp_lt(ret, 0)
def all(self):
"""Test whether all element not equal zero.
@@ -725,10 +658,10 @@ def all(self):
bool: True if all elements are not equal zero, False otherwise.
"""
- ret = ti.cmp_ne(self.entries[0], 0)
+ ret = ops_mod.cmp_ne(self.entries[0], 0)
for i in range(1, len(self.entries)):
- ret = ret + ti.cmp_ne(self.entries[i], 0)
- return -ti.cmp_eq(ret, -len(self.entries))
+ ret = ret + ops_mod.cmp_ne(self.entries[i], 0)
+ return -ops_mod.cmp_eq(ret, -len(self.entries))
@taichi_scope
def fill(self, val):
@@ -738,9 +671,9 @@ def fill(self, val):
val (Union[int, float]): Value to fill.
"""
def assign_renamed(x, y):
- return ti.assign(x, y)
+ return ops_mod.assign(x, y)
- return self.element_wise_writeback_binary(assign_renamed, val)
+ return self._element_wise_writeback_binary(assign_renamed, val)
@python_scope
def to_numpy(self, keep_dims=False):
@@ -755,7 +688,7 @@ def to_numpy(self, keep_dims=False):
"""
as_vector = self.m == 1 and not keep_dims
shape_ext = (self.n, ) if as_vector else (self.n, self.m)
- return np.array(self.value).reshape(shape_ext)
+ return np.array(self.to_list()).reshape(shape_ext)
@taichi_scope
def __ti_repr__(self):
@@ -788,8 +721,7 @@ def __str__(self):
So we have to make it happy with a dummy string...
'''
return f'<{self.n}x{self.m} ti.Matrix>'
- else:
- return str(self.to_numpy())
+ return str(self.to_numpy())
def __repr__(self):
return str(self.to_numpy())
@@ -809,12 +741,9 @@ def zero(dt, n, m=None):
"""
if m is None:
- return Vector([ti.cast(0, dt) for _ in range(n)],
- disable_local_tensor=True)
- else:
- return Matrix([[ti.cast(0, dt) for _ in range(m)]
- for _ in range(n)],
- disable_local_tensor=True)
+ return Vector([ops_mod.cast(0, dt) for _ in range(n)])
+ return Matrix([[ops_mod.cast(0, dt) for _ in range(m)]
+ for _ in range(n)])
@staticmethod
@taichi_scope
@@ -831,12 +760,9 @@ def one(dt, n, m=None):
"""
if m is None:
- return Vector([ti.cast(1, dt) for _ in range(n)],
- disable_local_tensor=True)
- else:
- return Matrix([[ti.cast(1, dt) for _ in range(m)]
- for _ in range(n)],
- disable_local_tensor=True)
+ return Vector([ops_mod.cast(1, dt) for _ in range(n)])
+ return Matrix([[ops_mod.cast(1, dt) for _ in range(m)]
+ for _ in range(n)])
@staticmethod
@taichi_scope
@@ -855,8 +781,7 @@ def unit(n, i, dt=None):
if dt is None:
dt = int
assert 0 <= i < n
- return Matrix([ti.cast(int(j == i), dt) for j in range(n)],
- disable_local_tensor=True)
+ return Vector([ops_mod.cast(int(j == i), dt) for j in range(n)])
@staticmethod
@taichi_scope
@@ -871,14 +796,14 @@ def identity(dt, n):
:class:`~taichi.lang.matrix.Matrix`: A n x n identity :class:`~taichi.lang.matrix.Matrix` instance.
"""
- return Matrix([[ti.cast(int(i == j), dt) for j in range(n)]
- for i in range(n)],
- disable_local_tensor=True)
+ return Matrix([[ops_mod.cast(int(i == j), dt) for j in range(n)]
+ for i in range(n)])
@staticmethod
def rotation2d(alpha):
- return Matrix([[ti.cos(alpha), -ti.sin(alpha)],
- [ti.sin(alpha), ti.cos(alpha)]])
+ return Matrix([[ops_mod.cos(alpha), -ops_mod.sin(alpha)],
+ [ops_mod.sin(alpha),
+ ops_mod.cos(alpha)]])
@classmethod
@python_scope
@@ -932,7 +857,8 @@ def field(cls,
entries, entries_grad = zip(*entries)
entries, entries_grad = MatrixField(entries, n, m), MatrixField(
entries_grad, n, m)
- entries.set_grad(entries_grad)
+ entries._set_grad(entries_grad)
+ impl.get_runtime().matrix_fields.append(entries)
if shape is None:
assert offset is None, "shape cannot be None when offset is being set"
@@ -950,43 +876,27 @@ def field(cls,
dim = len(shape)
if layout == Layout.SOA:
- for e in entries.get_field_members():
- ti.root.dense(impl.index_nd(dim),
- shape).place(ScalarField(e), offset=offset)
+ for e in entries._get_field_members():
+ impl.root.dense(impl.index_nd(dim),
+ shape).place(ScalarField(e), offset=offset)
if needs_grad:
- for e in entries_grad.get_field_members():
- ti.root.dense(impl.index_nd(dim),
- shape).place(ScalarField(e),
- offset=offset)
+ for e in entries_grad._get_field_members():
+ impl.root.dense(impl.index_nd(dim),
+ shape).place(ScalarField(e),
+ offset=offset)
else:
- ti.root.dense(impl.index_nd(dim), shape).place(entries,
- offset=offset)
+ impl.root.dense(impl.index_nd(dim), shape).place(entries,
+ offset=offset)
if needs_grad:
- ti.root.dense(impl.index_nd(dim),
- shape).place(entries_grad, offset=offset)
+ impl.root.dense(impl.index_nd(dim),
+ shape).place(entries_grad, offset=offset)
return entries
- @classmethod
- @python_scope
- @deprecated('ti.Matrix.var', 'ti.Matrix.field')
- def var(cls, n, m, dt, *args, **kwargs):
- """ti.Matrix.var"""
- _taichi_skip_traceback = 1
- return cls.field(n, m, dt, *args, **kwargs)
-
@classmethod
def _Vector_field(cls, n, dtype, *args, **kwargs):
"""ti.Vector.field"""
- _taichi_skip_traceback = 1
return cls.field(n, 1, dtype, *args, **kwargs)
- @classmethod
- @deprecated('ti.Vector.var', 'ti.Vector.field')
- def _Vector_var(cls, n, dt, *args, **kwargs):
- """ti.Vector.var"""
- _taichi_skip_traceback = 1
- return cls._Vector_field(n, dt, *args, **kwargs)
-
@classmethod
@python_scope
def ndarray(cls, n, m, dtype, shape, layout=Layout.AOS):
@@ -1030,7 +940,7 @@ def _Vector_ndarray(cls, n, dtype, shape, layout=Layout.AOS):
@staticmethod
def rows(rows):
- """Construct a Matrix instance by concactinating Vectors/lists row by row.
+ """Construct a Matrix instance by concatenating Vectors/lists row by row.
Args:
rows (List): A list of Vector (1-D Matrix) or a list of list.
@@ -1063,7 +973,7 @@ def rows(rows):
@staticmethod
def cols(cols):
- """Construct a Matrix instance by concactinating Vectors/lists column by column.
+ """Construct a Matrix instance by concatenating Vectors/lists column by column.
Args:
cols (List): A list of Vector (1-D Matrix) or a list of list.
@@ -1074,51 +984,12 @@ def cols(cols):
"""
return Matrix.rows(cols).transpose()
- @classmethod
- def empty(cls, n, m):
- """Clear the matrix and fill None.
-
- Args:
- n (int): The number of the row of the matrix.
- m (int): The number of the column of the matrix.
-
- Returns:
- :class:`~taichi.lang.matrix.Matrix`: A :class:`~taichi.lang.matrix.Matrix` instance filled with None.
-
- """
- return cls([[None] * m for _ in range(n)], disable_local_tensor=True)
-
- @classmethod
- def with_entries(cls, n, m, entries):
- """Construct a Matrix instance by giving all entries.
-
- Args:
- n (int): Number of rows of the matrix.
- m (int): Number of columns of the matrix.
- entries (List[Any]): Given entries.
-
- Returns:
- Matrix: A :class:`~taichi.lang.matrix.Matrix` instance filled with given entries.
- """
- assert n * m == len(entries), "Number of entries doesn't match n * m"
- mat = cls.empty(n, m)
- mat.entries = entries
- return mat
-
- @classmethod
- def new(cls, n, m):
- if impl.inside_kernel():
- return cls(n, m)
- else:
- return cls.empty(n, m)
-
def __hash__(self):
# TODO: refactor KernelTemplateMapper
# If not, we get `unhashable type: Matrix` when
# using matrices as template arguments.
return id(self)
- @kern_mod.pyfunc
def dot(self, other):
"""Perform the dot product with the input Vector (1-D Matrix).
@@ -1135,20 +1006,13 @@ def dot(self, other):
impl.static_assert(other.m == 1, "rhs for dot is not a vector"))
return (self * other).sum()
- @kern_mod.pyfunc
def _cross3d(self, other):
- ret = Matrix([
- self[1] * other[2] - self[2] * other[1],
- self[2] * other[0] - self[0] * other[2],
- self[0] * other[1] - self[1] * other[0],
- ],
- disable_local_tensor=True)
- return ret
+ from taichi._funcs import _matrix_cross3d # pylint: disable=C0415
+ return _matrix_cross3d(self, other)
- @kern_mod.pyfunc
def _cross2d(self, other):
- ret = self[0] * other[1] - self[1] * other[0]
- return ret
+ from taichi._funcs import _matrix_cross2d # pylint: disable=C0415
+ return _matrix_cross2d(self, other)
def cross(self, other):
"""Perform the cross product with the input Vector (1-D Matrix).
@@ -1163,15 +1027,12 @@ def cross(self, other):
if self.n == 3 and self.m == 1 and other.n == 3 and other.m == 1:
return self._cross3d(other)
- elif self.n == 2 and self.m == 1 and other.n == 2 and other.m == 1:
+ if self.n == 2 and self.m == 1 and other.n == 2 and other.m == 1:
return self._cross2d(other)
- else:
- raise ValueError(
- "Cross product is only supported between pairs of 2D/3D vectors"
- )
+ raise ValueError(
+ "Cross product is only supported between pairs of 2D/3D vectors")
- @kern_mod.pyfunc
def outer_product(self, other):
"""Perform the outer product with the input Vector (1-D Matrix).
@@ -1182,23 +1043,16 @@ def outer_product(self, other):
:class:`~taichi.lang.matrix.Matrix`: The outer product result (Matrix) of the two Vectors.
"""
- impl.static(
- impl.static_assert(self.m == 1,
- "lhs for outer_product is not a vector"))
- impl.static(
- impl.static_assert(other.m == 1,
- "rhs for outer_product is not a vector"))
- ret = Matrix([[self[i] * other[j] for j in range(other.n)]
- for i in range(self.n)],
- disable_local_tensor=True)
- return ret
+ from taichi._funcs import \
+ _matrix_outer_product # pylint: disable=C0415
+ return _matrix_outer_product(self, other)
def Vector(n, dt=None, **kwargs):
"""Construct a `Vector` instance i.e. 1-D Matrix.
Args:
- n (int): The desired number of entries of the Vector.
+ n (Union[int, list, tuple], np.ndarray): The desired number of entries of the Vector.
dt (DataType, optional): The desired data type of the Vector.
Returns:
@@ -1208,7 +1062,6 @@ def Vector(n, dt=None, **kwargs):
return Matrix(n, 1, dt=dt, **kwargs)
-Vector.var = Matrix._Vector_var
Vector.field = Matrix._Vector_field
Vector.ndarray = Matrix._Vector_ndarray
Vector.zero = Matrix.zero
@@ -1220,6 +1073,41 @@ def Vector(n, dt=None, **kwargs):
Vector.normalized = Matrix.normalized
+class _IntermediateMatrix(Matrix):
+ """Intermediate matrix class for compiler internal use only.
+
+ Args:
+ n (int): Number of rows of the matrix.
+ m (int): Number of columns of the matrix.
+ entries (List[Expr]): All entries of the matrix.
+ """
+ def __init__(self, n, m, entries):
+ assert isinstance(entries, list)
+ assert n * m == len(entries), "Number of entries doesn't match n * m"
+ self.n = n
+ self.m = m
+ self.entries = entries
+ self.local_tensor_proxy = None
+ self.any_array_access = None
+ self.grad = None
+ self.dynamic_index_stride = None
+
+
+class _MatrixFieldElement(_IntermediateMatrix):
+ """Matrix field element class for compiler internal use only.
+
+ Args:
+ field (MatrixField): The matrix field.
+ indices (taichi_core.ExprGroup): Indices of the element.
+ """
+ def __init__(self, field, indices):
+ super().__init__(field.n, field.m, [
+ expr.Expr(ti_core.subscript(e.ptr, indices))
+ for e in field._get_field_members()
+ ])
+ self.dynamic_index_stride = field.dynamic_index_stride
+
+
class MatrixField(Field):
"""Taichi matrix field with SNode implementation.
@@ -1228,15 +1116,12 @@ class MatrixField(Field):
n (Int): Number of rows.
m (Int): Number of columns.
"""
- def __init__(self, vars, n, m):
- assert len(vars) == n * m
- super().__init__(vars)
+ def __init__(self, _vars, n, m):
+ assert len(_vars) == n * m
+ super().__init__(_vars)
self.n = n
self.m = m
-
- @deprecated('x(i, j)', 'x.get_scalar_field(i, j)')
- def __call__(self, *indices):
- return self.get_scalar_field(*indices)
+ self.dynamic_index_stride = None
def get_scalar_field(self, *indices):
"""Creates a ScalarField using a specific field member. Only used for quant.
@@ -1252,6 +1137,37 @@ def get_scalar_field(self, *indices):
j = 0 if len(indices) == 1 else indices[1]
return ScalarField(self.vars[i * self.m + j])
+ def _calc_dynamic_index_stride(self):
+ # Algorithm: https://github.com/taichi-dev/taichi/issues/3810
+ paths = [ScalarField(var).snode._path_from_root() for var in self.vars]
+ num_members = len(paths)
+ if num_members == 1:
+ self.dynamic_index_stride = 0
+ return
+ length = len(paths[0])
+ if any(
+ len(path) != length or ti_core.is_custom_type(path[length -
+ 1]._dtype)
+ for path in paths):
+ return
+ for i in range(length):
+ if any(path[i] != paths[0][i] for path in paths):
+ depth_below_lca = i
+ break
+ for i in range(depth_below_lca, length - 1):
+ if any(path[i].ptr.type != ti_core.SNodeType.dense
+ or path[i]._cell_size_bytes != paths[0][i]._cell_size_bytes
+ or path[i + 1]._offset_bytes_in_parent_cell != paths[0][
+ i + 1]._offset_bytes_in_parent_cell for path in paths):
+ return
+ stride = paths[1][depth_below_lca]._offset_bytes_in_parent_cell - \
+ paths[0][depth_below_lca]._offset_bytes_in_parent_cell
+ for i in range(2, num_members):
+ if stride != paths[i][depth_below_lca]._offset_bytes_in_parent_cell \
+ - paths[i - 1][depth_below_lca]._offset_bytes_in_parent_cell:
+ return
+ self.dynamic_index_stride = stride
+
@python_scope
def fill(self, val):
"""Fills `self` with specific values.
@@ -1266,7 +1182,7 @@ def fill(self, val):
(list, tuple)) and isinstance(val[0], numbers.Number):
assert self.m == 1
val = tuple([(v, ) for v in val])
- elif isinstance(val, ti.Matrix):
+ elif isinstance(val, Matrix):
val_tuple = []
for i in range(val.n):
row = []
@@ -1277,10 +1193,11 @@ def fill(self, val):
val = tuple(val_tuple)
assert len(val) == self.n
assert len(val[0]) == self.m
- taichi.lang.meta.fill_matrix(self, val)
+ from taichi._kernels import fill_matrix # pylint: disable=C0415
+ fill_matrix(self, val)
@python_scope
- def to_numpy(self, keep_dims=False, as_vector=None, dtype=None):
+ def to_numpy(self, keep_dims=False, dtype=None):
"""Converts the field instance to a NumPy array.
Args:
@@ -1288,28 +1205,19 @@ def to_numpy(self, keep_dims=False, as_vector=None, dtype=None):
When keep_dims=True, on an n-D matrix field, the numpy array always has n+2 dims, even for 1x1, 1xn, nx1 matrix fields.
When keep_dims=False, the resulting numpy array should skip the matrix dims with size 1.
For example, a 4x1 or 1x4 matrix field with 5x6x7 elements results in an array of shape 5x6x7x4.
- as_vector (bool, deprecated): Whether to make the returned numpy array as a vector, i.e., with shape (n,) rather than (n, 1).
- Note that this argument has been deprecated.
- More discussion about `as_vector`: https://github.com/taichi-dev/taichi/pull/1046#issuecomment-633548858.
dtype (DataType, optional): The desired data type of returned numpy array.
Returns:
numpy.ndarray: The result NumPy array.
"""
- if as_vector is not None:
- warning(
- 'v.to_numpy(as_vector=True) is deprecated, '
- 'please use v.to_numpy() directly instead',
- DeprecationWarning,
- stacklevel=3)
if dtype is None:
dtype = to_numpy_type(self.dtype)
as_vector = self.m == 1 and not keep_dims
shape_ext = (self.n, ) if as_vector else (self.n, self.m)
- import numpy as np # pylint: disable=C0415
arr = np.zeros(self.shape + shape_ext, dtype=dtype)
- taichi.lang.meta.matrix_to_ext_arr(self, arr, as_vector)
- ti.sync()
+ from taichi._kernels import matrix_to_ext_arr # pylint: disable=C0415
+ matrix_to_ext_arr(self, arr, as_vector)
+ runtime_ops.sync()
return arr
def to_torch(self, device=None, keep_dims=False):
@@ -1326,11 +1234,13 @@ def to_torch(self, device=None, keep_dims=False):
import torch # pylint: disable=C0415
as_vector = self.m == 1 and not keep_dims
shape_ext = (self.n, ) if as_vector else (self.n, self.m)
+ # pylint: disable=E1101
arr = torch.empty(self.shape + shape_ext,
dtype=to_pytorch_type(self.dtype),
device=device)
- taichi.lang.meta.matrix_to_ext_arr(self, arr, as_vector)
- ti.sync()
+ from taichi._kernels import matrix_to_ext_arr # pylint: disable=C0415
+ matrix_to_ext_arr(self, arr, as_vector)
+ runtime_ops.sync()
return arr
@python_scope
@@ -1343,19 +1253,22 @@ def from_numpy(self, arr):
assert len(arr.shape) == len(self.shape) + 2
dim_ext = 1 if as_vector else 2
assert len(arr.shape) == len(self.shape) + dim_ext
- taichi.lang.meta.ext_arr_to_matrix(arr, self, as_vector)
- ti.sync()
+ from taichi._kernels import ext_arr_to_matrix # pylint: disable=C0415
+ ext_arr_to_matrix(arr, self, as_vector)
+ runtime_ops.sync()
@python_scope
def __setitem__(self, key, value):
- self.initialize_host_accessors()
- self[key].set_entries(value)
+ self._initialize_host_accessors()
+ self[key]._set_entries(value)
@python_scope
def __getitem__(self, key):
- self.initialize_host_accessors()
- key = self.pad_key(key)
- return Matrix.with_entries(self.n, self.m, self.host_access(key))
+ self._initialize_host_accessors()
+ key = self._pad_key(key)
+ _host_access = self._host_access(key)
+ return Matrix([[_host_access[i * self.m + j] for j in range(self.m)]
+ for i in range(self.n)])
def __repr__(self):
# make interactive shell happy, prevent materialization
@@ -1376,7 +1289,7 @@ def __call__(self, *args):
elif len(args) == 1:
# fill a single scalar
if isinstance(args[0], (numbers.Number, expr.Expr)):
- return self.scalar_filled(args[0])
+ return self.filled_with_scalar(args[0])
# fill a single vector or matrix
entries = args[0]
else:
@@ -1396,36 +1309,28 @@ def __call__(self, *args):
mat = self.cast(Matrix(entries, dt=self.dtype))
return mat
- def cast(self, mat, in_place=False):
- if not in_place:
- mat = mat.copy()
+ def cast(self, mat):
# sanity check shape
if self.m != mat.m or self.n != mat.n:
raise TaichiSyntaxError(
f"Incompatible arguments for the custom vector/matrix type: ({self.n}, {self.m}), ({mat.n}, {mat.m})"
)
if in_python_scope():
- mat.entries = [
- int(x) if self.dtype in ti.integer_types else x
- for x in mat.entries
- ]
- else:
- # only performs casting in Taichi scope
- mat.entries = [cast(x, self.dtype) for x in mat.entries]
- return mat
+ return Matrix([[
+ int(mat(i, j)) if self.dtype in primitive_types.integer_types
+ else float(mat(i, j)) for j in range(self.m)
+ ] for i in range(self.n)])
+ return mat.cast(self.dtype)
- def empty(self):
- """
- Create an empty instance of the given compound type.
- """
- return Matrix.empty(self.n, self.m)
+ def filled_with_scalar(self, value):
+ return Matrix([[value for _ in range(self.m)] for _ in range(self.n)])
def field(self, **kwargs):
return Matrix.field(self.n, self.m, dtype=self.dtype, **kwargs)
class MatrixNdarray(Ndarray):
- """Taichi ndarray with matrix elements implemented with a torch tensor.
+ """Taichi ndarray with matrix elements.
Args:
n (int): Number of rows of the matrix.
@@ -1436,21 +1341,16 @@ class MatrixNdarray(Ndarray):
"""
def __init__(self, n, m, dtype, shape, layout):
self.layout = layout
+ self.shape = shape
+ self.n = n
+ self.m = m
arr_shape = (n, m) + shape if layout == Layout.SOA else shape + (n, m)
super().__init__(dtype, arr_shape)
@property
- def n(self):
- return self.arr.shape[0 if self.layout == Layout.SOA else -2]
-
- @property
- def m(self):
- return self.arr.shape[1 if self.layout == Layout.SOA else -1]
-
- @property
- def shape(self):
+ def element_shape(self):
arr_shape = tuple(self.arr.shape)
- return arr_shape[2:] if self.layout == Layout.SOA else arr_shape[:-2]
+ return arr_shape[:2] if self.layout == Layout.SOA else arr_shape[-2:]
@python_scope
def __setitem__(self, key, value):
@@ -1466,17 +1366,35 @@ def __setitem__(self, key, value):
def __getitem__(self, key):
key = () if key is None else (
key, ) if isinstance(key, numbers.Number) else tuple(key)
- return Matrix.with_entries(self.n, self.m, [
- NdarrayHostAccess(self, key, (i, j)) for i in range(self.n)
- for j in range(self.m)
- ])
+ return Matrix(
+ [[NdarrayHostAccess(self, key, (i, j)) for j in range(self.m)]
+ for i in range(self.n)])
+
+ @python_scope
+ def to_numpy(self):
+ return self._ndarray_matrix_to_numpy(as_vector=0)
+
+ @python_scope
+ def from_numpy(self, arr):
+ self._ndarray_matrix_from_numpy(arr, as_vector=0)
+
+ def __deepcopy__(self, memo=None):
+ ret_arr = MatrixNdarray(self.n, self.m, self.dtype, self.shape,
+ self.layout)
+ ret_arr.copy_from(self)
+ return ret_arr
+
+ def _fill_by_kernel(self, val):
+ from taichi._kernels import \
+ fill_ndarray_matrix # pylint: disable=C0415
+ fill_ndarray_matrix(self, val)
def __repr__(self):
return f'<{self.n}x{self.m} {self.layout} ti.Matrix.ndarray>'
class VectorNdarray(Ndarray):
- """Taichi ndarray with vector elements implemented with a torch tensor.
+ """Taichi ndarray with vector elements.
Args:
n (int): Size of the vector.
@@ -1486,17 +1404,15 @@ class VectorNdarray(Ndarray):
"""
def __init__(self, n, dtype, shape, layout):
self.layout = layout
+ self.shape = shape
+ self.n = n
arr_shape = (n, ) + shape if layout == Layout.SOA else shape + (n, )
super().__init__(dtype, arr_shape)
@property
- def n(self):
- return self.arr.shape[0 if self.layout == Layout.SOA else -1]
-
- @property
- def shape(self):
+ def element_shape(self):
arr_shape = tuple(self.arr.shape)
- return arr_shape[1:] if self.layout == Layout.SOA else arr_shape[:-1]
+ return arr_shape[:1] if self.layout == Layout.SOA else arr_shape[-1:]
@python_scope
def __setitem__(self, key, value):
@@ -1509,9 +1425,29 @@ def __setitem__(self, key, value):
def __getitem__(self, key):
key = () if key is None else (
key, ) if isinstance(key, numbers.Number) else tuple(key)
- return Matrix.with_entries(
- self.n, 1,
+ return Vector(
[NdarrayHostAccess(self, key, (i, )) for i in range(self.n)])
+ @python_scope
+ def to_numpy(self):
+ return self._ndarray_matrix_to_numpy(as_vector=1)
+
+ @python_scope
+ def from_numpy(self, arr):
+ self._ndarray_matrix_from_numpy(arr, as_vector=1)
+
+ def __deepcopy__(self, memo=None):
+ ret_arr = VectorNdarray(self.n, self.dtype, self.shape, self.layout)
+ ret_arr.copy_from(self)
+ return ret_arr
+
+ def _fill_by_kernel(self, val):
+ from taichi._kernels import \
+ fill_ndarray_matrix # pylint: disable=C0415
+ fill_ndarray_matrix(self, val)
+
def __repr__(self):
return f'<{self.n} {self.layout} ti.Vector.ndarray>'
+
+
+__all__ = ["Matrix", "Vector", "MatrixField", "MatrixNdarray", "VectorNdarray"]
diff --git a/python/taichi/lang/mesh.py b/python/taichi/lang/mesh.py
new file mode 100644
index 0000000000000..850ece441b5cf
--- /dev/null
+++ b/python/taichi/lang/mesh.py
@@ -0,0 +1,537 @@
+import json
+
+import numpy as np
+from taichi._lib import core as _ti_core
+from taichi.lang import impl
+from taichi.lang.enums import Layout
+from taichi.lang.exception import TaichiSyntaxError
+from taichi.lang.field import Field, ScalarField
+from taichi.lang.matrix import (MatrixField, _IntermediateMatrix,
+ _MatrixFieldElement)
+from taichi.lang.struct import StructField
+from taichi.lang.util import python_scope
+from taichi.types import i32
+from taichi.types.compound_types import CompoundType
+
+from taichi import lang
+
+MeshTopology = _ti_core.MeshTopology
+MeshElementType = _ti_core.MeshElementType
+MeshRelationType = _ti_core.MeshRelationType
+ConvType = _ti_core.ConvType
+element_order = _ti_core.element_order
+from_end_element_order = _ti_core.from_end_element_order
+to_end_element_order = _ti_core.to_end_element_order
+relation_by_orders = _ti_core.relation_by_orders
+inverse_relation = _ti_core.inverse_relation
+element_type_name = _ti_core.element_type_name
+
+
+class MeshAttrType:
+ def __init__(self, name, dtype, reorder, needs_grad):
+ self.name = name
+ self.dtype = dtype
+ self.reorder = reorder
+ self.needs_grad = needs_grad
+
+
+class MeshReorderedScalarFieldProxy(ScalarField):
+ def __init__(self, field: ScalarField, mesh_ptr: _ti_core.MeshPtr,
+ element_type: MeshElementType, g2r_field: ScalarField):
+ self.vars = field.vars
+ self.host_accessors = field.host_accessors
+ self.grad = field.grad
+
+ self.mesh_ptr = mesh_ptr
+ self.element_type = element_type
+ self.g2r_field = g2r_field
+
+ @python_scope
+ def __setitem__(self, key, value):
+ self._initialize_host_accessors()
+ key = self.g2r_field[key]
+ self.host_accessors[0].setter(value, *self._pad_key(key))
+
+ @python_scope
+ def __getitem__(self, key):
+ self._initialize_host_accessors()
+ key = self.g2r_field[key]
+ return self.host_accessors[0].getter(*self._pad_key(key))
+
+
+class MeshReorderedMatrixFieldProxy(MatrixField):
+ def __init__(self, field: MatrixField, mesh_ptr: _ti_core.MeshPtr,
+ element_type: MeshElementType, g2r_field: ScalarField):
+ self.vars = field.vars
+ self.host_accessors = field.host_accessors
+ self.grad = field.grad
+ self.n = field.n
+ self.m = field.m
+ self.dynamic_index_stride = field.dynamic_index_stride
+
+ self.mesh_ptr = mesh_ptr
+ self.element_type = element_type
+ self.g2r_field = g2r_field
+
+ @python_scope
+ def __setitem__(self, key, value):
+ self._initialize_host_accessors()
+ self[key]._set_entries(value)
+
+ @python_scope
+ def __getitem__(self, key):
+ self._initialize_host_accessors()
+ key = self.g2r_field[key]
+ key = self._pad_key(key)
+ return _IntermediateMatrix(self.n, self.m, self._host_access(key))
+
+
+class MeshElementField:
+ def __init__(self, mesh_instance, _type, attr_dict, field_dict, g2r_field):
+ self.mesh = mesh_instance
+ self._type = _type
+ self.attr_dict = attr_dict
+ self.field_dict = field_dict
+ self.g2r_field = g2r_field
+
+ self._register_fields()
+
+ @property
+ def keys(self):
+ return list(self.field_dict.keys())
+
+ @property
+ def _members(self):
+ return list(self.field_dict.values())
+
+ @property
+ def _items(self):
+ return self.field_dict.items()
+
+ @staticmethod
+ def _make_getter(key):
+ def getter(self):
+ if key not in self.getter_dict:
+ if self.attr_dict[key].reorder:
+ if isinstance(self.field_dict[key], ScalarField):
+ self.getter_dict[key] = MeshReorderedScalarFieldProxy(
+ self.field_dict[key], self.mesh.mesh_ptr,
+ self._type, self.g2r_field)
+ elif isinstance(self.field_dict[key], MatrixField):
+ self.getter_dict[key] = MeshReorderedMatrixFieldProxy(
+ self.field_dict[key], self.mesh.mesh_ptr,
+ self._type, self.g2r_field)
+ else:
+ self.getter_dict[key] = self.field_dict[key]
+ """Get an entry from custom struct by name."""
+ return self.getter_dict[key]
+
+ return getter
+
+ def _register_fields(self):
+ self.getter_dict = {}
+ for k in self.keys:
+ setattr(MeshElementField, k,
+ property(fget=MeshElementField._make_getter(k)))
+
+ def _get_field_members(self):
+ field_members = []
+ for m in self._members:
+ assert isinstance(m, Field)
+ field_members += m._get_field_members()
+ return field_members
+
+ @python_scope
+ def copy_from(self, other):
+ assert isinstance(other, Field)
+ assert set(self.keys) == set(other.keys)
+ for k in self.keys:
+ self.field_dict[k].copy_from(other[k])
+
+ @python_scope
+ def fill(self, val):
+ for v in self._members:
+ v.fill(val)
+
+ def _initialize_host_accessors(self):
+ for v in self._members:
+ v._initialize_host_accessors()
+
+ def get_member_field(self, key):
+ return self.field_dict[key]
+
+ @python_scope
+ def from_numpy(self, array_dict):
+ for k, v in self._items:
+ v.from_numpy(array_dict[k])
+
+ @python_scope
+ def from_torch(self, array_dict):
+ for k, v in self._items:
+ v.from_torch(array_dict[k])
+
+ @python_scope
+ def to_numpy(self):
+ return {k: v.to_numpy() for k, v in self._items}
+
+ @python_scope
+ def to_torch(self, device=None):
+ return {k: v.to_torch(device=device) for k, v in self._items}
+
+ @python_scope
+ def __len__(self):
+ return _ti_core.get_num_elements(self.mesh.mesh_ptr, self._type)
+
+
+class MeshElement:
+ def __init__(self, _type, builder):
+ self.builder = builder
+ self._type = _type
+ self.layout = Layout.SOA
+ self.attr_dict = {}
+
+ def _SOA(self, soa=True): # AOS/SOA
+ self.layout = Layout.SOA if soa else Layout.AOS
+
+ def _AOS(self, aos=True):
+ self.layout = Layout.AOS if aos else Layout.SOA
+
+ SOA = property(fget=_SOA)
+ AOS = property(fget=_AOS)
+
+ def place(
+ self,
+ members,
+ reorder=False,
+ needs_grad=False,
+ ):
+ self.builder.elements.add(self._type)
+ for key, dtype in members.items():
+ if key in {'verts', 'edges', 'faces', 'cells'}:
+ raise TaichiSyntaxError(
+ f"'{key}' cannot use as attribute name. It has been reserved as ti.Mesh's keyword."
+ )
+ self.attr_dict[key] = MeshAttrType(key, dtype, reorder, needs_grad)
+
+ def build(self, mesh_instance, size, g2r_field):
+ field_dict = {}
+
+ for key, attr in self.attr_dict.items():
+ if isinstance(attr.dtype, CompoundType):
+ field_dict[key] = attr.dtype.field(shape=None,
+ needs_grad=attr.needs_grad)
+ else:
+ field_dict[key] = impl.field(attr.dtype,
+ shape=None,
+ needs_grad=attr.needs_grad)
+
+ if self.layout == Layout.SOA:
+ for key, field in field_dict.items():
+ impl.root.dense(impl.axes(0), size).place(field)
+ if self.attr_dict[key].needs_grad:
+ impl.root.dense(impl.axes(0), size).place(field.grad)
+ elif len(field_dict) > 0:
+ impl.root.dense(impl.axes(0),
+ size).place(*tuple(field_dict.values()))
+ grads = []
+ for key, field in field_dict.items():
+ if self.attr_dict[key].needs_grad:
+ grads.append(field.grad)
+ if len(grads) > 0:
+ impl.root.dense(impl.axes(0), size).place(*grads)
+
+ return MeshElementField(mesh_instance, self._type, self.attr_dict,
+ field_dict, g2r_field)
+
+ def link(self, element):
+ assert isinstance(element, MeshElement)
+ assert element.builder == self.builder
+ self.builder.relations.add(tuple([self._type, element._type]))
+ self.builder.elements.add(self._type)
+ self.builder.elements.add(element._type)
+
+
+# Define the instance of the Mesh Type, stores the field (type and data) info
+class MeshInstance:
+ def __init__(self, _type):
+ self._type = _type
+ self.mesh_ptr = _ti_core.create_mesh()
+
+ def set_owned_offset(self, element_type: MeshElementType,
+ owned_offset: ScalarField):
+ _ti_core.set_owned_offset(self.mesh_ptr, element_type,
+ owned_offset.vars[0].ptr.snode())
+
+ def set_total_offset(self, element_type: MeshElementType,
+ total_offset: ScalarField):
+ _ti_core.set_total_offset(self.mesh_ptr, element_type,
+ total_offset.vars[0].ptr.snode())
+
+ def set_index_mapping(self, element_type: MeshElementType,
+ conv_type: ConvType, mapping: ScalarField):
+ _ti_core.set_index_mapping(self.mesh_ptr, element_type, conv_type,
+ mapping.vars[0].ptr.snode())
+
+ def set_num_patches(self, num_patches: int):
+ _ti_core.set_num_patches(self.mesh_ptr, num_patches)
+
+ def set_patch_max_element_num(self, element_type: MeshElementType,
+ max_element_num: int):
+ _ti_core.set_patch_max_element_num(self.mesh_ptr, element_type,
+ max_element_num)
+
+ def set_relation_fixed(self, rel_type: MeshRelationType,
+ value: ScalarField):
+ _ti_core.set_relation_fixed(self.mesh_ptr, rel_type,
+ value.vars[0].ptr.snode())
+
+ def set_relation_dynamic(self, rel_type: MeshRelationType,
+ value: ScalarField, offset: ScalarField):
+ _ti_core.set_relation_dynamic(self.mesh_ptr, rel_type,
+ value.vars[0].ptr.snode(),
+ offset.vars[0].ptr.snode())
+
+ def add_mesh_attribute(self, element_type, snode, reorder_type):
+ _ti_core.add_mesh_attribute(self.mesh_ptr, element_type, snode,
+ reorder_type)
+
+
+class MeshMetadata:
+ def __init__(self, data):
+ self.num_patches = data["num_patches"]
+
+ self.element_fields = {}
+ self.relation_fields = {}
+ self.num_elements = {}
+ self.max_num_per_patch = {}
+
+ for element in data["elements"]:
+ element_type = MeshElementType(element["order"])
+ self.num_elements[element_type] = element["num"]
+ self.max_num_per_patch[element_type] = element["max_num_per_patch"]
+
+ element["l2g_mapping"] = np.array(element["l2g_mapping"])
+ element["l2r_mapping"] = np.array(element["l2r_mapping"])
+ element["g2r_mapping"] = np.array(element["g2r_mapping"])
+ self.element_fields[element_type] = {}
+ self.element_fields[element_type]["owned"] = impl.field(
+ dtype=i32, shape=self.num_patches + 1)
+ self.element_fields[element_type]["total"] = impl.field(
+ dtype=i32, shape=self.num_patches + 1)
+ self.element_fields[element_type]["l2g"] = impl.field(
+ dtype=i32, shape=element["l2g_mapping"].shape[0])
+ self.element_fields[element_type]["l2r"] = impl.field(
+ dtype=i32, shape=element["l2r_mapping"].shape[0])
+ self.element_fields[element_type]["g2r"] = impl.field(
+ dtype=i32, shape=element["g2r_mapping"].shape[0])
+
+ for relation in data["relations"]:
+ from_order = relation["from_order"]
+ to_order = relation["to_order"]
+ rel_type = MeshRelationType(
+ relation_by_orders(from_order, to_order))
+ self.relation_fields[rel_type] = {}
+ self.relation_fields[rel_type]["value"] = impl.field(
+ dtype=i32, shape=len(relation["value"]))
+ if from_order <= to_order:
+ self.relation_fields[rel_type]["offset"] = impl.field(
+ dtype=i32, shape=len(relation["offset"]))
+
+ for element in data["elements"]:
+ element_type = MeshElementType(element["order"])
+ self.element_fields[element_type]["owned"].from_numpy(
+ np.array(element["owned_offsets"]))
+ self.element_fields[element_type]["total"].from_numpy(
+ np.array(element["total_offsets"]))
+ self.element_fields[element_type]["l2g"].from_numpy(
+ element["l2g_mapping"])
+ self.element_fields[element_type]["l2r"].from_numpy(
+ element["l2r_mapping"])
+ self.element_fields[element_type]["g2r"].from_numpy(
+ element["g2r_mapping"])
+
+ for relation in data["relations"]:
+ from_order = relation["from_order"]
+ to_order = relation["to_order"]
+ rel_type = MeshRelationType(
+ relation_by_orders(from_order, to_order))
+ self.relation_fields[rel_type]["value"].from_numpy(
+ np.array(relation["value"]))
+ if from_order <= to_order:
+ self.relation_fields[rel_type]["offset"].from_numpy(
+ np.array(relation["offset"]))
+
+ self.attrs = {}
+ self.attrs["x"] = np.array(data["attrs"]["x"]).reshape(-1, 3)
+
+
+# Define the Mesh Type, stores the field type info
+class MeshBuilder:
+ def __init__(self, topology):
+ if not lang.misc.is_extension_supported(impl.current_cfg().arch,
+ lang.extension.mesh):
+ raise Exception('Backend ' + str(impl.current_cfg().arch) +
+ ' doesn\'t support MeshTaichi extension')
+
+ self.topology = topology
+ self.verts = MeshElement(MeshElementType.Vertex, self)
+ self.edges = MeshElement(MeshElementType.Edge, self)
+ self.faces = MeshElement(MeshElementType.Face, self)
+ if topology == MeshTopology.Tetrahedron:
+ self.cells = MeshElement(MeshElementType.Cell, self)
+
+ self.elements = set()
+ self.relations = set()
+
+ def build(self, metadata: MeshMetadata):
+ instance = MeshInstance(self)
+ instance.fields = {}
+
+ instance.set_num_patches(metadata.num_patches)
+
+ for element in self.elements:
+ _ti_core.set_num_elements(instance.mesh_ptr, element,
+ metadata.num_elements[element])
+ instance.set_patch_max_element_num(
+ element, metadata.max_num_per_patch[element])
+
+ element_name = element_type_name(element)
+ setattr(
+ instance, element_name,
+ getattr(self, element_name).build(
+ instance, metadata.num_elements[element],
+ metadata.element_fields[element]["g2r"]))
+ instance.fields[element] = getattr(instance, element_name)
+
+ instance.set_owned_offset(
+ element, metadata.element_fields[element]["owned"])
+ instance.set_total_offset(
+ element, metadata.element_fields[element]["total"])
+ instance.set_index_mapping(element, ConvType.l2g,
+ metadata.element_fields[element]["l2g"])
+ instance.set_index_mapping(element, ConvType.l2r,
+ metadata.element_fields[element]["l2r"])
+ instance.set_index_mapping(element, ConvType.g2r,
+ metadata.element_fields[element]["g2r"])
+
+ for relation in self.relations:
+ from_order = element_order(relation[0])
+ to_order = element_order(relation[1])
+ rel_type = MeshRelationType(
+ relation_by_orders(from_order, to_order))
+ if from_order <= to_order:
+ instance.set_relation_dynamic(
+ rel_type, metadata.relation_fields[rel_type]["value"],
+ metadata.relation_fields[rel_type]["offset"])
+ else:
+ instance.set_relation_fixed(
+ rel_type, metadata.relation_fields[rel_type]["value"])
+
+ if "x" in instance.verts.attr_dict: # pylint: disable=E1101
+ instance.verts.x.from_numpy(metadata.attrs["x"]) # pylint: disable=E1101
+
+ return instance
+
+
+# Mesh First Class
+class Mesh:
+ def __init__(self):
+ pass
+
+ @staticmethod
+ def Tet():
+ return MeshBuilder(MeshTopology.Tetrahedron)
+
+ @staticmethod
+ def Tri():
+ return MeshBuilder(MeshTopology.Triangle)
+
+ @staticmethod
+ def load_meta(filename):
+ with open(filename, "r") as fi:
+ data = json.loads(fi.read())
+ return MeshMetadata(data)
+
+ @staticmethod
+ def generate_meta(data):
+ return MeshMetadata(data)
+
+
+def TriMesh():
+ return Mesh.Tri()
+
+
+def TetMesh():
+ return Mesh.Tet()
+
+
+class MeshElementFieldProxy:
+ def __init__(self, mesh: MeshInstance, element_type: MeshElementType,
+ entry_expr: impl.Expr):
+ self.mesh = mesh
+ self.element_type = element_type
+ self.entry_expr = entry_expr
+
+ element_field = self.mesh.fields[self.element_type]
+ for key, attr in element_field.field_dict.items():
+ global_entry_expr = impl.Expr(
+ _ti_core.get_index_conversion(
+ self.mesh.mesh_ptr, element_type, entry_expr,
+ ConvType.l2r if element_field.attr_dict[key].reorder else
+ ConvType.l2g)) # transform index space
+ global_entry_expr_group = impl.make_expr_group(
+ *tuple([global_entry_expr]))
+ if isinstance(attr, MatrixField):
+ setattr(self, key,
+ _MatrixFieldElement(attr, global_entry_expr_group))
+ elif isinstance(attr, StructField):
+ raise RuntimeError('ti.Mesh has not support StructField yet')
+ else: # isinstance(attr, Field)
+ var = attr._get_field_members()[0].ptr
+ setattr(
+ self, key,
+ impl.Expr(_ti_core.subscript(var,
+ global_entry_expr_group)))
+
+ for element_type in self.mesh._type.elements:
+ setattr(self, element_type_name(element_type),
+ impl.mesh_relation_access(self.mesh, self, element_type))
+
+ @property
+ def ptr(self):
+ return self.entry_expr
+
+ @property
+ def id(self): # return the global non-reordered index
+ l2g_expr = impl.Expr(
+ _ti_core.get_index_conversion(self.mesh.mesh_ptr,
+ self.element_type, self.entry_expr,
+ ConvType.l2g))
+ return l2g_expr
+
+
+class MeshRelationAccessProxy:
+ def __init__(self, mesh: MeshInstance, from_index: impl.Expr,
+ to_element_type: MeshElementType):
+ self.mesh = mesh
+ self.from_index = from_index
+ self.to_element_type = to_element_type
+
+ @property
+ def size(self):
+ return impl.Expr(
+ _ti_core.get_relation_size(self.mesh.mesh_ptr, self.from_index.ptr,
+ self.to_element_type))
+
+ def subscript(self, *indices):
+ assert len(indices) == 1
+ entry_expr = _ti_core.get_relation_access(self.mesh.mesh_ptr,
+ self.from_index.ptr,
+ self.to_element_type,
+ impl.Expr(indices[0]).ptr)
+ entry_expr.type_check(impl.get_runtime().prog.config)
+ return MeshElementFieldProxy(self.mesh, self.to_element_type,
+ entry_expr)
+
+
+__all__ = ["Mesh", "TetMesh", "TriMesh"]
diff --git a/python/taichi/lang/meta.py b/python/taichi/lang/meta.py
deleted file mode 100644
index 5dc3a080d9348..0000000000000
--- a/python/taichi/lang/meta.py
+++ /dev/null
@@ -1,135 +0,0 @@
-from taichi.core import get_os_name
-from taichi.lang import impl
-from taichi.lang.expr import Expr
-from taichi.lang.field import ScalarField
-from taichi.lang.kernel_impl import kernel
-from taichi.type.annotations import ext_arr, template
-
-import taichi as ti
-
-# A set of helper (meta)functions
-
-
-@kernel
-def fill_tensor(tensor: template(), val: template()):
- for I in ti.grouped(tensor):
- tensor[I] = val
-
-
-@kernel
-def tensor_to_ext_arr(tensor: template(), arr: ext_arr()):
- for I in ti.grouped(tensor):
- arr[I] = tensor[I]
-
-
-@kernel
-def vector_to_fast_image(img: template(), out: ext_arr()):
- # FIXME: Why is ``for i, j in img:`` slower than:
- for i, j in ti.ndrange(*img.shape):
- r, g, b = 0, 0, 0
- color = img[i, img.shape[1] - 1 - j]
- if ti.static(img.dtype in [ti.f32, ti.f64]):
- r, g, b = min(255, max(0, int(color * 255)))
- else:
- impl.static_assert(img.dtype == ti.u8)
- r, g, b = color
- idx = j * img.shape[0] + i
- # We use i32 for |out| since OpenGL and Metal doesn't support u8 types
- if ti.static(get_os_name() != 'osx'):
- out[idx] = (r << 16) + (g << 8) + b
- else:
- # What's -16777216?
- #
- # On Mac, we need to set the alpha channel to 0xff. Since Mac's GUI
- # is big-endian, the color is stored in ABGR order, and we need to
- # add 0xff000000, which is -16777216 in I32's legit range. (Albeit
- # the clarity, adding 0xff000000 doesn't work.)
- alpha = -16777216
- out[idx] = (b << 16) + (g << 8) + r + alpha
-
-
-@kernel
-def tensor_to_image(tensor: template(), arr: ext_arr()):
- for I in ti.grouped(tensor):
- t = ti.cast(tensor[I], ti.f32)
- arr[I, 0] = t
- arr[I, 1] = t
- arr[I, 2] = t
-
-
-@kernel
-def vector_to_image(mat: template(), arr: ext_arr()):
- for I in ti.grouped(mat):
- for p in ti.static(range(mat.n)):
- arr[I, p] = ti.cast(mat[I][p], ti.f32)
- if ti.static(mat.n <= 2):
- arr[I, 2] = 0
-
-
-@kernel
-def tensor_to_tensor(tensor: template(), other: template()):
- for I in ti.grouped(tensor):
- tensor[I] = other[I]
-
-
-@kernel
-def ext_arr_to_tensor(arr: ext_arr(), tensor: template()):
- for I in ti.grouped(tensor):
- tensor[I] = arr[I]
-
-
-@kernel
-def matrix_to_ext_arr(mat: template(), arr: ext_arr(), as_vector: template()):
- for I in ti.grouped(mat):
- for p in ti.static(range(mat.n)):
- for q in ti.static(range(mat.m)):
- if ti.static(as_vector):
- arr[I, p] = mat[I][p]
- else:
- arr[I, p, q] = mat[I][p, q]
-
-
-@kernel
-def ext_arr_to_matrix(arr: ext_arr(), mat: template(), as_vector: template()):
- for I in ti.grouped(mat):
- for p in ti.static(range(mat.n)):
- for q in ti.static(range(mat.m)):
- if ti.static(as_vector):
- mat[I][p] = arr[I, p]
- else:
- mat[I][p, q] = arr[I, p, q]
-
-
-@kernel
-def clear_gradients(vars: template()):
- for I in ti.grouped(ScalarField(Expr(vars[0]))):
- for s in ti.static(vars):
- ScalarField(Expr(s))[I] = 0
-
-
-@kernel
-def clear_loss(l: template()):
- # Using SNode writers would result in a forced sync, therefore we wrap these
- # writes into a kernel.
- l[None] = 0
- l.grad[None] = 1
-
-
-@kernel
-def fill_matrix(mat: template(), vals: template()):
- for I in ti.grouped(mat):
- for p in ti.static(range(mat.n)):
- for q in ti.static(range(mat.m)):
- mat[I][p, q] = vals[p][q]
-
-
-@kernel
-def snode_deactivate(b: template()):
- for I in ti.grouped(b):
- ti.deactivate(b, I)
-
-
-@kernel
-def snode_deactivate_dynamic(b: template()):
- for I in ti.grouped(b.parent()):
- ti.deactivate(b, I)
diff --git a/python/taichi/lang/misc.py b/python/taichi/lang/misc.py
new file mode 100644
index 0000000000000..8d54644f427ba
--- /dev/null
+++ b/python/taichi/lang/misc.py
@@ -0,0 +1,667 @@
+import atexit
+import functools
+import os
+import shutil
+import tempfile
+import warnings
+from copy import deepcopy as _deepcopy
+
+from taichi._lib import core as _ti_core
+from taichi._lib.utils import locale_encode
+from taichi.lang import impl
+from taichi.lang.expr import Expr
+from taichi.lang.impl import axes, get_runtime
+from taichi.lang.snode import SNode
+from taichi.profiler.kernel_profiler import get_default_kernel_profiler
+from taichi.types.primitive_types import f32, f64, i32, i64
+
+from taichi import _logging, _snode, _version_check
+
+warnings.filterwarnings("once", category=DeprecationWarning, module="taichi")
+
+# ----------------------
+i = axes(0)
+"""Axis 0. For multi-dimensional arrays it's the direction downward the rows.
+For a 1d array it's the direction along this array.
+"""
+# ----------------------
+
+j = axes(1)
+"""Axis 1. For multi-dimensional arrays it's the direction across the columns.
+"""
+# ----------------------
+
+k = axes(2)
+"""Axis 2. For arrays of dimension `d` >= 3, view each cell as an array of
+lower dimension d-2, it's the first axis of this cell.
+"""
+# ----------------------
+
+l = axes(3)
+"""Axis 3. For arrays of dimension `d` >= 4, view each cell as an array of
+lower dimension d-2, it's the second axis of this cell.
+"""
+# ----------------------
+
+ij = axes(0, 1)
+"""Axes (0, 1).
+"""
+# ----------------------
+
+ik = axes(0, 2)
+"""Axes (0, 2).
+"""
+# ----------------------
+
+il = axes(0, 3)
+"""Axes (0, 3).
+"""
+# ----------------------
+
+jk = axes(1, 2)
+"""Axes (1, 2).
+"""
+# ----------------------
+
+jl = axes(1, 3)
+"""Axes (1, 3).
+"""
+# ----------------------
+
+kl = axes(2, 3)
+"""Axes (2, 3).
+"""
+# ----------------------
+
+ijk = axes(0, 1, 2)
+"""Axes (0, 1, 2).
+"""
+# ----------------------
+
+ijl = axes(0, 1, 3)
+"""Axes (0, 1, 3).
+"""
+# ----------------------
+
+ikl = axes(0, 2, 3)
+"""Axes (0, 2, 3).
+"""
+# ----------------------
+
+jkl = axes(1, 2, 3)
+"""Axes (1, 2, 3).
+"""
+# ----------------------
+
+ijkl = axes(0, 1, 2, 3)
+"""Axes (0, 1, 2, 3).
+"""
+# ----------------------
+
+# ----------------------
+
+x86_64 = _ti_core.x64
+"""The x64 CPU backend.
+"""
+# ----------------------
+
+x64 = _ti_core.x64
+"""The X64 CPU backend.
+"""
+# ----------------------
+
+arm64 = _ti_core.arm64
+"""The ARM CPU backend.
+"""
+# ----------------------
+
+cuda = _ti_core.cuda
+"""The CUDA backend.
+"""
+# ----------------------
+
+metal = _ti_core.metal
+"""The Apple Metal backend.
+"""
+# ----------------------
+
+opengl = _ti_core.opengl
+"""The OpenGL backend. OpenGL 4.3 required.
+"""
+# ----------------------
+
+# Skip annotating this one because it is barely maintained.
+cc = _ti_core.cc
+
+# ----------------------
+
+wasm = _ti_core.wasm
+"""The WebAssembly backend.
+"""
+# ----------------------
+
+vulkan = _ti_core.vulkan
+"""The Vulkan backend.
+"""
+# ----------------------
+
+dx11 = _ti_core.dx11
+"""The DX11 backend.
+"""
+# ----------------------
+
+gpu = [cuda, metal, opengl, vulkan, dx11]
+"""A list of GPU backends supported on the current system.
+
+When this is used, Taichi automatically picks the matching GPU backend. If no
+GPU is detected, Taichi falls back to the CPU backend.
+"""
+# ----------------------
+
+cpu = _ti_core.host_arch()
+"""A list of CPU backends supported on the current system.
+
+When this is used, Taichi automatically picks the matching CPU backend.
+"""
+# ----------------------
+
+timeline_clear = lambda: impl.get_runtime().prog.timeline_clear() # pylint: disable=unnecessary-lambda
+timeline_save = lambda fn: impl.get_runtime().prog.timeline_save(fn) # pylint: disable=unnecessary-lambda
+
+# Legacy API
+type_factory_ = _ti_core.get_type_factory_instance()
+
+extension = _ti_core.Extension
+"""An instance of Taichi extension.
+
+The list of currently available extensions is ['sparse', 'async_mode', 'quant', \
+ 'mesh', 'quant_basic', 'data64', 'adstack', 'bls', 'assertion', \
+ 'extfunc', 'packed', 'dynamic_index'].
+"""
+
+
+def is_extension_supported(arch, ext):
+ """Checks whether an extension is supported on an arch.
+
+ Args:
+ arch (taichi_core.Arch): Specified arch.
+ ext (taichi_core.Extension): Specified extension.
+
+ Returns:
+ bool: Whether `ext` is supported on `arch`.
+ """
+ return _ti_core.is_extension_supported(arch, ext)
+
+
+def reset():
+ """Resets Taichi to its initial state.
+ This will destroy all the allocated fields and kernels, and restore
+ the runtime to its default configuration.
+
+ Example::
+
+ >>> a = ti.field(ti.i32, shape=())
+ >>> a[None] = 1
+ >>> print("before reset: ", a)
+ before rest: 1
+ >>>
+ >>> ti.reset()
+ >>> print("after reset: ", a)
+ # will raise error because a is unavailable after reset.
+ """
+ impl.reset()
+ global runtime
+ runtime = impl.get_runtime()
+
+
+class _EnvironmentConfigurator:
+ def __init__(self, kwargs, _cfg):
+ self.cfg = _cfg
+ self.kwargs = kwargs
+ self.keys = []
+
+ def add(self, key, _cast=None):
+ _cast = _cast or self.bool_int
+
+ self.keys.append(key)
+
+ # TI_ASYNC= : no effect
+ # TI_ASYNC=0 : False
+ # TI_ASYNC=1 : True
+ name = 'TI_' + key.upper()
+ value = os.environ.get(name, '')
+ if len(value):
+ self[key] = _cast(value)
+ if key in self.kwargs:
+ _ti_core.warn(
+ f'ti.init argument "{key}" overridden by environment variable {name}={value}'
+ )
+ del self.kwargs[key] # mark as recognized
+ elif key in self.kwargs:
+ self[key] = self.kwargs[key]
+ del self.kwargs[key] # mark as recognized
+
+ def __getitem__(self, key):
+ return getattr(self.cfg, key)
+
+ def __setitem__(self, key, value):
+ setattr(self.cfg, key, value)
+
+ @staticmethod
+ def bool_int(x):
+ return bool(int(x))
+
+
+class _SpecialConfig:
+ # like CompileConfig in C++, this is the configurations that belong to other submodules
+ def __init__(self):
+ self.log_level = 'info'
+ self.gdb_trigger = False
+ self.short_circuit_operators = False
+
+
+def prepare_sandbox():
+ '''
+ Returns a temporary directory, which will be automatically deleted on exit.
+ It may contain the taichi_core shared object or some misc. files.
+ '''
+ tmp_dir = tempfile.mkdtemp(prefix='taichi-')
+ atexit.register(shutil.rmtree, tmp_dir)
+ print(f'[Taichi] preparing sandbox at {tmp_dir}')
+ os.mkdir(os.path.join(tmp_dir, 'runtime/'))
+ return tmp_dir
+
+
+def check_require_version(require_version):
+ '''
+ Check if installed version meets the requirements.
+ Allow to specify ....
+ . is optional. If not match, raise an exception.
+ '''
+ # Extract version number part (i.e. toss any revision / hash parts).
+ version_number_str = require_version
+ for c_idx, c in enumerate(require_version):
+ if not (c.isdigit() or c == "."):
+ version_number_str = require_version[:c_idx]
+ break
+ # Get required version.
+ try:
+ version_number_tuple = tuple(
+ [int(n) for n in version_number_str.split(".")])
+ major = version_number_tuple[0]
+ minor = version_number_tuple[1]
+ patch = 0
+ if len(version_number_tuple) > 2:
+ patch = version_number_tuple[2]
+ except:
+ raise Exception("The require_version should be formatted following PEP 440, " \
+ "and inlucdes major, minor, and patch number, " \
+ "e.g., major.minor.patch.") from None
+ # Get installed version
+ versions = [
+ int(_ti_core.get_version_major()),
+ int(_ti_core.get_version_minor()),
+ int(_ti_core.get_version_patch()),
+ ]
+ # Match installed version and required version.
+ match = major == versions[0] and (
+ minor < versions[1] or minor == versions[1] and patch <= versions[2])
+
+ if not match:
+ raise Exception(
+ f"Taichi version mismatch. Required version >= {major}.{minor}.{patch}, installed version = {_ti_core.get_version_string()}."
+ )
+
+
+def init(arch=None,
+ default_fp=None,
+ default_ip=None,
+ _test_mode=False,
+ enable_fallback=True,
+ require_version=None,
+ **kwargs):
+ """Initializes the Taichi runtime.
+
+ This should always be the entry point of your Taichi program. Most
+ importantly, it sets the backend used throughout the program.
+
+ Args:
+ arch: Backend to use. This is usually :const:`~taichi.lang.cpu` or :const:`~taichi.lang.gpu`.
+ default_fp (Optional[type]): Default floating-point type.
+ default_ip (Optional[type]): Default integral type.
+ require_version (Optional[string]): A version string.
+ **kwargs: Taichi provides highly customizable compilation through
+ ``kwargs``, which allows for fine grained control of Taichi compiler
+ behavior. Below we list some of the most frequently used ones. For a
+ complete list, please check out
+ https://github.com/taichi-dev/taichi/blob/master/taichi/program/compile_config.h.
+
+ * ``cpu_max_num_threads`` (int): Sets the number of threads used by the CPU thread pool.
+ * ``debug`` (bool): Enables the debug mode, under which Taichi does a few more things like boundary checks.
+ * ``print_ir`` (bool): Prints the CHI IR of the Taichi kernels.
+ * ``packed`` (bool): Enables the packed memory layout. See https://docs.taichi.graphics/lang/articles/advanced/layout.
+ """
+ # Check version for users every 7 days if not disabled by users.
+ _version_check.start_version_check_thread()
+
+ cfg = impl.default_cfg()
+ # Check if installed version meets the requirements.
+ if require_version is not None:
+ check_require_version(require_version)
+
+ # Make a deepcopy in case these args reference to items from ti.cfg, which are
+ # actually references. If no copy is made and the args are indeed references,
+ # ti.reset() could override the args to their default values.
+ default_fp = _deepcopy(default_fp)
+ default_ip = _deepcopy(default_ip)
+ kwargs = _deepcopy(kwargs)
+ reset()
+
+ spec_cfg = _SpecialConfig()
+ env_comp = _EnvironmentConfigurator(kwargs, cfg)
+ env_spec = _EnvironmentConfigurator(kwargs, spec_cfg)
+
+ # configure default_fp/ip:
+ # TODO: move these stuff to _SpecialConfig too:
+ env_default_fp = os.environ.get("TI_DEFAULT_FP")
+ if env_default_fp:
+ if default_fp is not None:
+ _ti_core.warn(
+ f'ti.init argument "default_fp" overridden by environment variable TI_DEFAULT_FP={env_default_fp}'
+ )
+ if env_default_fp == '32':
+ default_fp = f32
+ elif env_default_fp == '64':
+ default_fp = f64
+ elif env_default_fp is not None:
+ raise ValueError(
+ f'Invalid TI_DEFAULT_FP={env_default_fp}, should be 32 or 64')
+
+ env_default_ip = os.environ.get("TI_DEFAULT_IP")
+ if env_default_ip:
+ if default_ip is not None:
+ _ti_core.warn(
+ f'ti.init argument "default_ip" overridden by environment variable TI_DEFAULT_IP={env_default_ip}'
+ )
+ if env_default_ip == '32':
+ default_ip = i32
+ elif env_default_ip == '64':
+ default_ip = i64
+ elif env_default_ip is not None:
+ raise ValueError(
+ f'Invalid TI_DEFAULT_IP={env_default_ip}, should be 32 or 64')
+
+ if default_fp is not None:
+ impl.get_runtime().set_default_fp(default_fp)
+ if default_ip is not None:
+ impl.get_runtime().set_default_ip(default_ip)
+
+ # submodule configurations (spec_cfg):
+ env_spec.add('log_level', str)
+ env_spec.add('gdb_trigger')
+ env_spec.add('short_circuit_operators')
+
+ # compiler configurations (ti.cfg):
+ for key in dir(cfg):
+ if key in ['arch', 'default_fp', 'default_ip']:
+ continue
+ _cast = type(getattr(cfg, key))
+ if _cast is bool:
+ _cast = None
+ env_comp.add(key, _cast)
+
+ unexpected_keys = kwargs.keys()
+
+ if len(unexpected_keys):
+ raise KeyError(
+ f'Unrecognized keyword argument(s) for ti.init: {", ".join(unexpected_keys)}'
+ )
+
+ # dispatch configurations that are not in ti.cfg:
+ if not _test_mode:
+ _ti_core.set_core_trigger_gdb_when_crash(spec_cfg.gdb_trigger)
+ impl.get_runtime().short_circuit_operators = \
+ spec_cfg.short_circuit_operators
+ _logging.set_logging_level(spec_cfg.log_level.lower())
+
+ # select arch (backend):
+ env_arch = os.environ.get('TI_ARCH')
+ if env_arch is not None:
+ _logging.info(f'Following TI_ARCH setting up for arch={env_arch}')
+ arch = _ti_core.arch_from_name(env_arch)
+ cfg.arch = adaptive_arch_select(arch, enable_fallback, cfg.use_gles)
+ if cfg.arch == cc:
+ _ti_core.set_tmp_dir(locale_encode(prepare_sandbox()))
+ print(f'[Taichi] Starting on arch={_ti_core.arch_name(cfg.arch)}')
+
+ # user selected visible device
+ visible_device = os.environ.get("TI_VISIBLE_DEVICE")
+ if visible_device and (cfg.arch == vulkan or _ti_core.GGUI_AVAILABLE):
+ _ti_core.set_vulkan_visible_device(visible_device)
+
+ if _test_mode:
+ return spec_cfg
+
+ get_default_kernel_profiler().set_kernel_profiler_mode(cfg.kernel_profiler)
+
+ # create a new program:
+ impl.get_runtime().create_program()
+
+ _logging.trace('Materializing runtime...')
+ impl.get_runtime().prog.materialize_runtime()
+
+ impl._root_fb = _snode.FieldsBuilder()
+
+ if not os.environ.get("TI_DISABLE_SIGNAL_HANDLERS", False):
+ impl.get_runtime()._register_signal_handlers()
+
+ return None
+
+
+def no_activate(*args):
+ for v in args:
+ get_runtime().prog.no_activate(v._snode.ptr)
+
+
+def block_local(*args):
+ """Hints Taichi to cache the fields and to enable the BLS optimization.
+
+ Please visit https://docs.taichi.graphics/lang/articles/advanced/performance
+ for how BLS is used.
+
+ Args:
+ *args (List[Field]): A list of sparse Taichi fields.
+ """
+ if impl.current_cfg().opt_level == 0:
+ _logging.warn("""opt_level = 1 is enforced to enable bls analysis.""")
+ impl.current_cfg().opt_level = 1
+ for a in args:
+ for v in a._get_field_members():
+ get_runtime().prog.current_ast_builder().insert_snode_access_flag(
+ _ti_core.SNodeAccessFlag.block_local, v.ptr)
+
+
+def mesh_local(*args):
+ for a in args:
+ for v in a._get_field_members():
+ get_runtime().prog.current_ast_builder().insert_snode_access_flag(
+ _ti_core.SNodeAccessFlag.mesh_local, v.ptr)
+
+
+def cache_read_only(*args):
+ for a in args:
+ for v in a._get_field_members():
+ get_runtime().prog.current_ast_builder().insert_snode_access_flag(
+ _ti_core.SNodeAccessFlag.read_only, v.ptr)
+
+
+def assume_in_range(val, base, low, high):
+ return _ti_core.expr_assume_in_range(
+ Expr(val).ptr,
+ Expr(base).ptr, low, high)
+
+
+def loop_unique(val, covers=None):
+ if covers is None:
+ covers = []
+ if not isinstance(covers, (list, tuple)):
+ covers = [covers]
+ covers = [x.snode.ptr if isinstance(x, Expr) else x.ptr for x in covers]
+ return _ti_core.expr_loop_unique(Expr(val).ptr, covers)
+
+
+def parallelize(v):
+ get_runtime().prog.current_ast_builder().parallelize(v)
+
+
+serialize = lambda: parallelize(1)
+
+
+def block_dim(dim):
+ """Set the number of threads in a block to `dim`.
+ """
+ get_runtime().prog.current_ast_builder().block_dim(dim)
+
+
+def global_thread_idx():
+ return impl.get_runtime().prog.current_ast_builder(
+ ).insert_thread_idx_expr()
+
+
+def mesh_patch_idx():
+ return impl.get_runtime().prog.current_ast_builder().insert_patch_idx_expr(
+ )
+
+
+def Tape(loss, clear_gradients=True):
+ """Return a context manager of :class:`~taichi.lang.tape.TapeImpl`. The
+ context manager would catching all of the callings of functions that
+ decorated by :func:`~taichi.lang.kernel_impl.kernel` or
+ :func:`~taichi.ad.grad_replaced` under `with` statement, and calculate
+ all the partial gradients of a given loss variable by calling all of the
+ gradient function of the callings caught in reverse order while `with`
+ statement ended.
+
+ See also :func:`~taichi.lang.kernel_impl.kernel` and
+ :func:`~taichi.ad.grad_replaced` for gradient functions.
+
+ Args:
+ loss(:class:`~taichi.lang.expr.Expr`): The loss field, which shape should be ().
+ clear_gradients(Bool): Before `with` body start, clear all gradients or not.
+
+ Returns:
+ :class:`~taichi.lang.tape.TapeImpl`: The context manager.
+
+ Example::
+
+ >>> @ti.kernel
+ >>> def sum(a: ti.float32):
+ >>> for I in ti.grouped(x):
+ >>> y[None] += x[I] ** a
+ >>>
+ >>> with ti.Tape(loss = y):
+ >>> sum(2)
+ """
+ impl.get_runtime().materialize()
+ if len(loss.shape) != 0:
+ raise RuntimeError(
+ 'The loss of `Tape` must be a 0-D field, i.e. scalar')
+ if not loss.snode.ptr.has_grad():
+ raise RuntimeError(
+ 'Gradients of loss are not allocated, please use ti.field(..., needs_grad=True)'
+ ' for all fields that are required by autodiff.')
+ if clear_gradients:
+ clear_all_gradients()
+
+ from taichi._kernels import clear_loss # pylint: disable=C0415
+ clear_loss(loss)
+
+ return impl.get_runtime().get_tape(loss)
+
+
+def clear_all_gradients():
+ """Set the gradients of all fields to zero.
+ """
+ impl.get_runtime().materialize()
+
+ def visit(node):
+ places = []
+ for _i in range(node.ptr.get_num_ch()):
+ ch = node.ptr.get_ch(_i)
+ if not ch.is_place():
+ visit(SNode(ch))
+ else:
+ if not ch.is_primal():
+ places.append(ch.get_expr())
+
+ places = tuple(places)
+ if places:
+ from taichi._kernels import \
+ clear_gradients # pylint: disable=C0415
+ clear_gradients(places)
+
+ for root_fb in _snode.FieldsBuilder._finalized_roots():
+ visit(root_fb)
+
+
+def is_arch_supported(arch, use_gles=False):
+ """Checks whether an arch is supported on the machine.
+
+ Args:
+ arch (taichi_core.Arch): Specified arch.
+ use_gles (bool): If True, check is GLES is available otherwise
+ check if GLSL is available. Only effective when `arch` is `ti.opengl`.
+ Default is `False`.
+
+ Returns:
+ bool: Whether `arch` is supported on the machine.
+ """
+
+ arch_table = {
+ cuda: _ti_core.with_cuda,
+ metal: _ti_core.with_metal,
+ opengl: functools.partial(_ti_core.with_opengl, use_gles),
+ cc: _ti_core.with_cc,
+ vulkan: _ti_core.with_vulkan,
+ dx11: _ti_core.with_dx11,
+ wasm: lambda: True,
+ cpu: lambda: True,
+ }
+ with_arch = arch_table.get(arch, lambda: False)
+ try:
+ return with_arch()
+ except Exception as e:
+ arch = _ti_core.arch_name(arch)
+ _ti_core.warn(
+ f"{e.__class__.__name__}: '{e}' occurred when detecting "
+ f"{arch}, consider adding `TI_ENABLE_{arch.upper()}=0` "
+ f" to environment variables to suppress this warning message.")
+ return False
+
+
+def adaptive_arch_select(arch, enable_fallback, use_gles):
+ if arch is None:
+ return cpu
+ if not isinstance(arch, (list, tuple)):
+ arch = [arch]
+ for a in arch:
+ if is_arch_supported(a, use_gles):
+ return a
+ if not enable_fallback:
+ raise RuntimeError(f'Arch={arch} is not supported')
+ _logging.warn(f'Arch={arch} is not supported, falling back to CPU')
+ return cpu
+
+
+def get_host_arch_list():
+ return [_ti_core.host_arch()]
+
+
+__all__ = [
+ 'i', 'ij', 'ijk', 'ijkl', 'ijl', 'ik', 'ikl', 'il', 'j', 'jk', 'jkl', 'jl',
+ 'k', 'kl', 'l', 'x86_64', 'x64', 'dx11', 'wasm', 'arm64', 'cc', 'cpu',
+ 'cuda', 'gpu', 'metal', 'opengl', 'vulkan', 'extension', 'parallelize',
+ 'block_dim', 'global_thread_idx', 'Tape', 'assume_in_range', 'block_local',
+ 'cache_read_only', 'clear_all_gradients', 'init', 'mesh_local',
+ 'no_activate', 'reset', 'mesh_patch_idx'
+]
diff --git a/python/taichi/lang/ops.py b/python/taichi/lang/ops.py
index e8635e31efc04..b0a67774bf848 100644
--- a/python/taichi/lang/ops.py
+++ b/python/taichi/lang/ops.py
@@ -1,16 +1,12 @@
import builtins
-import ctypes
import functools
import math
import operator as _bt_ops_mod # bt for builtin
import traceback
-from taichi.core.util import ti_core as _ti_core
-from taichi.lang import impl, matrix
+from taichi._lib import core as _ti_core
+from taichi.lang import expr, impl
from taichi.lang.exception import TaichiSyntaxError
-from taichi.lang.expr import Expr, make_expr_group
-from taichi.lang.field import Field
-from taichi.lang.snode import SNode
from taichi.lang.util import cook_dtype, is_taichi_class, taichi_scope
unary_ops = []
@@ -28,27 +24,23 @@ def stack_info():
def is_taichi_expr(a):
- return isinstance(a, Expr)
+ return isinstance(a, expr.Expr)
def wrap_if_not_expr(a):
- _taichi_skip_traceback = 1
- return Expr(a) if not is_taichi_expr(a) else a
+ return expr.Expr(a) if not is_taichi_expr(a) else a
def unary(foo):
@functools.wraps(foo)
def imp_foo(x):
- _taichi_skip_traceback = 2
return foo(x)
@functools.wraps(foo)
def wrapped(a):
- _taichi_skip_traceback = 1
if is_taichi_class(a):
- return a.element_wise_unary(imp_foo)
- else:
- return imp_foo(a)
+ return a._element_wise_unary(imp_foo)
+ return imp_foo(a)
return wrapped
@@ -59,23 +51,19 @@ def wrapped(a):
def binary(foo):
@functools.wraps(foo)
def imp_foo(x, y):
- _taichi_skip_traceback = 2
return foo(x, y)
@functools.wraps(foo)
def rev_foo(x, y):
- _taichi_skip_traceback = 2
return foo(y, x)
@functools.wraps(foo)
def wrapped(a, b):
- _taichi_skip_traceback = 1
if is_taichi_class(a):
- return a.element_wise_binary(imp_foo, b)
- elif is_taichi_class(b):
- return b.element_wise_binary(rev_foo, a)
- else:
- return imp_foo(a, b)
+ return a._element_wise_binary(imp_foo, b)
+ if is_taichi_class(b):
+ return b._element_wise_binary(rev_foo, a)
+ return imp_foo(a, b)
binary_ops.append(wrapped)
return wrapped
@@ -87,30 +75,25 @@ def wrapped(a, b):
def ternary(foo):
@functools.wraps(foo)
def abc_foo(a, b, c):
- _taichi_skip_traceback = 2
return foo(a, b, c)
@functools.wraps(foo)
def bac_foo(b, a, c):
- _taichi_skip_traceback = 2
return foo(a, b, c)
@functools.wraps(foo)
def cab_foo(c, a, b):
- _taichi_skip_traceback = 2
return foo(a, b, c)
@functools.wraps(foo)
def wrapped(a, b, c):
- _taichi_skip_traceback = 1
if is_taichi_class(a):
- return a.element_wise_ternary(abc_foo, b, c)
- elif is_taichi_class(b):
- return b.element_wise_ternary(bac_foo, a, c)
- elif is_taichi_class(c):
- return c.element_wise_ternary(cab_foo, a, b)
- else:
- return abc_foo(a, b, c)
+ return a._element_wise_ternary(abc_foo, b, c)
+ if is_taichi_class(b):
+ return b._element_wise_ternary(bac_foo, a, c)
+ if is_taichi_class(c):
+ return c._element_wise_ternary(cab_foo, a, b)
+ return abc_foo(a, b, c)
ternary_ops.append(wrapped)
return wrapped
@@ -122,15 +105,13 @@ def wrapped(a, b, c):
def writeback_binary(foo):
@functools.wraps(foo)
def imp_foo(x, y):
- _taichi_skip_traceback = 2
return foo(x, wrap_if_not_expr(y))
@functools.wraps(foo)
def wrapped(a, b):
- _taichi_skip_traceback = 1
if is_taichi_class(a):
- return a.element_wise_writeback_binary(imp_foo, b)
- elif is_taichi_class(b):
+ return a._element_wise_writeback_binary(imp_foo, b)
+ if is_taichi_class(b):
raise TaichiSyntaxError(
f'cannot augassign taichi class {type(b)} to scalar expr')
else:
@@ -141,233 +122,441 @@ def wrapped(a, b):
def cast(obj, dtype):
- _taichi_skip_traceback = 1
+ """Copy and cast a scalar or a matrix to a specified data type.
+ Must be called in Taichi scope.
+
+ Args:
+ obj (Union[:mod:`~taichi.types.primitive_types`, :class:`~taichi.Matrix`]): \
+ Input scalar or matrix.
+
+ dtype (:mod:`~taichi.types.primitive_types`): A primitive type defined in :mod:`~taichi.types.primitive_types`.
+
+ Returns:
+ A copy of `obj`, casted to the specified data type `dtype`.
+
+ Example::
+
+ >>> @ti.kernel
+ >>> def test():
+ >>> x = ti.Matrix([0, 1, 2], ti.i32)
+ >>> y = ti.cast(x, ti.f32)
+ >>> print(y)
+ >>>
+ >>> test()
+ [0.0, 1.0, 2.0]
+ """
dtype = cook_dtype(dtype)
if is_taichi_class(obj):
# TODO: unify with element_wise_unary
return obj.cast(dtype)
- else:
- return Expr(_ti_core.value_cast(Expr(obj).ptr, dtype))
+ return expr.Expr(_ti_core.value_cast(expr.Expr(obj).ptr, dtype))
def bit_cast(obj, dtype):
- _taichi_skip_traceback = 1
+ """Copy and cast a scalar to a specified data type with its underlying
+ bits preserved. Must be called in taichi scope.
+
+ This function is equivalent to `reinterpret_cast` in C++.
+
+ Args:
+ obj (:mod:`~taichi.types.primitive_types`): Input scalar.
+
+ dtype (:mod:`~taichi.types.primitive_types`): Target data type, must have \
+ the same precision bits as the input (hence `f32` -> `f64` is not allowed).
+
+ Returns:
+ A copy of `obj`, casted to the specified data type `dtype`.
+
+ Example::
+
+ >>> @ti.kernel
+ >>> def test():
+ >>> x = 3.14
+ >>> y = ti.bit_cast(x, ti.i32)
+ >>> print(y) # 1078523331
+ >>>
+ >>> z = ti.bit_cast(y, ti.f32)
+ >>> print(z) # 3.14
+ """
dtype = cook_dtype(dtype)
if is_taichi_class(obj):
raise ValueError('Cannot apply bit_cast on Taichi classes')
else:
- return Expr(_ti_core.bits_cast(Expr(obj).ptr, dtype))
+ return expr.Expr(_ti_core.bits_cast(expr.Expr(obj).ptr, dtype))
def _unary_operation(taichi_op, python_op, a):
- _taichi_skip_traceback = 1
if is_taichi_expr(a):
- return Expr(taichi_op(a.ptr), tb=stack_info())
- else:
- return python_op(a)
+ return expr.Expr(taichi_op(a.ptr), tb=stack_info())
+ return python_op(a)
def _binary_operation(taichi_op, python_op, a, b):
- _taichi_skip_traceback = 1
if is_taichi_expr(a) or is_taichi_expr(b):
a, b = wrap_if_not_expr(a), wrap_if_not_expr(b)
- return Expr(taichi_op(a.ptr, b.ptr), tb=stack_info())
- else:
- return python_op(a, b)
+ return expr.Expr(taichi_op(a.ptr, b.ptr), tb=stack_info())
+ return python_op(a, b)
def _ternary_operation(taichi_op, python_op, a, b, c):
- _taichi_skip_traceback = 1
if is_taichi_expr(a) or is_taichi_expr(b) or is_taichi_expr(c):
a, b, c = wrap_if_not_expr(a), wrap_if_not_expr(b), wrap_if_not_expr(c)
- return Expr(taichi_op(a.ptr, b.ptr, c.ptr), tb=stack_info())
- else:
- return python_op(a, b, c)
+ return expr.Expr(taichi_op(a.ptr, b.ptr, c.ptr), tb=stack_info())
+ return python_op(a, b, c)
@unary
-def neg(a):
- """The negate function.
+def neg(x):
+ """Numerical negative, element-wise.
Args:
- a (Union[:class:`~taichi.lang.expr.Expr`, :class:`~taichi.lang.matrix.Matrix`]): A number or a matrix.
+ x (Union[:mod:`~taichi.types.primitive_types`, :class:`~taichi.Matrix`]): \
+ Input scalar or matrix.
Returns:
- The negative value of `a`.
+ Matrix or scalar `y`, so that `y = -x`. `y` has the same type as `x`.
+
+ Example::
+ >>> x = ti.Matrix([1, -1])
+ >>> y = ti.neg(a)
+ >>> y
+ [-1, 1]
"""
- return _unary_operation(_ti_core.expr_neg, _bt_ops_mod.neg, a)
+ return _unary_operation(_ti_core.expr_neg, _bt_ops_mod.neg, x)
@unary
-def sin(a):
- """The sine function.
+def sin(x):
+ """Trigonometric sine, element-wise.
Args:
- a (Union[:class:`~taichi.lang.expr.Expr`, :class:`~taichi.lang.matrix.Matrix`]): A number or a matrix.
+ x (Union[:mod:`~taichi.types.primitive_types`, :class:`~taichi.Matrix`]): \
+ Angle, in radians.
Returns:
- Sine of `a`.
+ The sine of each element of `x`.
+
+ Example::
+
+ >>> from math import pi
+ >>> x = ti.Matrix([-pi/2., 0, pi/2.])
+ >>> ti.sin(x)
+ [-1., 0., 1.]
"""
- return _unary_operation(_ti_core.expr_sin, math.sin, a)
+ return _unary_operation(_ti_core.expr_sin, math.sin, x)
@unary
-def cos(a):
- """The cosine function.
+def cos(x):
+ """Trigonometric cosine, element-wise.
Args:
- a (Union[:class:`~taichi.lang.expr.Expr`, :class:`~taichi.lang.matrix.Matrix`]): A number or a matrix.
+ x (Union[:mod:`~taichi.type.primitive_types`, :class:`~taichi.Matrix`]): \
+ Angle, in radians.
Returns:
- Cosine of `a`.
+ The cosine of each element of `x`.
+
+ Example::
+
+ >>> from math import pi
+ >>> x = ti.Matrix([-pi, 0, pi/2.])
+ >>> ti.cos(x)
+ [-1., 1., 0.]
"""
- return _unary_operation(_ti_core.expr_cos, math.cos, a)
+ return _unary_operation(_ti_core.expr_cos, math.cos, x)
@unary
-def asin(a):
- """The inverses function of sine.
+def asin(x):
+ """Trigonometric inverse sine, element-wise.
+
+ The inverse of `sin` so that, if `y = sin(x)`, then `x = asin(y)`.
+
+ For input `x` not in the domain `[-1, 1]`, this function returns `nan` if \
+ it's called in taichi scope, or raises exception if it's called in python scope.
Args:
- a (Union[:class:`~taichi.lang.expr.Expr`, :class:`~taichi.lang.matrix.Matrix`]): A number or a matrix with elements in [-1,1].
+ x (Union[:mod:`~taichi.types.primitive_types`, :class:`~taichi.Matrix`]): \
+ A scalar or a matrix with elements in [-1, 1].
Returns:
- The inverses function of sine of `a`.
+ The inverse sine of each element in `x`, in radians and in the closed \
+ interval `[-pi/2, pi/2]`.
+
+ Example::
+
+ >>> from math import pi
+ >>> ti.asin(ti.Matrix([-1.0, 0.0, 1.0])) * 180 / pi
+ [-90., 0., 90.]
"""
- return _unary_operation(_ti_core.expr_asin, math.asin, a)
+ return _unary_operation(_ti_core.expr_asin, math.asin, x)
@unary
-def acos(a):
- """The inverses function of cosine.
+def acos(x):
+ """Trigonometric inverse cosine, element-wise.
+
+ The inverse of `cos` so that, if `y = cos(x)`, then `x = acos(y)`.
+
+ For input `x` not in the domain `[-1, 1]`, this function returns `nan` if \
+ it's called in taichi scope, or raises exception if it's called in python scope.
Args:
- a (Union[:class:`~taichi.lang.expr.Expr`, :class:`~taichi.lang.matrix.Matrix`]): A number or a matrix with elements in [-1,1].
+ x (Union[:mod:`~taichi.types.primitive_types`, :class:`~taichi.Matrix`]): \
+ A scalar or a matrix with elements in [-1, 1].
Returns:
- The inverses function of cosine of `a`.
+ The inverse cosine of each element in `x`, in radians and in the closed \
+ interval `[0, pi]`. This is a scalar if `x` is a scalar.
+
+ Example::
+
+ >>> from math import pi
+ >>> ti.acos(ti.Matrix([-1.0, 0.0, 1.0])) * 180 / pi
+ [180., 90., 0.]
"""
- return _unary_operation(_ti_core.expr_acos, math.acos, a)
+ return _unary_operation(_ti_core.expr_acos, math.acos, x)
@unary
-def sqrt(a):
- """The square root function.
+def sqrt(x):
+ """Return the non-negative square-root of a scalar or a matrix,
+ element wise. If `x < 0` an exception is raised.
Args:
- a (Union[:class:`~taichi.lang.expr.Expr`, :class:`~taichi.lang.matrix.Matrix`]): A number or a matrix with elements not less than zero.
+ x (Union[:mod:`~taichi.types.primitive_types`, :class:`~taichi.Matrix`]): \
+ The scalar or matrix whose square-roots are required.
Returns:
- `x` such that `x>=0` and `x^2=a`.
+ The square-root `y` so that `y >= 0` and `y^2 = x`. `y` has the same type as `x`.
+
+ Example::
+
+ >>> x = ti.Matrix([1., 4., 9.])
+ >>> y = ti.sqrt(x)
+ >>> y
+ [1.0, 2.0, 3.0]
"""
- return _unary_operation(_ti_core.expr_sqrt, math.sqrt, a)
+ return _unary_operation(_ti_core.expr_sqrt, math.sqrt, x)
@unary
-def rsqrt(a):
+def rsqrt(x):
"""The reciprocal of the square root function.
Args:
- a (Union[:class:`~taichi.lang.expr.Expr`, :class:`~taichi.lang.matrix.Matrix`]): A number or a matrix.
+ x (Union[:mod:`~taichi.types.primitive_types`, :class:`~taichi.Matrix`]): \
+ A scalar or a matrix.
Returns:
- The reciprocal of `sqrt(a)`.
+ The reciprocal of `sqrt(x)`.
"""
- def _rsqrt(a):
- return 1 / math.sqrt(a)
+ def _rsqrt(x):
+ return 1 / math.sqrt(x)
- return _unary_operation(_ti_core.expr_rsqrt, _rsqrt, a)
+ return _unary_operation(_ti_core.expr_rsqrt, _rsqrt, x)
@unary
-def floor(a):
- """The floor function.
+def round(x): # pylint: disable=redefined-builtin
+ """Round to the nearest integer, element-wise.
Args:
- a (Union[:class:`~taichi.lang.expr.Expr`, :class:`~taichi.lang.matrix.Matrix`]): A number or a matrix.
+ x (Union[:mod:`~taichi.types.primitive_types`, :class:`~taichi.Matrix`]): \
+ A scalar or a matrix.
Returns:
- The greatest integer less than or equal to `a`.
+ The nearest integer of `x`.
+
+ Example::
+
+ >>> @ti.kernel
+ >>> def test():
+ >>> x = ti.Vector([-1.5, 1.2, 2.7])
+ >>> print(ti.round(x))
+ [-2., 1., 3.]
"""
- return _unary_operation(_ti_core.expr_floor, math.floor, a)
+ return _unary_operation(_ti_core.expr_round, builtins.round, x)
@unary
-def ceil(a):
- """The ceil function.
+def floor(x):
+ """Return the floor of the input, element-wise.
+
+ The floor of the scalar `x` is the largest integer `k`, such that `k <= x`.
Args:
- a (Union[:class:`~taichi.lang.expr.Expr`, :class:`~taichi.lang.matrix.Matrix`]): A number or a matrix.
+ x (Union[:mod:`~taichi.types.primitive_types`, :class:`~taichi.Matrix`]): \
+ Input scalar or matrix.
Returns:
- The least integer greater than or equal to `a`.
+ The floor of each element in `x`, with float type.
+
+ Example::
+
+ >>> @ti.kernel
+ >>> def test():
+ >>> x = ti.Matrix([3.14, -1.5])
+ >>> y = ti.floor(x)
+ >>> print(y) # [3.0, -2.0]
"""
- return _unary_operation(_ti_core.expr_ceil, math.ceil, a)
+ return _unary_operation(_ti_core.expr_floor, math.floor, x)
@unary
-def tan(a):
- """The tangent function.
+def ceil(x):
+ """Return the ceiling of the input, element-wise.
+
+ The ceil of the scalar `x` is the smallest integer `k`, such that `k >= x`.
Args:
- a (Union[:class:`~taichi.lang.expr.Expr`, :class:`~taichi.lang.matrix.Matrix`]): A number or a matrix.
+ x (Union[:mod:`~taichi.types.primitive_types`, :class:`~taichi.Matrix`]): \
+ Input scalar or matrix.
Returns:
- Tangent of `a`.
+ The ceiling of each element in `x`, with float dtype.
+
+ Example::
+
+ >>> @ti.kernel
+ >>> def test():
+ >>> x = ti.Matrix([3.14, -1.5])
+ >>> y = ti.ceil(x)
+ >>> print(y) # [4.0, -1.0]
"""
- return _unary_operation(_ti_core.expr_tan, math.tan, a)
+ return _unary_operation(_ti_core.expr_ceil, math.ceil, x)
@unary
-def tanh(a):
- """The hyperbolic tangent function.
+def tan(x):
+ """Trigonometric tangent function, element-wise.
+
+ Equivalent to `ti.sin(x)/ti.cos(x)` element-wise.
Args:
- a (Union[:class:`~taichi.lang.expr.Expr`, :class:`~taichi.lang.matrix.Matrix`]): A number or a matrix.
+ x (Union[:mod:`~taichi.types.primitive_types`, :class:`~taichi.Matrix`]): \
+ Input scalar or matrix.
Returns:
- `(e**x - e**(-x)) / (e**x + e**(-x))`.
+ The tangent values of `x`.
+
+ Example::
+
+ >>> from math import pi
+ >>> @ti.kernel
+ >>> def test():
+ >>> x = ti.Matrix([-pi, pi/2, pi])
+ >>> y = ti.tan(x)
+ >>> print(y)
+ >>>
+ >>> test()
+ [-0.0, -22877334.0, 0.0]
"""
- return _unary_operation(_ti_core.expr_tanh, math.tanh, a)
+ return _unary_operation(_ti_core.expr_tan, math.tan, x)
@unary
-def exp(a):
- """The exp function.
+def tanh(x):
+ """Compute the hyperbolic tangent of `x`, element-wise.
Args:
- a (Union[:class:`~taichi.lang.expr.Expr`, :class:`~taichi.lang.matrix.Matrix`]): A number or a matrix.
+ x (Union[:mod:`~taichi.types.primitive_types`, :class:`~taichi.Matrix`]): \
+ Input scalar or matrix.
+
+ Returns:
+ The corresponding hyperbolic tangent values.
+
+ Example::
+
+ >>> @ti.kernel
+ >>> def test():
+ >>> x = ti.Matrix([-1.0, 0.0, 1.0])
+ >>> y = ti.tanh(x)
+ >>> print(y)
+ >>>
+ >>> test()
+ [-0.761594, 0.000000, 0.761594]
+ """
+ return _unary_operation(_ti_core.expr_tanh, math.tanh, x)
+
+
+@unary
+def exp(x):
+ """Compute the exponential of all elements in `x`, element-wise.
+
+ Args:
+ x (Union[:mod:`~taichi.types.primitive_types`, :class:`~taichi.Matrix`]): \
+ Input scalar or matrix.
Returns:
- `e` to the `a`.
+ Element-wise exponential of `x`.
+
+ Example::
+
+ >>> @ti.kernel
+ >>> def test():
+ >>> x = ti.Matrix([-1.0, 0.0, 1.0])
+ >>> y = ti.exp(x)
+ >>> print(y)
+ >>>
+ >>> test()
+ [0.367879, 1.000000, 2.718282]
"""
- return _unary_operation(_ti_core.expr_exp, math.exp, a)
+ return _unary_operation(_ti_core.expr_exp, math.exp, x)
@unary
-def log(a):
- """The natural logarithm function.
+def log(x):
+ """Compute the natural logarithm, element-wise.
+
+ The natural logarithm `log` is the inverse of the exponential function,
+ so that `log(exp(x)) = x`. The natural logarithm is logarithm in base `e`.
Args:
- a (Union[:class:`~taichi.lang.expr.Expr`, :class:`~taichi.lang.matrix.Matrix`]): A number or a matrix with elements greater than zero.
+ x (Union[:mod:`~taichi.types.primitive_types`, :class:`~taichi.Matrix`]): \
+ Input scalar or matrix.
Returns:
- The natural logarithm of `a`.
+ The natural logarithm of `x`, element-wise.
+
+ Example::
+
+ >>> @ti.kernel
+ >>> def test():
+ >>> x = ti.Vector([-1.0, 0.0, 1.0])
+ >>> y = ti.log(x)
+ >>> print(y)
+ >>>
+ >>> test()
+ [-nan, -inf, 0.000000]
"""
- return _unary_operation(_ti_core.expr_log, math.log, a)
+ return _unary_operation(_ti_core.expr_log, math.log, x)
@unary
-def abs(a):
- """The absolute value function.
+def abs(x): # pylint: disable=W0622
+ """Compute the absolute value :math:`|x|` of `x`, element-wise.
Args:
- a (Union[:class:`~taichi.lang.expr.Expr`, :class:`~taichi.lang.matrix.Matrix`]): A number or a matrix.
+ x (Union[:mod:`~taichi.types.primitive_types`, :class:`~taichi.Matrix`]): \
+ Input scalar or matrix.
Returns:
- The absolute value of `a`.
+ The absolute value of each element in `x`.
+
+ Example::
+
+ >>> @ti.kernel
+ >>> def test():
+ >>> x = ti.Vector([-1.0, 0.0, 1.0])
+ >>> y = ti.abs(x)
+ >>> print(y)
+ >>>
+ >>> test()
+ [1.0, 0.0, 1.0]
"""
- return _unary_operation(_ti_core.expr_abs, builtins.abs, a)
+ return _unary_operation(_ti_core.expr_abs, builtins.abs, x)
@unary
@@ -397,16 +586,41 @@ def logical_not(a):
def random(dtype=float):
- """The random function.
+ """Return a single random float/integer according to the specified data type.
+ Must be called in taichi scope.
+
+ If the required `dtype` is float type, this function returns a random number
+ sampled from the uniform distribution in the half-open interval [0, 1).
+
+ For integer types this function returns a random integer in the
+ half-open interval [0, 2^32) if a 32-bit integer is required,
+ or a random integer in the half-open interval [0, 2^64) if a
+ 64-bit integer is required.
Args:
- dtype (DataType): Type of the random variable.
+ dtype (:mod:`~taichi.types.primitive_types`): Type of the required random value.
Returns:
- A random variable whose type is `dtype`.
+ A random value with type `dtype`.
+
+ Example::
+
+ >>> @ti.kernel
+ >>> def test():
+ >>> x = ti.random(float)
+ >>> print(x) # 0.090257
+ >>>
+ >>> y = ti.random(ti.f64)
+ >>> print(y) # 0.716101627301
+ >>>
+ >>> i = ti.random(ti.i32)
+ >>> print(i) # -963722261
+ >>>
+ >>> j = ti.random(ti.i64)
+ >>> print(j) # 73412986184350777
"""
dtype = cook_dtype(dtype)
- x = Expr(_ti_core.make_rand_expr(dtype))
+ x = expr.Expr(_ti_core.make_rand_expr(dtype))
return impl.expr_init(x)
@@ -456,37 +670,74 @@ def mul(a, b):
@binary
-def mod(a, b):
- """The remainder function.
+def mod(x1, x2):
+ """Returns the element-wise remainder of division.
+
+ This is equivalent to the Python modulus operator `x1 % x2` and
+ has the same sign as the divisor x2.
Args:
- a (Union[:class:`~taichi.lang.expr.Expr`, :class:`~taichi.lang.matrix.Matrix`]): A number or a matrix.
- b (Union[:class:`~taichi.lang.expr.Expr`, :class:`~taichi.lang.matrix.Matrix`]): A number or a matrix with elements not equal to zero.
+ x1 (Union[:mod:`~taichi.types.primitive_types`, :class:`~taichi.Matrix`]): \
+ Dividend scalar or matrix.
+
+ x2 (Union[:mod:`~taichi.types.primitive_types`, :class:`~taichi.Matrix`]): \
+ Divisor scalar or matrix. When both `x1` and `x2` are matrices they must have the same shape.
Returns:
- The remainder of `a` divided by `b`.
+ The element-wise remainder of the quotient `floordiv(x1, x2)`. This is a scalar \
+ if both `x1` and `x2` are scalars.
+
+ Example::
+
+ >>> @ti.kernel
+ >>> def test():
+ >>> x = ti.Matrix([3.0, 4.0, 5.0])
+ >>> y = 3
+ >>> z = ti.mod(y, x)
+ >>> print(z)
+ >>>
+ >>> test()
+ [1.0, 0.0, 4.0]
"""
def expr_python_mod(a, b):
# a % b = a - (a // b) * b
- quotient = Expr(_ti_core.expr_floordiv(a, b))
- multiply = Expr(_ti_core.expr_mul(b, quotient.ptr))
+ quotient = expr.Expr(_ti_core.expr_floordiv(a, b))
+ multiply = expr.Expr(_ti_core.expr_mul(b, quotient.ptr))
return _ti_core.expr_sub(a, multiply.ptr)
- return _binary_operation(expr_python_mod, _bt_ops_mod.mod, a, b)
+ return _binary_operation(expr_python_mod, _bt_ops_mod.mod, x1, x2)
@binary
-def pow(a, b):
- """The power function.
+def pow(x, a): # pylint: disable=W0622
+ """First array elements raised to powers from second array :math:`x^a`, element-wise.
+
+ Negative values raised to a non-integral value will return `nan`.
+ A zero value raised to a negative value will return `inf`.
Args:
- a (Union[:class:`~taichi.lang.expr.Expr`, :class:`~taichi.lang.matrix.Matrix`]): A number or a matrix.
- b (Union[:class:`~taichi.lang.expr.Expr`, :class:`~taichi.lang.matrix.Matrix`]): A number or a matrix.
+ x (Union[:mod:`~taichi.types.primitive_types`, :class:`~taichi.Matrix`]): \
+ The bases.
+ a (Union[:mod:`~taichi.types.primitive_types`, :class:`~taichi.Matrix`]): \
+ The exponents.
Returns:
- `a` to the `b`.
+ The bases in `x1` raised to the exponents in `x2`. This is a scalar if both \
+ `x1` and `x2` are scalars.
+
+ Example::
+
+ >>> @ti.kernel
+ >>> def test():
+ >>> x = ti.Matrix([-2.0, 0.0, 2.0])
+ >>> y = -2.2
+ >>> z = ti.pow(x, y)
+ >>> print(z)
+ >>>
+ >>> test()
+ [-nan, inf, 0.217638]
"""
- return _binary_operation(_ti_core.expr_pow, _bt_ops_mod.pow, a, b)
+ return _binary_operation(_ti_core.expr_pow, _bt_ops_mod.pow, x, a)
@binary
@@ -519,7 +770,7 @@ def truediv(a, b):
@binary
-def max(a, b):
+def max_impl(a, b):
"""The maxnimum function.
Args:
@@ -533,7 +784,7 @@ def max(a, b):
@binary
-def min(a, b):
+def min_impl(a, b):
"""The minimum function.
Args:
@@ -547,54 +798,89 @@ def min(a, b):
@binary
-def atan2(a, b):
- """The inverses of the tangent function.
+def atan2(x1, x2):
+ """Element-wise arc tangent of `x1/x2`.
Args:
- a (Union[:class:`~taichi.lang.expr.Expr`, :class:`~taichi.lang.matrix.Matrix`]): A number or a matrix.
- b (Union[:class:`~taichi.lang.expr.Expr`, :class:`~taichi.lang.matrix.Matrix`]): A number or a matrix with elements not equal to zero.
+ x1 (Union[:mod:`~taichi.types.primitive_types`, :class:`~taichi.Matrix`]): \
+ y-coordinates.
+ x2 (Union[:mod:`~taichi.types.primitive_types`, :class:`~taichi.Matrix`]): \
+ x-coordinates.
Returns:
- The inverses function of tangent of `b/a`.
+ Angles in radians, in the range `[-pi, pi]`.
+ This is a scalar if both `x1` and `x2` are scalars.
+
+ Example::
+
+ >>> from math import pi
+ >>> @ti.kernel
+ >>> def test():
+ >>> x = ti.Matrix([-1.0, 1.0, -1.0, 1.0])
+ >>> y = ti.Matrix([-1.0, -1.0, 1.0, 1.0])
+ >>> z = ti.atan2(y, x) * 180 / pi
+ >>> print(z)
+ >>>
+ >>> test()
+ [-135.0, -45.0, 135.0, 45.0]
"""
- return _binary_operation(_ti_core.expr_atan2, math.atan2, a, b)
+ return _binary_operation(_ti_core.expr_atan2, math.atan2, x1, x2)
@binary
-def raw_div(a, b):
- """Raw_div function.
+def raw_div(x1, x2):
+ """Return `x1 // x2` if both `x1`, `x2` are integers, otherwise return `x1/x2`.
Args:
- a (Union[:class:`~taichi.lang.expr.Expr`, :class:`~taichi.lang.matrix.Matrix`]): A number or a matrix.
- b (Union[:class:`~taichi.lang.expr.Expr`, :class:`~taichi.lang.matrix.Matrix`]): A number or a matrix with elements not equal to zero.
+ x1 (Union[:mod:`~taichi.types.primitive_types`, :class:`~taichi.Matrix`]): Dividend.
+ x2 (Union[:mod:`~taichi.types.primitive_types`, :class:`~taichi.Matrix`]): Divisor.
Returns:
- If `a` is a `int` and `b` is a `int`, then return `a//b`. Else return `a/b`.
+ Return `x1 // x2` if both `x1`, `x2` are integers, otherwise return `x1/x2`.
+
+ Example::
+
+ >>> @ti.kernel
+ >>> def main():
+ >>> x = 5
+ >>> y = 3
+ >>> print(raw_div(x, y)) # 1
+ >>> z = 4.0
+ >>> print(raw_div(x, z)) # 1.25
"""
def c_div(a, b):
if isinstance(a, int) and isinstance(b, int):
return a // b
- else:
- return a / b
+ return a / b
- return _binary_operation(_ti_core.expr_div, c_div, a, b)
+ return _binary_operation(_ti_core.expr_div, c_div, x1, x2)
@binary
-def raw_mod(a, b):
- """Raw_mod function. Both `a` and `b` can be `float`.
+def raw_mod(x1, x2):
+ """Return the remainder of `x1/x2`, element-wise.
+ This is the C-style `mod` function.
Args:
- a (Union[:class:`~taichi.lang.expr.Expr`, :class:`~taichi.lang.matrix.Matrix`]): A number or a matrix.
- b (Union[:class:`~taichi.lang.expr.Expr`, :class:`~taichi.lang.matrix.Matrix`]): A number or a matrix with elements not equal to zero.
+ x1 (Union[:mod:`~taichi.types.primitive_types`, :class:`~taichi.Matrix`]): \
+ The dividend.
+ x2 (Union[:mod:`~taichi.types.primitive_types`, :class:`~taichi.Matrix`]): \
+ The divisor.
Returns:
- The remainder of `a` divided by `b`.
+ The remainder of `x1` divided by `x2`.
+
+ Example::
+
+ >>> @ti.kernel
+ >>> def main():
+ >>> print(ti.mod(-4, 3)) # 2
+ >>> print(ti.raw_mod(-4, 3)) # -1
"""
- def c_mod(a, b):
- return a - b * int(float(a) / b)
+ def c_mod(x, y):
+ return x - y * int(float(x) / y)
- return _binary_operation(_ti_core.expr_mod, c_mod, a, b)
+ return _binary_operation(_ti_core.expr_mod, c_mod, x1, x2)
@binary
@@ -770,18 +1056,31 @@ def bit_sar(a, b):
@taichi_scope
@binary
-def bit_shr(a, b):
- """Compute bitwise shift right (in taichi scope)
+def bit_shr(x1, x2):
+ """Elements in `x1` shifted to the right by number of bits in `x2`.
+ Both `x1`, `x2` must have integer type.
Args:
- a (Union[:class:`~taichi.lang.expr.Expr`, :class:`~taichi.lang.matrix.Matrix`]): value LHS
- b (Union[:class:`~taichi.lang.expr.Expr`, :class:`~taichi.lang.matrix.Matrix`]): value RHS
+ x1 (Union[:mod:`~taichi.types.primitive_types`, :class:`~taichi.Matrix`]): \
+ Input data.
+ x2 (Union[:mod:`~taichi.types.primitive_types`, :class:`~taichi.Matrix`]): \
+ Number of bits to remove at the right of `x1`.
Returns:
- Union[:class:`~taichi.lang.expr.Expr`, int]: LHS >> RHS
-
+ Return `x1` with bits shifted `x2` times to the right.
+ This is a scalar if both `x1` and `x2` are scalars.
+
+ Example::
+ >>> @ti.kernel
+ >>> def main():
+ >>> x = ti.Matrix([7, 8])
+ >>> y = ti.Matrix([1, 2])
+ >>> print(ti.bit_shr(x, y))
+ >>>
+ >>> main()
+ [3, 2]
"""
- return _binary_operation(_ti_core.expr_bit_shr, _bt_ops_mod.rshift, a, b)
+ return _binary_operation(_ti_core.expr_bit_shr, _bt_ops_mod.rshift, x1, x2)
# We don't have logic_and/or instructions yet:
@@ -790,177 +1089,341 @@ def bit_shr(a, b):
@ternary
-def select(cond, a, b):
+def select(cond, x1, x2):
+ """Return an array drawn from elements in `x1` or `x2`,
+ depending on the conditions in `cond`.
+
+ Args:
+ cond (Union[:mod:`~taichi.types.primitive_types`, :class:`~taichi.Matrix`]): \
+ The array of conditions.
+ x1, x2 (Union[:mod:`~taichi.types.primitive_types`, :class:`~taichi.Matrix`]): \
+ The arrays where the output elements are taken from.
+
+ Returns:
+ The output at position `k` is the k-th element of `x1` if the k-th element
+ in `cond` is `True`, otherwise it's the k-th element of `x2`.
+
+ Example::
+
+ >>> @ti.kernel
+ >>> def main():
+ >>> cond = ti.Matrix([0, 1, 0, 1])
+ >>> x = ti.Matrix([1, 2, 3, 4])
+ >>> y = ti.Matrix([-1, -2, -3, -4])
+ >>> print(ti.select(cond, x, y))
+ >>>
+ >>> main()
+ [-1, 2, -3, 4]
+ """
# TODO: systematically resolve `-1 = True` problem by introducing u1:
cond = logical_not(logical_not(cond))
- def py_select(cond, a, b):
- return a * cond + b * (1 - cond)
+ def py_select(cond, x1, x2):
+ return x1 * cond + x2 * (1 - cond)
- return _ternary_operation(_ti_core.expr_select, py_select, cond, a, b)
+ return _ternary_operation(_ti_core.expr_select, py_select, cond, x1, x2)
@writeback_binary
-def atomic_add(a, b):
- return impl.expr_init(
- Expr(_ti_core.expr_atomic_add(a.ptr, b.ptr), tb=stack_info()))
+def atomic_add(x, y):
+ """Atomically compute `x + y`, store the result in `x`,
+ and return the old value of `x`.
+ `x` must be a writable target, constant expressions or scalars
+ are not allowed.
-@writeback_binary
-def atomic_sub(a, b):
- return impl.expr_init(
- Expr(_ti_core.expr_atomic_sub(a.ptr, b.ptr), tb=stack_info()))
-
+ Args:
+ x, y (Union[:mod:`~taichi.types.primitive_types`, :class:`~taichi.Matrix`]): \
+ The input.
-@writeback_binary
-def atomic_min(a, b):
+ Returns:
+ The old value of `x`.
+
+ Example::
+
+ >>> @ti.kernel
+ >>> def test():
+ >>> x = ti.Vector([0, 0, 0])
+ >>> y = ti.Vector([1, 2, 3])
+ >>> z = ti.atomic_add(x, y)
+ >>> print(x) # [1, 2, 3] the new value of x
+ >>> print(z) # [0, 0, 0], the old value of x
+ >>>
+ >>> ti.atomic_add(1, x) # will raise TaichiSyntaxError
+ """
return impl.expr_init(
- Expr(_ti_core.expr_atomic_min(a.ptr, b.ptr), tb=stack_info()))
+ expr.Expr(_ti_core.expr_atomic_add(x.ptr, y.ptr), tb=stack_info()))
@writeback_binary
-def atomic_max(a, b):
- return impl.expr_init(
- Expr(_ti_core.expr_atomic_max(a.ptr, b.ptr), tb=stack_info()))
+def atomic_sub(x, y):
+ """Atomically subtract `x` by `y`, store the result in `x`,
+ and return the old value of `x`.
+ `x` must be a writable target, constant expressions or scalars
+ are not allowed.
-@writeback_binary
-def atomic_and(a, b):
+ Args:
+ x, y (Union[:mod:`~taichi.types.primitive_types`, :class:`~taichi.Matrix`]): \
+ The input.
+
+ Returns:
+ The old value of `x`.
+
+ Example::
+
+ >>> @ti.kernel
+ >>> def test():
+ >>> x = ti.Vector([0, 0, 0])
+ >>> y = ti.Vector([1, 2, 3])
+ >>> z = ti.atomic_sub(x, y)
+ >>> print(x) # [-1, -2, -3] the new value of x
+ >>> print(z) # [0, 0, 0], the old value of x
+ >>>
+ >>> ti.atomic_sub(1, x) # will raise TaichiSyntaxError
+ """
return impl.expr_init(
- Expr(_ti_core.expr_atomic_bit_and(a.ptr, b.ptr), tb=stack_info()))
+ expr.Expr(_ti_core.expr_atomic_sub(x.ptr, y.ptr), tb=stack_info()))
@writeback_binary
-def atomic_or(a, b):
- return impl.expr_init(
- Expr(_ti_core.expr_atomic_bit_or(a.ptr, b.ptr), tb=stack_info()))
+def atomic_min(x, y):
+ """Atomically compute the minimum of `x` and `y`, element-wise.
+ Store the result in `x`, and return the old value of `x`.
+ `x` must be a writable target, constant expressions or scalars
+ are not allowed.
-@writeback_binary
-def atomic_xor(a, b):
+ Args:
+ x, y (Union[:mod:`~taichi.types.primitive_types`, :class:`~taichi.Matrix`]): \
+ The input.
+
+ Returns:
+ The old value of `x`.
+
+ Example::
+
+ >>> @ti.kernel
+ >>> def test():
+ >>> x = 2
+ >>> y = 1
+ >>> z = ti.atomic_min(x, y)
+ >>> print(x) # 1 the new value of x
+ >>> print(z) # 2, the old value of x
+ >>>
+ >>> ti.atomic_min(1, x) # will raise TaichiSyntaxError
+ """
return impl.expr_init(
- Expr(_ti_core.expr_atomic_bit_xor(a.ptr, b.ptr), tb=stack_info()))
+ expr.Expr(_ti_core.expr_atomic_min(x.ptr, y.ptr), tb=stack_info()))
@writeback_binary
-def assign(a, b):
- _ti_core.expr_assign(a.ptr, b.ptr, stack_info())
- return a
+def atomic_max(x, y):
+ """Atomically compute the maximum of `x` and `y`, element-wise.
+ Store the result in `x`, and return the old value of `x`.
+ `x` must be a writable target, constant expressions or scalars
+ are not allowed.
-def ti_max(*args):
- num_args = len(args)
- assert num_args >= 1
- if num_args == 1:
- return args[0]
- elif num_args == 2:
- return max(args[0], args[1])
- else:
- return max(args[0], ti_max(*args[1:]))
+ Args:
+ x, y (Union[:mod:`~taichi.types.primitive_types`, :class:`~taichi.Matrix`]): \
+ The input.
+ Returns:
+ The old value of `x`.
+
+ Example::
+
+ >>> @ti.kernel
+ >>> def test():
+ >>> x = 1
+ >>> y = 2
+ >>> z = ti.atomic_max(x, y)
+ >>> print(x) # 2 the new value of x
+ >>> print(z) # 1, the old value of x
+ >>>
+ >>> ti.atomic_max(1, x) # will raise TaichiSyntaxError
+ """
+ return impl.expr_init(
+ expr.Expr(_ti_core.expr_atomic_max(x.ptr, y.ptr), tb=stack_info()))
-def ti_min(*args):
- num_args = len(args)
- assert num_args >= 1
- if num_args == 1:
- return args[0]
- elif num_args == 2:
- return min(args[0], args[1])
- else:
- return min(args[0], ti_min(*args[1:]))
+@writeback_binary
+def atomic_and(x, y):
+ """Atomically compute the bit-wise AND of `x` and `y`, element-wise.
+ Store the result in `x`, and return the old value of `x`.
-def ti_any(a):
- return a.any()
+ `x` must be a writable target, constant expressions or scalars
+ are not allowed.
+ Args:
+ x, y (Union[:mod:`~taichi.types.primitive_types`, :class:`~taichi.Matrix`]): \
+ The input. When both are matrices they must have the same shape.
-def ti_all(a):
- return a.all()
+ Returns:
+ The old value of `x`.
+
+ Example::
+
+ >>> @ti.kernel
+ >>> def test():
+ >>> x = ti.Vector([-1, 0, 1])
+ >>> y = ti.Vector([1, 2, 3])
+ >>> z = ti.atomic_and(x, y)
+ >>> print(x) # [1, 0, 1] the new value of x
+ >>> print(z) # [-1, 0, 1], the old value of x
+ >>>
+ >>> ti.atomic_and(1, x) # will raise TaichiSyntaxError
+ """
+ return impl.expr_init(
+ expr.Expr(_ti_core.expr_atomic_bit_and(x.ptr, y.ptr), tb=stack_info()))
-def append(l, indices, val):
- a = impl.expr_init(
- _ti_core.insert_append(l.snode.ptr, make_expr_group(indices),
- Expr(val).ptr))
- return a
+@writeback_binary
+def atomic_or(x, y):
+ """Atomically compute the bit-wise OR of `x` and `y`, element-wise.
+ Store the result in `x`, and return the old value of `x`.
+ `x` must be a writable target, constant expressions or scalars
+ are not allowed.
-def external_func_call(func, args=[], outputs=[]):
- func_addr = ctypes.cast(func, ctypes.c_void_p).value
- _ti_core.insert_external_func_call(func_addr, '', make_expr_group(args),
- make_expr_group(outputs))
+ Args:
+ x, y (Union[:mod:`~taichi.types.primitive_types`, :class:`~taichi.Matrix`]): \
+ The input. When both are matrices they must have the same shape.
+ Returns:
+ The old value of `x`.
+
+ Example::
+
+ >>> @ti.kernel
+ >>> def test():
+ >>> x = ti.Vector([-1, 0, 1])
+ >>> y = ti.Vector([1, 2, 3])
+ >>> z = ti.atomic_or(x, y)
+ >>> print(x) # [-1, 2, 3] the new value of x
+ >>> print(z) # [-1, 0, 1], the old value of x
+ >>>
+ >>> ti.atomic_or(1, x) # will raise TaichiSyntaxError
+ """
+ return impl.expr_init(
+ expr.Expr(_ti_core.expr_atomic_bit_or(x.ptr, y.ptr), tb=stack_info()))
-def asm(source, inputs=[], outputs=[]):
- _ti_core.insert_external_func_call(0, source, make_expr_group(inputs),
- make_expr_group(outputs))
+@writeback_binary
+def atomic_xor(x, y):
+ """Atomically compute the bit-wise XOR of `x` and `y`, element-wise.
+ Store the result in `x`, and return the old value of `x`.
-def is_active(l, indices):
- return Expr(
- _ti_core.insert_is_active(l.snode.ptr, make_expr_group(indices)))
+ `x` must be a writable target, constant expressions or scalars
+ are not allowed.
+ Args:
+ x, y (Union[:mod:`~taichi.types.primitive_types`, :class:`~taichi.Matrix`]): \
+ The input. When both are matrices they must have the same shape.
-def activate(l, indices):
- _ti_core.insert_activate(l.snode.ptr, make_expr_group(indices))
+ Returns:
+ The old value of `x`.
+
+ Example::
+
+ >>> @ti.kernel
+ >>> def test():
+ >>> x = ti.Vector([-1, 0, 1])
+ >>> y = ti.Vector([1, 2, 3])
+ >>> z = ti.atomic_xor(x, y)
+ >>> print(x) # [-2, 2, 2] the new value of x
+ >>> print(z) # [-1, 0, 1], the old value of x
+ >>>
+ >>> ti.atomic_xor(1, x) # will raise TaichiSyntaxError
+ """
+ return impl.expr_init(
+ expr.Expr(_ti_core.expr_atomic_bit_xor(x.ptr, y.ptr), tb=stack_info()))
-def deactivate(l, indices):
- _ti_core.insert_deactivate(l.snode.ptr, make_expr_group(indices))
+@writeback_binary
+def assign(a, b):
+ impl.get_runtime().prog.current_ast_builder().expr_assign(
+ a.ptr, b.ptr, stack_info())
+ return a
-def length(l, indices):
- return Expr(_ti_core.insert_len(l.snode.ptr, make_expr_group(indices)))
+def max(*args): # pylint: disable=W0622
+ """Compute the maximum of the arguments, element-wise.
+ This function takes no effect on a single argument, even it's array-like.
+ When there are both scalar and matrix arguments in `args`, the matrices
+ must have the same shape, and scalars will be broadcasted to the same shape as the matrix.
-def rescale_index(a, b, I):
- """Rescales the index 'I' of field (or SNode) 'a' to match the shape of SNode 'b'
+ Args:
+ args: (List[:mod:`~taichi.types.primitive_types`, :class:`~taichi.Matrix`]): \
+ The input.
- Parameters
- ----------
- a: ti.field(), ti.Vector.field, ti.Matrix.field()
- input taichi field or snode
- b: ti.field(), ti.Vector.field, ti.Matrix.field()
- output taichi field or snode
- I: ti.Vector()
- grouped loop index
+ Returns:
+ Maximum of the inputs.
- Returns
- -------
- Ib: ti.Vector()
- rescaled grouped loop index
+ Example::
+ >>> @ti.kernel
+ >>> def foo():
+ >>> x = ti.Vector([0, 1, 2])
+ >>> y = ti.Vector([3, 4, 5])
+ >>> z = ti.max(x, y, 4)
+ >>> print(z) # [4, 4, 5]
"""
- assert isinstance(
- a, (Field, SNode)), "The first argument must be a field or an SNode"
- assert isinstance(
- b, (Field, SNode)), "The second argument must be a field or an SNode"
- if isinstance(I, list):
- I = matrix.Vector(I)
- else:
- assert isinstance(
- I, matrix.Matrix
- ), f"The third argument must be an index (list or ti.Vector)"
- Ib = I.copy()
- for n in range(min(I.n, min(len(a.shape), len(b.shape)))):
- if a.shape[n] > b.shape[n]:
- Ib.entries[n] = I.entries[n] // (a.shape[n] // b.shape[n])
- if a.shape[n] < b.shape[n]:
- Ib.entries[n] = I.entries[n] * (b.shape[n] // a.shape[n])
- return Ib
+ num_args = len(args)
+ assert num_args >= 1
+ if num_args == 1:
+ return args[0]
+ if num_args == 2:
+ return max_impl(args[0], args[1])
+ return max_impl(args[0], max(*args[1:]))
-def get_addr(f, indices):
- """Query the memory address (on CUDA/x64) of field `f` at index `indices`.
+def min(*args): # pylint: disable=W0622
+ """Compute the minimum of the arguments, element-wise.
- Currently, this function can only be called inside a taichi kernel.
+ This function takes no effect on a single argument, even it's array-like.
+ When there are both scalar and matrix arguments in `args`, the matrices
+ must have the same shape, and scalars will be broadcasted to the same shape as the matrix.
Args:
- f (Union[ti.field, ti.Vector.field, ti.Matrix.field]): Input taichi field for memory address query.
- indices (Union[int, ti.Vector()]): The specified field indices of the query.
+ args: (List[:mod:`~taichi.types.primitive_types`, :class:`~taichi.Matrix`]): \
+ The input.
Returns:
- ti.u64: The memory address of `f[indices]`.
+ Minimum of the inputs.
+ Example::
+
+ >>> @ti.kernel
+ >>> def foo():
+ >>> x = ti.Vector([0, 1, 2])
+ >>> y = ti.Vector([3, 4, 5])
+ >>> z = ti.min(x, y, 1)
+ >>> print(z) # [0, 1, 1]
"""
- return Expr(_ti_core.expr_get_addr(f.snode.ptr, make_expr_group(indices)))
+ num_args = len(args)
+ assert num_args >= 1
+ if num_args == 1:
+ return args[0]
+ if num_args == 2:
+ return min_impl(args[0], args[1])
+ return min_impl(args[0], min(*args[1:]))
+
+
+def ti_any(a):
+ return a.any()
+
+
+def ti_all(a):
+ return a.all()
+
+
+__all__ = [
+ "acos", "asin", "atan2", "atomic_and", "atomic_or", "atomic_xor",
+ "atomic_max", "atomic_sub", "atomic_min", "atomic_add", "bit_cast",
+ "bit_shr", "cast", "ceil", "cos", "exp", "floor", "log", "random",
+ "raw_mod", "raw_div", "round", "rsqrt", "sin", "sqrt", "tan", "tanh",
+ "max", "min", "select", "abs", "pow"
+]
diff --git a/python/taichi/lang/quant_impl.py b/python/taichi/lang/quant_impl.py
deleted file mode 100644
index 2b369d975f230..0000000000000
--- a/python/taichi/lang/quant_impl.py
+++ /dev/null
@@ -1,77 +0,0 @@
-from taichi.lang import impl
-from taichi.lang import type_factory_impl as tf_impl
-
-import taichi as ti
-
-
-class Quant:
- """Generator of quantized types.
-
- For more details, read https://yuanming.taichi.graphics/publication/2021-quantaichi/quantaichi.pdf.
- """
- @staticmethod
- def int(bits, signed=False, compute=None):
- """Generates a quantized type for integers.
-
- Args:
- bits (int): Number of bits.
- signed (bool): Signed or unsigned.
- compute (DataType): Type for computation.
-
- Returns:
- DataType: The specified type.
- """
- if compute is None:
- compute = impl.get_runtime().default_ip
- return tf_impl.type_factory.custom_int(bits, signed, compute)
-
- @staticmethod
- def fixed(frac, signed=True, range=1.0, compute=None):
- """Generates a quantized type for fixed-point real numbers.
-
- Args:
- frac (int): Number of bits.
- signed (bool): Signed or unsigned.
- range (float): Range of the number.
- compute (DataType): Type for computation.
-
- Returns:
- DataType: The specified type.
- """
- # TODO: handle cases with frac > 32
- frac_type = Quant.int(bits=frac, signed=signed, compute=ti.i32)
- if signed:
- scale = range / 2**(frac - 1)
- else:
- scale = range / 2**frac
- if compute is None:
- compute = impl.get_runtime().default_fp
- return tf_impl.type_factory.custom_float(frac_type, None, compute,
- scale)
-
- @staticmethod
- def float(exp, frac, signed=True, compute=None):
- """Generates a quantized type for floating-point real numbers.
-
- Args:
- exp (int): Number of exponent bits.
- frac (int): Number of fraction bits.
- signed (bool): Signed or unsigned.
- compute (DataType): Type for computation.
-
- Returns:
- DataType: The specified type.
- """
- # Exponent is always unsigned
- exp_type = Quant.int(bits=exp, signed=False, compute=ti.i32)
- # TODO: handle cases with frac > 32
- frac_type = Quant.int(bits=frac, signed=signed, compute=ti.i32)
- if compute is None:
- compute = impl.get_runtime().default_fp
- return tf_impl.type_factory.custom_float(significand_type=frac_type,
- exponent_type=exp_type,
- compute_type=compute)
-
-
-# Unstable API
-quant = Quant
diff --git a/python/taichi/lang/runtime_ops.py b/python/taichi/lang/runtime_ops.py
index aff86c95eeac2..a1103bc3439dd 100644
--- a/python/taichi/lang/runtime_ops.py
+++ b/python/taichi/lang/runtime_ops.py
@@ -7,3 +7,6 @@ def sync():
def async_flush():
impl.get_runtime().prog.async_flush()
+
+
+__all__ = ['sync']
diff --git a/python/taichi/lang/shell.py b/python/taichi/lang/shell.py
index c21262917edaa..87ea33eafcf21 100644
--- a/python/taichi/lang/shell.py
+++ b/python/taichi/lang/shell.py
@@ -1,20 +1,19 @@
-import atexit
import functools
import os
import sys
+from taichi._lib import core as _ti_core
from taichi._logging import info, warn
-from taichi.core.util import ti_core as _ti_core
try:
- import sourceinspect as oinspect
+ import sourceinspect as oinspect # pylint: disable=unused-import
except ImportError:
warn('`sourceinspect` not installed!')
warn(
'Without this package Taichi may not function well in Python IDLE interactive shell, '
'Blender scripting module and Python native shell.')
warn('Please run `python3 -m pip install sourceinspect` to install.')
- import inspect as oinspect
+ import inspect as oinspect # pylint: disable=unused-import
pybuf_enabled = False
_env_enable_pybuf = os.environ.get('TI_ENABLE_PYBUF', '1')
@@ -35,7 +34,6 @@ def _shell_pop_print(old_call):
@functools.wraps(old_call)
def new_call(*args, **kwargs):
- _taichi_skip_traceback = 1
ret = old_call(*args, **kwargs)
# print's in kernel won't take effect until ti.sync(), discussion:
# https://github.com/taichi-dev/taichi/pull/1303#discussion_r444897102
diff --git a/python/taichi/lang/snode.py b/python/taichi/lang/snode.py
index e0b9b31510dd5..f457c8428a7bd 100644
--- a/python/taichi/lang/snode.py
+++ b/python/taichi/lang/snode.py
@@ -1,15 +1,8 @@
import numbers
-# The reason we import just the taichi.core.util module, instead of the ti_core
-# object within it, is that ti_core is stateful. While in practice ti_core is
-# loaded during the import procedure, it's probably still good to delay the
-# access to it.
-import taichi.lang
-from taichi.core.util import ti_core as _ti_core
-from taichi.lang import impl
-from taichi.lang.expr import Expr
+from taichi._lib import core as _ti_core
+from taichi.lang import expr, impl, matrix
from taichi.lang.field import Field
-from taichi.misc.util import deprecated
class SNode:
@@ -59,13 +52,15 @@ def pointer(self, axes, dimensions):
self.ptr.pointer(axes, dimensions,
impl.current_cfg().packed))
- def hash(self, axes, dimensions):
+ @staticmethod
+ def _hash(axes, dimensions):
+ # original code is #def hash(self,axes, dimensions) without #@staticmethod before fix pylint R0201
"""Not supported."""
raise RuntimeError('hash not yet supported')
- if isinstance(dimensions, int):
- dimensions = [dimensions] * len(axes)
- return SNode(self.ptr.hash(axes, dimensions,
- impl.current_cfg().packed))
+ # if isinstance(dimensions, int):
+ # dimensions = [dimensions] * len(axes)
+ # return SNode(self.ptr.hash(axes, dimensions,
+ # impl.current_cfg().packed))
def dynamic(self, axis, dimension, chunk_size=None):
"""Adds a dynamic SNode as a child component of `self`.
@@ -101,10 +96,6 @@ def bitmasked(self, axes, dimensions):
self.ptr.bitmasked(axes, dimensions,
impl.current_cfg().packed))
- @deprecated('_bit_struct', 'bit_struct')
- def _bit_struct(self, num_bits):
- return self.bit_struct(num_bits)
-
def bit_struct(self, num_bits: int):
"""Adds a bit_struct SNode as a child component of `self`.
@@ -116,10 +107,6 @@ def bit_struct(self, num_bits: int):
"""
return SNode(self.ptr.bit_struct(num_bits, impl.current_cfg().packed))
- @deprecated('_bit_array', 'bit_array')
- def _bit_array(self, axes, dimensions, num_bits):
- return self.bit_array(axes, dimensions, num_bits)
-
def bit_array(self, axes, dimensions, num_bits):
"""Adds a bit_array SNode as a child component of `self`.
@@ -157,7 +144,7 @@ def place(self, *args, offset=None, shared_exponent=False):
for arg in args:
if isinstance(arg, Field):
- for var in arg.get_field_members():
+ for var in arg._get_field_members():
self.ptr.place(var.ptr, offset)
elif isinstance(arg, list):
for x in arg:
@@ -199,8 +186,22 @@ def parent(self, n=1):
return impl.root
return SNode(p)
+ def _path_from_root(self):
+ """Gets the path from root to `self` in the SNode tree.
+
+ Returns:
+ List[Union[_Root, SNode]]: The list of SNodes on the path from root to `self`.
+ """
+ p = self
+ res = [p]
+ while p != impl.root:
+ p = p.parent()
+ res.append(p)
+ res.reverse()
+ return res
+
@property
- def dtype(self):
+ def _dtype(self):
"""Gets the data type of `self`.
Returns:
@@ -208,16 +209,8 @@ def dtype(self):
"""
return self.ptr.data_type()
- @deprecated('x.data_type()', 'x.dtype')
- def data_type(self):
- return self.dtype
-
- @deprecated('x.dim()', 'len(x.shape)')
- def dim(self):
- return len(self.shape)
-
@property
- def id(self):
+ def _id(self):
"""Gets the id of `self`.
Returns:
@@ -233,21 +226,11 @@ def shape(self):
Tuple[int]: The number of elements from root in each axis of `self`.
"""
dim = self.ptr.num_active_indices()
- ret = [self.ptr.get_shape_along_axis(i) for i in range(dim)]
-
- class callable_tuple(tuple):
- @deprecated('x.shape()', 'x.shape')
- def __call__(self):
- return self
+ ret = tuple(self.ptr.get_shape_along_axis(i) for i in range(dim))
- ret = callable_tuple(ret)
return ret
- @deprecated('x.get_shape(i)', 'x.shape[i]')
- def get_shape(self, i):
- return self.shape[i]
-
- def loop_range(self):
+ def _loop_range(self):
"""Gets the taichi_core.Expr wrapping the taichi_core.GlobalVariableExpression corresponding to `self` to serve as loop range.
Returns:
@@ -256,7 +239,7 @@ def loop_range(self):
return _ti_core.global_var_expr_from_snode(self.ptr)
@property
- def name(self):
+ def _name(self):
"""Gets the name of `self`.
Returns:
@@ -265,24 +248,14 @@ def name(self):
return self.ptr.name()
@property
- def snode(self):
+ def _snode(self):
"""Gets `self`.
-
Returns:
SNode: `self`.
"""
return self
- @property
- def needs_grad(self):
- """Checks whether `self` has a corresponding gradient :class:`~taichi.lang.SNode`.
-
- Returns:
- bool: Whether `self` has a corresponding gradient :class:`~taichi.lang.SNode`.
- """
- return self.ptr.has_grad()
-
- def get_children(self):
+ def _get_children(self):
"""Gets all children components of `self`.
Returns:
@@ -294,30 +267,38 @@ def get_children(self):
return children
@property
- def num_dynamically_allocated(self):
+ def _num_dynamically_allocated(self):
runtime = impl.get_runtime()
- runtime.materialize()
+ runtime.materialize_root_fb(False)
return runtime.prog.get_snode_num_dynamically_allocated(self.ptr)
@property
- def cell_size_bytes(self):
- runtime = impl.get_runtime()
- runtime.materialize()
+ def _cell_size_bytes(self):
+ impl.get_runtime().materialize_root_fb(False)
return self.ptr.cell_size_bytes
+ @property
+ def _offset_bytes_in_parent_cell(self):
+ impl.get_runtime().materialize_root_fb(False)
+ return self.ptr.offset_bytes_in_parent_cell
+
def deactivate_all(self):
"""Recursively deactivate all children components of `self`."""
- ch = self.get_children()
+ ch = self._get_children()
for c in ch:
c.deactivate_all()
SNodeType = _ti_core.SNodeType
if self.ptr.type == SNodeType.pointer or self.ptr.type == SNodeType.bitmasked:
- taichi.lang.meta.snode_deactivate(self)
+ from taichi._kernels import \
+ snode_deactivate # pylint: disable=C0415
+ snode_deactivate(self)
if self.ptr.type == SNodeType.dynamic:
# Note that dynamic nodes are different from other sparse nodes:
# instead of deactivating each element, we only need to deactivate
# its parent, whose linked list of chunks of elements will be deleted.
- taichi.lang.meta.snode_deactivate_dynamic(self)
+ from taichi._kernels import \
+ snode_deactivate_dynamic # pylint: disable=C0415
+ snode_deactivate_dynamic(self)
def __repr__(self):
type_ = str(self.ptr.type)[len('SNodeType.'):]
@@ -334,7 +315,7 @@ def __str__(self):
def __eq__(self, other):
return self.ptr == other.ptr
- def physical_index_position(self):
+ def _physical_index_position(self):
"""Gets mappings from virtual axes to physical axes.
Returns:
@@ -346,3 +327,90 @@ def physical_index_position(self):
if physical != -1:
ret[virtual] = physical
return ret
+
+
+def rescale_index(a, b, I):
+ """Rescales the index 'I' of field (or SNode) 'a' to match the shape of SNode 'b'
+
+ Parameters
+ ----------
+ a: ti.field(), ti.Vector.field, ti.Matrix.field()
+ input taichi field or snode
+ b: ti.field(), ti.Vector.field, ti.Matrix.field()
+ output taichi field or snode
+ I: ti.Vector()
+ grouped loop index
+
+ Returns
+ -------
+ Ib: ti.Vector()
+ rescaled grouped loop index
+
+ """
+ assert isinstance(
+ a, (Field, SNode)), "The first argument must be a field or an SNode"
+ assert isinstance(
+ b, (Field, SNode)), "The second argument must be a field or an SNode"
+ if isinstance(I, list):
+ I = matrix.Vector(I)
+ else:
+ assert isinstance(
+ I, matrix.Matrix
+ ), "The third argument must be an index (list or ti.Vector)"
+ entries = [I(i) for i in range(I.n)]
+ for n in range(min(I.n, min(len(a.shape), len(b.shape)))):
+ if a.shape[n] > b.shape[n]:
+ entries[n] = I(n) // (a.shape[n] // b.shape[n])
+ if a.shape[n] < b.shape[n]:
+ entries[n] = I(n) * (b.shape[n] // a.shape[n])
+ return matrix.Vector(entries)
+
+
+def append(l, indices, val):
+ a = impl.expr_init(
+ _ti_core.insert_append(l._snode.ptr, expr.make_expr_group(indices),
+ expr.Expr(val).ptr))
+ return a
+
+
+def is_active(l, indices):
+ return expr.Expr(
+ _ti_core.insert_is_active(l._snode.ptr, expr.make_expr_group(indices)))
+
+
+def activate(l, indices):
+ impl.get_runtime().prog.current_ast_builder().insert_activate(
+ l._snode.ptr, expr.make_expr_group(indices))
+
+
+def deactivate(l, indices):
+ impl.get_runtime().prog.current_ast_builder().insert_deactivate(
+ l._snode.ptr, expr.make_expr_group(indices))
+
+
+def length(l, indices):
+ return expr.Expr(
+ _ti_core.insert_len(l._snode.ptr, expr.make_expr_group(indices)))
+
+
+def get_addr(f, indices):
+ """Query the memory address (on CUDA/x64) of field `f` at index `indices`.
+
+ Currently, this function can only be called inside a taichi kernel.
+
+ Args:
+ f (Union[ti.field, ti.Vector.field, ti.Matrix.field]): Input taichi field for memory address query.
+ indices (Union[int, ti.Vector()]): The specified field indices of the query.
+
+ Returns:
+ ti.u64: The memory address of `f[indices]`.
+
+ """
+ return expr.Expr(
+ _ti_core.expr_get_addr(f._snode.ptr, expr.make_expr_group(indices)))
+
+
+__all__ = [
+ 'activate', 'append', 'deactivate', 'get_addr', 'is_active', 'length',
+ 'rescale_index', "SNode"
+]
diff --git a/python/taichi/lang/source_builder.py b/python/taichi/lang/source_builder.py
new file mode 100644
index 0000000000000..5a151ef8149d1
--- /dev/null
+++ b/python/taichi/lang/source_builder.py
@@ -0,0 +1,148 @@
+import atexit
+import ctypes
+import os
+import shutil
+import subprocess
+import tempfile
+
+from taichi._lib import core as _ti_core
+from taichi.lang import impl
+from taichi.lang.exception import TaichiSyntaxError
+from taichi.lang.expr import make_expr_group
+from taichi.lang.util import get_clangpp
+
+
+class SourceBuilder:
+ def __init__(self):
+ self.bc = None
+ self.so = None
+ self.mode = None
+ self.td = None
+
+ def cleanup():
+ if self.td is not None:
+ shutil.rmtree(self.td)
+
+ atexit.register(cleanup)
+
+ @classmethod
+ def from_file(cls, filename, compile_fn=None, _temp_dir=None):
+ self = cls()
+ self.td = _temp_dir
+ if self.td is None:
+ self.td = tempfile.mkdtemp()
+
+ if filename.endswith((".cpp", ".c", ".cc")):
+ if impl.current_cfg().arch not in [
+ _ti_core.Arch.x64, _ti_core.Arch.cuda
+ ]:
+ raise TaichiSyntaxError(
+ "Unsupported arch for external function call")
+ if compile_fn is None:
+
+ def compile_fn_impl(filename):
+ if impl.current_cfg().arch == _ti_core.Arch.x64:
+ subprocess.call(get_clangpp() + ' -flto -c ' +
+ filename + ' -o ' +
+ os.path.join(self.td, 'source.bc'),
+ shell=True)
+ else:
+ subprocess.call(get_clangpp() + ' -flto -c ' +
+ filename + ' -o ' +
+ os.path.join(self.td, 'source.bc') +
+ ' -target nvptx64-nvidia-cuda',
+ shell=True)
+ return os.path.join(self.td, 'source.bc')
+
+ compile_fn = compile_fn_impl
+ self.bc = compile_fn(filename)
+ self.mode = 'bc'
+ elif filename.endswith(".cu"):
+ if impl.current_cfg().arch not in [_ti_core.Arch.cuda]:
+ raise TaichiSyntaxError(
+ "Unsupported arch for external function call")
+ if compile_fn is None:
+ shutil.copy(filename, os.path.join(self.td, 'source.cu'))
+
+ def compile_fn_impl(filename):
+ # Cannot use -o to specify multiple output files
+ subprocess.call(
+ get_clangpp() + ' ' +
+ os.path.join(self.td, 'source.cu') +
+ ' -c -emit-llvm -std=c++17 --cuda-gpu-arch=sm_50 -nocudalib',
+ cwd=self.td,
+ shell=True)
+ return os.path.join(
+ self.td, 'source-cuda-nvptx64-nvidia-cuda-sm_50.bc')
+
+ compile_fn = compile_fn_impl
+ self.bc = compile_fn(filename)
+ self.mode = 'bc'
+ elif filename.endswith((".so", ".dylib", ".dll")):
+ if impl.current_cfg().arch not in [_ti_core.Arch.x64]:
+ raise TaichiSyntaxError(
+ "Unsupported arch for external function call")
+ self.so = ctypes.CDLL(filename)
+ self.mode = 'so'
+ elif filename.endswith(".ll"):
+ if impl.current_cfg().arch not in [
+ _ti_core.Arch.x64, _ti_core.Arch.cuda
+ ]:
+ raise TaichiSyntaxError(
+ "Unsupported arch for external function call")
+ subprocess.call('llvm-as ' + filename + ' -o ' +
+ os.path.join(self.td, 'source.bc'),
+ shell=True)
+ self.bc = os.path.join(self.td, 'source.bc')
+ self.mode = 'bc'
+ elif filename.endswith(".bc"):
+ if impl.current_cfg().arch not in [
+ _ti_core.Arch.x64, _ti_core.Arch.cuda
+ ]:
+ raise TaichiSyntaxError(
+ "Unsupported arch for external function call")
+ self.bc = filename
+ self.mode = 'bc'
+ else:
+ raise TaichiSyntaxError(
+ 'Unsupported file type for external function call.')
+ return self
+
+ @classmethod
+ def from_source(cls, source_code, compile_fn=None):
+ if impl.current_cfg().arch not in [
+ _ti_core.Arch.x64, _ti_core.Arch.cuda
+ ]:
+ raise TaichiSyntaxError(
+ "Unsupported arch for external function call")
+ _temp_dir = tempfile.mkdtemp()
+ _temp_source = os.path.join(_temp_dir, '_temp_source.cpp')
+ with open(_temp_source, 'w') as f:
+ f.write(source_code)
+ return SourceBuilder.from_file(_temp_source, compile_fn, _temp_dir)
+
+ def __getattr__(self, item):
+ def bitcode_func_call_wrapper(*args):
+ impl.get_runtime().prog.current_ast_builder(
+ ).insert_external_func_call(0, '', self.bc, item,
+ make_expr_group(args),
+ make_expr_group([]))
+
+ if self.mode == 'bc':
+ return bitcode_func_call_wrapper
+
+ def external_func_call_wrapper(args=[], outputs=[]):
+ func_addr = ctypes.cast(self.so.__getattr__(item),
+ ctypes.c_void_p).value
+ impl.get_runtime().prog.current_ast_builder(
+ ).insert_external_func_call(func_addr, '', '', '',
+ make_expr_group(args),
+ make_expr_group(outputs))
+
+ if self.mode == 'so':
+ return external_func_call_wrapper
+
+ raise TaichiSyntaxError('Error occurs when calling external function.')
+
+
+__all__ = []
diff --git a/python/taichi/lang/stmt_builder.py b/python/taichi/lang/stmt_builder.py
deleted file mode 100644
index fdc34d875c79c..0000000000000
--- a/python/taichi/lang/stmt_builder.py
+++ /dev/null
@@ -1,717 +0,0 @@
-import ast
-import copy
-
-import astor
-from taichi.lang import impl
-from taichi.lang.ast.symbol_resolver import ASTResolver
-from taichi.lang.ast_builder_utils import *
-from taichi.lang.exception import TaichiSyntaxError
-from taichi.lang.expr_builder import build_expr, build_exprs
-from taichi.lang.util import to_taichi_type
-
-import taichi as ti
-
-
-class StmtBuilder(Builder):
- @staticmethod
- def set_subscript_index(node, value):
- assert isinstance(node, ast.Subscript), type(node)
- if isinstance(node.slice, ast.Index):
- node.slice.value = value
- else:
- node.slice = value
-
- @staticmethod
- def make_single_statement(stmts):
- template = 'if 1: pass'
- t = ast.parse(template).body[0]
- t.body = stmts
- return t
-
- @staticmethod
- def make_constant(value):
- # Do not use ast.Constant which does not exist in python3.5
- node = parse_expr('0')
- node.value = value
- return node
-
- @staticmethod
- def build_AugAssign(ctx, node):
- node.target = build_expr(ctx, node.target)
- node.value = build_expr(ctx, node.value)
- template = 'x.augassign(0, 0)'
- t = ast.parse(template).body[0]
- t.value.func.value = node.target
- t.value.func.value.ctx = ast.Load()
- t.value.args[0] = node.value
- t.value.args[1] = ast.Str(s=type(node.op).__name__,
- ctx=ast.Load(),
- kind=None)
- return ast.copy_location(t, node)
-
- @staticmethod
- def _is_string_mod_args(msg):
- # 1. str % (a, b, c, ...)
- # 2. str % single_item
- # Note that |msg.right| may not be a tuple.
- return isinstance(msg, ast.BinOp) and isinstance(
- msg.left, ast.Str) and isinstance(msg.op, ast.Mod)
-
- @staticmethod
- def _handle_string_mod_args(ctx, msg):
- assert StmtBuilder._is_string_mod_args(msg)
- s = msg.left.s
- t = None
- if isinstance(msg.right, ast.Tuple):
- t = msg.right
- else:
- # assuming the format is `str % single_item`
- t = ast.Tuple(elts=[msg.right], ctx=ast.Load())
- t = build_expr(ctx, t)
- return s, t
-
- @staticmethod
- def build_Assert(ctx, node):
- extra_args = ast.List(elts=[], ctx=ast.Load())
- if node.msg is not None:
- if isinstance(node.msg, ast.Constant):
- msg = node.msg.value
- elif isinstance(node.msg, ast.Str):
- msg = node.msg.s
- elif StmtBuilder._is_string_mod_args(node.msg):
- msg = build_expr(ctx, node.msg)
- msg, extra_args = StmtBuilder._handle_string_mod_args(ctx, msg)
- else:
- raise ValueError(
- f"assert info must be constant, not {ast.dump(node.msg)}")
- else:
- msg = astor.to_source(node.test)
- node.test = build_expr(ctx, node.test)
-
- new_node = parse_stmt('ti.ti_assert(0, 0, [])')
- new_node.value.args[0] = node.test
- new_node.value.args[1] = parse_expr("'{}'".format(msg.strip()))
- new_node.value.args[2] = extra_args
- new_node = ast.copy_location(new_node, node)
- return new_node
-
- @staticmethod
- def build_Assign(ctx, node):
- node.value = build_expr(ctx, node.value)
- node.targets = build_exprs(ctx, node.targets)
-
- is_static_assign = isinstance(
- node.value, ast.Call) and ASTResolver.resolve_to(
- node.value.func, ti.static, globals())
- if is_static_assign:
- return node
-
- # Keep all generated assign statements and compose single one at last.
- # The variable is introduced to support chained assignments.
- # Ref https://github.com/taichi-dev/taichi/issues/2659.
- assign_stmts = []
- for node_target in node.targets:
- if isinstance(node_target, ast.Tuple):
- assign_stmts.append(
- StmtBuilder.build_assign_unpack(ctx, node, node_target))
- else:
- assign_stmts.append(
- StmtBuilder.build_assign_basic(ctx, node, node_target,
- node.value))
- return StmtBuilder.make_single_statement(assign_stmts)
-
- @staticmethod
- def build_assign_unpack(ctx, node, node_target):
- """Build the unpack assignments like this: (target1, target2) = (value1, value2).
- The function should be called only if the node target is a tuple.
-
- Args:
- ctx (ast_builder_utils.BuilderContext): The builder context.
- node (ast.Assign): An assignment. targets is a list of nodes,
- and value is a single node.
- node_target (ast.Tuple): A list or tuple object. elts holds a
- list of nodes representing the elements.
- """
-
- targets = node_target.elts
-
- # Create
- stmts = []
-
- # Create a temp list and keep values in it, delete it after the initialization is finished.
- holder = parse_stmt('__tmp_tuple = ti.expr_init_list(0, '
- f'{len(targets)})')
- holder.value.args[0] = node.value
-
- stmts.append(holder)
-
- def tuple_indexed(i):
- indexing = parse_stmt('__tmp_tuple[0]')
- StmtBuilder.set_subscript_index(indexing.value, parse_expr(f"{i}"))
- return indexing.value
-
- # Generate assign statements for every target, then merge them into one.
- for i, target in enumerate(targets):
- stmts.append(
- StmtBuilder.build_assign_basic(ctx, node, target,
- tuple_indexed(i)))
- stmts.append(parse_stmt('del __tmp_tuple'))
- return StmtBuilder.make_single_statement(stmts)
-
- @staticmethod
- def build_assign_basic(ctx, node, target, value):
- """Build basic assginment like this: target = value.
-
- Args:
- ctx (ast_builder_utils.BuilderContext): The builder context.
- node (ast.Assign): An assignment. targets is a list of nodes,
- and value is a single node.
- target (ast.Name): A variable name. id holds the name as
- a string.
- value: A node representing the value.
- """
- is_local = isinstance(target, ast.Name)
- if is_local and ctx.is_creation(target.id):
- var_name = target.id
- target.ctx = ast.Store()
- # Create, no AST resolution needed
- init = ast.Attribute(value=ast.Name(id='ti', ctx=ast.Load()),
- attr='expr_init',
- ctx=ast.Load())
- rhs = ast.Call(
- func=init,
- args=[value],
- keywords=[],
- )
- ctx.create_variable(var_name)
- return ast.copy_location(
- ast.Assign(targets=[target], value=rhs, type_comment=None),
- node)
- else:
- # Assign
- target.ctx = ast.Load()
- func = ast.Attribute(value=target, attr='assign', ctx=ast.Load())
- call = ast.Call(func=func, args=[value], keywords=[])
- return ast.copy_location(ast.Expr(value=call), node)
-
- @staticmethod
- def build_Try(ctx, node):
- raise TaichiSyntaxError(
- "Keyword 'try' not supported in Taichi kernels")
-
- @staticmethod
- def build_While(ctx, node):
- if node.orelse:
- raise TaichiSyntaxError(
- "'else' clause for 'while' not supported in Taichi kernels")
-
- with ctx.control_scope():
- ctx.current_control_scope().append('while')
-
- template = '''
-if 1:
- ti.core.begin_frontend_while(ti.Expr(1).ptr)
- __while_cond = 0
- if __while_cond:
- pass
- else:
- break
- ti.core.pop_scope()
-'''
- cond = node.test
- t = ast.parse(template).body[0]
- t.body[1].value = cond
- t.body = t.body[:3] + node.body + t.body[3:]
-
- t.body = build_stmts(ctx, t.body)
- return ast.copy_location(t, node)
-
- @staticmethod
- def build_If(ctx, node):
- node.test = build_expr(ctx, node.test)
- node.body = build_stmts(ctx, node.body)
- node.orelse = build_stmts(ctx, node.orelse)
-
- is_static_if = isinstance(node.test, ast.Call) and isinstance(
- node.test.func, ast.Attribute)
- if is_static_if:
- attr = node.test.func
- if attr.attr == 'static':
- is_static_if = True
- else:
- is_static_if = False
-
- if is_static_if:
- # Do nothing
- return node
-
- template = '''
-if 1:
- __cond = 0
- ti.begin_frontend_if(__cond)
- ti.core.begin_frontend_if_true()
- ti.core.pop_scope()
- ti.core.begin_frontend_if_false()
- ti.core.pop_scope()
-'''
- t = ast.parse(template).body[0]
- cond = node.test
- t.body[0].value = cond
- t.body = t.body[:5] + node.orelse + t.body[5:]
- t.body = t.body[:3] + node.body + t.body[3:]
- return ast.copy_location(t, node)
-
- @staticmethod
- def get_for_loop_targets(node):
- """
- Returns the list of indices of the for loop |node|.
- See also: https://docs.python.org/3/library/ast.html#ast.For
- """
- if isinstance(node.target, ast.Name):
- return [node.target.id]
- else:
- assert isinstance(node.target, ast.Tuple)
- return [name.id for name in node.target.elts]
-
- @staticmethod
- def get_decorator(node):
- if not isinstance(node, ast.Call):
- return ''
- for wanted, name in [
- (ti.static, 'static'),
- (ti.grouped, 'grouped'),
- (ti.ndrange, 'ndrange'),
- ]:
- if ASTResolver.resolve_to(node.func, wanted, globals()):
- return name
- return ''
-
- @staticmethod
- def build_static_for(ctx, node, is_grouped):
- # for i in ti.static(range(n))
- # for i, j in ti.static(ti.ndrange(n))
- # for I in ti.static(ti.grouped(ti.ndrange(n, m)))
-
- ctx.current_control_scope().append('static')
- node.body = build_stmts(ctx, node.body)
- if is_grouped:
- assert len(node.iter.args[0].args) == 1
- template = '''
-if 1:
- __ndrange_arg = 0
- from taichi.lang.exception import TaichiSyntaxError
- if not isinstance(__ndrange_arg, ti.ndrange):
- raise TaichiSyntaxError("Only 'ti.ndrange' is allowed in 'ti.static(ti.grouped(...))'.")
- pass
- del a
- '''
- t = ast.parse(template).body[0]
- t.body[0].value = node.iter.args[0].args[0]
- t.body[3] = node
- t.body[3].iter.args[0].args[0] = parse_expr('__ndrange_arg')
- else:
- t = parse_stmt('if 1: pass; del a')
- t.body[0] = node
- target = copy.deepcopy(node.target)
- target.ctx = ast.Del()
- if isinstance(target, ast.Tuple):
- for tar in target.elts:
- tar.ctx = ast.Del()
- t.body[-1].targets = [target]
- return t
-
- @staticmethod
- def build_range_for(ctx, node):
- # for i in range(n)
- node.body = build_stmts(ctx, node.body)
- loop_var = node.target.id
- ctx.check_loop_var(loop_var)
- template = '''
-if 1:
- {} = ti.Expr(ti.core.make_id_expr(''))
- ___begin = ti.Expr(0)
- ___end = ti.Expr(0)
- ___begin = ti.cast(___begin, ti.i32)
- ___end = ti.cast(___end, ti.i32)
- ti.core.begin_frontend_range_for({}.ptr, ___begin.ptr, ___end.ptr)
- ti.core.end_frontend_range_for()
- '''.format(loop_var, loop_var)
- t = ast.parse(template).body[0]
-
- assert len(node.iter.args) in [1, 2]
- if len(node.iter.args) == 2:
- bgn = build_expr(ctx, node.iter.args[0])
- end = build_expr(ctx, node.iter.args[1])
- else:
- bgn = StmtBuilder.make_constant(value=0)
- end = build_expr(ctx, node.iter.args[0])
-
- t.body[1].value.args[0] = bgn
- t.body[2].value.args[0] = end
- t.body = t.body[:6] + node.body + t.body[6:]
- t.body.append(parse_stmt('del {}'.format(loop_var)))
- return ast.copy_location(t, node)
-
- @staticmethod
- def build_ndrange_for(ctx, node):
- # for i, j in ti.ndrange(n)
- template = f'''
-if ti.static(1):
- __ndrange{id(node)} = 0
- for __ndrange_I{id(node)} in range(0):
- __I = __ndrange_I{id(node)}
- '''
- t = ast.parse(template).body[0]
- t.body[0].value = node.iter
- t_loop = t.body[1]
- t_loop.iter.args[0] = parse_expr(
- f'__ndrange{id(node)}.acc_dimensions[0]')
- targets = StmtBuilder.get_for_loop_targets(node)
- targets_tmp = ['__' + name for name in targets]
- loop_body = t_loop.body
- for i in range(len(targets)):
- if i + 1 < len(targets):
- stmt = '{} = __I // __ndrange{}.acc_dimensions[{}]'.format(
- targets_tmp[i], id(node), i + 1)
- else:
- stmt = '{} = __I'.format(targets_tmp[i])
- loop_body.append(parse_stmt(stmt))
- stmt = '{} = {} + __ndrange{}.bounds[{}][0]'.format(
- targets[i], targets_tmp[i], id(node), i)
- loop_body.append(parse_stmt(stmt))
- if i + 1 < len(targets):
- stmt = '__I = __I - {} * __ndrange{}.acc_dimensions[{}]'.format(
- targets_tmp[i], id(node), i + 1)
- loop_body.append(parse_stmt(stmt))
- loop_body += node.body
-
- node = ast.copy_location(t, node)
- return build_stmt(ctx, node) # further translate as a range for
-
- @staticmethod
- def build_grouped_ndrange_for(ctx, node):
- # for I in ti.grouped(ti.ndrange(n, m))
- node.body = build_stmts(ctx, node.body)
- target = node.target.id
- template = '''
-if ti.static(1):
- __ndrange = 0
- {} = ti.expr_init(ti.Vector([0] * len(__ndrange.dimensions), disable_local_tensor=True))
- ___begin = ti.Expr(0)
- ___end = __ndrange.acc_dimensions[0]
- ___begin = ti.cast(___begin, ti.i32)
- ___end = ti.cast(___end, ti.i32)
- __ndrange_I = ti.Expr(ti.core.make_id_expr(''))
- ti.core.begin_frontend_range_for(__ndrange_I.ptr, ___begin.ptr, ___end.ptr)
- __I = __ndrange_I
- for __grouped_I in range(len(__ndrange.dimensions)):
- __grouped_I_tmp = 0
- if __grouped_I + 1 < len(__ndrange.dimensions):
- __grouped_I_tmp = __I // __ndrange.acc_dimensions[__grouped_I + 1]
- else:
- __grouped_I_tmp = __I
- ti.subscript({}, __grouped_I).assign(__grouped_I_tmp + __ndrange.bounds[__grouped_I][0])
- if __grouped_I + 1 < len(__ndrange.dimensions):
- __I = __I - __grouped_I_tmp * __ndrange.acc_dimensions[__grouped_I + 1]
- ti.core.end_frontend_range_for()
- '''.format(target, target)
- t = ast.parse(template).body[0]
- node.iter.args[0].args = build_exprs(ctx, node.iter.args[0].args)
- t.body[0].value = node.iter.args[0]
- cut = len(t.body) - 1
- t.body = t.body[:cut] + node.body + t.body[cut:]
- return ast.copy_location(t, node)
-
- @staticmethod
- def build_struct_for(ctx, node, is_grouped):
- # for i, j in x
- # for I in ti.grouped(x)
- node.body = build_stmts(ctx, node.body)
- targets = StmtBuilder.get_for_loop_targets(node)
-
- for loop_var in targets:
- ctx.check_loop_var(loop_var)
-
- var_decl = ''.join(
- ' {} = ti.Expr(ti.core.make_id_expr(""))\n'.format(name)
- for name in targets) # indent: 4 spaces
- vars = ', '.join(targets)
- if is_grouped:
- template = '''
-if 1:
- ___loop_var = 0
- {} = ti.lang.expr.make_var_vector(size=len(___loop_var.shape))
- ___expr_group = ti.lang.expr.make_expr_group({})
- ti.begin_frontend_struct_for(___expr_group, ___loop_var)
- ti.core.end_frontend_range_for()
- '''.format(vars, vars)
- t = ast.parse(template).body[0]
- cut = 4
- t.body[0].value = node.iter
- t.body = t.body[:cut] + node.body + t.body[cut:]
- else:
- template = '''
-if 1:
-{}
- ___loop_var = 0
- ___expr_group = ti.lang.expr.make_expr_group({})
- ti.begin_frontend_struct_for(___expr_group, ___loop_var)
- ti.core.end_frontend_range_for()
- '''.format(var_decl, vars)
- t = ast.parse(template).body[0]
- cut = len(targets) + 3
- t.body[cut - 3].value = node.iter
- t.body = t.body[:cut] + node.body + t.body[cut:]
- for loop_var in reversed(targets):
- t.body.append(parse_stmt('del {}'.format(loop_var)))
- return ast.copy_location(t, node)
-
- @staticmethod
- def build_For(ctx, node):
- if node.orelse:
- raise TaichiSyntaxError(
- "'else' clause for 'for' not supported in Taichi kernels")
-
- with ctx.control_scope():
- ctx.current_control_scope().append('for')
-
- decorator = StmtBuilder.get_decorator(node.iter)
- double_decorator = ''
- if decorator != '' and len(node.iter.args) == 1:
- double_decorator = StmtBuilder.get_decorator(node.iter.args[0])
- ast.fix_missing_locations(node)
-
- if decorator == 'static':
- if double_decorator == 'static':
- raise TaichiSyntaxError("'ti.static' cannot be nested")
- return StmtBuilder.build_static_for(
- ctx, node, double_decorator == 'grouped')
- elif decorator == 'ndrange':
- if double_decorator != '':
- raise TaichiSyntaxError(
- "No decorator is allowed inside 'ti.ndrange")
- return StmtBuilder.build_ndrange_for(ctx, node)
- elif decorator == 'grouped':
- if double_decorator == 'static':
- raise TaichiSyntaxError(
- "'ti.static' is not allowed inside 'ti.grouped'")
- elif double_decorator == 'ndrange':
- return StmtBuilder.build_grouped_ndrange_for(ctx, node)
- elif double_decorator == 'grouped':
- raise TaichiSyntaxError("'ti.grouped' cannot be nested")
- else:
- return StmtBuilder.build_struct_for(ctx,
- node,
- is_grouped=True)
- elif isinstance(node.iter, ast.Call) and isinstance(
- node.iter.func, ast.Name) and node.iter.func.id == 'range':
- return StmtBuilder.build_range_for(ctx, node)
- else: # Struct for
- return StmtBuilder.build_struct_for(ctx,
- node,
- is_grouped=False)
-
- @staticmethod
- def build_Break(ctx, node):
- if 'static' in ctx.current_control_scope():
- return node
- else:
- return parse_stmt('ti.core.insert_break_stmt()')
-
- @staticmethod
- def build_Continue(ctx, node):
- if 'static' in ctx.current_control_scope():
- return node
- else:
- return parse_stmt('ti.core.insert_continue_stmt()')
-
- @staticmethod
- def build_FunctionDef(ctx, node):
- args = node.args
- assert args.vararg is None
- assert args.kwonlyargs == []
- assert args.kw_defaults == []
- assert args.kwarg is None
-
- arg_decls = []
-
- def transform_as_kernel():
- # Treat return type
- if node.returns is not None:
- ret_init = parse_stmt(
- 'ti.lang.kernel_arguments.decl_scalar_ret(0)')
- ret_init.value.args[0] = node.returns
- ctx.returns = node.returns
- arg_decls.append(ret_init)
- node.returns = None
-
- for i, arg in enumerate(args.args):
- # Directly pass in template arguments,
- # such as class instances ("self"), fields, SNodes, etc.
- if isinstance(ctx.func.argument_annotations[i], ti.template):
- continue
- if isinstance(ctx.func.argument_annotations[i],
- ti.linalg.sparse_matrix_builder):
- arg_init = parse_stmt(
- 'x = ti.lang.kernel_arguments.decl_sparse_matrix()')
- arg_init.targets[0].id = arg.arg
- ctx.create_variable(arg.arg)
- arg_decls.append(arg_init)
- elif isinstance(ctx.func.argument_annotations[i], ti.any_arr):
- arg_init = parse_stmt(
- 'x = ti.lang.kernel_arguments.decl_any_arr_arg(0, 0, 0, 0)'
- )
- arg_init.targets[0].id = arg.arg
- ctx.create_variable(arg.arg)
- array_dt = ctx.arg_features[i][0]
- array_dim = ctx.arg_features[i][1]
- array_element_shape = ctx.arg_features[i][2]
- array_layout = ctx.arg_features[i][3]
- array_dt = to_taichi_type(array_dt)
- dt_expr = 'ti.' + ti.core.data_type_name(array_dt)
- dt = parse_expr(dt_expr)
- arg_init.value.args[0] = dt
- arg_init.value.args[1] = parse_expr("{}".format(array_dim))
- arg_init.value.args[2] = parse_expr(
- "{}".format(array_element_shape))
- arg_init.value.args[3] = parse_expr(
- "ti.{}".format(array_layout))
- arg_decls.append(arg_init)
- else:
- arg_init = parse_stmt(
- 'x = ti.lang.kernel_arguments.decl_scalar_arg(0)')
- arg_init.targets[0].id = arg.arg
- dt = arg.annotation
- arg_init.value.args[0] = dt
- arg_decls.append(arg_init)
- # remove original args
- node.args.args = []
-
- if ctx.is_kernel: # ti.kernel
- for decorator in node.decorator_list:
- if ASTResolver.resolve_to(decorator, ti.func, globals()):
- raise TaichiSyntaxError(
- "Function definition not allowed in 'ti.kernel'.")
- transform_as_kernel()
-
- else: # ti.func
- for decorator in node.decorator_list:
- if ASTResolver.resolve_to(decorator, ti.func, globals()):
- raise TaichiSyntaxError(
- "Function definition not allowed in 'ti.func'.")
- if impl.get_runtime().experimental_real_function:
- transform_as_kernel()
- else:
- # Transform as force-inlined func
- arg_decls = []
- for i, arg in enumerate(args.args):
- # Remove annotations because they are not used.
- args.args[i].annotation = None
- # Template arguments are passed by reference.
- if isinstance(ctx.func.argument_annotations[i],
- ti.template):
- ctx.create_variable(ctx.func.argument_names[i])
- continue
- # Create a copy for non-template arguments,
- # so that they are passed by value.
- arg_init = parse_stmt('x = ti.expr_init_func(0)')
- arg_init.targets[0].id = arg.arg
- ctx.create_variable(arg.arg)
- arg_init.value.args[0] = parse_expr(arg.arg +
- '_by_value__')
- args.args[i].arg += '_by_value__'
- arg_decls.append(arg_init)
-
- with ctx.variable_scope():
- node.body = build_stmts(ctx, node.body)
-
- node.body = arg_decls + node.body
- node.body = [parse_stmt('import taichi as ti')] + node.body
- return node
-
- @staticmethod
- def build_Return(ctx, node):
- node.value = build_expr(ctx, node.value)
- if ctx.is_kernel or impl.get_runtime().experimental_real_function:
- # TODO: check if it's at the end of a kernel, throw TaichiSyntaxError if not
- if node.value is not None:
- if ctx.returns is None:
- raise TaichiSyntaxError(
- f'A {"kernel" if ctx.is_kernel else "function"} '
- 'with a return value must be annotated '
- 'with a return type, e.g. def func() -> ti.f32')
- ret_expr = parse_expr('ti.cast(ti.Expr(0), 0)')
- ret_expr.args[0].args[0] = node.value
- ret_expr.args[1] = ctx.returns
- ret_stmt = parse_stmt('ti.core.create_kernel_return(ret.ptr)')
- # For args[0], it is an ast.Attribute, because it loads the
- # attribute, |ptr|, of the expression |ret_expr|. Therefore we
- # only need to replace the object part, i.e. args[0].value
- ret_stmt.value.args[0].value = ret_expr
- return ret_stmt
- return node
-
- @staticmethod
- def build_Module(ctx, node):
- with ctx.variable_scope():
- # Do NOT use |build_stmts| which inserts 'del' statements to the
- # end and deletes parameters passed into the module
- node.body = [build_stmt(ctx, stmt) for stmt in list(node.body)]
- return node
-
- @staticmethod
- def build_Global(ctx, node):
- raise TaichiSyntaxError(
- "Keyword 'global' not supported in Taichi kernels")
-
- @staticmethod
- def build_Nonlocal(ctx, node):
- raise TaichiSyntaxError(
- "Keyword 'nonlocal' not supported in Taichi kernels")
-
- @staticmethod
- def build_Raise(ctx, node):
- node.exc = build_expr(ctx, node.exc)
- return node
-
- @staticmethod
- def build_Expr(ctx, node):
- if not isinstance(node.value, ast.Call):
- # A statement with a single expression.
- return node
-
- # A function call.
- node.value = build_expr(ctx, node.value)
- # Note that we can only return an ast.Expr instead of an ast.Call.
-
- if impl.get_runtime().experimental_real_function:
- # Generates code that inserts a FrontendExprStmt if the function
- # called is a Taichi function.
- # We cannot insert the FrontendExprStmt here because we do not
- # know if the function is a Taichi function now.
- node.value.args = [node.value.func] + node.value.args
- node.value.func = parse_expr('ti.insert_expr_stmt_if_ti_func')
- return node
-
- @staticmethod
- def build_Import(ctx, node):
- return node
-
- @staticmethod
- def build_ImportFrom(ctx, node):
- return node
-
- @staticmethod
- def build_Pass(ctx, node):
- return node
-
-
-build_stmt = StmtBuilder()
-
-
-def build_stmts(ctx, stmts):
- result = []
- with ctx.variable_scope(result):
- for stmt in list(stmts):
- result.append(build_stmt(ctx, stmt))
- return result
diff --git a/python/taichi/lang/struct.py b/python/taichi/lang/struct.py
index 9dd7a6cadb93a..f4f582f57b733 100644
--- a/python/taichi/lang/struct.py
+++ b/python/taichi/lang/struct.py
@@ -1,29 +1,46 @@
-import copy
import numbers
-from numpy import broadcast
-from taichi.lang import expr, impl
+from taichi.lang import expr, impl, ops
from taichi.lang.common_ops import TaichiOperations
from taichi.lang.enums import Layout
from taichi.lang.exception import TaichiSyntaxError
from taichi.lang.field import Field, ScalarField, SNodeHostAccess
from taichi.lang.matrix import Matrix
-from taichi.lang.ops import cast
-from taichi.lang.types import CompoundType
from taichi.lang.util import (cook_dtype, in_python_scope, is_taichi_class,
python_scope, taichi_scope)
-
-import taichi as ti
+from taichi.types import primitive_types
+from taichi.types.compound_types import CompoundType
class Struct(TaichiOperations):
"""The Struct type class.
- Args:
- entries (Dict[str, Union[Dict, Expr, Matrix, Struct]]): keys and values for struct members.
+
+ A struct is a dictionary-like data structure that stores members as
+ (key, value) pairs. Valid data members of a struct can be scalars,
+ matrices or other dictionary-like stuctures.
"""
- is_taichi_class = True
+ _is_taichi_class = True
def __init__(self, *args, **kwargs):
+ """
+ Args:
+ entries (Dict[str, Union[Dict, Expr, Matrix, Struct]]): \
+ keys and values for struct members.
+
+ Returns:
+ An instance of this struct.
+
+ Example::
+
+ >>> vec3 = ti.types.vector(3, ti.f32)
+ >>> a = ti.Struct(v=vec3([0, 0, 0]), t=1.0)
+ >>> print(a.items)
+ dict_items([('v', [0. 0. 0.]), ('t', 1.0)])
+ >>>
+ >>> B = ti.Struct(v=vec3([0., 0., 0.]), t=1.0, A=a)
+ >>> print(B.items)
+ dict_items([('v', [0. 0. 0.]), ('t', 1.0), ('A', {'v': [[0.], [0.], [0.]], 't': 1.0})])
+ """
# converts lists to matrices and dicts to structs
if len(args) == 1 and kwargs == {} and isinstance(args[0], dict):
self.entries = args[0]
@@ -38,47 +55,43 @@ def __init__(self, *args, **kwargs):
v = Matrix(v)
if isinstance(v, dict):
v = Struct(v)
- self.entries[k] = v
- self.register_members()
- self.local_tensor_proxy = None
- self.any_array_access = None
+ self.entries[k] = v if in_python_scope() else impl.expr_init(v)
+ self._register_members()
@property
def keys(self):
return list(self.entries.keys())
@property
- def members(self):
+ def _members(self):
return list(self.entries.values())
@property
def items(self):
return self.entries.items()
- def register_members(self):
+ def _register_members(self):
for k in self.keys:
setattr(Struct, k,
property(
- Struct.make_getter(k),
- Struct.make_setter(k),
+ Struct._make_getter(k),
+ Struct._make_setter(k),
))
def __getitem__(self, key):
- _taichi_skip_traceback = 1
ret = self.entries[key]
if isinstance(ret, SNodeHostAccess):
ret = ret.accessor.getter(*ret.key)
return ret
def __setitem__(self, key, value):
- _taichi_skip_traceback = 1
if isinstance(self.entries[key], SNodeHostAccess):
self.entries[key].accessor.setter(value, *self.entries[key].key)
else:
if in_python_scope():
if isinstance(self.entries[key], Struct) or isinstance(
self.entries[key], Matrix):
- self.entries[key].set_entries(value)
+ self.entries[key]._set_entries(value)
else:
if isinstance(value, numbers.Number):
self.entries[key] = value
@@ -89,121 +102,89 @@ def __setitem__(self, key, value):
else:
self.entries[key] = value
- def set_entries(self, value):
+ def _set_entries(self, value):
if isinstance(value, dict):
value = Struct(value)
for k in self.keys:
self[k] = value[k]
@staticmethod
- def make_getter(key):
+ def _make_getter(key):
def getter(self):
"""Get an entry from custom struct by name."""
- _taichi_skip_traceback = 1
return self[key]
return getter
@staticmethod
- def make_setter(key):
+ def _make_setter(key):
@python_scope
def setter(self, value):
- _taichi_skip_traceback = 1
self[key] = value
return setter
- def element_wise_unary(self, foo):
- _taichi_skip_traceback = 1
- ret = self.empty_copy()
+ def _element_wise_unary(self, foo):
+ entries = {}
for k, v in self.items:
- if isinstance(v, expr.Expr):
- ret.entries[k] = foo(v)
+ if is_taichi_class(v):
+ entries[k] = v._element_wise_unary(foo)
else:
- ret.entries[k] = v.element_wise_unary(foo)
- return ret
+ entries[k] = foo(v)
+ return Struct(entries)
- def element_wise_binary(self, foo, other):
- _taichi_skip_traceback = 1
- ret = self.empty_copy()
- if isinstance(other, (dict)):
- other = Struct(other)
- if isinstance(other, Struct):
- if self.entries.keys() != other.entries.keys():
- raise TypeError(
- f"Member mismatch between structs {self.keys}, {other.keys}"
- )
- for k, v in self.items:
- if isinstance(v, expr.Expr):
- ret.entries[k] = foo(v, other.entries[k])
- else:
- ret.entries[k] = v.element_wise_binary(
- foo, other.entries[k])
- else: # assumed to be scalar
- for k, v in self.items:
- if isinstance(v, expr.Expr):
- ret.entries[k] = foo(v, other)
- else:
- ret.entries[k] = v.element_wise_binary(foo, other)
- return ret
+ def _element_wise_binary(self, foo, other):
+ other = self._broadcast_copy(other)
+ entries = {}
+ for k, v in self.items:
+ if is_taichi_class(v):
+ entries[k] = v._element_wise_binary(foo, other.entries[k])
+ else:
+ entries[k] = foo(v, other.entries[k])
+ return Struct(entries)
- def broadcast_copy(self, other):
+ def _broadcast_copy(self, other):
if isinstance(other, dict):
other = Struct(other)
if not isinstance(other, Struct):
- ret = self.empty_copy()
- for k, v in ret.items:
- if isinstance(v, (Matrix, Struct)):
- ret.entries[k] = v.broadcast_copy(other)
+ entries = {}
+ for k, v in self.items:
+ if is_taichi_class(v):
+ entries[k] = v._broadcast_copy(other)
else:
- ret.entries[k] = other
- other = ret
+ entries[k] = other
+ other = Struct(entries)
if self.entries.keys() != other.entries.keys():
raise TypeError(
f"Member mismatch between structs {self.keys}, {other.keys}")
return other
- def element_wise_writeback_binary(self, foo, other):
- ret = self.empty_copy()
- if isinstance(other, (dict)):
- other = Struct(other)
- if is_taichi_class(other):
- other = other.variable()
- if foo.__name__ == 'assign' and not isinstance(other, Struct):
+ def _element_wise_writeback_binary(self, foo, other):
+ if foo.__name__ == 'assign' and not isinstance(other, (dict, Struct)):
raise TaichiSyntaxError(
'cannot assign scalar expr to '
f'taichi class {type(self)}, maybe you want to use `a.fill(b)` instead?'
)
- if isinstance(other, Struct):
- if self.entries.keys() != other.entries.keys():
- raise TypeError(
- f"Member mismatch between structs {self.keys}, {other.keys}"
- )
- for k, v in self.items:
- if isinstance(v, expr.Expr):
- ret.entries[k] = foo(v, other.entries[k])
- else:
- ret.entries[k] = v.element_wise_binary(
- foo, other.entries[k])
- else: # assumed to be scalar
- for k, v in self.items:
- if isinstance(v, expr.Expr):
- ret.entries[k] = foo(v, other)
- else:
- ret.entries[k] = v.element_wise_binary(foo, other)
- return ret
+ other = self._broadcast_copy(other)
+ entries = {}
+ for k, v in self.items:
+ if is_taichi_class(v):
+ entries[k] = v._element_wise_binary(foo, other.entries[k])
+ else:
+ entries[k] = foo(v, other.entries[k])
+ return self if foo.__name__ == 'assign' else Struct(entries)
- def element_wise_ternary(self, foo, other, extra):
- ret = self.empty_copy()
- other = self.broadcast_copy(other)
- extra = self.broadcast_copy(extra)
+ def _element_wise_ternary(self, foo, other, extra):
+ other = self._broadcast_copy(other)
+ extra = self._broadcast_copy(extra)
+ entries = {}
for k, v in self.items:
- if isinstance(v, expr.Expr):
- ret.entries[k] = foo(v, other.entries[k], extra.entries[k])
+ if is_taichi_class(v):
+ entries[k] = v._element_wise_ternary(foo, other.entries[k],
+ extra.entries[k])
else:
- ret.entries[k] = v.element_wise_ternary(
- foo, other.entries[k], extra.entries[k])
- return ret
+ entries[k] = foo(v, other.entries[k], extra.entries[k])
+ return Struct(entries)
@taichi_scope
def fill(self, val):
@@ -213,35 +194,9 @@ def fill(self, val):
val (Union[int, float]): Value to fill.
"""
def assign_renamed(x, y):
- return ti.assign(x, y)
-
- return self.element_wise_writeback_binary(assign_renamed, val)
+ return ops.assign(x, y)
- def empty_copy(self):
- """
- Nested structs and matrices need to be recursively handled.
- """
- struct = Struct.empty(self.keys)
- for k, v in self.items:
- if isinstance(v, (Struct, Matrix)):
- struct.entries[k] = v.empty_copy()
- return struct
-
- def copy(self):
- ret = self.empty_copy()
- ret.entries = copy.copy(self.entries)
- return ret
-
- @taichi_scope
- def variable(self):
- ret = self.copy()
- ret.entries = {
- k: impl.expr_init(v) if isinstance(v,
- (numbers.Number,
- expr.Expr)) else v.variable()
- for k, v in ret.items
- }
- return ret
+ return self._element_wise_writeback_binary(assign_renamed, val)
def __len__(self):
"""Get the number of entries in a custom struct"""
@@ -256,13 +211,11 @@ def __str__(self):
item_str = ", ".join(
[str(k) + "=" + str(v) for k, v in self.items])
return f''
- else:
- return str(self.to_dict())
+ return str(self.to_dict())
def __repr__(self):
return str(self.to_dict())
- @python_scope
def to_dict(self):
"""Converts the Struct to a dictionary.
@@ -271,19 +224,11 @@ def to_dict(self):
Returns:
Dict: The result dictionary.
"""
- return self.entries
-
- @classmethod
- def empty(cls, entries):
- """Clear the struct and fill None.
-
- Args:
- members (Dict[str, DataType]): the names and data types for struct members.
- Returns:
- :class:`~taichi.lang.struct.Struct`: A :class:`~taichi.lang.struct.Struct` instance filled with None.
-
- """
- return cls({k: None for k in entries})
+ return {
+ k: v.to_dict() if isinstance(v, Struct) else
+ v.to_list() if isinstance(v, Matrix) else v
+ for k, v in self.entries.items()
+ }
@classmethod
@python_scope
@@ -328,23 +273,35 @@ def field(cls,
dim = len(shape)
if layout == Layout.SOA:
for e in field_dict.values():
- ti.root.dense(impl.index_nd(dim),
- shape).place(e, offset=offset)
+ impl.root.dense(impl.index_nd(dim),
+ shape).place(e, offset=offset)
if needs_grad:
for e in field_dict.values():
- ti.root.dense(impl.index_nd(dim),
- shape).place(e.grad, offset=offset)
+ impl.root.dense(impl.index_nd(dim),
+ shape).place(e.grad, offset=offset)
else:
- ti.root.dense(impl.index_nd(dim),
- shape).place(*tuple(field_dict.values()),
- offset=offset)
+ impl.root.dense(impl.index_nd(dim),
+ shape).place(*tuple(field_dict.values()),
+ offset=offset)
if needs_grad:
grads = tuple(e.grad for e in field_dict.values())
- ti.root.dense(impl.index_nd(dim),
- shape).place(*grads, offset=offset)
+ impl.root.dense(impl.index_nd(dim),
+ shape).place(*grads, offset=offset)
return StructField(field_dict, name=name)
+class _IntermediateStruct(Struct):
+ """Intermediate struct class for compiler internal use only.
+
+ Args:
+ entries (Dict[str, Union[Expr, Matrix, Struct]]): keys and values for struct members.
+ """
+ def __init__(self, entries):
+ assert isinstance(entries, dict)
+ self.entries = entries
+ self._register_members()
+
+
class StructField(Field):
"""Taichi struct field with SNode implementation.
Instead of directly contraining Expr entries, the StructField object
@@ -357,80 +314,74 @@ class StructField(Field):
def __init__(self, field_dict, name=None):
# will not call Field initializer
self.field_dict = field_dict
- self._name = name
- self.register_fields()
-
- @property
- def name(self):
- return self._name
+ self.name = name
+ self._register_fields()
@property
def keys(self):
return list(self.field_dict.keys())
@property
- def members(self):
+ def _members(self):
return list(self.field_dict.values())
@property
- def items(self):
+ def _items(self):
return self.field_dict.items()
@staticmethod
- def make_getter(key):
+ def _make_getter(key):
def getter(self):
"""Get an entry from custom struct by name."""
- _taichi_skip_traceback = 1
return self.field_dict[key]
return getter
@staticmethod
- def make_setter(key):
+ def _make_setter(key):
@python_scope
def setter(self, value):
- _taichi_skip_traceback = 1
self.field_dict[key] = value
return setter
- def register_fields(self):
+ def _register_fields(self):
for k in self.keys:
setattr(
StructField, k,
property(
- StructField.make_getter(k),
- StructField.make_setter(k),
+ StructField._make_getter(k),
+ StructField._make_setter(k),
))
- def get_field_members(self):
+ def _get_field_members(self):
"""Get A flattened list of all struct elements.
Returns:
A list of struct elements.
"""
field_members = []
- for m in self.members:
+ for m in self._members:
assert isinstance(m, Field)
- field_members += m.get_field_members()
+ field_members += m._get_field_members()
return field_members
@property
- def snode(self):
+ def _snode(self):
"""Gets representative SNode for info purposes.
Returns:
SNode: Representative SNode (SNode of first field member).
"""
- return self.members[0].snode
+ return self._members[0]._snode
- def loop_range(self):
+ def _loop_range(self):
"""Gets representative field member for loop range info.
Returns:
taichi_core.Expr: Representative (first) field member.
"""
- return self.members[0].loop_range()
+ return self._members[0]._loop_range()
@python_scope
def copy_from(self, other):
@@ -453,12 +404,12 @@ def fill(self, val):
Args:
val (Union[int, float]): Value to fill.
"""
- for v in self.members:
+ for v in self._members:
v.fill(val)
- def initialize_host_accessors(self):
- for v in self.members:
- v.initialize_host_accessors()
+ def _initialize_host_accessors(self):
+ for v in self._members:
+ v._initialize_host_accessors()
def get_member_field(self, key):
"""Creates a ScalarField using a specific field member. Only used for quant.
@@ -473,12 +424,12 @@ def get_member_field(self, key):
@python_scope
def from_numpy(self, array_dict):
- for k, v in self.items:
+ for k, v in self._items:
v.from_numpy(array_dict[k])
@python_scope
def from_torch(self, array_dict):
- for k, v in self.items:
+ for k, v in self._items:
v.from_torch(array_dict[k])
@python_scope
@@ -490,7 +441,7 @@ def to_numpy(self):
Returns:
Dict[str, Union[numpy.ndarray, Dict]]: The result NumPy array.
"""
- return {k: v.to_numpy() for k, v in self.items}
+ return {k: v.to_numpy() for k, v in self._items}
@python_scope
def to_torch(self, device=None):
@@ -502,21 +453,21 @@ def to_torch(self, device=None):
Returns:
Dict[str, Union[torch.Tensor, Dict]]: The result PyTorch tensor.
"""
- return {k: v.to_torch(device=device) for k, v in self.items}
+ return {k: v.to_torch(device=device) for k, v in self._items}
@python_scope
def __setitem__(self, indices, element):
- self.initialize_host_accessors()
- self[indices].set_entries(element)
+ self._initialize_host_accessors()
+ self[indices]._set_entries(element)
@python_scope
def __getitem__(self, indices):
- self.initialize_host_accessors()
+ self._initialize_host_accessors()
# scalar fields does not instantiate SNodeHostAccess by default
entries = {
- k: v.host_access(self.pad_key(indices))[0] if isinstance(
+ k: v._host_access(self._pad_key(indices))[0] if isinstance(
v, ScalarField) else v[indices]
- for k, v in self.items
+ for k, v in self._items
}
return Struct(entries)
@@ -542,43 +493,43 @@ def __call__(self, *args, **kwargs):
elif len(args) == 1:
# fill a single scalar
if isinstance(args[0], (numbers.Number, expr.Expr)):
- entries = self.scalar_filled(args[0])
+ entries = self.filled_with_scalar(args[0])
else:
- # fill a single vector or matrix
# initialize struct members by dictionary
entries = Struct(args[0])
struct = self.cast(entries)
return struct
- def cast(self, struct, in_place=False):
- if not in_place:
- struct = struct.copy()
+ def cast(self, struct):
# sanity check members
if self.members.keys() != struct.entries.keys():
raise TaichiSyntaxError(
"Incompatible arguments for custom struct members!")
+ entries = {}
for k, dtype in self.members.items():
if isinstance(dtype, CompoundType):
- struct.entries[k] = dtype.cast(struct.entries[k])
+ entries[k] = dtype.cast(struct.entries[k])
else:
if in_python_scope():
v = struct.entries[k]
- struct.entries[k] = int(
- v) if dtype in ti.integer_types else float(v)
+ entries[k] = int(
+ v
+ ) if dtype in primitive_types.integer_types else float(v)
else:
- struct.entries[k] = cast(struct.entries[k], dtype)
- return struct
+ entries[k] = ops.cast(struct.entries[k], dtype)
+ return Struct(entries)
- def empty(self):
- """
- Create an empty instance of the given compound type.
- Nested structs and matrices need to be recursively handled.
- """
- struct = Struct.empty(self.members.keys())
+ def filled_with_scalar(self, value):
+ entries = {}
for k, dtype in self.members.items():
if isinstance(dtype, CompoundType):
- struct.entries[k] = dtype.empty()
- return struct
+ entries[k] = dtype.filled_with_scalar(value)
+ else:
+ entries[k] = value
+ return Struct(entries)
def field(self, **kwargs):
return Struct.field(self.members, **kwargs)
+
+
+__all__ = ["Struct", "StructField"]
diff --git a/python/taichi/lang/tape.py b/python/taichi/lang/tape.py
index 74be931ba3f41..c9101d306cbd6 100644
--- a/python/taichi/lang/tape.py
+++ b/python/taichi/lang/tape.py
@@ -11,7 +11,7 @@ def __enter__(self):
assert not self.entered, "Tape can be entered only once."
self.entered = True
- def __exit__(self, type, value, tb):
+ def __exit__(self, _type, value, tb):
# print('# kernel calls', len(self.calls))
self.runtime.target_tape = None
if self.eval_on_exit:
diff --git a/python/taichi/lang/type_factory_impl.py b/python/taichi/lang/type_factory_impl.py
deleted file mode 100644
index cd26e1433a540..0000000000000
--- a/python/taichi/lang/type_factory_impl.py
+++ /dev/null
@@ -1,54 +0,0 @@
-from taichi.core.util import ti_core as _ti_core
-from taichi.lang import impl
-
-
-class TypeFactory:
- """A Python-side TypeFactory wrapper."""
- def __init__(self):
- self.core = _ti_core.get_type_factory_instance()
-
- def custom_int(self, bits, signed=True, compute_type=None):
- """Generates a custom int type.
-
- Args:
- bits (int): Number of bits.
- signed (bool): Signed or unsigned.
- compute_type (DataType): Type for computation.
-
- Returns:
- DataType: The specified type.
- """
- if compute_type is None:
- compute_type = impl.get_runtime().default_ip
- if isinstance(compute_type, _ti_core.DataType):
- compute_type = compute_type.get_ptr()
- return self.core.get_custom_int_type(bits, signed, compute_type)
-
- def custom_float(self,
- significand_type,
- exponent_type=None,
- compute_type=None,
- scale=1.0):
- """Generates a custom float type.
-
- Args:
- significand_type (DataType): Type of significand.
- exponent_type (DataType): Type of exponent.
- compute_type (DataType): Type for computation.
- scale (float): Scaling factor.
-
- Returns:
- DataType: The specified type.
- """
- if compute_type is None:
- compute_type = impl.get_runtime().default_fp
- if isinstance(compute_type, _ti_core.DataType):
- compute_type = compute_type.get_ptr()
- return self.core.get_custom_float_type(significand_type,
- exponent_type,
- compute_type,
- scale=scale)
-
-
-# Unstable API
-type_factory = TypeFactory()
diff --git a/python/taichi/lang/types.py b/python/taichi/lang/types.py
deleted file mode 100644
index 1299e378b43f5..0000000000000
--- a/python/taichi/lang/types.py
+++ /dev/null
@@ -1,31 +0,0 @@
-import numbers
-
-import taichi.lang.matrix
-from taichi.lang.exception import TaichiSyntaxError
-
-
-class CompoundType:
- def empty(self):
- """
- Create an empty instance of the given compound type.
- """
- raise NotImplementedError
-
- def scalar_filled(self, value):
- instance = self.empty()
- return instance.broadcast_copy(value)
-
- def field(self, **kwargs):
- raise NotImplementedError
-
-
-def matrix(m, n, dtype=None):
- return taichi.lang.matrix.MatrixType(m, n, dtype=dtype)
-
-
-def vector(m, dtype=None):
- return taichi.lang.matrix.MatrixType(m, 1, dtype=dtype)
-
-
-def struct(**kwargs):
- return taichi.lang.struct.StructType(**kwargs)
diff --git a/python/taichi/lang/util.py b/python/taichi/lang/util.py
index 6d9fd2d578d8c..feb18831d91cb 100644
--- a/python/taichi/lang/util.py
+++ b/python/taichi/lang/util.py
@@ -1,11 +1,13 @@
import functools
import os
+import traceback
import numpy as np
-from taichi.core.util import ti_core as _ti_core
+from colorama import Fore, Style
+from taichi._lib import core as _ti_core
from taichi.lang import impl
-
-import taichi as ti
+from taichi.types.primitive_types import (f16, f32, f64, i8, i16, i32, i64, u8,
+ u16, u32, u64)
_has_pytorch = False
@@ -28,10 +30,29 @@ def has_pytorch():
return _has_pytorch
+from distutils.spawn import find_executable
+
+# Taichi itself uses llvm-10.0.0 to compile.
+# There will be some issues compiling CUDA with other clang++ version.
+_clangpp_candidates = ['clang++-10']
+_clangpp_presence = None
+for c in _clangpp_candidates:
+ if find_executable(c) is not None:
+ _clangpp_presence = find_executable(c)
+
+
+def has_clangpp():
+ return _clangpp_presence is not None
+
+
+def get_clangpp():
+ return _clangpp_presence
+
+
def is_taichi_class(rhs):
taichi_class = False
try:
- if rhs.is_taichi_class:
+ if rhs._is_taichi_class:
taichi_class = True
except:
pass
@@ -48,28 +69,29 @@ def to_numpy_type(dt):
DataType: The counterpart data type in numpy.
"""
- if dt == ti.f32:
+ if dt == f32:
return np.float32
- elif dt == ti.f64:
+ if dt == f64:
return np.float64
- elif dt == ti.i32:
+ if dt == i32:
return np.int32
- elif dt == ti.i64:
+ if dt == i64:
return np.int64
- elif dt == ti.i8:
+ if dt == i8:
return np.int8
- elif dt == ti.i16:
+ if dt == i16:
return np.int16
- elif dt == ti.u8:
+ if dt == u8:
return np.uint8
- elif dt == ti.u16:
+ if dt == u16:
return np.uint16
- elif dt == ti.u32:
+ if dt == u32:
return np.uint32
- elif dt == ti.u64:
+ if dt == u64:
return np.uint64
- else:
- assert False
+ if dt == f16:
+ return np.half
+ assert False
def to_pytorch_type(dt):
@@ -82,28 +104,27 @@ def to_pytorch_type(dt):
DataType: The counterpart data type in torch.
"""
- if dt == ti.f32:
+ # pylint: disable=E1101
+ if dt == f32:
return torch.float32
- elif dt == ti.f64:
+ if dt == f64:
return torch.float64
- elif dt == ti.i32:
+ if dt == i32:
return torch.int32
- elif dt == ti.i64:
+ if dt == i64:
return torch.int64
- elif dt == ti.i8:
+ if dt == i8:
return torch.int8
- elif dt == ti.i16:
+ if dt == i16:
return torch.int16
- elif dt == ti.u8:
+ if dt == u8:
return torch.uint8
- elif dt == ti.u16:
- return torch.uint16
- elif dt == ti.u32:
- return torch.uint32
- elif dt == ti.u64:
- return torch.uint64
- else:
- assert False
+ if dt == f16:
+ return torch.float16
+ if dt in (u16, u32, u64):
+ raise RuntimeError(
+ f'PyTorch doesn\'t support {dt.to_string()} data type.')
+ assert False
def to_taichi_type(dt):
@@ -120,63 +141,63 @@ def to_taichi_type(dt):
return dt
if dt == np.float32:
- return ti.f32
- elif dt == np.float64:
- return ti.f64
- elif dt == np.int32:
- return ti.i32
- elif dt == np.int64:
- return ti.i64
- elif dt == np.int8:
- return ti.i8
- elif dt == np.int16:
- return ti.i16
- elif dt == np.uint8:
- return ti.u8
- elif dt == np.uint16:
- return ti.u16
- elif dt == np.uint32:
- return ti.u32
- elif dt == np.uint64:
- return ti.u64
+ return f32
+ if dt == np.float64:
+ return f64
+ if dt == np.int32:
+ return i32
+ if dt == np.int64:
+ return i64
+ if dt == np.int8:
+ return i8
+ if dt == np.int16:
+ return i16
+ if dt == np.uint8:
+ return u8
+ if dt == np.uint16:
+ return u16
+ if dt == np.uint32:
+ return u32
+ if dt == np.uint64:
+ return u64
+ if dt == np.half:
+ return f16
if has_pytorch():
+ # pylint: disable=E1101
if dt == torch.float32:
- return ti.f32
- elif dt == torch.float64:
- return ti.f64
- elif dt == torch.int32:
- return ti.i32
- elif dt == torch.int64:
- return ti.i64
- elif dt == torch.int8:
- return ti.i8
- elif dt == torch.int16:
- return ti.i16
- elif dt == torch.uint8:
- return ti.u8
- elif dt == torch.uint16:
- return ti.u16
- elif dt == torch.uint32:
- return ti.u32
- elif dt == torch.uint64:
- return ti.u64
-
- raise AssertionError("Unknown type {}".format(dt))
+ return f32
+ if dt == torch.float64:
+ return f64
+ if dt == torch.int32:
+ return i32
+ if dt == torch.int64:
+ return i64
+ if dt == torch.int8:
+ return i8
+ if dt == torch.int16:
+ return i16
+ if dt == torch.uint8:
+ return u8
+ if dt == torch.float16:
+ return f16
+ if dt in (u16, u32, u64):
+ raise RuntimeError(
+ f'PyTorch doesn\'t support {dt.to_string()} data type.')
+
+ raise AssertionError(f"Unknown type {dt}")
def cook_dtype(dtype):
- _taichi_skip_traceback = 1
if isinstance(dtype, _ti_core.DataType):
return dtype
- elif isinstance(dtype, _ti_core.Type):
+ if isinstance(dtype, _ti_core.Type):
return _ti_core.DataType(dtype)
- elif dtype is float:
+ if dtype is float:
return impl.get_runtime().default_fp
- elif dtype is int:
+ if dtype is int:
return impl.get_runtime().default_ip
- else:
- raise ValueError(f'Invalid data type {dtype}')
+ raise ValueError(f'Invalid data type {dtype}')
def in_taichi_scope():
@@ -190,7 +211,6 @@ def in_python_scope():
def taichi_scope(func):
@functools.wraps(func)
def wrapped(*args, **kwargs):
- _taichi_skip_traceback = 1
assert in_taichi_scope(), \
f'{func.__name__} cannot be called in Python-scope'
return func(*args, **kwargs)
@@ -201,9 +221,32 @@ def wrapped(*args, **kwargs):
def python_scope(func):
@functools.wraps(func)
def wrapped(*args, **kwargs):
- _taichi_skip_traceback = 1
assert in_python_scope(), \
f'{func.__name__} cannot be called in Taichi-scope'
return func(*args, **kwargs)
return wrapped
+
+
+def warning(msg, warning_type=UserWarning, stacklevel=1, print_stack=True):
+ """Print a warning message. Note that the builtin `warnings` module is
+ unreliable since it may be suppressed by other packages such as IPython.
+
+ Args:
+ msg (str): message to print.
+ warning_type (Warning): type of warning.
+ stacklevel (int): warning stack level from the caller.
+ print_stack (bool): whether to print the stack
+ """
+ msg = f'{warning_type.__name__}: {msg}'
+ if print_stack:
+ msg += f'\n{get_traceback(stacklevel)}'
+ print(Fore.YELLOW + Style.BRIGHT + msg + Style.RESET_ALL)
+
+
+def get_traceback(stacklevel=1):
+ s = traceback.extract_stack()[:-1 - stacklevel]
+ return ''.join(traceback.format_list(s))
+
+
+__all__ = []
diff --git a/python/taichi/linalg/__init__.py b/python/taichi/linalg/__init__.py
index a56d77fb4846d..469762c2ee0cd 100644
--- a/python/taichi/linalg/__init__.py
+++ b/python/taichi/linalg/__init__.py
@@ -1,3 +1,2 @@
-from taichi.linalg.sparse_matrix import (SparseMatrix, SparseMatrixBuilder,
- sparse_matrix_builder)
+from taichi.linalg.sparse_matrix import *
from taichi.linalg.sparse_solver import SparseSolver
diff --git a/python/taichi/linalg/sparse_matrix.py b/python/taichi/linalg/sparse_matrix.py
index 9e0c21a935d9c..928aec91d1382 100644
--- a/python/taichi/linalg/sparse_matrix.py
+++ b/python/taichi/linalg/sparse_matrix.py
@@ -1,7 +1,8 @@
import numpy as np
-from taichi.core.util import ti_core as _ti_core
from taichi.lang.field import Field
-from taichi.type.primitive_types import f32
+from taichi.lang.impl import get_runtime
+from taichi.lang.util import warning
+from taichi.types import annotations, f32
class SparseMatrix:
@@ -18,7 +19,7 @@ def __init__(self, n=None, m=None, sm=None, dtype=f32):
if sm is None:
self.n = n
self.m = m if m else n
- self.matrix = _ti_core.create_sparse_matrix(n, m)
+ self.matrix = get_runtime().prog.create_sparse_matrix(n, m)
else:
self.n = sm.num_rows()
self.m = sm.num_cols()
@@ -55,11 +56,13 @@ def __mul__(self, other):
if isinstance(other, float):
sm = self.matrix * other
return SparseMatrix(sm=sm)
- elif isinstance(other, SparseMatrix):
+ if isinstance(other, SparseMatrix):
assert self.n == other.n and self.m == other.m, f"Dimension mismatch between sparse matrices ({self.n}, {self.m}) and ({other.n}, {other.m})"
sm = self.matrix * other.matrix
return SparseMatrix(sm=sm)
+ return None
+
def __rmul__(self, other):
"""Right scalar multiplication for sparse matrix.
@@ -72,6 +75,8 @@ def __rmul__(self, other):
sm = other * self.matrix
return SparseMatrix(sm=sm)
+ return None
+
def transpose(self):
"""Sparse Matrix transpose.
@@ -93,16 +98,15 @@ def __matmul__(self, other):
assert self.m == other.n, f"Dimension mismatch between sparse matrices ({self.n}, {self.m}) and ({other.n}, {other.m})"
sm = self.matrix.matmul(other.matrix)
return SparseMatrix(sm=sm)
- elif isinstance(other, Field):
+ if isinstance(other, Field):
assert self.m == other.shape[
0], f"Dimension mismatch between sparse matrix ({self.n}, {self.m}) and vector ({other.shape})"
return self.matrix.mat_vec_mul(other.to_numpy())
- elif isinstance(other, np.ndarray):
+ if isinstance(other, np.ndarray):
assert self.m == other.shape[
0], f"Dimension mismatch between sparse matrix ({self.n}, {self.m}) and vector ({other.shape})"
return self.matrix.mat_vec_mul(other)
- else:
- assert False, f"Sparse matrix-matrix/vector multiplication does not support {type(other)} for now. Supported types are SparseMatrix, ti.field, and numpy.ndarray."
+ assert False, f"Sparse matrix-matrix/vector multiplication does not support {type(other)} for now. Supported types are SparseMatrix, ti.field, and numpy.ndarray."
def __getitem__(self, indices):
return self.matrix.get_element(indices[0], indices[1])
@@ -117,6 +121,10 @@ def __str__(self):
def __repr__(self):
return self.matrix.to_string()
+ def shape(self):
+ """The shape of the sparse matrix."""
+ return (self.n, self.m)
+
class SparseMatrixBuilder:
"""A python wrap around sparse matrix builder.
@@ -135,11 +143,12 @@ def __init__(self,
dtype=f32):
self.num_rows = num_rows
self.num_cols = num_cols if num_cols else num_rows
+ self.dtype = dtype
if num_rows is not None:
- self.ptr = _ti_core.create_sparse_matrix_builder(
- num_rows, num_cols, max_num_triplets)
+ self.ptr = get_runtime().prog.create_sparse_matrix_builder(
+ num_rows, num_cols, max_num_triplets, dtype)
- def get_addr(self):
+ def _get_addr(self):
"""Get the address of the sparse matrix"""
return self.ptr.get_addr()
@@ -147,11 +156,18 @@ def print_triplets(self):
"""Print the triplets stored in the builder"""
self.ptr.print_triplets()
- def build(self, dtype=f32, format='CSR'):
+ def build(self, dtype=f32, _format='CSR'):
"""Create a sparse matrix using the triplets"""
sm = self.ptr.build()
return SparseMatrix(sm=sm)
-sparse_matrix_builder = SparseMatrixBuilder
-# Alias for :class:`SparseMatrixBuilder`
+# TODO: remove this in 1.0 release
+class sparse_matrix_builder(annotations.sparse_matrix_builder):
+ def __init__(self):
+ warning(
+ 'ti.linalg.sparse_matrix_builder is deprecated. Please use ti.types.sparse_matrix_builder instead.',
+ DeprecationWarning)
+
+
+__all__ = ['SparseMatrix', 'SparseMatrixBuilder', 'sparse_matrix_builder']
diff --git a/python/taichi/linalg/sparse_solver.py b/python/taichi/linalg/sparse_solver.py
index 4b886fe96fa3d..67fca2e5f2ab3 100644
--- a/python/taichi/linalg/sparse_solver.py
+++ b/python/taichi/linalg/sparse_solver.py
@@ -1,8 +1,9 @@
import numpy as np
import taichi.lang
-from taichi.core.util import ti_core as _ti_core
+from taichi._lib import core as _ti_core
+from taichi.lang.field import Field
from taichi.linalg import SparseMatrix
-from taichi.type.primitive_types import f32
+from taichi.types.primitive_types import f32
class SparseSolver:
@@ -20,12 +21,13 @@ def __init__(self, dtype=f32, solver_type="LLT", ordering="AMD"):
if solver_type in solver_type_list and ordering in solver_ordering:
taichi_arch = taichi.lang.impl.get_runtime().prog.config.arch
assert taichi_arch == _ti_core.Arch.x64 or taichi_arch == _ti_core.Arch.arm64, "SparseSolver only supports CPU for now."
- self.solver = _ti_core.make_sparse_solver(solver_type, ordering)
+ self.solver = _ti_core.make_sparse_solver(dtype, solver_type,
+ ordering)
else:
assert False, f"The solver type {solver_type} with {ordering} is not supported for now. Only {solver_type_list} with {solver_ordering} are supported."
@staticmethod
- def type_assert(sparse_matrix):
+ def _type_assert(sparse_matrix):
assert False, f"The parameter type: {type(sparse_matrix)} is not supported in linear solvers for now."
def compute(self, sparse_matrix):
@@ -37,7 +39,7 @@ def compute(self, sparse_matrix):
if isinstance(sparse_matrix, SparseMatrix):
self.solver.compute(sparse_matrix.matrix)
else:
- self.type_assert(sparse_matrix)
+ self._type_assert(sparse_matrix)
def analyze_pattern(self, sparse_matrix):
"""Reorder the nonzero elements of the matrix, such that the factorization step creates less fill-in.
@@ -48,7 +50,7 @@ def analyze_pattern(self, sparse_matrix):
if isinstance(sparse_matrix, SparseMatrix):
self.solver.analyze_pattern(sparse_matrix.matrix)
else:
- self.type_assert(sparse_matrix)
+ self._type_assert(sparse_matrix)
def factorize(self, sparse_matrix):
"""Do the factorization step
@@ -59,7 +61,7 @@ def factorize(self, sparse_matrix):
if isinstance(sparse_matrix, SparseMatrix):
self.solver.factorize(sparse_matrix.matrix)
else:
- self.type_assert(sparse_matrix)
+ self._type_assert(sparse_matrix)
def solve(self, b):
"""Computes the solution of the linear systems.
@@ -69,12 +71,11 @@ def solve(self, b):
Returns:
numpy.array: The solution of linear systems.
"""
- if isinstance(b, taichi.lang.Field):
+ if isinstance(b, Field):
return self.solver.solve(b.to_numpy())
- elif isinstance(b, np.ndarray):
+ if isinstance(b, np.ndarray):
return self.solver.solve(b)
- else:
- assert False, f"The parameter type: {type(b)} is not supported in linear solvers for now."
+ assert False, f"The parameter type: {type(b)} is not supported in linear solvers for now."
def info(self):
"""Check if the linear systems are solved successfully.
diff --git a/python/taichi/misc/__init__.py b/python/taichi/misc/__init__.py
deleted file mode 100644
index b629d01d1ee2c..0000000000000
--- a/python/taichi/misc/__init__.py
+++ /dev/null
@@ -1,7 +0,0 @@
-from .error import *
-from .gui import *
-from .image import *
-from .task import Task
-from .util import *
-
-__all__ = [s for s in dir() if not s.startswith('_')]
diff --git a/python/taichi/misc/error.py b/python/taichi/misc/error.py
deleted file mode 100644
index 7f88c853c2b58..0000000000000
--- a/python/taichi/misc/error.py
+++ /dev/null
@@ -1,48 +0,0 @@
-import functools
-import sys
-import traceback
-
-from colorama import Fore, Style
-
-
-def enable_excepthook():
- def excepthook(exctype, value, tb):
- skip = 0
- back = 4
- forward = 2
- bar = f'{Fore.LIGHTBLACK_EX}{"-"*44}{Fore.RESET}'
- print(
- f'{Fore.LIGHTBLACK_EX}========== Taichi Stack Traceback =========={Fore.RESET}'
- )
- for frame, lineno in traceback.walk_tb(tb):
- name = frame.f_code.co_name
- filename = frame.f_code.co_filename
- if '_taichi_skip_traceback' in frame.f_locals:
- skip = frame.f_locals['_taichi_skip_traceback']
- if skip > 0:
- skip -= 1
- continue
- print(
- f'In {Fore.LIGHTYELLOW_EX}{name}{Fore.RESET}() at {Fore.LIGHTMAGENTA_EX}{filename}{Fore.RESET}:{Fore.LIGHTCYAN_EX}{lineno}{Fore.RESET}:\n{bar}'
- )
- with open(filename) as f:
- lines = [''] + f.readlines()
- if lines[lineno][-1] == '\n':
- lines[lineno] = lines[lineno][:-1]
- lines[lineno] = f'{Fore.LIGHTRED_EX}' + lines[
- lineno] + f' {Fore.LIGHTYELLOW_EX}<--{Fore.LIGHTBLACK_EX}\n'
- line = ''.join(lines[max(1, lineno -
- back):min(len(lines), lineno +
- forward + 1)])
- if line[-1] != '\n':
- line += '\n'
- print(f'{Fore.LIGHTWHITE_EX}{line}{bar}')
- value = str(value)
- if len(value):
- value = f': {Fore.LIGHTRED_EX}{value}'
- print(
- f'{Fore.LIGHTGREEN_EX}{exctype.__name__}{Fore.RESET}{value}{Fore.RESET}'
- )
-
- if sys.excepthook is not excepthook:
- sys.excepthook = excepthook
diff --git a/python/taichi/misc/gui.py b/python/taichi/misc/gui.py
deleted file mode 100644
index e921f3006043c..0000000000000
--- a/python/taichi/misc/gui.py
+++ /dev/null
@@ -1,859 +0,0 @@
-import math
-import numbers
-import os
-
-import numpy as np
-import taichi.lang
-from taichi.core import ti_core as _ti_core
-from taichi.lang.field import Field, ScalarField
-
-import taichi as ti
-
-from .util import core_veci, deprecated
-
-
-class GUI:
- """Taichi Graphical User Interface class.
-
- Args:
- name (str, optional): The name of the GUI to be constructed.
- Default is 'Taichi'.
- res (Union[int, List[int]], optional): The resolution of created
- GUI. Default is 512*512. If `res` is scalar, then width will be equal to height.
- background_color (int, optional): The background color of created GUI.
- Default is 0x000000.
- show_gui (bool, optional): Specify whether to render the GUI. Default is True.
- fullscreen (bool, optional): Specify whether to render the GUI in
- fullscreen mode. Default is False.
- fast_gui (bool, optional): Specify whether to use fast gui mode of
- Taichi. Default is False.
-
- Returns:
- :class:`~taichi.misc.gui.GUI` :The created taichi GUI object.
-
- """
- class Event:
- pass
-
- # Event keys
- SHIFT = 'Shift'
- ALT = 'Alt'
- CTRL = 'Control'
- ESCAPE = 'Escape'
- RETURN = 'Return'
- TAB = 'Tab'
- BACKSPACE = 'BackSpace'
- SPACE = ' '
- UP = 'Up'
- DOWN = 'Down'
- LEFT = 'Left'
- RIGHT = 'Right'
- CAPSLOCK = 'Caps_Lock'
- LMB = 'LMB'
- MMB = 'MMB'
- RMB = 'RMB'
- EXIT = 'WMClose'
- WHEEL = 'Wheel'
- MOVE = 'Motion'
-
- # Event types
- MOTION = _ti_core.KeyEvent.EType.Move
- PRESS = _ti_core.KeyEvent.EType.Press
- RELEASE = _ti_core.KeyEvent.EType.Release
-
- def __init__(self,
- name='Taichi',
- res=512,
- background_color=0x0,
- show_gui=True,
- fullscreen=False,
- fast_gui=False):
- show_gui = self.get_bool_environ('TI_GUI_SHOW', show_gui)
- fullscreen = self.get_bool_environ('TI_GUI_FULLSCREEN', fullscreen)
- fast_gui = self.get_bool_environ('TI_GUI_FAST', fast_gui)
-
- self.name = name
- if isinstance(res, numbers.Number):
- res = (res, res)
- self.res = res
- self.fast_gui = fast_gui
- if fast_gui:
- self.img = np.ascontiguousarray(
- np.zeros(self.res[0] * self.res[1], dtype=np.uint32))
- fast_buf = self.img.ctypes.data
- else:
- # The GUI canvas uses RGBA for storage, therefore we need NxMx4 for an image.
- self.img = np.ascontiguousarray(
- np.zeros(self.res + (4, ), np.float32))
- fast_buf = 0
- self.core = _ti_core.GUI(name, core_veci(*res), show_gui, fullscreen,
- fast_gui, fast_buf)
- self.canvas = self.core.get_canvas()
- self.background_color = background_color
- self.key_pressed = set()
- self.event = None
- self.frame = 0
- self.clear()
-
- def __enter__(self):
- return self
-
- def __exit__(self, type, val, tb):
- self.close()
-
- def __del__(self):
- self.close()
-
- def close(self):
- self.core = None # dereference to call GUI::~GUI()
-
- ## Widget system
-
- class WidgetValue:
- def __init__(self, gui, wid):
- self.gui = gui
- self.wid = wid
-
- @property
- def value(self):
- return self.gui.core.get_widget_value(self.wid)
-
- @value.setter
- def value(self, value):
- self.gui.core.set_widget_value(self.wid, value)
-
- def get_bool_environ(self, key, default):
- """Get an environment variable and cast to bool.
- Args:
- key (str): The environment variable key.
- default (bool): The default value.
- Return:
- The environment variable value cast to bool. If the value is not found, directly return argument 'default'.
- """
- if key not in os.environ:
- return default
- return bool(int(os.environ[key]))
-
- def slider(self, text, minimum, maximum, step=1):
- """Create a slider object on canvas to be manipulated with.
-
- Args:
- text (str): The title of slider.
- minimum (Number): The minimum value of slider.
- maximum (Number): The maximum value of slider.
- step (Number, optional): The changing step of slider. Default is 1.
-
- Return:
- :class:`~taichi.misc.gui.GUI.WidgetValue` :The created slider object.
-
- """
- wid = self.core.make_slider(text, minimum, minimum, maximum, step)
- return GUI.WidgetValue(self, wid)
-
- def label(self, text):
- """Create a label object on canvas.
-
- Args:
- text (str): The title of label.
-
- Return:
- :class:`~taichi.misc.gui.GUI.WidgetValue` :The created label object.
-
- """
- wid = self.core.make_label(text, 0)
- return GUI.WidgetValue(self, wid)
-
- def button(self, text, event_name=None):
- """Create a button object on canvas to be manipulated with.
-
- Args:
- text (str): The title of button.
- event_name (str, optional): The event name associated with button.
- Default is WidgetButton_{text}
-
- Return:
- The event name associated with created button.
-
- """
- event_name = event_name or f'WidgetButton_{text}'
- self.core.make_button(text, event_name)
- return event_name
-
- ## Drawing system
-
- def clear(self, color=None):
- """Clear the canvas with the color provided.
-
- Args:
- color (int, optional): Specify the color to clear the canvas. Default
- is the background color of GUI.
-
- """
- if color is None:
- color = self.background_color
- self.canvas.clear(color)
-
- def cook_image(self, img):
- if img.dtype in [np.uint8, np.uint16, np.uint32, np.uint64]:
- img = img.astype(np.float32) * (1 / np.iinfo(img.dtype).max)
- elif img.dtype in [np.float32, np.float64]:
- img = img.astype(np.float32)
- else:
- raise ValueError(
- f'Data type {img.dtype} not supported in GUI.set_image')
-
- if len(img.shape) == 2:
- img = img[..., None]
-
- if img.shape[2] == 1:
- img = img + np.zeros((1, 1, 4), np.float32)
- if img.shape[2] == 3:
- zeros = np.zeros((img.shape[0], img.shape[1], 1), np.float32)
- img = np.concatenate([img, zeros], axis=2)
- if img.shape[2] == 2:
- zeros = np.zeros((img.shape[0], img.shape[1], 2), np.float32)
- img = np.concatenate([img, zeros], axis=2)
-
- assert img.shape[2] == 4, "Image must be grayscale, RG, RGB or RGBA"
-
- res = img.shape[:2]
- assert res == self.res, "Image resolution does not match GUI resolution"
- return np.ascontiguousarray(img)
-
- def get_image(self):
- """Get the image data.
-
- Returns:
- :class:`numpy.array` :The image data in numpy contiguous array type.
-
- """
- self.img = np.ascontiguousarray(self.img)
- self.core.get_img(self.img.ctypes.data)
- return self.img
-
- def set_image(self, img):
- """Draw an image on canvas.
-
- Args:
- img (Union[ti.field, numpy.array]): The color array representing the
- image to be drawn. Support greyscale, RG, RGB, and RGBA color
- representations. Its shape must match GUI resolution.
-
- """
-
- if self.fast_gui:
- assert isinstance(img, taichi.lang.matrix.MatrixField), \
- "Only ti.Vector.field is supported in GUI.set_image when fast_gui=True"
- assert img.shape == self.res, \
- "Image resolution does not match GUI resolution"
- assert img.n in [3, 4] and img.m == 1, \
- "Only RGB images are supported in GUI.set_image when fast_gui=True"
- assert img.dtype in [ti.f32, ti.f64, ti.u8], \
- "Only f32, f64, u8 are supported in GUI.set_image when fast_gui=True"
-
- taichi.lang.meta.vector_to_fast_image(img, self.img)
- return
-
- if isinstance(img, ScalarField):
- if _ti_core.is_integral(img.dtype) or len(img.shape) != 2:
- # Images of uint is not optimized by xxx_to_image
- self.img = self.cook_image(img.to_numpy())
- else:
- # Type matched! We can use an optimized copy kernel.
- assert img.shape \
- == self.res, "Image resolution does not match GUI resolution"
- taichi.lang.meta.tensor_to_image(img, self.img)
- ti.sync()
-
- elif isinstance(img, taichi.lang.matrix.MatrixField):
- if _ti_core.is_integral(img.dtype):
- self.img = self.cook_image(img.to_numpy())
- else:
- # Type matched! We can use an optimized copy kernel.
- assert img.shape == self.res, \
- "Image resolution does not match GUI resolution"
- assert img.n in [2, 3, 4] and img.m == 1, \
- "Only greyscale, RG, RGB or RGBA images are supported in GUI.set_image"
-
- taichi.lang.meta.vector_to_image(img, self.img)
- ti.sync()
-
- elif isinstance(img, np.ndarray):
- self.img = self.cook_image(img)
-
- else:
- raise ValueError(
- f"GUI.set_image only takes a Taichi field or NumPy array, not {type(img)}"
- )
-
- self.core.set_img(self.img.ctypes.data)
-
- def circle(self, pos, color=0xFFFFFF, radius=1):
- """Draw a single circle on canvas.
-
- Args:
- pos (Union[List[int], numpy.array]): The position of the circle.
- color (int, Optional): The color of the circle. Default is 0xFFFFFF.
- radius (Number, Optional): The radius of the circle. Default is 1.
-
- """
- self.canvas.circle_single(pos[0], pos[1], color, radius)
-
- def circles(self,
- pos,
- radius=1,
- color=0xFFFFFF,
- palette=None,
- palette_indices=None):
- """Draw a list of circles on canvas.
-
- Args:
- pos (numpy.array): The positions of the circles.
- radius (Number, optional): The radius of the circles. Default is 1.
- color (int, optional): The color of the circles. Default is 0xFFFFFF.
- palette (list[int], optional): The List of colors from which to
- choose to draw. Default is None.
- palette_indices (Union[list[int], ti.field, numpy.array], optional):
- The List of indices that choose color from palette for each
- circle. Shape must match pos. Default is None.
-
- """
- n = pos.shape[0]
- if len(pos.shape) == 3:
- assert pos.shape[2] == 1
- pos = pos[:, :, 0]
-
- assert pos.shape == (n, 2)
- pos = np.ascontiguousarray(pos.astype(np.float32))
- # Note: do not use "pos = int(pos.ctypes.data)" here
- # Otherwise pos will get garbage collected by Python
- # and the pointer to its data becomes invalid
- pos_ptr = int(pos.ctypes.data)
-
- if isinstance(color, np.ndarray):
- assert color.shape == (n, )
- color = np.ascontiguousarray(color.astype(np.uint32))
- color_array = int(color.ctypes.data)
- color_single = 0
- elif isinstance(color, int):
- color_array = 0
- color_single = color
- else:
- raise ValueError(
- 'Color must be an ndarray or int (e.g., 0x956333)')
-
- if palette is not None:
- assert palette_indices is not None, 'palette must be used together with palette_indices'
-
- if isinstance(palette_indices, Field):
- ind_int = palette_indices.to_numpy().astype(np.uint32)
- elif isinstance(palette_indices, list) or isinstance(
- palette_indices, np.ndarray):
- ind_int = np.array(palette_indices).astype(np.uint32)
- else:
- try:
- ind_int = np.array(palette_indices)
- except:
- raise TypeError(
- 'palette_indices must be a type that can be converted to numpy.ndarray'
- )
-
- assert issubclass(
- ind_int.dtype.type,
- np.integer), 'palette_indices must be an integer array'
- assert ind_int.shape == (
- n,
- ), 'palette_indices must be in 1-d shape with shape (num_particles, )'
- assert min(
- ind_int
- ) >= 0, 'the min of palette_indices must not be less than zero'
- assert max(ind_int) < len(
- palette
- ), 'the max of palette_indices must not exceed the length of palette'
- color_array = np.array(palette, dtype=np.uint32)[ind_int]
- color_array = np.ascontiguousarray(color_array)
- color_array = color_array.ctypes.data
-
- if isinstance(radius, np.ndarray):
- assert radius.shape == (n, )
- radius = np.ascontiguousarray(radius.astype(np.float32))
- radius_array = int(radius.ctypes.data)
- radius_single = 0
- elif isinstance(radius, numbers.Number):
- radius_array = 0
- radius_single = radius
- else:
- raise ValueError('Radius must be an ndarray or float (e.g., 0.4)')
-
- self.canvas.circles_batched(n, pos_ptr, color_single, color_array,
- radius_single, radius_array)
-
- def triangles(self, a, b, c, color=0xFFFFFF):
- """Draw a list of triangles on canvas.
-
- Args:
- a (numpy.array): The positions of the first points of triangles.
- b (numpy.array): The positions of the second points of triangles.
- c (numpy.array): The positions of the thrid points of triangles.
- color (Union[int, numpy.array], optional): The color or colors of triangles.
- Can be either a single color or a list of colors whose shape matches
- the shape of a & b & c. Default is 0xFFFFFF.
-
- """
- assert a.shape == b.shape
- assert a.shape == c.shape
- n = a.shape[0]
- if len(a.shape) == 3:
- assert a.shape[2] == 1
- a = a[:, :, 0]
- b = b[:, :, 0]
- c = c[:, :, 0]
-
- assert a.shape == (n, 2)
- a = np.ascontiguousarray(a.astype(np.float32))
- b = np.ascontiguousarray(b.astype(np.float32))
- c = np.ascontiguousarray(c.astype(np.float32))
- # Note: do not use "a = int(a.ctypes.data)" here
- # Otherwise a will get garbage collected by Python
- # and the pointer to its data becomes invalid
- a_ptr = int(a.ctypes.data)
- b_ptr = int(b.ctypes.data)
- c_ptr = int(c.ctypes.data)
-
- if isinstance(color, np.ndarray):
- assert color.shape == (n, )
- color = np.ascontiguousarray(color.astype(np.uint32))
- color_array = int(color.ctypes.data)
- color_single = 0
- elif isinstance(color, int):
- color_array = 0
- color_single = color
- else:
- raise ValueError(
- '"color" must be an ndarray or int (e.g., 0x956333)')
-
- self.canvas.triangles_batched(n, a_ptr, b_ptr, c_ptr, color_single,
- color_array)
-
- def triangle(self, a, b, c, color=0xFFFFFF):
- """Draw a single triangle on canvas.
-
- Args:
- a (List[Number]): The position of the first point of triangle. Shape must be 2.
- b (List[Number]): The position of the second point of triangle. Shape must be 2.
- c (List[Number]): The position of the third point of triangle. Shape must be 2.
- color (int, optional): The color of the triangle. Default is 0xFFFFFF.
-
- """
- self.canvas.triangle_single(a[0], a[1], b[0], b[1], c[0], c[1], color)
-
- def lines(self, begin, end, radius=1, color=0xFFFFFF):
- """Draw a list of lines on canvas.
-
- Args:
- begin (numpy.array): The positions of one end of lines.
- end (numpy.array): The positions of the other end of lines.
- radius (Union[Number, numpy.array], optional): The width of lines.
- Can be either a single width or a list of width whose shape matches
- the shape of begin & end. Default is 1.
- color (Union[int, numpy.array], optional): The color or colors of lines.
- Can be either a single color or a list of colors whose shape matches
- the shape of begin & end. Default is 0xFFFFFF.
-
- """
- assert begin.shape == end.shape
- n = begin.shape[0]
- if len(begin.shape) == 3:
- assert begin.shape[2] == 1
- begin = begin[:, :, 0]
- end = end[:, :, 0]
-
- assert begin.shape == (n, 2)
- begin = np.ascontiguousarray(begin.astype(np.float32))
- end = np.ascontiguousarray(end.astype(np.float32))
- # Note: do not use "begin = int(begin.ctypes.data)" here
- # Otherwise begin will get garbage collected by Python
- # and the pointer to its data becomes invalid
- begin_ptr = int(begin.ctypes.data)
- end_ptr = int(end.ctypes.data)
-
- if isinstance(color, np.ndarray):
- assert color.shape == (n, )
- color = np.ascontiguousarray(color.astype(np.uint32))
- color_array = int(color.ctypes.data)
- color_single = 0
- elif isinstance(color, int):
- color_array = 0
- color_single = color
- else:
- raise ValueError(
- 'Color must be an ndarray or int (e.g., 0x956333)')
-
- if isinstance(radius, np.ndarray):
- assert radius.shape == (n, )
- radius = np.ascontiguousarray(radius.astype(np.float32))
- radius_array = int(radius.ctypes.data)
- radius_single = 0
- elif isinstance(radius, numbers.Number):
- radius_array = 0
- radius_single = radius
- else:
- raise ValueError('Radius must be an ndarray or float (e.g., 0.4)')
-
- self.canvas.paths_batched(n, begin_ptr, end_ptr, color_single,
- color_array, radius_single, radius_array)
-
- def line(self, begin, end, radius=1, color=0xFFFFFF):
- """Draw a single line on canvas.
-
- Args:
- begin (List[Number]): The position of one end of line. Shape must be 2.
- end (List[Number]): The position of the other end of line. Shape must be 2.
- radius (Number, optional): The width of line. Default is 1.
- color (int, optional): The color of line. Default is 0xFFFFFF.
-
- """
- self.canvas.path_single(begin[0], begin[1], end[0], end[1], color,
- radius)
-
- @staticmethod
- def _arrow_to_lines(orig, major, tip_scale=0.2, angle=45):
- angle = math.radians(180 - angle)
- c, s = math.cos(angle), math.sin(angle)
- minor1 = np.array([
- major[:, 0] * c - major[:, 1] * s,
- major[:, 0] * s + major[:, 1] * c
- ]).swapaxes(0, 1)
- minor2 = np.array([
- major[:, 0] * c + major[:, 1] * s,
- -major[:, 0] * s + major[:, 1] * c
- ]).swapaxes(0, 1)
- end = orig + major
- return [(orig, end), (end, end + minor1 * tip_scale),
- (end, end + minor2 * tip_scale)]
-
- def arrows(self, orig, dir, radius=1, color=0xffffff, **kwargs):
- """Draw a list arrows on canvas.
-
- Args:
- orig (numpy.array): The positions where arrows start.
- dir (numpy.array): The directions where arrows point to.
- radius (Union[Number, np.array], optional): The width of arrows. Default is 1.
- color (Union[int, np.array], optional): The color or colors of arrows. Default is 0xffffff.
-
- """
- for begin, end in self._arrow_to_lines(orig, dir, **kwargs):
- self.lines(begin, end, radius, color)
-
- def arrow(self, orig, dir, radius=1, color=0xffffff, **kwargs):
- """Draw a single arrow on canvas.
-
- Args:
- orig (List[Number]): The position where arrow starts. Shape must be 2.
- dir (List[Number]): The direction where arrow points to. Shape must be 2.
- radius (Number, optional): The width of arrow. Default is 1.
- color (int, optional): The color of arrow. Default is 0xFFFFFF.
-
- """
- orig = np.array([orig])
- dir = np.array([dir])
- for begin, end in self._arrow_to_lines(orig, dir, **kwargs):
- self.line(begin[0], end[0], radius, color)
-
- def rect(self, topleft, bottomright, radius=1, color=0xFFFFFF):
- """Draw a single rectangle on canvas.
-
- Args:
- topleft (List[Number]): The position of the topleft corner of rectangle.
- Shape must be 2.
- bottomright (List[Number]): The position of the bottomright corner
- of rectangle. Shape must be 2.
- radius (Number, optional): The width of rectangle's sides. Default is 1.
- color (int, optional): The color of rectangle. Default is 0xFFFFFF.
-
- """
- a = topleft[0], topleft[1]
- b = bottomright[0], topleft[1]
- c = bottomright[0], bottomright[1]
- d = topleft[0], bottomright[1]
- self.line(a, b, radius, color)
- self.line(b, c, radius, color)
- self.line(c, d, radius, color)
- self.line(d, a, radius, color)
-
- def text(self, content, pos, font_size=15, color=0xFFFFFF):
- """Draw texts on canvas.
-
- Args:
- content (str): The text to be drawn on canvas.
- pos (List[Number]): The position where the text is to be put.
- font_size (Number, optional): The font size of the text.
- color (int, optional): The color of the text. Default is 0xFFFFFF.
-
- """
-
- # TODO: refactor Canvas::text
- font_size = float(font_size)
- pos = ti.core_vec(*pos)
- r, g, b = hex_to_rgb(color)
- color = ti.core_vec(r, g, b, 1)
- self.canvas.text(content, pos, font_size, color)
-
- @staticmethod
- def _make_field_base(w, h, bound):
- x = np.linspace(bound / w, 1 - bound / w, w)
- y = np.linspace(bound / h, 1 - bound / h, h)
- base = np.array(np.meshgrid(x, y))
- base = base.swapaxes(0, 1).swapaxes(1, 2).swapaxes(0, 1)
- return base.reshape(w * h, 2)
-
- def point_field(self, radius, color=0xffffff, bound=0.5):
- """Draw a field of points on canvas.
-
- Args:
- radius (np.array): The pattern and radius of the field of points.
- color (Union[int, np.array], optional): The color or colors of points.
- Default is 0xFFFFFF.
- bound (Number, optional): The boundary of the field. Default is 0.5.
-
- """
- assert len(radius.shape) == 2
- base = self._make_field_base(radius.shape[0], radius.shape[1], bound)
- radius = radius.reshape(radius.shape[0] * radius.shape[1])
- self.circles(base, radius=radius, color=color)
-
- def arrow_field(self, dir, radius=1, color=0xffffff, bound=0.5, **kwargs):
- """Draw a field of arrows on canvas.
-
- Args:
- dir (np.array): The pattern and direction of the field of arrows.
- color (Union[int, np.array], optional): The color or colors of arrows.
- Default is 0xFFFFFF.
- bound (Number, optional): The boundary of the field. Default is 0.5.
-
- """
- assert len(dir.shape) == 3
- assert dir.shape[2] == 2
- base = self._make_field_base(dir.shape[0], dir.shape[1], bound)
- dir = dir.reshape(dir.shape[0] * dir.shape[1], 2)
- self.arrows(base, dir, radius=radius, color=color, **kwargs)
-
- def show(self, file=None):
- """Show the frame or save current frame as a picture.
-
- Args:
- file (str, optional): The path & name of the picture to be saved.
- Default is None.
-
- """
- self.core.update()
- if file:
- self.core.screenshot(file)
- self.frame += 1
- self.clear()
-
- ## Event system
-
- class EventFilter:
- def __init__(self, *filter):
- self.filter = set()
- for ent in filter:
- if isinstance(ent, (list, tuple)):
- type, key = ent
- ent = (type, key)
- self.filter.add(ent)
-
- def match(self, e):
- if (e.type, e.key) in self.filter:
- return True
- if e.type in self.filter:
- return True
- if e.key in self.filter:
- return True
- return False
-
- def has_key_event(self):
- """Check if there are any key event registered.
-
- Returns:
- Bool to indicate whether there is any key event registered.
-
- """
- return self.core.has_key_event()
-
- def get_event(self, *filter):
- """Check if the specific event is triggered.
-
- Args:
- *filter (ti.GUI.EVENT): The specific event to be checked.
-
- Returns:
- Bool to indicate whether the specific event is triggered.
-
- """
- for e in self.get_events(*filter):
- self.event = e
- return True
- else:
- return False
-
- def get_events(self, *filter):
- """Get a list of events that are triggered.
-
- Args:
- *filter (List[ti.GUI.EVENT]): The type of events to be filtered.
-
- Returns:
- :class:`~taichi.misc.gui.GUI.EVENT` :A list of events that are triggered.
-
- """
- filter = filter and GUI.EventFilter(*filter) or None
-
- while True:
- if not self.has_key_event():
- break
- e = self.get_key_event()
- if filter is None or filter.match(e):
- yield e
-
- def get_key_event(self):
- """Get keyboard triggered event.
-
- Returns:
- :class:`~taichi.misc.gui.GUI.EVENT` :The keyboard triggered event.
-
- """
- self.core.wait_key_event()
-
- e = GUI.Event()
- event = self.core.get_key_event_head()
-
- e.type = event.type
- e.key = event.key
- e.pos = self.core.canvas_untransform(event.pos)
- e.pos = (e.pos[0], e.pos[1])
- e.modifier = []
-
- if e.key == GUI.WHEEL:
- e.delta = event.delta
- else:
- e.delta = (0, 0)
-
- for mod in ['Shift', 'Alt', 'Control']:
- if self.is_pressed(mod):
- e.modifier.append(mod)
-
- if e.type == GUI.PRESS:
- self.key_pressed.add(e.key)
- else:
- self.key_pressed.discard(e.key)
-
- self.core.pop_key_event_head()
- return e
-
- def is_pressed(self, *keys):
- """Check if the specific key or keys are pressed.
-
- Args:
- *keys (Union[str, List[str]]): The string that stands for keys in keyboard.
-
- Returns:
- Bool to indicate whether the key or keys are pressed.
-
- """
- for key in keys:
- if key in ['Shift', 'Alt', 'Control']:
- if key + '_L' in self.key_pressed or key + '_R' in self.key_pressed:
- return True
- if key in self.key_pressed:
- return True
- else:
- return False
-
- def get_cursor_pos(self):
- """Get the current position of mouse.
-
- Returns:
- The current position of mouse.
-
- """
- pos = self.core.get_cursor_pos()
- return pos[0], pos[1]
-
- @deprecated('gui.has_key_pressed()', 'gui.get_event()')
- def has_key_pressed(self):
- if self.has_key_event():
- self.get_key_event() # pop to update self.key_pressed
- return len(self.key_pressed) != 0
-
- @property
- def running(self):
- """Get the property of whether the gui is running.
-
- Returns:
- The running property of gui(bool).
-
- """
- return not self.core.should_close
-
- @running.setter
- def running(self, value):
- if value:
- self.core.should_close = 0
- elif not self.core.should_close:
- self.core.should_close = 1
-
- @property
- def fps_limit(self):
- """Get the property of fps limit.
-
- Returns:
- The property of fps limit of gui.
-
- """
- if self.core.frame_delta_limit == 0:
- return None
- else:
- return 1 / self.core.frame_delta_limit
-
- @fps_limit.setter
- def fps_limit(self, value):
- if value is None:
- self.core.frame_delta_limit = 0
- else:
- self.core.frame_delta_limit = 1 / value
-
-
-def rgb_to_hex(c):
- """Convert rgb color format to hex color format.
-
- Args:
- c (List[int]): The rgb representation of color.
-
- Returns:
- The hex representation of color.
-
- """
- to255 = lambda x: np.clip(np.int32(x * 255), 0, 255)
- return (to255(c[0]) << 16) + (to255(c[1]) << 8) + to255(c[2])
-
-
-def hex_to_rgb(color):
- """Convert hex color format to rgb color format.
-
- Args:
- color (int): The hex representation of color.
-
- Returns:
- The rgb representation of color.
-
- """
- r, g, b = (color >> 16) & 0xff, (color >> 8) & 0xff, color & 0xff
- return r / 255, g / 255, b / 255
-
-
-__all__ = [
- 'GUI',
- 'rgb_to_hex',
- 'hex_to_rgb',
-]
diff --git a/python/taichi/misc/task.py b/python/taichi/misc/task.py
deleted file mode 100644
index e06231ec69d30..0000000000000
--- a/python/taichi/misc/task.py
+++ /dev/null
@@ -1,36 +0,0 @@
-from taichi.core import ti_core as _ti_core
-from taichi.misc.util import config_from_dict
-
-
-def _unit(unit_name):
- def decorator(target_class):
- if target_class.__init__ != object.__init__:
- original_init = target_class.__init__
- else:
-
- def dummy_init(*args, **kwargs):
- pass
-
- original_init = dummy_init
-
- def new_init(self, name, *args, **kwargs):
- self.c = getattr(_ti_core, 'create_' + unit_name)(name)
- self.c.initialize(config_from_dict(kwargs))
- original_init(self, *args, **kwargs)
-
- target_class.__init__ = new_init
-
- def new_getattr_(self, key):
- return self.c.__getattribute__(key)
-
- target_class.__getattr__ = new_getattr_
-
- return target_class
-
- return decorator
-
-
-@_unit('task')
-class Task:
- def run(self, *args):
- return self.c.run(args)
diff --git a/python/taichi/misc/util.py b/python/taichi/misc/util.py
deleted file mode 100644
index 21e81253c3028..0000000000000
--- a/python/taichi/misc/util.py
+++ /dev/null
@@ -1,273 +0,0 @@
-import copy
-import functools
-import subprocess
-import sys
-import traceback
-
-from colorama import Fore, Style
-from taichi.core import ti_core as _ti_core
-
-import taichi as ti
-
-
-def config_from_dict(args):
- d = copy.copy(args)
- for k in d:
- if isinstance(d[k], _ti_core.Vector2f):
- d[k] = '({}, {})'.format(d[k].x, d[k].y)
- if isinstance(d[k], _ti_core.Vector3f):
- d[k] = '({}, {}, {})'.format(d[k].x, d[k].y, d[k].z)
- d[k] = str(d[k])
- return _ti_core.config_from_dict(d)
-
-
-def core_veci(*args):
- if isinstance(args[0], _ti_core.Vector2i):
- return args[0]
- if isinstance(args[0], _ti_core.Vector3i):
- return args[0]
- if isinstance(args[0], tuple):
- args = tuple(*args)
- if len(args) == 2:
- return _ti_core.Vector2i(int(args[0]), int(args[1]))
- elif len(args) == 3:
- return _ti_core.Vector3i(int(args[0]), int(args[1]), int(args[2]))
- elif len(args) == 4:
- return _ti_core.Vector4i(int(args[0]), int(args[1]), int(args[2]),
- int(args[3]))
- else:
- assert False, type(args[0])
-
-
-def core_vec(*args):
- if isinstance(args[0], _ti_core.Vector2f):
- return args[0]
- if isinstance(args[0], _ti_core.Vector3f):
- return args[0]
- if isinstance(args[0], _ti_core.Vector4f):
- return args[0]
- if isinstance(args[0], _ti_core.Vector2d):
- return args[0]
- if isinstance(args[0], _ti_core.Vector3d):
- return args[0]
- if isinstance(args[0], _ti_core.Vector4d):
- return args[0]
- if isinstance(args[0], tuple):
- args = tuple(*args)
- if _ti_core.get_default_float_size() == 4:
- if len(args) == 2:
- return _ti_core.Vector2f(float(args[0]), float(args[1]))
- elif len(args) == 3:
- return _ti_core.Vector3f(float(args[0]), float(args[1]),
- float(args[2]))
- elif len(args) == 4:
- return _ti_core.Vector4f(float(args[0]), float(args[1]),
- float(args[2]), float(args[3]))
- else:
- assert False, type(args[0])
- else:
- if len(args) == 2:
- return _ti_core.Vector2d(float(args[0]), float(args[1]))
- elif len(args) == 3:
- return _ti_core.Vector3d(float(args[0]), float(args[1]),
- float(args[2]))
- elif len(args) == 4:
- return _ti_core.Vector4d(float(args[0]), float(args[1]),
- float(args[2]), float(args[3]))
- else:
- assert False, type(args[0])
-
-
-class Tee():
- def __init__(self, name):
- self.file = open(name, 'w')
- self.stdout = sys.stdout
- self.stderr = sys.stderr
- sys.stdout = self
- sys.stderr = self
-
- def __del__(self):
- self.file.close()
-
- def write(self, data):
- self.file.write(data)
- self.stdout.write(data)
- self.file.flush()
- self.stdout.flush()
-
- def write_to_file(self, data):
- self.file.write(data)
-
-
-# The builtin `warnings` module is unreliable since it may be suppressed
-# by other packages such as IPython.
-def warning(msg, type=UserWarning, stacklevel=1):
- """Print warning message
-
- Args:
- msg (str): massage to print.
- type (builtin warning type): type of warning.
- stacklevel (int): warning stack level from the caller.
- """
- s = traceback.extract_stack()[:-stacklevel]
- raw = ''.join(traceback.format_list(s))
- print(Fore.YELLOW + Style.BRIGHT, end='')
- print(f'{type.__name__}: {msg}')
- print(f'\n{raw}')
- print(Style.RESET_ALL, end='')
-
-
-def deprecated(old, new, warning_type=DeprecationWarning):
- """Mark an API as deprecated.
-
- Args:
- old (str): old method.
- new (str): new method.
- warning_type (builtin warning type): type of warning.
-
- Example::
-
- >>> @deprecated('ti.sqr(x)', 'x**2')
- >>> def sqr(x):
- >>> return x**2
-
- Returns:
- Decorated fuction with warning message
- """
- def decorator(foo):
- @functools.wraps(foo)
- def wrapped(*args, **kwargs):
- _taichi_skip_traceback = 1
- msg = f'{old} is deprecated. Please use {new} instead.'
- warning(msg, warning_type, stacklevel=2)
- return foo(*args, **kwargs)
-
- return wrapped
-
- return decorator
-
-
-def obsolete(old, new):
- """
- Mark an API as obsolete. Usage:
-
- sqr = obsolete('ti.sqr(x)', 'x**2')
- """
- def wrapped(*args, **kwargs):
- _taichi_skip_traceback = 1
- msg = f'{old} is obsolete. Please use {new} instead.'
- raise SyntaxError(msg)
-
- return wrapped
-
-
-def get_traceback(stacklevel=1):
- s = traceback.extract_stack()[:-1 - stacklevel]
- return ''.join(traceback.format_list(s))
-
-
-def duplicate_stdout_to_file(fn):
- _ti_core.duplicate_stdout_to_file(fn)
-
-
-def set_gdb_trigger(on=True):
- _ti_core.set_core_trigger_gdb_when_crash(on)
-
-
-def print_profile_info():
- """Print time elapsed on the host tasks in a hierarchical format.
-
- This profiler is automatically on.
-
- Call function imports from C++ : _ti_core.print_profile_info()
-
- Example::
-
- >>> import taichi as ti
- >>> ti.init(arch=ti.cpu)
- >>> var = ti.field(ti.f32, shape=1)
- >>> @ti.kernel
- >>> def compute():
- >>> var[0] = 1.0
- >>> print("Setting var[0] =", var[0])
- >>> compute()
- >>> ti.print_profile_info()
- """
- _ti_core.print_profile_info()
-
-
-def clear_profile_info():
- """Clear profiler's records about time elapsed on the host tasks.
-
- Call function imports from C++ : _ti_core.clear_profile_info()
- """
- _ti_core.clear_profile_info()
-
-
-@deprecated('ti.vec(x, y)', 'ti.core_vec(x, y)')
-def vec(*args, **kwargs):
- return core_vec(*args, **kwargs)
-
-
-@deprecated('ti.veci(x, y)', 'ti.core_veci(x, y)')
-def veci(*args, **kwargs):
- return core_veci(*args, **kwargs)
-
-
-def dump_dot(filepath=None, rankdir=None, embed_states_threshold=0):
- d = _ti_core.dump_dot(rankdir, embed_states_threshold)
- if filepath is not None:
- with open(filepath, 'w') as fh:
- fh.write(d)
- return d
-
-
-def dot_to_pdf(dot, filepath):
- assert filepath.endswith('.pdf')
- p = subprocess.Popen(['dot', '-Tpdf'],
- stdin=subprocess.PIPE,
- stdout=subprocess.PIPE)
- pdf_contents = p.communicate(input=dot.encode())[0]
- with open(filepath, 'wb') as fh:
- fh.write(pdf_contents)
-
-
-def get_kernel_stats():
- return _ti_core.get_kernel_stats()
-
-
-def print_async_stats(include_kernel_profiler=False):
- if include_kernel_profiler:
- ti.print_kernel_profile_info()
- print()
- stat = ti.get_kernel_stats()
- counters = stat.get_counters()
- print('=======================')
- print('Async benchmark metrics')
- print('-----------------------')
- print(f'Async mode: {ti.current_cfg().async_mode}')
- print(f'Kernel time: {ti.kernel_profiler_total_time():.3f} s')
- print(f'Tasks launched: {int(counters["launched_tasks"])}')
- print(f'Instructions emitted: {int(counters["codegen_statements"])}')
- print(f'Tasks compiled: {int(counters["codegen_offloaded_tasks"])}')
- NUM_FUSED_TASKS_KEY = 'num_fused_tasks'
- if NUM_FUSED_TASKS_KEY in counters:
- print(f'Tasks fused: {int(counters["num_fused_tasks"])}')
- print('=======================')
-
-
-__all__ = [
- 'vec',
- 'veci',
- 'core_vec',
- 'core_veci',
- 'deprecated',
- 'dump_dot',
- 'dot_to_pdf',
- 'obsolete',
- 'get_kernel_stats',
- 'get_traceback',
- 'set_gdb_trigger',
- 'print_profile_info',
- 'clear_profile_info',
-]
diff --git a/python/taichi/profiler/__init__.py b/python/taichi/profiler/__init__.py
index 0aad976a6cfb5..780b2a937f507 100644
--- a/python/taichi/profiler/__init__.py
+++ b/python/taichi/profiler/__init__.py
@@ -1,3 +1,4 @@
-from taichi.profiler.kernelprofiler import \
- KernelProfiler # import for docstring-gen
-from taichi.profiler.kernelprofiler import get_default_kernel_profiler
+from taichi.profiler.kernel_metrics import *
+from taichi.profiler.kernel_profiler import *
+from taichi.profiler.memory_profiler import *
+from taichi.profiler.scoped_profiler import *
diff --git a/python/taichi/profiler/kernelmetrics.py b/python/taichi/profiler/kernel_metrics.py
similarity index 71%
rename from python/taichi/profiler/kernelmetrics.py
rename to python/taichi/profiler/kernel_metrics.py
index 909f1f06874fd..a8a837d3b0fcc 100644
--- a/python/taichi/profiler/kernelmetrics.py
+++ b/python/taichi/profiler/kernel_metrics.py
@@ -1,21 +1,18 @@
-from dataclasses import dataclass
+from taichi._lib import core as _ti_core
-from taichi.core import ti_core as _ti_core
-
-@dataclass
class CuptiMetric:
- """A data class to add CUPTI metric for :class:`~taichi.lang.KernelProfiler`.
+ """A class to add CUPTI metric for :class:`~taichi.profiler.kernel_profiler.KernelProfiler`.
- This data class is designed to add user selected CUPTI metrics.
+ This class is designed to add user selected CUPTI metrics.
Only available for the CUDA backend now, i.e. you need ``ti.init(kernel_profiler=True, arch=ti.cuda)``.
- For usage of this class, see examples in func :func:`~taichi.lang.set_kernel_profile_metrics` and :func:`~taichi.lang.collect_kernel_profile_metrics`.
+ For usage of this class, see examples in func :func:`~taichi.profiler.set_kernel_profiler_metrics` and :func:`~taichi.profiler.collect_kernel_profiler_metrics`.
Args:
- name (str): name of metric that collected by CUPTI toolkit. used by :func:`~taichi.lang.set_kernel_profile_metrics` and :func:`~taichi.lang.collect_kernel_profile_metrics`.
- header (str): column header of this metric, used by :func:`~taichi.lang.print_kernel_profile_info`.
- format (str): format for print metric value (and unit of this value), used by :func:`~taichi.lang.print_kernel_profile_info`.
- scale (float): scale of metric value, used by :func:`~taichi.lang.print_kernel_profile_info`.
+ name (str): name of metric that collected by CUPTI toolkit. used by :func:`~taichi.profiler.set_kernel_profiler_metrics` and :func:`~taichi.profiler.collect_kernel_profiler_metrics`.
+ header (str): column header of this metric, used by :func:`~taichi.profiler.print_kernel_profiler_info`.
+ val_format (str): format for print metric value (and unit of this value), used by :func:`~taichi.profiler.print_kernel_profiler_info`.
+ scale (float): scale of metric value, used by :func:`~taichi.profiler.print_kernel_profiler_info`.
Example::
@@ -33,62 +30,67 @@ class CuptiMetric:
>>> for i in x:
>>> y[None] += x[i]
- >>> global_op_atom = ti.CuptiMetric(
+ >>> global_op_atom = ti.profiler.CuptiMetric(
>>> name='l1tex__t_set_accesses_pipe_lsu_mem_global_op_atom.sum',
>>> header=' global.atom ',
- >>> format=' {:8.0f} ')
+ >>> val_format=' {:8.0f} ')
>>> # add and set user defined metrics
- >>> profiling_metrics = ti.get_predefined_cupti_metrics('global_access') + [global_op_atom]
- >>> ti.set_kernel_profile_metrics(profiling_metrics)
+ >>> profiling_metrics = ti.profiler.get_predefined_cupti_metrics('global_access') + [global_op_atom]
+ >>> ti.profiler.set_kernel_profile_metrics(profiling_metrics)
>>> for i in range(16):
>>> reduction()
- >>> ti.print_kernel_profile_info('trace')
+ >>> ti.profiler.print_kernel_profiler_info('trace')
Note:
For details about using CUPTI in Taichi, please visit https://docs.taichi.graphics/docs/lang/articles/misc/profiler#advanced-mode.
"""
- name: str = ''
- header: str = ''
- format: str = ''
- scale: float = 1.0
+ def __init__(self,
+ name='',
+ header='unnamed_header',
+ val_format=' {:8.0f} ',
+ scale=1.0):
+ self.name = name
+ self.header = header
+ self.val_format = val_format
+ self.scale = scale
# Global Memory Metrics
dram_utilization = CuptiMetric(
name='dram__throughput.avg.pct_of_peak_sustained_elapsed',
header=' global.uti ',
- format=' {:6.2f} % ')
+ val_format=' {:6.2f} % ')
dram_bytes_sum = CuptiMetric(name='dram__bytes.sum',
header=' global.R&W ',
- format='{:9.3f} MB ',
+ val_format='{:9.3f} MB ',
scale=1.0 / 1024 / 1024)
dram_bytes_throughput = CuptiMetric(name='dram__bytes.sum.per_second',
header=' global.R&W/s ',
- format='{:8.3f} GB/s ',
+ val_format='{:8.3f} GB/s ',
scale=1.0 / 1024 / 1024 / 1024)
dram_bytes_read = CuptiMetric(name='dram__bytes_read.sum',
header=' global.R ',
- format='{:8.3f} MB ',
+ val_format='{:8.3f} MB ',
scale=1.0 / 1024 / 1024)
dram_read_throughput = CuptiMetric(name='dram__bytes_read.sum.per_second',
header=' global.R/s ',
- format='{:8.3f} GB/s ',
+ val_format='{:8.3f} GB/s ',
scale=1.0 / 1024 / 1024 / 1024)
dram_bytes_write = CuptiMetric(name='dram__bytes_write.sum',
header=' global.W ',
- format='{:8.3f} MB ',
+ val_format='{:8.3f} MB ',
scale=1.0 / 1024 / 1024)
dram_write_throughput = CuptiMetric(name='dram__bytes_write.sum.per_second',
header=' global.W/s ',
- format='{:8.3f} GB/s ',
+ val_format='{:8.3f} GB/s ',
scale=1.0 / 1024 / 1024 / 1024)
# Shared Memory Metrics
@@ -96,73 +98,73 @@ class CuptiMetric:
name=
'l1tex__data_pipe_lsu_wavefronts_mem_shared.avg.pct_of_peak_sustained_elapsed',
header=' uti.shared ',
- format=' {:6.2f} % ')
+ val_format=' {:6.2f} % ')
shared_transactions_load = CuptiMetric(
name='l1tex__data_pipe_lsu_wavefronts_mem_shared_op_ld.sum',
header=' shared.trans.W ',
- format=' {:10.0f} ')
+ val_format=' {:10.0f} ')
shared_transactions_store = CuptiMetric(
name='l1tex__data_pipe_lsu_wavefronts_mem_shared_op_st.sum',
header=' shared.trans.R ',
- format=' {:10.0f} ')
+ val_format=' {:10.0f} ')
shared_bank_conflicts_store = CuptiMetric(
name='l1tex__data_bank_conflicts_pipe_lsu_mem_shared_op_st.sum',
header=' bank.conflict.W ',
- format=' {:10.0f} ')
+ val_format=' {:10.0f} ')
shared_bank_conflicts_load = CuptiMetric(
name='l1tex__data_bank_conflicts_pipe_lsu_mem_shared_op_ld.sum',
header=' bank.conflict.R ',
- format=' {:10.0f} ')
+ val_format=' {:10.0f} ')
# Atomic Metrics
global_op_atom = CuptiMetric(
name='l1tex__t_set_accesses_pipe_lsu_mem_global_op_atom.sum',
header=' global.atom ',
- format=' {:8.0f} ')
+ val_format=' {:8.0f} ')
global_op_reduction = CuptiMetric(
name='l1tex__t_set_accesses_pipe_lsu_mem_global_op_red.sum',
header=' global.red ',
- format=' {:8.0f} ')
+ val_format=' {:8.0f} ')
# Hardware Utilization Metrics
sm_throughput = CuptiMetric(
name='sm__throughput.avg.pct_of_peak_sustained_elapsed',
header=' core.uti ',
- format=' {:6.2f} % ')
+ val_format=' {:6.2f} % ')
dram_throughput = CuptiMetric(
name='gpu__dram_throughput.avg.pct_of_peak_sustained_elapsed',
header=' mem.uti ',
- format=' {:6.2f} % ')
+ val_format=' {:6.2f} % ')
l1tex_throughput = CuptiMetric(
name='l1tex__throughput.avg.pct_of_peak_sustained_elapsed',
header=' L1.uti ',
- format=' {:6.2f} % ')
+ val_format=' {:6.2f} % ')
l2_throughput = CuptiMetric(
name='lts__throughput.avg.pct_of_peak_sustained_elapsed',
header=' L2.uti ',
- format=' {:6.2f} % ')
+ val_format=' {:6.2f} % ')
# Misc Metrics
l1_hit_rate = CuptiMetric(name='l1tex__t_sector_hit_rate.pct',
header=' L1.hit ',
- format=' {:6.2f} % ')
+ val_format=' {:6.2f} % ')
l2_hit_rate = CuptiMetric(name='lts__t_sector_hit_rate.pct',
header=' L2.hit ',
- format=' {:6.2f} % ')
+ val_format=' {:6.2f} % ')
achieved_occupancy = CuptiMetric(
name='sm__warps_active.avg.pct_of_peak_sustained_active',
header=' occupancy',
- format=' {:6.0f} ')
+ val_format=' {:6.0f} ')
# metric suite: global load & store
global_access = [
@@ -219,9 +221,10 @@ def get_predefined_cupti_metrics(name=''):
for key in predefined_cupti_metrics:
_ti_core.warn(f" '{key}'")
return None
- else:
- return predefined_cupti_metrics[name]
+ return predefined_cupti_metrics[name]
# Default metrics list
default_cupti_metrics = [dram_bytes_sum]
+
+__all__ = ['CuptiMetric', 'get_predefined_cupti_metrics']
diff --git a/python/taichi/profiler/kernelprofiler.py b/python/taichi/profiler/kernel_profiler.py
similarity index 55%
rename from python/taichi/profiler/kernelprofiler.py
rename to python/taichi/profiler/kernel_profiler.py
index 558bb08856d14..21615b5f06dfc 100644
--- a/python/taichi/profiler/kernelprofiler.py
+++ b/python/taichi/profiler/kernel_profiler.py
@@ -1,10 +1,8 @@
from contextlib import contextmanager
-from taichi.core import ti_core as _ti_core
+from taichi._lib import core as _ti_core
from taichi.lang import impl
-from taichi.profiler.kernelmetrics import default_cupti_metrics
-
-import taichi as ti
+from taichi.profiler.kernel_metrics import default_cupti_metrics
class StatisticalResult:
@@ -42,7 +40,7 @@ class KernelProfiler:
"""Kernel profiler of Taichi.
Kernel profiler acquires kernel profiling records from backend, counts records in Python scope,
- and prints the results to the console by :func:`~taichi.profiler.kernelprofiler.KernelProfiler.print_info`.
+ and prints the results to the console by :func:`~taichi.profiler.kernel_profiler.KernelProfiler.print_info`.
``KernelProfiler`` now support detailed low-level performance metrics (such as memory bandwidth consumption) in its advanced mode.
This mode is only available for the CUDA backend with CUPTI toolkit, i.e. you need ``ti.init(kernel_profiler=True, arch=ti.cuda)``.
@@ -52,6 +50,7 @@ class KernelProfiler:
"""
def __init__(self):
self._profiling_mode = False
+ self._profiling_toolkit = 'default'
self._metric_list = [default_cupti_metrics]
self._total_time_ms = 0.0
self._traced_records = []
@@ -60,7 +59,7 @@ def __init__(self):
# public methods
def set_kernel_profiler_mode(self, mode=False):
- """Turn on or off :class:`~taichi.profiler.kernelprofiler.KernelProfiler`."""
+ """Turn on or off :class:`~taichi.profiler.kernel_profiler.KernelProfiler`."""
if type(mode) is bool:
self._profiling_mode = mode
else:
@@ -69,9 +68,22 @@ def set_kernel_profiler_mode(self, mode=False):
)
def get_kernel_profiler_mode(self):
- """Get status of :class:`~taichi.profiler.kernelprofiler.KernelProfiler`."""
+ """Get status of :class:`~taichi.profiler.kernel_profiler.KernelProfiler`."""
return self._profiling_mode
+ def set_toolkit(self, toolkit_name='default'):
+ if self._check_not_turned_on_with_warning_message():
+ return False
+ status = impl.get_runtime().prog.set_kernel_profiler_toolkit(
+ toolkit_name)
+ if status is True:
+ self._profiling_toolkit = toolkit_name
+ else:
+ _ti_core.warn(
+ f'Failed to set kernel profiler toolkit ({toolkit_name}) , keep using ({self._profiling_toolkit}).'
+ )
+ return status
+
def get_total_time(self):
"""Get elapsed time of all kernels recorded in KernelProfiler.
@@ -85,7 +97,7 @@ def get_total_time(self):
return self._total_time_ms / 1000 # ms to s
def clear_info(self):
- """Clear all records both in front-end :class:`~taichi.profiler.kernelprofiler.KernelProfiler` and back-end instance ``KernelProfilerBase``.
+ """Clear all records both in front-end :class:`~taichi.profiler.kernel_profiler.KernelProfiler` and back-end instance ``KernelProfilerBase``.
Note:
The values of ``self._profiling_mode`` and ``self._metric_list`` will not be cleared.
@@ -93,13 +105,15 @@ def clear_info(self):
if self._check_not_turned_on_with_warning_message():
return None
#sync first
- impl.get_runtime().sync()
+ impl.get_runtime().prog.sync_kernel_profiler()
#then clear backend & frontend info
impl.get_runtime().prog.clear_kernel_profile_info()
self._clear_frontend()
+ return None
+
def query_info(self, name):
- """For docsting of this function, see :func:`~taichi.lang.query_kernel_profile_info`."""
+ """For docstring of this function, see :func:`~taichi.profiler.query_kernel_profiler_info`."""
if self._check_not_turned_on_with_warning_message():
return None
self._update_records() # kernel records
@@ -108,7 +122,7 @@ def query_info(self, name):
return impl.get_runtime().prog.query_kernel_profile_info(name)
def set_metrics(self, metric_list=default_cupti_metrics):
- """For docsting of this function, see :func:`~taichi.lang.set_kernel_profile_metrics`."""
+ """For docstring of this function, see :func:`~taichi.profiler.set_kernel_profiler_metrics`."""
if self._check_not_turned_on_with_warning_message():
return None
self._metric_list = metric_list
@@ -117,11 +131,13 @@ def set_metrics(self, metric_list=default_cupti_metrics):
impl.get_runtime().prog.reinit_kernel_profiler_with_metrics(
metric_name_list)
+ return None
+
@contextmanager
def collect_metrics_in_context(self, metric_list=default_cupti_metrics):
"""This function is not exposed to user now.
- For usage of this function, see :func:`~taichi.lang.collect_kernel_profile_metrics`.
+ For usage of this function, see :func:`~taichi.profiler.collect_kernel_profiler_metrics`.
"""
if self._check_not_turned_on_with_warning_message():
return None
@@ -129,6 +145,8 @@ def collect_metrics_in_context(self, metric_list=default_cupti_metrics):
yield self
self.set_metrics() #back to default metric list
+ return None
+
# mode of print_info
COUNT = 'count' # print the statistical results (min,max,avg time) of Taichi kernels.
TRACE = 'trace' # print the records of launched Taichi kernels with specific profiling metrics (time, memory load/store and core utilization etc.)
@@ -136,7 +154,7 @@ def collect_metrics_in_context(self, metric_list=default_cupti_metrics):
def print_info(self, mode=COUNT):
"""Print the profiling results of Taichi kernels.
- For usage of this function, see :func:`~taichi.lang.print_kernel_profile_info`.
+ For usage of this function, see :func:`~taichi.profiler.print_kernel_profiler_info`.
Args:
mode (str): the way to print profiling results.
@@ -154,21 +172,22 @@ def print_info(self, mode=COUNT):
self._print_kernel_info()
else:
raise ValueError(
- f'Arg `mode` must be of type \'str\', and has the value \'count\' or \'trace\'.'
+ 'Arg `mode` must be of type \'str\', and has the value \'count\' or \'trace\'.'
)
+ return None
+
# private methods
def _check_not_turned_on_with_warning_message(self):
if self._profiling_mode is False:
_ti_core.warn(
- f'use \'ti.init(kernel_profiler = True)\' to turn on KernelProfiler.'
+ 'use \'ti.init(kernel_profiler = True)\' to turn on KernelProfiler.'
)
return True
- else:
- return False
+ return False
def _clear_frontend(self):
- """Clear member variables in :class:`~taichi.profiler.kernelprofiler.KernelProfiler`.
+ """Clear member variables in :class:`~taichi.profiler.kernel_profiler.KernelProfiler`.
Note:
The values of ``self._profiling_mode`` and ``self._metric_list`` will not be cleared.
@@ -179,7 +198,7 @@ def _clear_frontend(self):
def _update_records(self):
"""Acquires kernel records from a backend."""
- impl.get_runtime().sync()
+ impl.get_runtime().prog.sync_kernel_profiler()
self._clear_frontend()
self._traced_records = impl.get_runtime(
).prog.get_kernel_profiler_records()
@@ -204,8 +223,8 @@ def _count_statistics(self):
}
def _make_table_header(self, mode):
- header_str = f'Kernel Profiler({mode})'
- arch_name = f' @ {_ti_core.arch_name(ti.cfg.arch).upper()}'
+ header_str = f'Kernel Profiler({mode}, {self._profiling_toolkit})'
+ arch_name = f' @ {_ti_core.arch_name(impl.current_cfg().arch).upper()}'
device_name = impl.get_runtime().prog.get_kernel_profiler_device_name()
if len(device_name) > 1: # default device_name = ' '
device_name = ' on ' + device_name
@@ -334,3 +353,240 @@ def get_default_kernel_profiler():
For data retention purposes, multiple instances will be considered in the future.
"""
return _ti_kernel_profiler
+
+
+def print_kernel_profiler_info(mode='count'):
+ """Print the profiling results of Taichi kernels.
+
+ To enable this profiler, set ``kernel_profiler=True`` in ``ti.init()``.
+ ``'count'`` mode: print the statistics (min,max,avg time) of launched kernels,
+ ``'trace'`` mode: print the records of launched kernels with specific profiling metrics (time, memory load/store and core utilization etc.),
+ and defaults to ``'count'``.
+
+ Args:
+ mode (str): the way to print profiling results.
+
+ Example::
+
+ >>> import taichi as ti
+
+ >>> ti.init(ti.cpu, kernel_profiler=True)
+ >>> var = ti.field(ti.f32, shape=1)
+
+ >>> @ti.kernel
+ >>> def compute():
+ >>> var[0] = 1.0
+
+ >>> compute()
+ >>> ti.profiler.print_kernel_profiler_info()
+ >>> # equivalent calls :
+ >>> # ti.profiler.print_kernel_profiler_info('count')
+
+ >>> ti.profiler.print_kernel_profiler_info('trace')
+
+ Note:
+ Currently the result of `KernelProfiler` could be incorrect on OpenGL
+ backend due to its lack of support for `ti.sync()`.
+
+ For advanced mode of `KernelProfiler`, please visit https://docs.taichi.graphics/docs/lang/articles/misc/profiler#advanced-mode.
+ """
+ get_default_kernel_profiler().print_info(mode)
+
+
+def query_kernel_profiler_info(name):
+ """Query kernel elapsed time(min,avg,max) on devices using the kernel name.
+
+ To enable this profiler, set `kernel_profiler=True` in `ti.init`.
+
+ Args:
+ name (str): kernel name.
+
+ Returns:
+ KernelProfilerQueryResult (class): with member variables(counter, min, max, avg)
+
+ Example::
+
+ >>> import taichi as ti
+
+ >>> ti.init(ti.cpu, kernel_profiler=True)
+ >>> n = 1024*1024
+ >>> var = ti.field(ti.f32, shape=n)
+
+ >>> @ti.kernel
+ >>> def fill():
+ >>> for i in range(n):
+ >>> var[i] = 0.1
+
+ >>> fill()
+ >>> ti.profiler.clear_kernel_profiler_info() #[1]
+ >>> for i in range(100):
+ >>> fill()
+ >>> query_result = ti.profiler.query_kernel_profiler_info(fill.__name__) #[2]
+ >>> print("kernel excuted times =",query_result.counter)
+ >>> print("kernel elapsed time(min_in_ms) =",query_result.min)
+ >>> print("kernel elapsed time(max_in_ms) =",query_result.max)
+ >>> print("kernel elapsed time(avg_in_ms) =",query_result.avg)
+
+ Note:
+ [1] To get the correct result, query_kernel_profiler_info() must be used in conjunction with
+ clear_kernel_profiler_info().
+
+ [2] Currently the result of `KernelProfiler` could be incorrect on OpenGL
+ backend due to its lack of support for `ti.sync()`.
+ """
+ return get_default_kernel_profiler().query_info(name)
+
+
+def clear_kernel_profiler_info():
+ """Clear all KernelProfiler records."""
+ get_default_kernel_profiler().clear_info()
+
+
+def get_kernel_profiler_total_time():
+ """Get elapsed time of all kernels recorded in KernelProfiler.
+
+ Returns:
+ time (float): total time in second.
+ """
+ return get_default_kernel_profiler().get_total_time()
+
+
+def set_kernel_profiler_toolkit(toolkit_name='default'):
+ """Set the toolkit used by KernelProfiler.
+
+ Currently, we only support toolkits: ``'default'`` and ``'cupti'``.
+
+ Args:
+ toolkit_name (str): string of toolkit name.
+
+ Returns:
+ status (bool): whether the setting is successful or not.
+
+ Example::
+
+ >>> import taichi as ti
+
+ >>> ti.init(arch=ti.cuda, kernel_profiler=True)
+ >>> x = ti.field(ti.f32, shape=1024*1024)
+
+ >>> @ti.kernel
+ >>> def fill():
+ >>> for i in x:
+ >>> x[i] = i
+
+ >>> ti.profiler.set_kernel_profiler_toolkit('cupti')
+ >>> for i in range(100):
+ >>> fill()
+ >>> ti.profiler.print_kernel_profiler_info()
+
+ >>> ti.profiler.set_kernel_profiler_toolkit('default')
+ >>> for i in range(100):
+ >>> fill()
+ >>> ti.profiler.print_kernel_profiler_info()
+ """
+ return get_default_kernel_profiler().set_toolkit(toolkit_name)
+
+
+def set_kernel_profiler_metrics(metric_list=default_cupti_metrics):
+ """Set metrics that will be collected by the CUPTI toolkit.
+
+ Args:
+ metric_list (list): a list of :class:`~taichi.profiler.CuptiMetric()` instances, default value: :data:`~taichi.profiler.kernel_metrics.default_cupti_metrics`.
+
+ Example::
+
+ >>> import taichi as ti
+
+ >>> ti.init(kernel_profiler=True, arch=ti.cuda)
+ >>> ti.profiler.set_kernel_profiler_toolkit('cupti')
+ >>> num_elements = 128*1024*1024
+
+ >>> x = ti.field(ti.f32, shape=num_elements)
+ >>> y = ti.field(ti.f32, shape=())
+ >>> y[None] = 0
+
+ >>> @ti.kernel
+ >>> def reduction():
+ >>> for i in x:
+ >>> y[None] += x[i]
+
+ >>> # In the case of not pramater, Taichi will print its pre-defined metrics list
+ >>> ti.profiler.get_predefined_cupti_metrics()
+ >>> # get Taichi pre-defined metrics
+ >>> profiling_metrics = ti.profiler.get_predefined_cupti_metrics('shared_access')
+
+ >>> global_op_atom = ti.profiler.CuptiMetric(
+ >>> name='l1tex__t_set_accesses_pipe_lsu_mem_global_op_atom.sum',
+ >>> header=' global.atom ',
+ >>> format=' {:8.0f} ')
+ >>> # add user defined metrics
+ >>> profiling_metrics += [global_op_atom]
+
+ >>> # metrics setting will be retained until the next configuration
+ >>> ti.profiler.set_kernel_profile_metrics(profiling_metrics)
+ >>> for i in range(16):
+ >>> reduction()
+ >>> ti.profiler.print_kernel_profiler_info('trace')
+
+ Note:
+ Metrics setting will be retained until the next configuration.
+ """
+ get_default_kernel_profiler().set_metrics(metric_list)
+
+
+@contextmanager
+def collect_kernel_profiler_metrics(metric_list=default_cupti_metrics):
+ """Set temporary metrics that will be collected by the CUPTI toolkit within this context.
+
+ Args:
+ metric_list (list): a list of :class:`~taichi.profiler.CuptiMetric()` instances, default value: :data:`~taichi.profiler.kernel_metrics.default_cupti_metrics`.
+
+ Example::
+
+ >>> import taichi as ti
+
+ >>> ti.init(kernel_profiler=True, arch=ti.cuda)
+ >>> ti.profiler.set_kernel_profiler_toolkit('cupti')
+ >>> num_elements = 128*1024*1024
+
+ >>> x = ti.field(ti.f32, shape=num_elements)
+ >>> y = ti.field(ti.f32, shape=())
+ >>> y[None] = 0
+
+ >>> @ti.kernel
+ >>> def reduction():
+ >>> for i in x:
+ >>> y[None] += x[i]
+
+ >>> # In the case of not pramater, Taichi will print its pre-defined metrics list
+ >>> ti.profiler.get_predefined_cupti_metrics()
+ >>> # get Taichi pre-defined metrics
+ >>> profiling_metrics = ti.profiler.get_predefined_cupti_metrics('device_utilization')
+
+ >>> global_op_atom = ti.profiler.CuptiMetric(
+ >>> name='l1tex__t_set_accesses_pipe_lsu_mem_global_op_atom.sum',
+ >>> header=' global.atom ',
+ >>> format=' {:8.0f} ')
+ >>> # add user defined metrics
+ >>> profiling_metrics += [global_op_atom]
+
+ >>> # metrics setting is temporary, and will be clear when exit from this context.
+ >>> with ti.profiler.collect_kernel_profiler_metrics(profiling_metrics):
+ >>> for i in range(16):
+ >>> reduction()
+ >>> ti.profiler.print_kernel_profiler_info('trace')
+
+ Note:
+ The configuration of the ``metric_list`` will be clear when exit from this context.
+ """
+ get_default_kernel_profiler().set_metrics(metric_list)
+ yield get_default_kernel_profiler()
+ get_default_kernel_profiler().set_metrics()
+
+
+__all__ = [
+ 'clear_kernel_profiler_info', 'collect_kernel_profiler_metrics',
+ 'get_kernel_profiler_total_time', 'print_kernel_profiler_info',
+ 'query_kernel_profiler_info', 'set_kernel_profiler_metrics',
+ 'set_kernel_profiler_toolkit'
+]
diff --git a/python/taichi/profiler/memory_profiler.py b/python/taichi/profiler/memory_profiler.py
new file mode 100644
index 0000000000000..0c1030bf742aa
--- /dev/null
+++ b/python/taichi/profiler/memory_profiler.py
@@ -0,0 +1,13 @@
+from taichi.lang.impl import get_runtime
+
+
+def print_memory_profiler_info():
+ """Memory profiling tool for LLVM backends with full sparse support.
+
+ This profiler is automatically on.
+ """
+ get_runtime().materialize()
+ get_runtime().prog.print_memory_profiler_info()
+
+
+__all__ = ['print_memory_profiler_info']
diff --git a/python/taichi/profiler/scoped_profiler.py b/python/taichi/profiler/scoped_profiler.py
new file mode 100644
index 0000000000000..c9884993ebe08
--- /dev/null
+++ b/python/taichi/profiler/scoped_profiler.py
@@ -0,0 +1,34 @@
+from taichi._lib import core as _ti_core
+
+
+def print_scoped_profiler_info():
+ """Print time elapsed on the host tasks in a hierarchical format.
+
+ This profiler is automatically on.
+
+ Call function imports from C++ : _ti_core.print_profile_info()
+
+ Example::
+
+ >>> import taichi as ti
+ >>> ti.init(arch=ti.cpu)
+ >>> var = ti.field(ti.f32, shape=1)
+ >>> @ti.kernel
+ >>> def compute():
+ >>> var[0] = 1.0
+ >>> print("Setting var[0] =", var[0])
+ >>> compute()
+ >>> ti.profiler.print_scoped_profiler_info()
+ """
+ _ti_core.print_profile_info()
+
+
+def clear_scoped_profiler_info():
+ """Clear profiler's records about time elapsed on the host tasks.
+
+ Call function imports from C++ : _ti_core.clear_profile_info()
+ """
+ _ti_core.clear_profile_info()
+
+
+__all__ = ['print_scoped_profiler_info', 'clear_scoped_profiler_info']
diff --git a/python/taichi/shaders/Circles_vk.vert b/python/taichi/shaders/Circles_vk.vert
index c77358f4e7a08..712132ec4e495 100644
--- a/python/taichi/shaders/Circles_vk.vert
+++ b/python/taichi/shaders/Circles_vk.vert
@@ -3,7 +3,7 @@
layout(location = 0) in vec3 in_position;
layout(location = 1) in vec3 in_normal;
layout(location = 2) in vec2 in_texcoord;
-layout(location = 3) in vec3 in_color;
+layout(location = 3) in vec4 in_color;
layout(binding = 0) uniform UBO {
vec3 color;
@@ -25,6 +25,6 @@ void main() {
if (ubo.use_per_vertex_color == 0) {
selected_color = ubo.color;
} else {
- selected_color = in_color;
+ selected_color = in_color.rgb;
}
}
diff --git a/python/taichi/shaders/Circles_vk_frag.spv b/python/taichi/shaders/Circles_vk_frag.spv
index e18b1cd09d69d..51146c0b7358b 100644
Binary files a/python/taichi/shaders/Circles_vk_frag.spv and b/python/taichi/shaders/Circles_vk_frag.spv differ
diff --git a/python/taichi/shaders/Circles_vk_vert.spv b/python/taichi/shaders/Circles_vk_vert.spv
index fc0a8c566bb5b..f511e9dc7e74d 100644
Binary files a/python/taichi/shaders/Circles_vk_vert.spv and b/python/taichi/shaders/Circles_vk_vert.spv differ
diff --git a/python/taichi/shaders/Lines_vk.vert b/python/taichi/shaders/Lines_vk.vert
index e009eec432945..4176ecc3ea8cc 100644
--- a/python/taichi/shaders/Lines_vk.vert
+++ b/python/taichi/shaders/Lines_vk.vert
@@ -3,7 +3,7 @@
layout(location = 0) in vec3 in_position;
layout(location = 1) in vec3 in_normal;
layout(location = 2) in vec2 in_texcoord;
-layout(location = 3) in vec3 in_color;
+layout(location = 3) in vec4 in_color;
layout(location = 0) out vec2 frag_texcoord;
@@ -25,6 +25,6 @@ void main() {
if (ubo.use_per_vertex_color == 0) {
selected_color = ubo.color;
} else {
- selected_color = in_color;
+ selected_color = in_color.rgb;
}
}
diff --git a/python/taichi/shaders/Mesh_vk.frag b/python/taichi/shaders/Mesh_vk.frag
index 45f788ef1b4dc..7596e10647ef9 100644
--- a/python/taichi/shaders/Mesh_vk.frag
+++ b/python/taichi/shaders/Mesh_vk.frag
@@ -32,10 +32,10 @@ layout(binding = 1, std430) buffer SSBO {
}
ssbo;
-layout(location = 3) in vec3 selected_color;
+layout(location = 3) in vec4 selected_color;
vec3 lambertian() {
- vec3 ambient = ubo.scene.ambient_light * selected_color;
+ vec3 ambient = ubo.scene.ambient_light * selected_color.rgb;
vec3 result = ambient;
for (int i = 0; i < ubo.scene.point_light_count; ++i) {
@@ -50,7 +50,7 @@ vec3 lambertian() {
else{
factor = max(dot(light_dir, normal), 0);
}
- vec3 diffuse = factor * selected_color * light_color;
+ vec3 diffuse = factor * selected_color.rgb * light_color;
result += diffuse;
}
@@ -58,5 +58,5 @@ vec3 lambertian() {
}
void main() {
- out_color = vec4(lambertian(), 1);
+ out_color = vec4(lambertian(), selected_color.a);
}
diff --git a/python/taichi/shaders/Mesh_vk.vert b/python/taichi/shaders/Mesh_vk.vert
index d3b99471eb705..327d80280b62a 100644
--- a/python/taichi/shaders/Mesh_vk.vert
+++ b/python/taichi/shaders/Mesh_vk.vert
@@ -3,12 +3,12 @@
layout(location = 0) in vec3 in_position;
layout(location = 1) in vec3 in_normal;
layout(location = 2) in vec2 in_texcoord;
-layout(location = 3) in vec3 in_color;
+layout(location = 3) in vec4 in_color;
layout(location = 0) out vec3 frag_pos;
layout(location = 1) out vec3 frag_normal;
layout(location = 2) out vec2 frag_texcoord;
-layout(location = 3) out vec3 selected_color;
+layout(location = 3) out vec4 selected_color;
struct SceneUBO {
vec3 camera_pos;
@@ -39,7 +39,7 @@ void main() {
frag_normal = in_normal;
if (ubo.use_per_vertex_color == 0) {
- selected_color = ubo.color;
+ selected_color = vec4(ubo.color, 1.0);
} else {
selected_color = in_color;
}
diff --git a/python/taichi/shaders/Mesh_vk_frag.spv b/python/taichi/shaders/Mesh_vk_frag.spv
index 137dcc5a80052..3656b926bbfec 100644
Binary files a/python/taichi/shaders/Mesh_vk_frag.spv and b/python/taichi/shaders/Mesh_vk_frag.spv differ
diff --git a/python/taichi/shaders/Mesh_vk_vert.spv b/python/taichi/shaders/Mesh_vk_vert.spv
index 2b2893e3b8a67..9d12325b1199d 100644
Binary files a/python/taichi/shaders/Mesh_vk_vert.spv and b/python/taichi/shaders/Mesh_vk_vert.spv differ
diff --git a/python/taichi/shaders/Particles_vk.frag b/python/taichi/shaders/Particles_vk.frag
index 81221d730b6e2..751ee9aa616d5 100644
--- a/python/taichi/shaders/Particles_vk.frag
+++ b/python/taichi/shaders/Particles_vk.frag
@@ -32,7 +32,7 @@ ssbo;
layout(location = 0) out vec4 out_color;
layout(location = 0) in vec4 pos_camera_space;
-layout(location = 1) in vec3 selected_color;
+layout(location = 1) in vec4 selected_color;
float project_z(float view_z) {
vec4 projected = ubo.scene.projection * vec4(0, 0, view_z, 1);
@@ -46,7 +46,7 @@ vec3 to_camera_space(vec3 pos) {
// operates in camera space !!
vec3 lambertian(vec3 frag_pos, vec3 frag_normal) {
- vec3 ambient = ubo.scene.ambient_light * selected_color;
+ vec3 ambient = ubo.scene.ambient_light * selected_color.rgb;
vec3 result = ambient;
for (int i = 0; i < ubo.scene.point_light_count; ++i) {
@@ -56,7 +56,7 @@ vec3 lambertian(vec3 frag_pos, vec3 frag_normal) {
normalize(to_camera_space(ssbo.point_lights[i].pos) - frag_pos);
vec3 normal = normalize(frag_normal);
vec3 diffuse =
- max(dot(light_dir, normal), 0.0) * selected_color * light_color;
+ max(dot(light_dir, normal), 0.0) * selected_color.rgb * light_color;
result += diffuse;
}
@@ -80,7 +80,7 @@ void main() {
pos_camera_space.xyz / pos_camera_space.w + coord_in_sphere * ubo.radius;
vec3 frag_normal = coord_in_sphere;
vec3 color = lambertian(frag_pos, frag_normal);
- out_color = vec4(color, 1.0);
+ out_color = vec4(color, selected_color.a);
float depth =
(pos_camera_space.z / pos_camera_space.w) + z_in_sphere * ubo.radius;
diff --git a/python/taichi/shaders/Particles_vk.vert b/python/taichi/shaders/Particles_vk.vert
index 5dbfeb7c79cc4..6a963fe846868 100644
--- a/python/taichi/shaders/Particles_vk.vert
+++ b/python/taichi/shaders/Particles_vk.vert
@@ -3,7 +3,7 @@
layout(location = 0) in vec3 in_position;
layout(location = 1) in vec3 in_normal;
layout(location = 2) in vec2 in_texcoord;
-layout(location = 3) in vec3 in_color;
+layout(location = 3) in vec4 in_color;
struct SceneUBO {
vec3 camera_pos;
@@ -25,7 +25,7 @@ layout(binding = 0) uniform UBO {
ubo;
layout(location = 0) out vec4 pos_camera_space;
-layout(location = 1) out vec3 selected_color;
+layout(location = 1) out vec4 selected_color;
void main() {
float distance = length(in_position - ubo.scene.camera_pos);
@@ -38,7 +38,7 @@ void main() {
gl_Position.y *= -1;
if (ubo.use_per_vertex_color == 0) {
- selected_color = ubo.color;
+ selected_color = vec4(ubo.color, 1.0);
} else {
selected_color = in_color;
}
diff --git a/python/taichi/shaders/Particles_vk_frag.spv b/python/taichi/shaders/Particles_vk_frag.spv
index cd3dabb8dd778..5038c2506c2cf 100644
Binary files a/python/taichi/shaders/Particles_vk_frag.spv and b/python/taichi/shaders/Particles_vk_frag.spv differ
diff --git a/python/taichi/shaders/Particles_vk_vert.spv b/python/taichi/shaders/Particles_vk_vert.spv
index c432761fbe874..d46476649e80a 100644
Binary files a/python/taichi/shaders/Particles_vk_vert.spv and b/python/taichi/shaders/Particles_vk_vert.spv differ
diff --git a/python/taichi/shaders/SetImage_vk.vert b/python/taichi/shaders/SetImage_vk.vert
index 9f40ffe1ec0ca..b8610c5486662 100644
--- a/python/taichi/shaders/SetImage_vk.vert
+++ b/python/taichi/shaders/SetImage_vk.vert
@@ -5,7 +5,7 @@
layout(location = 0) in vec3 in_position;
layout(location = 1) in vec3 in_normal;
layout(location = 2) in vec2 in_texcoord;
-layout(location = 3) in vec3 in_color;
+layout(location = 3) in vec4 in_color;
layout(location = 0) out vec2 frag_texcoord;
diff --git a/python/taichi/shaders/Triangles_vk.vert b/python/taichi/shaders/Triangles_vk.vert
index 1bc1f8ded4a03..914396bf003a5 100644
--- a/python/taichi/shaders/Triangles_vk.vert
+++ b/python/taichi/shaders/Triangles_vk.vert
@@ -3,7 +3,7 @@
layout(location = 0) in vec3 in_position;
layout(location = 1) in vec3 in_normal;
layout(location = 2) in vec2 in_texcoord;
-layout(location = 3) in vec3 in_color;
+layout(location = 3) in vec4 in_color;
layout(location = 0) out vec2 frag_texcoord;
layout(location = 1) out vec3 selected_color;
@@ -24,6 +24,6 @@ void main() {
if (ubo.use_per_vertex_color == 0) {
selected_color = ubo.color;
} else {
- selected_color = in_color;
+ selected_color = in_color.rgb;
}
}
diff --git a/python/taichi/snode/__init__.py b/python/taichi/snode/__init__.py
deleted file mode 100644
index e3b852b98d876..0000000000000
--- a/python/taichi/snode/__init__.py
+++ /dev/null
@@ -1 +0,0 @@
-from taichi.snode.fields_builder import FieldsBuilder
diff --git a/python/taichi/tools/__init__.py b/python/taichi/tools/__init__.py
index 98e9b2efbd0b3..c70d7f8f46cab 100644
--- a/python/taichi/tools/__init__.py
+++ b/python/taichi/tools/__init__.py
@@ -1,5 +1,6 @@
-from .np2ply import PLYWriter
-from .patterns import taichi_logo
-from .video import VideoManager
-
-__all__ = [s for s in dir() if not s.startswith('_')]
+from taichi.tools.async_utils import *
+from taichi.tools.cc_compose import *
+from taichi.tools.diagnose import *
+from taichi.tools.image import *
+from taichi.tools.np2ply import *
+from taichi.tools.video import *
diff --git a/python/taichi/tools/async_utils.py b/python/taichi/tools/async_utils.py
new file mode 100644
index 0000000000000..9cbec3391b044
--- /dev/null
+++ b/python/taichi/tools/async_utils.py
@@ -0,0 +1,29 @@
+import subprocess
+
+from taichi._lib import core as _ti_core
+from taichi.lang.impl import get_runtime
+
+
+def dump_dot(filepath=None, rankdir=None, embed_states_threshold=0):
+ d = get_runtime().prog.dump_dot(rankdir, embed_states_threshold)
+ if filepath is not None:
+ with open(filepath, 'w') as fh:
+ fh.write(d)
+ return d
+
+
+def dot_to_pdf(dot, filepath):
+ assert filepath.endswith('.pdf')
+ with subprocess.Popen(['dot', '-Tpdf'],
+ stdin=subprocess.PIPE,
+ stdout=subprocess.PIPE) as p:
+ pdf_contents = p.communicate(input=dot.encode())[0]
+ with open(filepath, 'wb') as fh:
+ fh.write(pdf_contents)
+
+
+def get_kernel_stats():
+ return _ti_core.get_kernel_stats()
+
+
+__all__ = []
diff --git a/python/taichi/cc_compose.py b/python/taichi/tools/cc_compose.py
similarity index 97%
rename from python/taichi/cc_compose.py
rename to python/taichi/tools/cc_compose.py
index 99bb960811e88..01bd803d22ba8 100644
--- a/python/taichi/cc_compose.py
+++ b/python/taichi/tools/cc_compose.py
@@ -24,7 +24,6 @@ def do_group_begin(self, e):
self.launches = []
def do_group_end(self, e):
- name = e['content']
self.groups[self.current_group] = list(self.launches)
self.current_group = None
self.launches = []
@@ -77,7 +76,6 @@ def do_compile_layout(self, e):
self.emit('')
def do_allocate_buffer(self, e):
- root_size = e['root_size']
gtmp_size = e['gtmp_size']
extr_size = 4 * 1024 * 1024 # pinpoint: 4 MB
@@ -106,7 +104,6 @@ def do_allocate_buffer(self, e):
self.emit('')
def do_compile_kernel(self, e):
- name = e['kernel_name']
source = e['kernel_source']
if self.emscripten:
@@ -137,7 +134,7 @@ def main(fin_name, fout_name, hdrout_name, emscripten=False):
import yaml # pylint: disable=C0415
with open(fin_name, 'r') as fin:
warnings.filterwarnings('ignore')
- obj = yaml.load(fin)
+ obj = yaml.load(fin, Loader=yaml.FullLoader)
with open(hdrout_name, 'w') as hdrout:
with open(fout_name, 'w') as fout:
@@ -147,3 +144,5 @@ def main(fin_name, fout_name, hdrout_name, emscripten=False):
if __name__ == '__main__':
main(sys.argv[1], sys.argv[2], sys.argv[3], len(sys.argv) > 4)
+
+__all__ = []
diff --git a/python/taichi/diagnose.py b/python/taichi/tools/diagnose.py
similarity index 95%
rename from python/taichi/diagnose.py
rename to python/taichi/tools/diagnose.py
index 184e508e06da0..a2d06fd84e297 100644
--- a/python/taichi/diagnose.py
+++ b/python/taichi/tools/diagnose.py
@@ -47,7 +47,7 @@ def try_print(tag, expr):
try_print('import', 'ti')
print('')
for arch in ['cc', 'cpu', 'metal', 'opengl', 'cuda', 'vulkan']:
- try_print(arch, f'ti.is_arch_supported(ti.{arch})')
+ try_print(arch, f'ti.lang.misc.is_arch_supported(ti.{arch})')
print('')
try:
@@ -112,7 +112,7 @@ def try_print(tag, expr):
ti_laplace = subprocess.check_output(
[executable, '-m', 'taichi', 'example', 'minimal'])
except Exception as e:
- print(f'`examples/laplace.py` failed: {e}')
+ print(f'`python/taichi/examples/algorithm/laplace.py` failed: {e}')
else:
print(f'{ti_laplace.decode()}')
@@ -123,3 +123,5 @@ def try_print(tag, expr):
if __name__ == '__main__':
main()
+
+__all__ = []
diff --git a/python/taichi/tools/file.py b/python/taichi/tools/file.py
deleted file mode 100644
index 22797a45afc13..0000000000000
--- a/python/taichi/tools/file.py
+++ /dev/null
@@ -1,9 +0,0 @@
-import os
-
-
-def clear_directory_with_suffix(directory, suffix):
- files = os.listdir(directory)
- assert suffix[0] != '.', "No '.' needed."
- for f in files:
- if f.endswith('.' + suffix):
- os.remove(os.path.join(directory, f))
diff --git a/python/taichi/misc/image.py b/python/taichi/tools/image.py
similarity index 75%
rename from python/taichi/misc/image.py
rename to python/taichi/tools/image.py
index 018043764324e..29e7d2d025e2b 100644
--- a/python/taichi/misc/image.py
+++ b/python/taichi/tools/image.py
@@ -1,7 +1,5 @@
-from io import BytesIO
-
import numpy as np
-from taichi.core import ti_core as _ti_core
+from taichi._lib import core as _ti_core
import taichi as ti
@@ -10,7 +8,7 @@ def cook_image_to_bytes(img):
"""
Takes a NumPy array or Taichi field of any type.
Returns a NumPy array of uint8.
- This is used by ti.imwrite and ti.imdisplay.
+ This is used by ti.imwrite.
"""
if not isinstance(img, np.ndarray):
img = img.to_numpy()
@@ -34,26 +32,6 @@ def cook_image_to_bytes(img):
return img.swapaxes(0, 1)[::-1, :]
-def imdisplay(img):
- """
- Try to display image in interactive shell.
-
- Args:
- img (Union[ti.field, np.ndarray]): A field of of array with shape `(width, height)` or `(height, width, 3)` or `(height, width, 4)`.
- """
- try:
- get_ipython()
- except:
- ti.imshow(img)
- else:
- import IPython.display # pylint: disable=C0415
- import PIL.Image # pylint: disable=C0415
- img = cook_image_to_bytes(img)
- with BytesIO() as f:
- PIL.Image.fromarray(img).save(f, 'png')
- IPython.display.display(IPython.display.Image(data=f.getvalue()))
-
-
def imresize(img, w, h=None):
"""Resize an image to a specific size.
@@ -113,32 +91,36 @@ def imread(filename, channels=0):
return img.swapaxes(0, 1)[:, ::-1, :]
-def imshow(img, window_name='imshow'):
- """Show image in a Taichi GUI.
+def imshow(img, title='imshow'):
+ """Display a taichi.field or a numpy.ndarray in a Taichi GUI window or an interactive Ipython notebook.
+ For an interactive Ipython environment, the image will be shown in the notebook.
Args:
img (Union[ti.field, np.ndarray]): A field of of array with shape `(width, height)` or `(height, width, 3)` or `(height, width, 4)`.
- window_name (str, optional): The title of GUI window. Default to `imshow`.
+ title (str, optional): The title of GUI window. Default to `imshow`.
"""
- if not isinstance(img, np.ndarray):
- img = img.to_numpy()
- assert len(img.shape) in [2,
- 3], "Image must be either RGB/RGBA or greyscale"
-
- with ti.GUI(window_name, res=img.shape[:2]) as gui:
- img = gui.cook_image(img)
- while gui.running:
- if gui.get_event(ti.GUI.ESCAPE):
- gui.running = False
-
- gui.set_image(img)
- gui.show()
+ try: # check if we are in Ipython environment
+ get_ipython()
+ except:
+ if not isinstance(img, np.ndarray):
+ img = img.to_numpy()
+ assert len(
+ img.shape) in [2,
+ 3], "Image must be either RGB/RGBA or greyscale"
+
+ with ti.GUI(title, res=img.shape[:2]) as gui:
+ img = gui.cook_image(img)
+ while gui.running:
+ if gui.get_event(ti.GUI.ESCAPE):
+ gui.running = False
+
+ gui.set_image(img)
+ gui.show()
+ else:
+ import IPython.display # pylint: disable=C0415
+ import PIL.Image # pylint: disable=C0415
+ img = cook_image_to_bytes(img)
+ IPython.display.display(PIL.Image.fromarray(img))
-__all__ = [
- 'imshow',
- 'imread',
- 'imwrite',
- 'imresize',
- 'imdisplay',
-]
+__all__ = ['imread', 'imresize', 'imshow', 'imwrite']
diff --git a/python/taichi/tools/messenger.py b/python/taichi/tools/messenger.py
deleted file mode 100644
index 906ad30c0868f..0000000000000
--- a/python/taichi/tools/messenger.py
+++ /dev/null
@@ -1,85 +0,0 @@
-import atexit
-import os
-import smtplib
-import socket
-
-import taichi as tc
-
-gmail_sender = 'taichi.messager@gmail.com'
-gmail_passwd = '6:L+XbNOp^'
-
-emailed = False
-
-
-def send_crash_report(message, receiver=None):
- global emailed
- if emailed:
- return
- emailed = True
- if receiver is None:
- receiver = os.environ.get('TI_MONITOR_EMAIL', None)
- if receiver is None:
- tc.warn('No receiver in $TI_MONITOR_EMAIL')
- return
- tc.warn('Emailing {}'.format(receiver))
- TO = receiver
- SUBJECT = 'Report'
- TEXT = message
-
- server = smtplib.SMTP('smtp.gmail.com', 587)
- server.ehlo()
- server.starttls()
- server.login(gmail_sender, gmail_passwd)
-
- BODY = '\r\n'.join([
- 'To: %s' % TO,
- 'From: %s' % gmail_sender,
- 'Subject: %s' % SUBJECT, '', TEXT
- ])
-
- try:
- server.sendmail(gmail_sender, [TO], BODY)
- except:
- print('Error sending mail')
- server.quit()
- print('Press enter or Ctrl + \\ to exit.')
-
-
-def enable(task_name):
- register_call_back(task_name)
-
-
-crashed = False
-keep = []
-
-
-def register_call_back(task_name):
- def at_exit():
- if not crashed:
- message = 'Congratulations! Your task [{}] at machine [{}] has finished.'.format(
- task_name, socket.gethostname())
- send_crash_report(message)
-
- def email_call_back(_):
- global crashed
- crashed = True
- tc.warn('Task has crashed.')
- message = 'Your task [{}] at machine [{}] has crashed.'.format(
- task_name, socket.gethostname())
- send_crash_report(message)
- atexit.unregister(at_exit)
- exit(-1)
-
- keep.append(email_call_back)
- # TODO: email_call_back should be passed to Taichi core (C++). It will then called by the signal handler when Taichi crashes
- # (std::function python_at_exit)
- # Simply register a callback in the Python scope will not work in cases when Taichi crashes
- # call_back = tc.function11(email_call_back)
- # tc.core.register_at_exit(call_back)
-
- atexit.register(at_exit)
-
-
-if __name__ == '__main__':
- register_call_back('test')
- tc.core.trigger_sig_fpe()
diff --git a/python/taichi/tools/np2ply.py b/python/taichi/tools/np2ply.py
index 5902fa4c7ea2a..e8ff38c75f076 100644
--- a/python/taichi/tools/np2ply.py
+++ b/python/taichi/tools/np2ply.py
@@ -3,8 +3,6 @@
import numpy as np
-import taichi as ti
-
class PLYWriter:
def __init__(self,
@@ -25,9 +23,8 @@ def __init__(self,
np.float32, np.float64
]
self.type_map = {}
- for i in range(len(self.ply_supported_types)):
- self.type_map[self.ply_supported_types[
- i]] = self.corresponding_numpy_types[i]
+ for i, ply_type in enumerate(self.ply_supported_types):
+ self.type_map[ply_type] = self.corresponding_numpy_types[i]
self.num_vertices = num_vertices
self.num_vertex_channels = 0
@@ -46,9 +43,10 @@ def __init__(self,
self.face_indices = -np.ones((self.num_faces, 4), dtype=np.int32)
self.comment = comment
- def add_vertex_channel(self, key: str, type: str, data: np.array):
- if type not in self.ply_supported_types:
- print("Unknown type " + type + " detected, skipping this channel")
+ def add_vertex_channel(self, key: str, data_type: str, data: np.array):
+ if data_type not in self.ply_supported_types:
+ print("Unknown type " + data_type +
+ " detected, skipping this channel")
return
if data.ndim == 1:
assert data.size == self.num_vertices, "The dimension of the vertex channel is not correct"
@@ -56,8 +54,8 @@ def add_vertex_channel(self, key: str, type: str, data: np.array):
if key in self.vertex_channels:
print("WARNING: duplicate key " + key + " detected")
self.vertex_channels.append(key)
- self.vertex_data_type.append(type)
- self.vertex_data.append(self.type_map[type](data))
+ self.vertex_data_type.append(data_type)
+ self.vertex_data.append(self.type_map[data_type](data))
else:
num_col = data.size // self.num_vertices
assert data.ndim == 2 and data.size == num_col * \
@@ -69,8 +67,8 @@ def add_vertex_channel(self, key: str, type: str, data: np.array):
if item_key in self.vertex_channels:
print("WARNING: duplicate key " + item_key + " detected")
self.vertex_channels.append(item_key)
- self.vertex_data_type.append(type)
- self.vertex_data.append(self.type_map[type](data[:, i]))
+ self.vertex_data_type.append(data_type)
+ self.vertex_data.append(self.type_map[data_type](data[:, i]))
def add_vertex_pos(self, x: np.array, y: np.array, z: np.array):
self.add_vertex_channel("x", "float", x)
@@ -164,9 +162,10 @@ def add_faces(self, indices: np.array):
(self.num_faces, vert_per_face))
self.face_indices = self.face_indices.astype(np.int32)
- def add_face_channel(self, key: str, type: str, data: np.array):
- if type not in self.ply_supported_types:
- print("Unknown type " + type + " detected, skipping this channel")
+ def add_face_channel(self, key: str, data_type: str, data: np.array):
+ if data_type not in self.ply_supported_types:
+ print("Unknown type " + data_type +
+ " detected, skipping this channel")
return
if data.ndim == 1:
assert data.size == self.num_faces, "The dimension of the face channel is not correct"
@@ -174,8 +173,8 @@ def add_face_channel(self, key: str, type: str, data: np.array):
if key in self.face_channels:
print("WARNING: duplicate key " + key + " detected")
self.face_channels.append(key)
- self.face_data_type.append(type)
- self.face_data.append(self.type_map[type](data))
+ self.face_data_type.append(data_type)
+ self.face_data.append(self.type_map[data_type](data))
else:
num_col = data.size // self.num_faces
assert data.ndim == 2 and data.size == num_col * \
@@ -187,8 +186,8 @@ def add_face_channel(self, key: str, type: str, data: np.array):
if item_key in self.face_channels:
print("WARNING: duplicate key " + item_key + " detected")
self.face_channels.append(item_key)
- self.face_data_type.append(type)
- self.face_data.append(self.type_map[type](data[:, i]))
+ self.face_data_type.append(data_type)
+ self.face_data.append(self.type_map[data_type](data[:, i]))
def add_face_id(self):
self.add_face_channel("id", "int", np.arange(self.num_faces))
@@ -204,17 +203,17 @@ def sanity_check(self):
for idx in self.face_indices.flatten():
assert idx >= 0 and idx < self.num_vertices, "The face indices are invalid"
- def print_header(self, path: str, format: str):
+ def print_header(self, path: str, _format: str):
with open(path, "w") as f:
f.writelines([
- "ply\n", "format " + format + " 1.0\n",
+ "ply\n", "format " + _format + " 1.0\n",
"comment " + self.comment + "\n"
])
f.write("element vertex " + str(self.num_vertices) + "\n")
for i in range(self.num_vertex_channels):
f.write("property " + self.vertex_data_type[i] + " " +
self.vertex_channels[i] + "\n")
- if (self.num_faces != 0):
+ if self.num_faces != 0:
f.write("element face " + str(self.num_faces) + "\n")
f.write("property list uchar int vertex_indices\n")
for i in range(self.num_face_channels):
@@ -267,7 +266,7 @@ def export_frame_ascii(self, series_num: int, path: str):
if last_4_char == ".ply":
path = path[:-4]
- real_path = path + "_" + "{0:0=6d}".format(series_num) + ".ply"
+ real_path = path + "_" + f"{series_num:0=6d}" + ".ply"
self.export_ascii(real_path)
def export_frame(self, series_num: int, path: str):
@@ -276,5 +275,8 @@ def export_frame(self, series_num: int, path: str):
if last_4_char == ".ply":
path = path[:-4]
- real_path = path + "_" + "{0:0=6d}".format(series_num) + ".ply"
+ real_path = path + "_" + f"{series_num:0=6d}" + ".ply"
self.export(real_path)
+
+
+__all__ = ['PLYWriter']
diff --git a/python/taichi/tools/video.py b/python/taichi/tools/video.py
index 846d92fd4e663..4a9e1bcf31297 100644
--- a/python/taichi/tools/video.py
+++ b/python/taichi/tools/video.py
@@ -1,8 +1,8 @@
import os
import shutil
-from taichi.core import get_os_name
-from taichi.misc.image import imwrite
+from taichi._lib.utils import get_os_name
+from taichi.tools.image import imwrite
FRAME_FN_TEMPLATE = '%06d.png'
FRAME_DIR = 'frames'
@@ -10,21 +10,22 @@
# Write the frames to the disk and then make videos (mp4 or gif) if necessary
-def scale_video(input, output, ratiow, ratioh):
- os.system('ffmpeg -i {} -vf "scale=iw*{:.4f}:ih*{:.4f}" {}'.format(
- input, ratiow, ratioh, output))
+def scale_video(input_fn, output_fn, ratiow, ratioh):
+ os.system(
+ f'ffmpeg -i {input_fn} -vf "scale=iw*{ratiow:.4f}:ih*{ratioh:.4f}" {output_fn}'
+ )
-def crop_video(input, output, x_begin, x_end, y_begin, y_end):
+def crop_video(input_fn, output_fn, x_begin, x_end, y_begin, y_end):
os.system(
- 'ffmpeg -i {} -filter:v "crop=iw*{:.4f}:ih*{:.4f}:iw*{:0.4f}:ih*{:0.4f}" {}'
- .format(input, x_end - x_begin, y_end - y_begin, x_begin, 1 - y_end,
- output))
+ f'ffmpeg -i {input_fn} -filter:v "crop=iw*{x_end - x_begin:.4f}:ih*{y_end - y_begin:.4f}:iw*{x_begin:0.4f}:ih*{1 - y_end:0.4f}" {output_fn}'
+ )
-def accelerate_video(input, output, speed):
- os.system('ffmpeg -i {} -filter:v "setpts={:.4f}*PTS" {}'.format(
- input, 1 / speed, output))
+def accelerate_video(input_fn, output_fn, speed):
+ os.system(
+ f'ffmpeg -i {input_fn} -filter:v "setpts={1 / speed:.4f}*PTS" {output_fn}'
+ )
def get_ffmpeg_path():
@@ -36,19 +37,17 @@ def mp4_to_gif(input_fn, output_fn, framerate):
palette_name = 'palette.png'
if get_os_name() == 'win':
command = get_ffmpeg_path(
- ) + " -loglevel panic -i %s -vf 'palettegen' -y %s" % (input_fn,
- palette_name)
+ ) + f" -loglevel panic -i {input_fn} -vf 'palettegen' -y {palette_name}"
else:
command = get_ffmpeg_path(
- ) + " -loglevel panic -i %s -vf 'fps=%d,scale=320:640:flags=lanczos,palettegen' -y %s" % (
- input_fn, framerate, palette_name)
+ ) + f" -loglevel panic -i {input_fn} -vf 'fps={framerate}," \
+ f"scale=320:640:flags=lanczos,palettegen' -y {palette_name}"
# print command
os.system(command)
# Generate the GIF
command = get_ffmpeg_path(
- ) + " -loglevel panic -i %s -i %s -lavfi paletteuse -y %s" % (
- input_fn, palette_name, output_fn)
+ ) + f" -loglevel panic -i {input_fn} -i {palette_name} -lavfi paletteuse -y {output_fn}"
# print command
os.system(command)
os.remove(palette_name)
@@ -117,7 +116,8 @@ def clean_frames(self):
def make_video(self, mp4=True, gif=True):
fn = self.get_output_filename('.mp4')
- command = (get_ffmpeg_path() + " -loglevel panic -framerate %d -i " % self.framerate) + os.path.join(self.frame_directory, FRAME_FN_TEMPLATE) + \
+ command = (get_ffmpeg_path() + f" -loglevel panic -framerate {self.framerate} -i ") + os.path.join(
+ self.frame_directory, FRAME_FN_TEMPLATE) + \
" -s:v " + str(self.width) + 'x' + str(self.height) + \
" -c:v libx264 -profile:v high -crf 1 -pix_fmt yuv420p -y " + fn
@@ -139,7 +139,7 @@ def interpolate_frames(frame_dir, mul=4):
images_interpolated = []
for f in sorted(files):
if f.endswith('png'):
- images.append(cv2.imread(f) / 255.0)
+ images.append(cv2.imread(f) / 255.0) # pylint: disable=E1101
for i in range(len(images) - 1):
images_interpolated.append(images[i])
@@ -152,12 +152,12 @@ def interpolate_frames(frame_dir, mul=4):
os.makedirs('interpolated', exist_ok=True)
for i, img in enumerate(images_interpolated):
- cv2.imwrite('interpolated/{:05d}.png'.format(i), img * 255.0)
+ cv2.imwrite(f'interpolated/{i:05d}.png', img * 255.0) # pylint: disable=E1101
-def ffmpeg_common_args(frame_rate, input, width, height, crf, output_path):
- return f"{get_ffmpeg_path()} -y -loglevel panic -framerate {frame_rate} -i {input} -s:v {width}x{height} " + \
- f"-c:v libx264 -profile:v high -crf {crf} -pix_fmt yuv420p {output_path}"
+def ffmpeg_common_args(frame_rate, input_fn, width, height, crf, output_path):
+ return f"{get_ffmpeg_path()} -y -loglevel panic -framerate {frame_rate} -i {input_fn} -s:v {width}x{height} " + \
+ f"-c:v libx264 -profile:v high -crf {crf} -pix_fmt yuv420p {output_path}"
def make_video(input_files,
@@ -173,20 +173,20 @@ def make_video(input_files,
tmp_dir = 'tmp_ffmpeg_dir'
os.mkdir(tmp_dir)
if width % 2 != 0:
- print("Width ({}) not divisible by 2".format(width))
+ print(f"Width ({width}) not divisible by 2")
width -= 1
if height % 2 != 0:
- print("Height ({}) not divisible by 2".format(width))
+ print(f"Height ({width}) not divisible by 2")
height -= 1
for i, inp in enumerate(input_files):
- shutil.copy(inp, os.path.join(tmp_dir, '%06d.png' % i))
+ shutil.copy(inp, os.path.join(tmp_dir, f'{i:06d}.png'))
inputs = f'{tmp_dir}/%06d.png'
command = ffmpeg_common_args(frame_rate, inputs, width, height, crf,
output_path)
ret = os.system(command)
assert ret == 0, "ffmpeg failed to generate video file."
for i in range(len(input_files)):
- os.remove(os.path.join(tmp_dir, '%06d.png' % i))
+ os.remove(os.path.join(tmp_dir, f'{i:06d}.png'))
os.rmdir(tmp_dir)
elif isinstance(input_files, str):
assert width != 0 and height != 0
@@ -196,3 +196,6 @@ def make_video(input_files,
assert ret == 0, "ffmpeg failed to generate video file."
else:
assert False, f'input_files should be list (of files) or str (of file template, e.g., "%04d.png") instead of {type(input_files)}'
+
+
+__all__ = ['VideoManager']
diff --git a/python/taichi/torch_io.py b/python/taichi/torch_io.py
deleted file mode 100644
index 819c4d9164572..0000000000000
--- a/python/taichi/torch_io.py
+++ /dev/null
@@ -1,26 +0,0 @@
-from taichi.lang.kernel_impl import kernel
-from taichi.type.annotations import ext_arr, template
-
-
-@kernel
-def from_torch_template(expr: template(), torch_tensor: ext_arr()):
- for i in expr:
- expr[i] = torch_tensor[i]
-
-
-@kernel
-def to_torch_template(expr: template(), torch_tensor: ext_arr()):
- for i in expr:
- torch_tensor[i] = expr[i]
-
-
-def from_torch(expr, torch_tensor):
- if not expr.from_torch_:
- expr.from_torch_ = lambda x: from_torch_template(expr, x.contiguous())
- expr.from_torch_(torch_tensor)
-
-
-def to_torch(expr, torch_tensor):
- if not expr.to_torch_:
- expr.to_torch_ = lambda x: to_torch_template(expr, x.contiguous())
- expr.to_torch_(torch_tensor)
diff --git a/python/taichi/type/__init__.py b/python/taichi/type/__init__.py
deleted file mode 100644
index cc82c2a60626c..0000000000000
--- a/python/taichi/type/__init__.py
+++ /dev/null
@@ -1,2 +0,0 @@
-from taichi.type.annotations import *
-from taichi.type.primitive_types import *
diff --git a/python/taichi/type/primitive_types.py b/python/taichi/type/primitive_types.py
deleted file mode 100644
index 09a47006529a9..0000000000000
--- a/python/taichi/type/primitive_types.py
+++ /dev/null
@@ -1,86 +0,0 @@
-from taichi.core import ti_core
-
-# Real types
-
-float32 = ti_core.DataType_f32
-"""32-bit single precision floating point data type.
-"""
-f32 = float32
-"""Alias for :const:`~taichi.type.primitive_types.float32`
-"""
-float64 = ti_core.DataType_f64
-"""64-bit double precision floating point data type.
-"""
-f64 = float64
-"""Alias for :const:`~taichi.type.primitive_types.float64`
-"""
-
-real_types = [f32, f64, float]
-real_type_ids = [id(t) for t in real_types]
-
-# Integer types
-
-int8 = ti_core.DataType_i8
-i8 = int8
-int16 = ti_core.DataType_i16
-i16 = int16
-int32 = ti_core.DataType_i32
-"""32-bit signed integer data type.
-"""
-i32 = int32
-"""Alias for :const:`~taichi.type.primitive_types.int32`
-"""
-int64 = ti_core.DataType_i64
-"""64-bit signed integer data type.
-"""
-i64 = int64
-"""Alias for :const:`~taichi.type.primitive_types.int64`
-"""
-
-uint8 = ti_core.DataType_u8
-u8 = uint8
-uint16 = ti_core.DataType_u16
-u16 = uint16
-uint32 = ti_core.DataType_u32
-"""32-bit unsigned integer data type.
-"""
-u32 = uint32
-"""Alias for :const:`~taichi.type.primitive_types.uint32`
-"""
-uint64 = ti_core.DataType_u64
-"""64-bit unsigned integer data type.
-"""
-u64 = uint64
-"""Alias for :const:`~taichi.type.primitive_types.uint64`
-"""
-
-integer_types = [i8, i16, i32, i64, u8, u16, u32, u64, int]
-integer_type_ids = [id(t) for t in integer_types]
-
-types = real_types + integer_types
-type_ids = [id(t) for t in types]
-
-__all__ = [
- 'float32',
- 'f32',
- 'float64',
- 'f64',
- 'int8',
- 'i8',
- 'int16',
- 'i16',
- 'int32',
- 'i32',
- 'int64',
- 'i64',
- 'uint8',
- 'u8',
- 'uint16',
- 'u16',
- 'uint32',
- 'u32',
- 'uint64',
- 'u64',
- 'real_types',
- 'integer_types',
-]
diff --git a/python/taichi/types/__init__.py b/python/taichi/types/__init__.py
new file mode 100644
index 0000000000000..ee913ba4f7ccd
--- /dev/null
+++ b/python/taichi/types/__init__.py
@@ -0,0 +1,5 @@
+from taichi.types.annotations import *
+from taichi.types.compound_types import *
+from taichi.types.primitive_types import *
+from taichi.types.quantized_types import *
+from taichi.types.utils import *
diff --git a/python/taichi/type/annotations.py b/python/taichi/types/annotations.py
similarity index 58%
rename from python/taichi/type/annotations.py
rename to python/taichi/types/annotations.py
index 3de93335c435b..4e98870c2dc7a 100644
--- a/python/taichi/type/annotations.py
+++ b/python/taichi/types/annotations.py
@@ -6,28 +6,54 @@ class ArgAnyArray:
Args:
element_dim (Union[Int, NoneType], optional): None if not specified (will be treated as 0 for external arrays), 0 if scalar elements, 1 if vector elements, and 2 if matrix elements.
+ element_shape (Union[Tuple[Int], NoneType]): None if not specified, shapes of each element. For example, element_shape must be 1d for vector and 2d tuple for matrix. This argument is ignored for external arrays for now.
+ field_dim (Union[Int, NoneType]): None if not specified, number of field dimensions. This argument is ignored for external arrays for now.
layout (Union[Layout, NoneType], optional): None if not specified (will be treated as Layout.AOS for external arrays), Layout.AOS or Layout.SOA.
"""
- def __init__(self, element_dim=None, layout=None):
+ def __init__(self,
+ element_dim=None,
+ element_shape=None,
+ field_dim=None,
+ layout=None):
if element_dim is not None and (element_dim < 0 or element_dim > 2):
raise ValueError(
"Only scalars, vectors, and matrices are allowed as elements of ti.any_arr()"
)
- self.element_dim = element_dim
+ if element_dim is not None and element_shape is not None and len(
+ element_shape) != element_dim:
+ raise ValueError(
+ f"Both element_shape and element_dim are specified, but shape doesn't match specified dim: {len(element_shape)}!={element_dim}"
+ )
+ self.element_shape = element_shape
+ self.element_dim = len(
+ element_shape) if element_shape is not None else element_dim
+ self.field_dim = field_dim
self.layout = layout
- def check_element_dim(self, arg, arg_dim):
+ def _check_element_dim(self, arg, arg_dim):
if self.element_dim is not None and self.element_dim != arg_dim:
raise ValueError(
f"Invalid argument into ti.any_arr() - required element_dim={self.element_dim}, but {arg} is provided"
)
- def check_layout(self, arg):
+ def _check_layout(self, arg):
if self.layout is not None and self.layout != arg.layout:
raise ValueError(
f"Invalid argument into ti.any_arr() - required layout={self.layout}, but {arg} is provided"
)
+ def _check_element_shape(self, shapes):
+ if self.element_shape is not None and shapes != self.element_shape:
+ raise ValueError(
+ f"Invalid argument into ti.any_arr() - required element_shape={self.element_shape}, but {shapes} is provided"
+ )
+
+ def _check_field_dim(self, field_dim):
+ if self.field_dim is not None and field_dim != self.field_dim:
+ raise ValueError(
+ f"Invalid argument into ti.any_arr() - required field_dim={self.field_dim}, but {field_dim} is provided"
+ )
+
def ext_arr():
"""Type annotation for external arrays.
@@ -49,7 +75,7 @@ def ext_arr():
any_arr = ArgAnyArray
-"""Alias for :class:`~taichi.type.annotations.ArgAnyArray`.
+"""Alias for :class:`~taichi.types.annotations.ArgAnyArray`.
Example::
@@ -80,7 +106,12 @@ def __init__(self, tensor=None, dim=None):
template = Template
-"""Alias for :class:`~taichi.type.annotations.Template`.
+"""Alias for :class:`~taichi.types.annotations.Template`.
"""
-__all__ = ['ext_arr', 'any_arr', 'template']
+
+class sparse_matrix_builder:
+ pass
+
+
+__all__ = ['ext_arr', 'any_arr', 'template', 'sparse_matrix_builder']
diff --git a/python/taichi/types/compound_types.py b/python/taichi/types/compound_types.py
new file mode 100644
index 0000000000000..cad8854969f16
--- /dev/null
+++ b/python/taichi/types/compound_types.py
@@ -0,0 +1,21 @@
+import taichi
+
+
+class CompoundType:
+ pass
+
+
+# TODO: maybe move MatrixType, StructType here to avoid the circular import?
+def matrix(n, m, dtype):
+ return taichi.lang.matrix.MatrixType(n, m, dtype)
+
+
+def vector(n, dtype):
+ return taichi.lang.matrix.MatrixType(n, 1, dtype)
+
+
+def struct(**kwargs):
+ return taichi.lang.struct.StructType(**kwargs)
+
+
+__all__ = ['matrix', 'vector', 'struct']
diff --git a/python/taichi/types/primitive_types.py b/python/taichi/types/primitive_types.py
new file mode 100644
index 0000000000000..726da1831a32a
--- /dev/null
+++ b/python/taichi/types/primitive_types.py
@@ -0,0 +1,176 @@
+from taichi._lib import core as ti_core
+
+# ========================================
+# real types
+
+# ----------------------------------------
+
+float16 = ti_core.DataType_f16
+"""16-bit precision floating point data type.
+"""
+
+# ----------------------------------------
+
+f16 = float16
+"""Alias for :const:`~taichi.types.primitive_types.float16`
+"""
+
+# ----------------------------------------
+
+float32 = ti_core.DataType_f32
+"""32-bit single precision floating point data type.
+"""
+
+# ----------------------------------------
+
+f32 = float32
+"""Alias for :const:`~taichi.types.primitive_types.float32`
+"""
+
+# ----------------------------------------
+
+float64 = ti_core.DataType_f64
+"""64-bit double precision floating point data type.
+"""
+
+# ----------------------------------------
+
+f64 = float64
+"""Alias for :const:`~taichi.types.primitive_types.float64`
+"""
+# ----------------------------------------
+
+# ========================================
+# Integer types
+
+# ----------------------------------------
+
+int8 = ti_core.DataType_i8
+"""8-bit signed integer data type.
+"""
+
+# ----------------------------------------
+
+i8 = int8
+"""Alias for :const:`~taichi.types.primitive_types.int8`
+"""
+
+# ----------------------------------------
+
+int16 = ti_core.DataType_i16
+"""16-bit signed integer data type.
+"""
+
+# ----------------------------------------
+
+i16 = int16
+"""Alias for :const:`~taichi.types.primitive_types.int16`
+"""
+
+# ----------------------------------------
+
+int32 = ti_core.DataType_i32
+"""32-bit signed integer data type.
+"""
+
+# ----------------------------------------
+
+i32 = int32
+"""Alias for :const:`~taichi.types.primitive_types.int32`
+"""
+
+# ----------------------------------------
+
+int64 = ti_core.DataType_i64
+"""64-bit signed integer data type.
+"""
+
+# ----------------------------------------
+
+i64 = int64
+"""Alias for :const:`~taichi.types.primitive_types.int64`
+"""
+
+# ----------------------------------------
+
+uint8 = ti_core.DataType_u8
+"""8-bit unsigned integer data type.
+"""
+
+# ----------------------------------------
+
+u8 = uint8
+"""Alias for :const:`~taichi.types.primitive_types.uint8`
+"""
+
+# ----------------------------------------
+
+uint16 = ti_core.DataType_u16
+"""16-bit unsigned integer data type.
+"""
+
+# ----------------------------------------
+
+u16 = uint16
+"""Alias for :const:`~taichi.types.primitive_types.uint16`
+"""
+
+# ----------------------------------------
+
+uint32 = ti_core.DataType_u32
+"""32-bit unsigned integer data type.
+"""
+
+# ----------------------------------------
+
+u32 = uint32
+"""Alias for :const:`~taichi.types.primitive_types.uint32`
+"""
+
+# ----------------------------------------
+
+uint64 = ti_core.DataType_u64
+"""64-bit unsigned integer data type.
+"""
+
+# ----------------------------------------
+
+u64 = uint64
+"""Alias for :const:`~taichi.types.primitive_types.uint64`
+"""
+
+# ----------------------------------------
+
+real_types = [f16, f32, f64, float]
+real_type_ids = [id(t) for t in real_types]
+
+integer_types = [i8, i16, i32, i64, u8, u16, u32, u64, int]
+integer_type_ids = [id(t) for t in integer_types]
+
+types = real_types + integer_types
+type_ids = [id(t) for t in types]
+
+__all__ = [
+ 'float32',
+ 'f32',
+ 'float64',
+ 'f64',
+ 'float16',
+ 'f16',
+ 'int8',
+ 'i8',
+ 'int16',
+ 'i16',
+ 'int32',
+ 'i32',
+ 'int64',
+ 'i64',
+ 'uint8',
+ 'u8',
+ 'uint16',
+ 'u16',
+ 'uint32',
+ 'u32',
+ 'uint64',
+ 'u64',
+]
diff --git a/python/taichi/types/quantized_types.py b/python/taichi/types/quantized_types.py
new file mode 100644
index 0000000000000..19a6dec44a2cb
--- /dev/null
+++ b/python/taichi/types/quantized_types.py
@@ -0,0 +1,129 @@
+from taichi._lib.utils import ti_core as _ti_core
+from taichi.lang import impl
+from taichi.types.primitive_types import i32
+
+
+class TypeFactory:
+ """A Python-side TypeFactory wrapper."""
+ def __init__(self):
+ self.core = _ti_core.get_type_factory_instance()
+
+ def custom_int(self, bits, signed=True, compute_type=None):
+ """Generates a custom int type.
+
+ Args:
+ bits (int): Number of bits.
+ signed (bool): Signed or unsigned.
+ compute_type (DataType): Type for computation.
+
+ Returns:
+ DataType: The specified type.
+ """
+ if compute_type is None:
+ compute_type = impl.get_runtime().default_ip
+ if isinstance(compute_type, _ti_core.DataType):
+ compute_type = compute_type.get_ptr()
+ return self.core.get_custom_int_type(bits, signed, compute_type)
+
+ def custom_float(self,
+ significand_type,
+ exponent_type=None,
+ compute_type=None,
+ scale=1.0):
+ """Generates a custom float type.
+
+ Args:
+ significand_type (DataType): Type of significand.
+ exponent_type (DataType): Type of exponent.
+ compute_type (DataType): Type for computation.
+ scale (float): Scaling factor.
+
+ Returns:
+ DataType: The specified type.
+ """
+ if compute_type is None:
+ compute_type = impl.get_runtime().default_fp
+ if isinstance(compute_type, _ti_core.DataType):
+ compute_type = compute_type.get_ptr()
+ return self.core.get_custom_float_type(significand_type,
+ exponent_type,
+ compute_type,
+ scale=scale)
+
+
+# Unstable API
+type_factory = TypeFactory()
+
+
+class Quant:
+ """Generator of quantized types.
+
+ For more details, read https://yuanming.taichi.graphics/publication/2021-quantaichi/quantaichi.pdf.
+ """
+ @staticmethod
+ def int(bits, signed=False, compute=None):
+ """Generates a quantized type for integers.
+
+ Args:
+ bits (int): Number of bits.
+ signed (bool): Signed or unsigned.
+ compute (DataType): Type for computation.
+
+ Returns:
+ DataType: The specified type.
+ """
+ if compute is None:
+ compute = impl.get_runtime().default_ip
+ return type_factory.custom_int(bits, signed, compute)
+
+ @staticmethod
+ def fixed(frac, signed=True, num_range=1.0, compute=None):
+ """Generates a quantized type for fixed-point real numbers.
+
+ Args:
+ frac (int): Number of bits.
+ signed (bool): Signed or unsigned.
+ num_range (float): Range of the number.
+ compute (DataType): Type for computation.
+
+ Returns:
+ DataType: The specified type.
+ """
+ # TODO: handle cases with frac > 32
+ frac_type = Quant.int(bits=frac, signed=signed, compute=i32)
+ if signed:
+ scale = num_range / 2**(frac - 1)
+ else:
+ scale = num_range / 2**frac
+ if compute is None:
+ compute = impl.get_runtime().default_fp
+ return type_factory.custom_float(frac_type, None, compute, scale)
+
+ @staticmethod
+ def float(exp, frac, signed=True, compute=None):
+ """Generates a quantized type for floating-point real numbers.
+
+ Args:
+ exp (int): Number of exponent bits.
+ frac (int): Number of fraction bits.
+ signed (bool): Signed or unsigned.
+ compute (DataType): Type for computation.
+
+ Returns:
+ DataType: The specified type.
+ """
+ # Exponent is always unsigned
+ exp_type = Quant.int(bits=exp, signed=False, compute=i32)
+ # TODO: handle cases with frac > 32
+ frac_type = Quant.int(bits=frac, signed=signed, compute=i32)
+ if compute is None:
+ compute = impl.get_runtime().default_fp
+ return type_factory.custom_float(significand_type=frac_type,
+ exponent_type=exp_type,
+ compute_type=compute)
+
+
+# Unstable API
+quant = Quant
+
+__all__ = []
diff --git a/python/taichi/types/utils.py b/python/taichi/types/utils.py
new file mode 100644
index 0000000000000..5f168cfc8d605
--- /dev/null
+++ b/python/taichi/types/utils.py
@@ -0,0 +1,7 @@
+from taichi._lib import core as ti_core
+
+is_signed = ti_core.is_signed
+
+is_integral = ti_core.is_integral
+
+__all__ = ['is_signed', 'is_integral']
diff --git a/python/taichi/ui/__init__.py b/python/taichi/ui/__init__.py
index 3023d814c3c81..972d3fc806635 100644
--- a/python/taichi/ui/__init__.py
+++ b/python/taichi/ui/__init__.py
@@ -1 +1,2 @@
+from .gui import *
from .ui import *
diff --git a/python/taichi/ui/camera.py b/python/taichi/ui/camera.py
index 09586945f39c1..e9c20f5b5849b 100644
--- a/python/taichi/ui/camera.py
+++ b/python/taichi/ui/camera.py
@@ -3,7 +3,6 @@
from taichi.lang.matrix import Vector
from .utils import euler_to_vec, vec_to_euler
-from .window import Window
class Camera:
@@ -51,7 +50,7 @@ def bottom(self, bottom):
def z_near(self, z_near):
self.ptr.z_near(z_near)
- def z_near(self, z_far):
+ def z_far(self, z_far):
self.ptr.z_far(z_far)
# move the camera according to user inputs, FPS game style.
@@ -64,6 +63,7 @@ def track_user_inputs(self,
front = (self.curr_lookat - self.curr_position).normalized()
position_change = Vector([0.0, 0.0, 0.0])
left = self.curr_up.cross(front)
+ up = self.curr_up
if window.is_pressed('w'):
position_change = front * movement_speed
if window.is_pressed('s'):
@@ -72,6 +72,10 @@ def track_user_inputs(self,
position_change = left * movement_speed
if window.is_pressed('d'):
position_change = -left * movement_speed
+ if window.is_pressed('e'):
+ position_change = up * movement_speed
+ if window.is_pressed('q'):
+ position_change = -up * movement_speed
self.position(*(self.curr_position + position_change))
self.lookat(*(self.curr_lookat + position_change))
diff --git a/python/taichi/ui/canvas.py b/python/taichi/ui/canvas.py
index eccf0fb6e5353..985b4ce3f28a1 100644
--- a/python/taichi/ui/canvas.py
+++ b/python/taichi/ui/canvas.py
@@ -1,12 +1,6 @@
-from taichi.core import ti_core as _ti_core
-from taichi.lang.impl import default_cfg
-from taichi.lang.kernel_impl import kernel
-from taichi.lang.ops import get_addr
-from taichi.type.annotations import ext_arr, template
-
from .staging_buffer import (copy_colors_to_vbo, copy_vertices_to_vbo,
get_vbo_field, to_u8_rgba)
-from .utils import *
+from .utils import get_field_info
class Canvas:
@@ -92,4 +86,4 @@ def circles(self,
def scene(self, scene):
"""Draw a 3D scene on the canvas"""
- self.canvas.scene(scene)
+ self.canvas.scene(scene.scene)
diff --git a/python/taichi/ui/gui.py b/python/taichi/ui/gui.py
index 404e120451db3..946d21d49ffef 100644
--- a/python/taichi/ui/gui.py
+++ b/python/taichi/ui/gui.py
@@ -1,99 +1,938 @@
-import pathlib
-from contextlib import contextmanager
+import math
+import numbers
+import os
-from taichi.core import ti_core as _ti_core
-from taichi.lang.impl import default_cfg
-from taichi.lang.kernel_impl import kernel
-from taichi.lang.ops import get_addr
-from taichi.type.annotations import ext_arr, template
+import numpy as np
+import taichi.lang
+from taichi._kernels import (tensor_to_image, vector_to_fast_image,
+ vector_to_image)
+from taichi._lib import core as _ti_core
+from taichi.lang.field import Field, ScalarField
-from .utils import *
+import taichi as ti
-class Gui:
- def __init__(self, gui) -> None:
- self.gui = gui #reference to a PyGui
+# For window creation and drawing in the original ti.GUI system.
+class GUI:
+ """Taichi Graphical User Interface class.
- @contextmanager
- def sub_window(self, name, x, y, width, height):
- """Creating a context manager for subwindow
+ Args:
+ name (str, optional): The name of the GUI to be constructed.
+ Default is 'Taichi'.
+ res (Union[int, List[int]], optional): The resolution of created
+ GUI. Default is 512*512. If `res` is scalar, then width will be equal to height.
+ background_color (int, optional): The background color of created GUI.
+ Default is 0x000000.
+ show_gui (bool, optional): Specify whether to render the GUI. Default is True.
+ fullscreen (bool, optional): Specify whether to render the GUI in
+ fullscreen mode. Default is False.
+ fast_gui (bool, optional): Specify whether to use fast gui mode of
+ Taichi. Default is False.
- Note:
- All args of this method should align with `begin`.
+ Returns:
+ :class:`~taichi.misc.gui.GUI` :The created taichi GUI object.
+
+ """
+ class Event:
+ def __init__(self):
+ self.type = None
+ self.modifier = None
+ self.pos = None
+ self.key = None
+ self.delta = None
+
+ # Event keys
+ SHIFT = 'Shift'
+ ALT = 'Alt'
+ CTRL = 'Control'
+ ESCAPE = 'Escape'
+ RETURN = 'Return'
+ TAB = 'Tab'
+ BACKSPACE = 'BackSpace'
+ SPACE = ' '
+ UP = 'Up'
+ DOWN = 'Down'
+ LEFT = 'Left'
+ RIGHT = 'Right'
+ CAPSLOCK = 'Caps_Lock'
+ LMB = 'LMB'
+ MMB = 'MMB'
+ RMB = 'RMB'
+ EXIT = 'WMClose'
+ WHEEL = 'Wheel'
+ MOVE = 'Motion'
+
+ # Event types
+ MOTION = _ti_core.KeyEvent.EType.Move
+ PRESS = _ti_core.KeyEvent.EType.Press
+ RELEASE = _ti_core.KeyEvent.EType.Release
+
+ def __init__(self,
+ name='Taichi',
+ res=512,
+ background_color=0x0,
+ show_gui=True,
+ fullscreen=False,
+ fast_gui=False):
+ show_gui = self.get_bool_environ('TI_GUI_SHOW', show_gui)
+ fullscreen = self.get_bool_environ('TI_GUI_FULLSCREEN', fullscreen)
+ fast_gui = self.get_bool_environ('TI_GUI_FAST', fast_gui)
+
+ self.name = name
+ if isinstance(res, numbers.Number):
+ res = (res, res)
+ self.res = res
+ self.fast_gui = fast_gui
+ if fast_gui:
+ self.img = np.ascontiguousarray(
+ np.zeros(self.res[0] * self.res[1], dtype=np.uint32))
+ fast_buf = self.img.ctypes.data
+ else:
+ # The GUI canvas uses RGBA for storage, therefore we need NxMx4 for an image.
+ self.img = np.ascontiguousarray(
+ np.zeros(self.res + (4, ), np.float32))
+ fast_buf = 0
+ self.core = _ti_core.GUI(name, core_veci(*res), show_gui, fullscreen,
+ fast_gui, fast_buf)
+ self.canvas = self.core.get_canvas()
+ self.background_color = background_color
+ self.key_pressed = set()
+ self.event = None
+ self.frame = 0
+ self.clear()
+
+ def __enter__(self):
+ return self
+
+ def __exit__(self, e_type, val, tb):
+ self.close()
+
+ def __del__(self):
+ self.close()
+
+ def close(self):
+ self.core = None # dereference to call GUI::~GUI()
+
+ # Widget system
+
+ class WidgetValue:
+ def __init__(self, gui, wid):
+ self.gui = gui
+ self.wid = wid
+
+ @property
+ def value(self):
+ return self.gui.core.get_widget_value(self.wid)
+
+ @value.setter
+ def value(self, value):
+ self.gui.core.set_widget_value(self.wid, value)
+
+ @staticmethod
+ def get_bool_environ(key, default):
+ """Get an environment variable and cast to bool.
+ Args:
+ key (str): The environment variable key.
+ default (bool): The default value.
+ Return:
+ The environment variable value cast to bool. If the value is not found, directly return argument 'default'.
+ """
+ if key not in os.environ:
+ return default
+ return bool(int(os.environ[key]))
+
+ def slider(self, text, minimum, maximum, step=1):
+ """Create a slider object on canvas to be manipulated with.
Args:
- x (float): The x-coordinate (between 0 and 1) of the top-left corner of the subwindow, relative to the full window.
- y (float): The y-coordinate (between 0 and 1) of the top-left corner of the subwindow, relative to the full window.
- width (float): The width of the subwindow relative to the full window.
- height (float): The height of the subwindow relative to the full window.
+ text (str): The title of slider.
+ minimum (Number): The minimum value of slider.
+ maximum (Number): The maximum value of slider.
+ step (Number, optional): The changing step of slider. Default is 1.
- Usage::
+ Return:
+ :class:`~taichi.misc.gui.GUI.WidgetValue` :The created slider object.
- >>> with gui.sub_window(name, x, y, width, height) as g:
- >>> g.text("Hello, World!")
"""
- self.begin(name, x, y, width, height)
- try:
- yield self
- finally:
- self.end()
+ wid = self.core.make_slider(text, minimum, minimum, maximum, step)
+ return GUI.WidgetValue(self, wid)
- def begin(self, name, x, y, width, height):
- """Creates a subwindow that holds imgui widgets.
+ def label(self, text):
+ """Create a label object on canvas.
- All widget function calls (e.g. `text`, `button`) after the `begin` and before the next `end` will describe the widgets within this subwindow.
+ Args:
+ text (str): The title of label.
+
+ Return:
+ :class:`~taichi.misc.gui.GUI.WidgetValue` :The created label object.
+
+ """
+ wid = self.core.make_label(text, 0)
+ return GUI.WidgetValue(self, wid)
+
+ def button(self, text, event_name=None):
+ """Create a button object on canvas to be manipulated with.
Args:
- x (float): The x-coordinate (between 0 and 1) of the top-left corner of the subwindow, relative to the full window.
- y (float): The y-coordinate (between 0 and 1) of the top-left corner of the subwindow, relative to the full window.
- width (float): The width of the subwindow relative to the full window.
- height (float): The height of the subwindow relative to the full window.
+ text (str): The title of button.
+ event_name (str, optional): The event name associated with button.
+ Default is WidgetButton_{text}
+
+ Return:
+ The event name associated with created button.
+
"""
- self.gui.begin(name, x, y, width, height)
+ event_name = event_name or f'WidgetButton_{text}'
+ self.core.make_button(text, event_name)
+ return event_name
+
+ # Drawing system
+
+ def clear(self, color=None):
+ """Clear the canvas with the color provided.
+
+ Args:
+ color (int, optional): Specify the color to clear the canvas. Default
+ is the background color of GUI.
- def end(self):
- """End the description of the current subwindow.
"""
- self.gui.end()
+ if color is None:
+ color = self.background_color
+ self.canvas.clear(color)
+
+ def cook_image(self, img):
+ if img.dtype in [np.uint8, np.uint16, np.uint32, np.uint64]:
+ img = img.astype(np.float32) * (1 / np.iinfo(img.dtype).max)
+ elif img.dtype in [np.float16, np.float32, np.float64]:
+ img = img.astype(np.float32)
+ else:
+ raise ValueError(
+ f'Data type {img.dtype} not supported in GUI.set_image')
+
+ if len(img.shape) == 2:
+ img = img[..., None]
+
+ if img.shape[2] == 1:
+ img = img + np.zeros((1, 1, 4), np.float32)
+ if img.shape[2] == 3:
+ zeros = np.zeros((img.shape[0], img.shape[1], 1), np.float32)
+ img = np.concatenate([img, zeros], axis=2)
+ if img.shape[2] == 2:
+ zeros = np.zeros((img.shape[0], img.shape[1], 2), np.float32)
+ img = np.concatenate([img, zeros], axis=2)
+
+ assert img.shape[2] == 4, "Image must be grayscale, RG, RGB or RGBA"
+
+ res = img.shape[:2]
+ assert res == self.res, "Image resolution does not match GUI resolution"
+ return np.ascontiguousarray(img)
+
+ def get_image(self):
+ """Get the image data.
+
+ Returns:
+ :class:`numpy.array` :The image data in numpy contiguous array type.
- def text(self, text):
- """Declares a line of text.
"""
- self.gui.text(text)
+ self.img = np.ascontiguousarray(self.img)
+ self.core.get_img(self.img.ctypes.data)
+ return self.img
- def checkbox(self, text, old_value):
- """Declares a checkbox, and returns whether or not it has been checked.
+ def set_image(self, img):
+ """Sets an image to display on the window.
+ The image pixels are set from the values of `img[i, j]`, where `i` indicates the horizontal coordinates (from left to right) and `j` the vertical coordinates (from bottom to top).
+ If the window size is `(x, y)`, then `img` must be one of:
+ - `ti.field(shape=(x, y))`, a gray-scale image
+ - `ti.field(shape=(x, y, 3))`, where `3` is for `(r, g, b)` channels
+ - `ti.field(shape=(x, y, 2))`, where `2` is for `(r, g)` channels
+ - `ti.Vector.field(3, shape=(x, y))` `(r, g, b)` channels on each component
+ - `ti.Vector.field(2, shape=(x, y))` `(r, g)` channels on each component
+ - `np.ndarray(shape=(x, y))`
+ - `np.ndarray(shape=(x, y, 3))`
+ - `np.ndarray(shape=(x, y, 2))`
+ The data type of `img` must be one of:
+ - `uint8`, range `[0, 255]`
+ - `uint16`, range `[0, 65535]`
+ - `uint32`, range `[0, 4294967295]`
+ - `float32`, range `[0, 1]`
+ - `float64`, range `[0, 1]`
Args:
- text (str): a line of text to be shown next to the checkbox
- old_value (bool): whether the checkbox is currently checked
+ img (Union[ti.field, numpy.array]): The color array representing the
+ image to be drawn. Support greyscale, RG, RGB, and RGBA color
+ representations. Its shape must match GUI resolution.
+
"""
- return self.gui.checkbox(text, old_value)
- def slider_float(self, text, old_value, minimum, maximum):
- """Declares a slider, and returns its newest value.
+ if self.fast_gui:
+ assert isinstance(img, taichi.lang.matrix.MatrixField), \
+ "Only ti.Vector.field is supported in GUI.set_image when fast_gui=True"
+ assert img.shape == self.res, \
+ "Image resolution does not match GUI resolution"
+ assert img.n in [3, 4] and img.m == 1, \
+ "Only RGB images are supported in GUI.set_image when fast_gui=True"
+ assert img.dtype in [ti.f32, ti.f64, ti.u8], \
+ "Only f32, f64, u8 are supported in GUI.set_image when fast_gui=True"
+
+ vector_to_fast_image(img, self.img)
+ return
+
+ if isinstance(img, ScalarField):
+ if _ti_core.is_integral(img.dtype) or len(img.shape) != 2:
+ # Images of uint is not optimized by xxx_to_image
+ self.img = self.cook_image(img.to_numpy())
+ else:
+ # Type matched! We can use an optimized copy kernel.
+ assert img.shape \
+ == self.res, "Image resolution does not match GUI resolution"
+ tensor_to_image(img, self.img)
+ ti.sync()
+
+ elif isinstance(img, taichi.lang.matrix.MatrixField):
+ if _ti_core.is_integral(img.dtype):
+ self.img = self.cook_image(img.to_numpy())
+ else:
+ # Type matched! We can use an optimized copy kernel.
+ assert img.shape == self.res, \
+ "Image resolution does not match GUI resolution"
+ assert img.n in [2, 3, 4] and img.m == 1, \
+ "Only greyscale, RG, RGB or RGBA images are supported in GUI.set_image"
+
+ vector_to_image(img, self.img)
+ ti.sync()
+
+ elif isinstance(img, np.ndarray):
+ self.img = self.cook_image(img)
+
+ else:
+ raise ValueError(
+ f"GUI.set_image only takes a Taichi field or NumPy array, not {type(img)}"
+ )
+
+ self.core.set_img(self.img.ctypes.data)
+
+ def circle(self, pos, color=0xFFFFFF, radius=1):
+ """Draw a single circle on canvas.
Args:
- text (str): a line of text to be shown next to the slider
- old_value (float): the current value of the slider.
- minimum (float): the minimum value of the slider.
- maximum (float): the maximum value of the slider.
+ pos (Union[List[int], numpy.array]): The position of the circle.
+ color (int, Optional): The color of the circle. Default is 0xFFFFFF.
+ radius (Number, Optional): The radius of the circle in pixel. Default is 1.
+
"""
- return self.gui.slider_float(text, old_value, minimum, maximum)
+ self.canvas.circle_single(pos[0], pos[1], color, radius)
- def color_edit_3(self, text, old_value):
- """Declares a color edit palate.
+ def circles(self,
+ pos,
+ radius=1,
+ color=0xFFFFFF,
+ palette=None,
+ palette_indices=None):
+ """Draw a list of circles on canvas.
Args:
- text (str): a line of text to be shown next to the palate
- old_value (Tuple[float]): the current value of the color, this should be a tuple of floats in [0,1] that indicates RGB values.
+ pos (numpy.array): The positions of the circles.
+ radius (Number, optional): The radius of the circles in pixel. Default is 1.
+ color (int, optional): The color of the circles. Default is 0xFFFFFF.
+ palette (list[int], optional): The List of colors from which to
+ choose to draw. Default is None.
+ palette_indices (Union[list[int], ti.field, numpy.array], optional):
+ The List of indices that choose color from palette for each
+ circle. Shape must match pos. Default is None.
+
"""
- return self.gui.color_edit_3(text, old_value)
+ n = pos.shape[0]
+ if len(pos.shape) == 3:
+ assert pos.shape[2] == 1
+ pos = pos[:, :, 0]
- def button(self, text):
- """Declares a button, and returns whether or not it had just been clicked.
+ assert pos.shape == (n, 2)
+ pos = np.ascontiguousarray(pos.astype(np.float32))
+ # Note: do not use "pos = int(pos.ctypes.data)" here
+ # Otherwise pos will get garbage collected by Python
+ # and the pointer to its data becomes invalid
+ pos_ptr = int(pos.ctypes.data)
+
+ if isinstance(color, np.ndarray):
+ assert color.shape == (n, )
+ color = np.ascontiguousarray(color.astype(np.uint32))
+ color_array = int(color.ctypes.data)
+ color_single = 0
+ elif isinstance(color, int):
+ color_array = 0
+ color_single = color
+ else:
+ raise ValueError(
+ 'Color must be an ndarray or int (e.g., 0x956333)')
+
+ if palette is not None:
+ assert palette_indices is not None, 'palette must be used together with palette_indices'
+
+ if isinstance(palette_indices, Field):
+ ind_int = palette_indices.to_numpy().astype(np.uint32)
+ elif isinstance(palette_indices, list) or isinstance(
+ palette_indices, np.ndarray):
+ ind_int = np.array(palette_indices).astype(np.uint32)
+ else:
+ try:
+ ind_int = np.array(palette_indices)
+ except:
+ raise TypeError(
+ 'palette_indices must be a type that can be converted to numpy.ndarray'
+ )
+
+ assert issubclass(
+ ind_int.dtype.type,
+ np.integer), 'palette_indices must be an integer array'
+ assert ind_int.shape == (
+ n,
+ ), 'palette_indices must be in 1-d shape with shape (num_particles, )'
+ assert min(
+ ind_int
+ ) >= 0, 'the min of palette_indices must not be less than zero'
+ assert max(ind_int) < len(
+ palette
+ ), 'the max of palette_indices must not exceed the length of palette'
+ color_array = np.array(palette, dtype=np.uint32)[ind_int]
+ color_array = np.ascontiguousarray(color_array)
+ color_array = color_array.ctypes.data
+
+ if isinstance(radius, np.ndarray):
+ assert radius.shape == (n, )
+ radius = np.ascontiguousarray(radius.astype(np.float32))
+ radius_array = int(radius.ctypes.data)
+ radius_single = 0
+ elif isinstance(radius, numbers.Number):
+ radius_array = 0
+ radius_single = radius
+ else:
+ raise ValueError('Radius must be an ndarray or float (e.g., 0.4)')
+
+ self.canvas.circles_batched(n, pos_ptr, color_single, color_array,
+ radius_single, radius_array)
+
+ def triangles(self, a, b, c, color=0xFFFFFF):
+ """Draw a list of triangles on canvas.
Args:
- text (str): a line of text to be shown next to the button
+ a (numpy.array): The positions of the first points of triangles.
+ b (numpy.array): The positions of the second points of triangles.
+ c (numpy.array): The positions of the thrid points of triangles.
+ color (Union[int, numpy.array], optional): The color or colors of triangles.
+ Can be either a single color or a list of colors whose shape matches
+ the shape of a & b & c. Default is 0xFFFFFF.
+
"""
- return self.gui.button(text)
+ assert a.shape == b.shape
+ assert a.shape == c.shape
+ n = a.shape[0]
+ if len(a.shape) == 3:
+ assert a.shape[2] == 1
+ a = a[:, :, 0]
+ b = b[:, :, 0]
+ c = c[:, :, 0]
+
+ assert a.shape == (n, 2)
+ a = np.ascontiguousarray(a.astype(np.float32))
+ b = np.ascontiguousarray(b.astype(np.float32))
+ c = np.ascontiguousarray(c.astype(np.float32))
+ # Note: do not use "a = int(a.ctypes.data)" here
+ # Otherwise a will get garbage collected by Python
+ # and the pointer to its data becomes invalid
+ a_ptr = int(a.ctypes.data)
+ b_ptr = int(b.ctypes.data)
+ c_ptr = int(c.ctypes.data)
+
+ if isinstance(color, np.ndarray):
+ assert color.shape == (n, )
+ color = np.ascontiguousarray(color.astype(np.uint32))
+ color_array = int(color.ctypes.data)
+ color_single = 0
+ elif isinstance(color, int):
+ color_array = 0
+ color_single = color
+ else:
+ raise ValueError(
+ '"color" must be an ndarray or int (e.g., 0x956333)')
+
+ self.canvas.triangles_batched(n, a_ptr, b_ptr, c_ptr, color_single,
+ color_array)
+
+ def triangle(self, a, b, c, color=0xFFFFFF):
+ """Draw a single triangle on canvas.
+
+ Args:
+ a (List[Number]): The position of the first point of triangle. Shape must be 2.
+ b (List[Number]): The position of the second point of triangle. Shape must be 2.
+ c (List[Number]): The position of the third point of triangle. Shape must be 2.
+ color (int, optional): The color of the triangle. Default is 0xFFFFFF.
+
+ """
+ self.canvas.triangle_single(a[0], a[1], b[0], b[1], c[0], c[1], color)
+
+ def lines(self, begin, end, radius=1, color=0xFFFFFF):
+ """Draw a list of lines on canvas.
+
+ Args:
+ begin (numpy.array): The positions of one end of lines.
+ end (numpy.array): The positions of the other end of lines.
+ radius (Union[Number, numpy.array], optional): The width of lines.
+ Can be either a single width or a list of width whose shape matches
+ the shape of begin & end. Default is 1.
+ color (Union[int, numpy.array], optional): The color or colors of lines.
+ Can be either a single color or a list of colors whose shape matches
+ the shape of begin & end. Default is 0xFFFFFF.
+
+ """
+ assert begin.shape == end.shape
+ n = begin.shape[0]
+ if len(begin.shape) == 3:
+ assert begin.shape[2] == 1
+ begin = begin[:, :, 0]
+ end = end[:, :, 0]
+
+ assert begin.shape == (n, 2)
+ begin = np.ascontiguousarray(begin.astype(np.float32))
+ end = np.ascontiguousarray(end.astype(np.float32))
+ # Note: do not use "begin = int(begin.ctypes.data)" here
+ # Otherwise begin will get garbage collected by Python
+ # and the pointer to its data becomes invalid
+ begin_ptr = int(begin.ctypes.data)
+ end_ptr = int(end.ctypes.data)
+
+ if isinstance(color, np.ndarray):
+ assert color.shape == (n, )
+ color = np.ascontiguousarray(color.astype(np.uint32))
+ color_array = int(color.ctypes.data)
+ color_single = 0
+ elif isinstance(color, int):
+ color_array = 0
+ color_single = color
+ else:
+ raise ValueError(
+ 'Color must be an ndarray or int (e.g., 0x956333)')
+
+ if isinstance(radius, np.ndarray):
+ assert radius.shape == (n, )
+ radius = np.ascontiguousarray(radius.astype(np.float32))
+ radius_array = int(radius.ctypes.data)
+ radius_single = 0
+ elif isinstance(radius, numbers.Number):
+ radius_array = 0
+ radius_single = radius
+ else:
+ raise ValueError('Radius must be an ndarray or float (e.g., 0.4)')
+
+ self.canvas.paths_batched(n, begin_ptr, end_ptr, color_single,
+ color_array, radius_single, radius_array)
+
+ def line(self, begin, end, radius=1, color=0xFFFFFF):
+ """Draw a single line on canvas.
+
+ Args:
+ begin (List[Number]): The position of one end of line. Shape must be 2.
+ end (List[Number]): The position of the other end of line. Shape must be 2.
+ radius (Number, optional): The width of line. Default is 1.
+ color (int, optional): The color of line. Default is 0xFFFFFF.
+
+ """
+ self.canvas.path_single(begin[0], begin[1], end[0], end[1], color,
+ radius)
+
+ @staticmethod
+ def _arrow_to_lines(orig, major, tip_scale=0.2, angle=45):
+ angle = math.radians(180 - angle)
+ c, s = math.cos(angle), math.sin(angle)
+ minor1 = np.array([
+ major[:, 0] * c - major[:, 1] * s,
+ major[:, 0] * s + major[:, 1] * c
+ ]).swapaxes(0, 1)
+ minor2 = np.array([
+ major[:, 0] * c + major[:, 1] * s,
+ -major[:, 0] * s + major[:, 1] * c
+ ]).swapaxes(0, 1)
+ end = orig + major
+ return [(orig, end), (end, end + minor1 * tip_scale),
+ (end, end + minor2 * tip_scale)]
+
+ def arrows(self, orig, direction, radius=1, color=0xffffff, **kwargs):
+ """Draw a list arrows on canvas.
+
+ Args:
+ orig (numpy.array): The positions where arrows start.
+ direction (numpy.array): The directions where arrows point to.
+ radius (Union[Number, np.array], optional): The width of arrows. Default is 1.
+ color (Union[int, np.array], optional): The color or colors of arrows. Default is 0xffffff.
+
+ """
+ for begin, end in self._arrow_to_lines(orig, direction, **kwargs):
+ self.lines(begin, end, radius, color)
+
+ def arrow(self, orig, direction, radius=1, color=0xffffff, **kwargs):
+ """Draw a single arrow on canvas.
+
+ Args:
+ orig (List[Number]): The position where arrow starts. Shape must be 2.
+ direction (List[Number]): The direction where arrow points to. Shape must be 2.
+ radius (Number, optional): The width of arrow. Default is 1.
+ color (int, optional): The color of arrow. Default is 0xFFFFFF.
+
+ """
+ orig = np.array([orig])
+ direction = np.array([direction])
+ for begin, end in self._arrow_to_lines(orig, direction, **kwargs):
+ self.line(begin[0], end[0], radius, color)
+
+ def rect(self, topleft, bottomright, radius=1, color=0xFFFFFF):
+ """Draw a single rectangle on canvas.
+
+ Args:
+ topleft (List[Number]): The position of the topleft corner of rectangle.
+ Shape must be 2.
+ bottomright (List[Number]): The position of the bottomright corner
+ of rectangle. Shape must be 2.
+ radius (Number, optional): The width of rectangle's sides. Default is 1.
+ color (int, optional): The color of rectangle. Default is 0xFFFFFF.
+
+ """
+ a = topleft[0], topleft[1]
+ b = bottomright[0], topleft[1]
+ c = bottomright[0], bottomright[1]
+ d = topleft[0], bottomright[1]
+ self.line(a, b, radius, color)
+ self.line(b, c, radius, color)
+ self.line(c, d, radius, color)
+ self.line(d, a, radius, color)
+
+ def text(self, content, pos, font_size=15, color=0xFFFFFF):
+ """Draw texts on canvas.
+
+ Args:
+ content (str): The text to be drawn on canvas.
+ pos (List[Number]): The position where the text is to be put.
+ font_size (Number, optional): The font size of the text.
+ color (int, optional): The color of the text. Default is 0xFFFFFF.
+
+ """
+
+ # TODO: refactor Canvas::text
+ font_size = float(font_size)
+ pos = core_vec(*pos)
+ r, g, b = hex_to_rgb(color)
+ color = core_vec(r, g, b, 1)
+ self.canvas.text(content, pos, font_size, color)
+
+ @staticmethod
+ def _make_field_base(w, h, bound):
+ x = np.linspace(bound / w, 1 - bound / w, w)
+ y = np.linspace(bound / h, 1 - bound / h, h)
+ base = np.array(np.meshgrid(x, y))
+ base = base.swapaxes(0, 1).swapaxes(1, 2).swapaxes(0, 1)
+ return base.reshape(w * h, 2)
+
+ def point_field(self, radius, color=0xffffff, bound=0.5):
+ """Draw a field of points on canvas.
+
+ Args:
+ radius (np.array): The pattern and radius of the field of points.
+ color (Union[int, np.array], optional): The color or colors of points.
+ Default is 0xFFFFFF.
+ bound (Number, optional): The boundary of the field. Default is 0.5.
+
+ """
+ assert len(radius.shape) == 2
+ base = self._make_field_base(radius.shape[0], radius.shape[1], bound)
+ radius = radius.reshape(radius.shape[0] * radius.shape[1])
+ self.circles(base, radius=radius, color=color)
+
+ def arrow_field(self,
+ direction,
+ radius=1,
+ color=0xffffff,
+ bound=0.5,
+ **kwargs):
+ """Draw a field of arrows on canvas.
+
+ Args:
+ direction (np.array): The pattern and direction of the field of arrows.
+ color (Union[int, np.array], optional): The color or colors of arrows.
+ Default is 0xFFFFFF.
+ bound (Number, optional): The boundary of the field. Default is 0.5.
+
+ """
+ assert len(direction.shape) == 3
+ assert direction.shape[2] == 2
+ base = self._make_field_base(direction.shape[0], direction.shape[1],
+ bound)
+ direction = direction.reshape(direction.shape[0] * direction.shape[1],
+ 2)
+ self.arrows(base, direction, radius=radius, color=color, **kwargs)
+
+ def show(self, file=None):
+ """Show the frame or save current frame as a picture.
+
+ Args:
+ file (str, optional): The path & name of the picture to be saved.
+ Default is None.
+
+ """
+ self.core.update()
+ if file:
+ self.core.screenshot(file)
+ self.frame += 1
+ self.clear()
+
+ # Event system
+
+ class EventFilter:
+ def __init__(self, *e_filter):
+ self.filter = set()
+ for ent in e_filter:
+ if isinstance(ent, (list, tuple)):
+ e_type, key = ent
+ ent = (e_type, key)
+ self.filter.add(ent)
+
+ def match(self, e):
+ if (e.type, e.key) in self.filter:
+ return True
+ if e.type in self.filter:
+ return True
+ if e.key in self.filter:
+ return True
+ return False
+
+ def has_key_event(self):
+ """Check if there are any key event registered.
+
+ Returns:
+ Bool to indicate whether there is any key event registered.
+
+ """
+ return self.core.has_key_event()
+
+ def get_event(self, *e_filter):
+ """Check if the specific event is triggered.
+
+ Args:
+ *e_filter (ti.GUI.EVENT): The specific event to be checked.
+
+ Returns:
+ Bool to indicate whether the specific event is triggered.
+
+ """
+ for e in self.get_events(*e_filter):
+ self.event = e
+ return True
+ else:
+ return False
+
+ def get_events(self, *e_filter):
+ """Get a list of events that are triggered.
+
+ Args:
+ *e_filter (List[ti.GUI.EVENT]): The type of events to be filtered.
+
+ Returns:
+ :class:`~taichi.misc.gui.GUI.EVENT` :A list of events that are triggered.
+
+ """
+ e_filter = e_filter and GUI.EventFilter(*e_filter) or None
+
+ while True:
+ if not self.has_key_event():
+ break
+ e = self.get_key_event()
+ if e_filter is None or e_filter.match(e): # pylint: disable=E1101
+ yield e
+
+ def get_key_event(self):
+ """Get keyboard triggered event.
+
+ Returns:
+ :class:`~taichi.misc.gui.GUI.EVENT` :The keyboard triggered event.
+
+ """
+ self.core.wait_key_event()
+
+ e = GUI.Event()
+ event = self.core.get_key_event_head()
+
+ e.type = event.type
+ e.key = event.key
+ e.pos = self.core.canvas_untransform(event.pos)
+ e.pos = (e.pos[0], e.pos[1])
+ e.modifier = []
+
+ if e.key == GUI.WHEEL:
+ e.delta = event.delta
+ else:
+ e.delta = (0, 0)
+
+ for mod in ['Shift', 'Alt', 'Control']:
+ if self.is_pressed(mod):
+ e.modifier.append(mod)
+
+ if e.type == GUI.PRESS:
+ self.key_pressed.add(e.key)
+ else:
+ self.key_pressed.discard(e.key)
+
+ self.core.pop_key_event_head()
+ return e
+
+ def is_pressed(self, *keys):
+ """Check if the specific key or keys are pressed.
+
+ Args:
+ *keys (Union[str, List[str]]): The string that stands for keys in keyboard.
+
+ Returns:
+ Bool to indicate whether the key or keys are pressed.
+
+ """
+ for key in keys:
+ if key in ['Shift', 'Alt', 'Control']:
+ if key + '_L' in self.key_pressed or key + '_R' in self.key_pressed:
+ return True
+ if key in self.key_pressed:
+ return True
+ else:
+ return False
+
+ def get_cursor_pos(self):
+ """Get the current position of mouse.
+
+ Returns:
+ The current position of mouse.
+
+ """
+ pos = self.core.get_cursor_pos()
+ return pos[0], pos[1]
+
+ @property
+ def running(self):
+ """Get the property of whether the gui is running.
+
+ Returns:
+ The running property of gui(bool).
+
+ """
+ return not self.core.should_close
+
+ @running.setter
+ def running(self, value):
+ if value:
+ self.core.should_close = 0
+ elif not self.core.should_close:
+ self.core.should_close = 1
+
+ @property
+ def fps_limit(self):
+ """Get the property of fps limit.
+
+ Returns:
+ The property of fps limit of gui.
+
+ """
+ if self.core.frame_delta_limit == 0:
+ return None
+ return 1 / self.core.frame_delta_limit
+
+ @fps_limit.setter
+ def fps_limit(self, value):
+ if value is None:
+ self.core.frame_delta_limit = 0
+ else:
+ self.core.frame_delta_limit = 1 / value
+
+
+def rgb_to_hex(c):
+ """Convert rgb color format to hex color format.
+
+ Args:
+ c (List[int]): The rgb representation of color.
+
+ Returns:
+ The hex representation of color.
+
+ """
+ def to255(x):
+ return np.clip(np.int32(x * 255), 0, 255)
+
+ return (to255(c[0]) << 16) + (to255(c[1]) << 8) + to255(c[2])
+
+
+def hex_to_rgb(color):
+ """Convert hex color format to rgb color format.
+
+ Args:
+ color (int): The hex representation of color.
+
+ Returns:
+ The rgb representation of color.
+
+ """
+ r, g, b = (color >> 16) & 0xff, (color >> 8) & 0xff, color & 0xff
+ return r / 255, g / 255, b / 255
+
+
+def core_veci(*args):
+ if isinstance(args[0], _ti_core.Vector2i):
+ return args[0]
+ if isinstance(args[0], _ti_core.Vector3i):
+ return args[0]
+ if isinstance(args[0], tuple):
+ args = tuple(*args)
+ if len(args) == 2:
+ return _ti_core.Vector2i(int(args[0]), int(args[1]))
+ if len(args) == 3:
+ return _ti_core.Vector3i(int(args[0]), int(args[1]), int(args[2]))
+ if len(args) == 4:
+ return _ti_core.Vector4i(int(args[0]), int(args[1]), int(args[2]),
+ int(args[3]))
+ assert False, type(args[0])
+
+
+def core_vec(*args):
+ if isinstance(args[0], _ti_core.Vector2f):
+ return args[0]
+ if isinstance(args[0], _ti_core.Vector3f):
+ return args[0]
+ if isinstance(args[0], _ti_core.Vector4f):
+ return args[0]
+ if isinstance(args[0], _ti_core.Vector2d):
+ return args[0]
+ if isinstance(args[0], _ti_core.Vector3d):
+ return args[0]
+ if isinstance(args[0], _ti_core.Vector4d):
+ return args[0]
+ if isinstance(args[0], tuple):
+ args = tuple(*args)
+ if _ti_core.get_default_float_size() == 4:
+ if len(args) == 2:
+ return _ti_core.Vector2f(float(args[0]), float(args[1]))
+ if len(args) == 3:
+ return _ti_core.Vector3f(float(args[0]), float(args[1]),
+ float(args[2]))
+ if len(args) == 4:
+ return _ti_core.Vector4f(float(args[0]), float(args[1]),
+ float(args[2]), float(args[3]))
+ assert False, type(args[0])
+ else:
+ if len(args) == 2:
+ return _ti_core.Vector2d(float(args[0]), float(args[1]))
+ if len(args) == 3:
+ return _ti_core.Vector3d(float(args[0]), float(args[1]),
+ float(args[2]))
+ if len(args) == 4:
+ return _ti_core.Vector4d(float(args[0]), float(args[1]),
+ float(args[2]), float(args[3]))
+ assert False, type(args[0])
+
+
+__all__ = [
+ 'GUI',
+ 'rgb_to_hex',
+ 'hex_to_rgb',
+]
diff --git a/python/taichi/ui/imgui.py b/python/taichi/ui/imgui.py
new file mode 100644
index 0000000000000..28b426d48246e
--- /dev/null
+++ b/python/taichi/ui/imgui.py
@@ -0,0 +1,91 @@
+from contextlib import contextmanager
+
+
+#For declaring IMGUI components in a ti.Window created by the GGUI system.
+class Gui:
+ def __init__(self, gui) -> None:
+ self.gui = gui #reference to a PyGui
+
+ @contextmanager
+ def sub_window(self, name, x, y, width, height):
+ """Creating a context manager for subwindow
+
+ Note:
+ All args of this method should align with `begin`.
+
+ Args:
+ x (float): The x-coordinate (between 0 and 1) of the top-left corner of the subwindow, relative to the full window.
+ y (float): The y-coordinate (between 0 and 1) of the top-left corner of the subwindow, relative to the full window.
+ width (float): The width of the subwindow relative to the full window.
+ height (float): The height of the subwindow relative to the full window.
+
+ Usage::
+
+ >>> with gui.sub_window(name, x, y, width, height) as g:
+ >>> g.text("Hello, World!")
+ """
+ self.begin(name, x, y, width, height)
+ try:
+ yield self
+ finally:
+ self.end()
+
+ def begin(self, name, x, y, width, height):
+ """Creates a subwindow that holds imgui widgets.
+
+ All widget function calls (e.g. `text`, `button`) after the `begin` and before the next `end` will describe the widgets within this subwindow.
+
+ Args:
+ x (float): The x-coordinate (between 0 and 1) of the top-left corner of the subwindow, relative to the full window.
+ y (float): The y-coordinate (between 0 and 1) of the top-left corner of the subwindow, relative to the full window.
+ width (float): The width of the subwindow relative to the full window.
+ height (float): The height of the subwindow relative to the full window.
+ """
+ self.gui.begin(name, x, y, width, height)
+
+ def end(self):
+ """End the description of the current subwindow.
+ """
+ self.gui.end()
+
+ def text(self, text):
+ """Declares a line of text.
+ """
+ self.gui.text(text)
+
+ def checkbox(self, text, old_value):
+ """Declares a checkbox, and returns whether or not it has been checked.
+
+ Args:
+ text (str): a line of text to be shown next to the checkbox
+ old_value (bool): whether the checkbox is currently checked
+ """
+ return self.gui.checkbox(text, old_value)
+
+ def slider_float(self, text, old_value, minimum, maximum):
+ """Declares a slider, and returns its newest value.
+
+ Args:
+ text (str): a line of text to be shown next to the slider
+ old_value (float): the current value of the slider.
+ minimum (float): the minimum value of the slider.
+ maximum (float): the maximum value of the slider.
+ """
+ return self.gui.slider_float(text, old_value, minimum, maximum)
+
+ def color_edit_3(self, text, old_value):
+ """Declares a color edit palate.
+
+ Args:
+ text (str): a line of text to be shown next to the palate
+ old_value (Tuple[float]): the current value of the color, this should be a tuple of floats in [0,1] that indicates RGB values.
+ """
+ return self.gui.color_edit_3(text, old_value)
+
+ def button(self, text):
+ """Declares a button, and returns whether or not it had just been clicked.
+
+ Args:
+ text (str): a line of text to be shown next to the button
+ """
+ return self.gui.button(text)
diff --git a/python/taichi/ui/scene.py b/python/taichi/ui/scene.py
index 8b5f5646aaa88..d7333f9630d26 100644
--- a/python/taichi/ui/scene.py
+++ b/python/taichi/ui/scene.py
@@ -1,17 +1,14 @@
-import pathlib
-
-from taichi.core import ti_core as _ti_core
-from taichi.lang.impl import default_cfg, field
+from taichi._lib import core as _ti_core
+from taichi.lang.impl import field
from taichi.lang.kernel_impl import kernel
from taichi.lang.matrix import Vector
-from taichi.lang.ops import atomic_add, get_addr
-from taichi.type.annotations import ext_arr, template
-from taichi.type.primitive_types import f32
+from taichi.lang.ops import atomic_add
+from taichi.types.annotations import template
+from taichi.types.primitive_types import f32
-from .camera import Camera
from .staging_buffer import (copy_colors_to_vbo, copy_normals_to_vbo,
copy_vertices_to_vbo, get_vbo_field)
-from .utils import get_field_info
+from .utils import check_ggui_availability, get_field_info
normals_field_cache = {}
@@ -23,8 +20,7 @@ def get_normals_field(vertices):
normal_weights = field(f32, shape=(N, ))
normals_field_cache[vertices] = (normals, normal_weights)
return (normals, normal_weights)
- else:
- return normals_field_cache[vertices]
+ return normals_field_cache[vertices]
@kernel
@@ -76,14 +72,15 @@ def gen_normals(vertices, indices):
return normals
-class Scene(_ti_core.PyScene):
+class Scene:
"""A 3D scene, which can contain meshes and particles, and can be rendered on a canvas
"""
def __init__(self):
- super().__init__()
+ check_ggui_availability()
+ self.scene = _ti_core.PyScene()
def set_camera(self, camera):
- super().set_camera(camera.ptr)
+ self.scene.set_camera(camera.ptr)
def mesh(self,
vertices,
@@ -113,8 +110,8 @@ def mesh(self,
vbo_info = get_field_info(vbo)
indices_info = get_field_info(indices)
- super().mesh(vbo_info, has_per_vertex_color, indices_info, color,
- two_sided)
+ self.scene.mesh(vbo_info, has_per_vertex_color, indices_info, color,
+ two_sided)
def particles(self,
centers,
@@ -135,10 +132,10 @@ def particles(self,
if has_per_vertex_color:
copy_colors_to_vbo(vbo, per_vertex_color)
vbo_info = get_field_info(vbo)
- super().particles(vbo_info, has_per_vertex_color, color, radius)
+ self.scene.particles(vbo_info, has_per_vertex_color, color, radius)
- def point_light(self, pos, color):
- super().point_light(pos, color)
+ def point_light(self, pos, color): # pylint: disable=W0235
+ self.scene.point_light(pos, color)
def ambient_light(self, color):
- super().ambient_light(color)
+ self.scene.ambient_light(color)
diff --git a/python/taichi/ui/staging_buffer.py b/python/taichi/ui/staging_buffer.py
index b66822a125b7c..f7377c9843002 100644
--- a/python/taichi/ui/staging_buffer.py
+++ b/python/taichi/ui/staging_buffer.py
@@ -1,16 +1,10 @@
-from taichi.core import ti_core as _ti_core
-from taichi.lang.impl import default_cfg, field, static
from taichi.lang.kernel_impl import kernel
from taichi.lang.matrix import Vector
-from taichi.lang.ndrange import ndrange
-from taichi.lang.ops import atomic_add, get_addr
-from taichi.type.annotations import ext_arr, template
-from taichi.type.primitive_types import f32, u8
+from taichi.types.annotations import template
+from taichi.types.primitive_types import f32, u8
import taichi as ti
-from .utils import get_field_info
-
vbo_field_cache = {}
@@ -20,13 +14,12 @@ def get_vbo_field(vertices):
pos = 3
normal = 3
tex_coord = 2
- color = 3
+ color = 4
vertex_stride = pos + normal + tex_coord + color
vbo = Vector.field(vertex_stride, f32, shape=(N, ))
vbo_field_cache[vertices] = vbo
return vbo
- else:
- return vbo_field_cache[vertices]
+ return vbo_field_cache[vertices]
@kernel
@@ -37,6 +30,14 @@ def copy_to_vbo(vbo: template(), src: template(), offset: template(),
vbo[i][offset + c] = src[i][c]
+@kernel
+def fill_vbo(vbo: template(), value: template(), offset: template(),
+ num_components: template()):
+ for i in vbo:
+ for c in ti.static(range(num_components)):
+ vbo[i][offset + c] = value
+
+
def validate_input_field(f, name):
if f.dtype != f32:
raise Exception(f"{name} needs to have dtype f32")
@@ -53,29 +54,31 @@ def validate_input_field(f, name):
def copy_vertices_to_vbo(vbo, vertices):
validate_input_field(vertices, "vertices")
if not 2 <= vertices.n <= 3:
- raise Exception(f'vertices can only be 2D or 3D vector fields')
+ raise Exception('vertices can only be 2D or 3D vector fields')
copy_to_vbo(vbo, vertices, 0, vertices.n)
def copy_normals_to_vbo(vbo, normals):
validate_input_field(normals, "normals")
if normals.n != 3:
- raise Exception(f'normals can only be 3D vector fields')
+ raise Exception('normals can only be 3D vector fields')
copy_to_vbo(vbo, normals, 3, normals.n)
def copy_texcoords_to_vbo(vbo, texcoords):
validate_input_field(texcoords, "texcoords")
if texcoords.n != 2:
- raise Exception(f'texcoords can only be 3D vector fields')
+ raise Exception('texcoords can only be 3D vector fields')
copy_to_vbo(vbo, texcoords, 6, texcoords.n)
def copy_colors_to_vbo(vbo, colors):
validate_input_field(colors, "colors")
- if colors.n != 3:
- raise Exception(f'colors can only be 3D vector fields')
+ if colors.n != 3 and colors.n != 4:
+ raise Exception('colors can only be 3D/4D vector fields')
copy_to_vbo(vbo, colors, 8, colors.n)
+ if colors.n == 3:
+ fill_vbo(vbo, 1.0, 11, 1)
@ti.kernel
@@ -87,6 +90,9 @@ def copy_image_f32_to_u8(src: ti.template(), dst: ti.template(),
c = max(0.0, min(1.0, c))
c = c * 255
dst[i, j][k] = int(c)
+ if num_components < 4:
+ # alpha channel
+ dst[i, j][3] = 255
@ti.kernel
@@ -95,6 +101,9 @@ def copy_image_u8_to_u8(src: ti.template(), dst: ti.template(),
for i, j in src:
for k in ti.static(range(num_components)):
dst[i, j][k] = src[i, j][k]
+ if num_components < 4:
+ # alpha channel
+ dst[i, j][3] = 255
# ggui renderer always assumes the input image to be u8 RGBA
@@ -105,11 +114,11 @@ def copy_image_u8_to_u8(src: ti.template(), dst: ti.template(),
def to_u8_rgba(image):
if not hasattr(image, 'n') or image.m != 1:
raise Exception(
- f'the input image needs to be a Vector field (matrix with 1 column)'
+ 'the input image needs to be a Vector field (matrix with 1 column)'
)
if len(image.shape) != 2:
raise Exception(
- f"the shape of the image must be of the form (width,height)")
+ "the shape of the image must be of the form (width,height)")
if image.dtype == u8 and image.n == 4:
# already in the desired format
diff --git a/python/taichi/ui/ui.py b/python/taichi/ui/ui.py
index 7026e4c6b01b8..d98c1d1215ae1 100644
--- a/python/taichi/ui/ui.py
+++ b/python/taichi/ui/ui.py
@@ -1,37 +1,17 @@
-import pathlib
+from taichi._lib import core as _ti_core
-from taichi.core import ti_core as _ti_core
-from taichi.lang.impl import default_cfg, field
-from taichi.lang.kernel_impl import kernel
-from taichi.lang.ops import get_addr
-from taichi.type.annotations import ext_arr, template
+from .camera import Camera
+from .canvas import Canvas # pylint: disable=unused-import
+from .constants import * # pylint: disable=unused-import,wildcard-import
+from .imgui import Gui # pylint: disable=unused-import
+from .scene import Scene # pylint: disable=unused-import
+from .utils import check_ggui_availability
+from .window import Window # pylint: disable=unused-import
-if _ti_core.GGUI_AVAILABLE:
- from .camera import Camera
- from .canvas import Canvas
- from .constants import *
- from .gui import Gui
- from .scene import Scene
- from .window import Window
+def make_camera():
+ check_ggui_availability()
+ return Camera(_ti_core.PyCamera())
- def make_camera():
- return Camera(_ti_core.PyCamera())
- ProjectionMode = _ti_core.ProjectionMode
-
-else:
-
- def err_no_ggui():
- raise Exception("GGUI Not Available")
-
- class Window:
- def __init__(self, name, res, vsync=False):
- err_no_ggui()
-
- class Scene:
- def __init__(self):
- err_no_ggui()
-
- def make_camera():
- err_no_ggui()
+ProjectionMode = _ti_core.ProjectionMode if _ti_core.GGUI_AVAILABLE else None
diff --git a/python/taichi/ui/utils.py b/python/taichi/ui/utils.py
index 0b18c726f81f1..3bc7aa1604450 100644
--- a/python/taichi/ui/utils.py
+++ b/python/taichi/ui/utils.py
@@ -1,13 +1,8 @@
-import pathlib
-from math import acos, asin, cos, pi, sin
+from math import acos, asin, cos, sin
-from taichi.core import ti_core as _ti_core
+from taichi._lib import core as _ti_core
from taichi.lang.impl import default_cfg
-from taichi.lang.kernel_impl import kernel
from taichi.lang.matrix import Vector
-from taichi.lang.ops import get_addr
-from taichi.type.annotations import ext_arr, template
-from taichi.type.primitive_types import u64
def get_field_info(field):
@@ -20,6 +15,8 @@ def get_field_info(field):
info.field_source = _ti_core.FieldSource.TaichiCuda
elif default_cfg().arch == _ti_core.x64:
info.field_source = _ti_core.FieldSource.TaichiX64
+ elif default_cfg().arch == _ti_core.vulkan:
+ info.field_source = _ti_core.FieldSource.TaichiVulkan
else:
raise Exception("unsupported taichi backend")
info.shape = [n for n in field.shape]
@@ -59,7 +56,12 @@ def vec_to_euler(v):
yaw = 0
else:
yaw = acos(cos_yaw)
- if (sin_yaw < 0):
+ if sin_yaw < 0:
yaw = -yaw
return yaw, pitch
+
+
+def check_ggui_availability():
+ if not _ti_core.GGUI_AVAILABLE:
+ raise Exception("GGUI Not Available")
diff --git a/python/taichi/ui/window.py b/python/taichi/ui/window.py
index fd442798bf007..7d6827e3d787e 100644
--- a/python/taichi/ui/window.py
+++ b/python/taichi/ui/window.py
@@ -1,18 +1,15 @@
import pathlib
-from taichi.core import ti_core as _ti_core
-from taichi.lang.impl import default_cfg
-from taichi.lang.kernel_impl import kernel
-from taichi.lang.ops import get_addr
-from taichi.type.annotations import ext_arr, template
+from taichi._lib import core as _ti_core
+from taichi.lang.impl import default_cfg, get_runtime
from .canvas import Canvas
-from .constants import *
-from .gui import Gui
-from .utils import get_field_info
+from .constants import PRESS, RELEASE
+from .imgui import Gui
+from .utils import check_ggui_availability
-class Window(_ti_core.PyWindow):
+class Window:
"""The window class.
Args:
@@ -20,20 +17,31 @@ class Window(_ti_core.PyWindow):
res (Tuple[Int]): resolution (width, height) of the window, in pixels.
layout (vsync): whether or not vertical sync should be enabled.
"""
- def __init__(self, name, res, vsync=False):
+ def __init__(self, name, res, vsync=False, show_window=True):
+ check_ggui_availability()
package_path = str(pathlib.Path(__file__).parent.parent)
ti_arch = default_cfg().arch
is_packed = default_cfg().packed
- super().__init__(name, res, vsync, package_path, ti_arch, is_packed)
+ self.window = _ti_core.PyWindow(get_runtime().prog, name, res, vsync,
+ show_window, package_path, ti_arch,
+ is_packed)
@property
def running(self):
- return self.is_running()
+ return self.window.is_running()
@running.setter
def running(self, value):
- self.set_is_running(value)
+ self.window.set_is_running(value)
+
+ @property
+ def event(self):
+ return self.window.get_current_event()
+
+ @event.setter
+ def event(self, value):
+ self.window.set_current_event(value)
def get_events(self, tag=None):
""" Obtain a list of unprocessed events.
@@ -42,11 +50,11 @@ def get_events(self, tag=None):
tag (str): A tag used for filtering events. If it is None, then all events are returned.
"""
if tag is None:
- return super().get_events(_ti_core.EventType.Any)
- elif tag is PRESS:
- return super().get_events(_ti_core.EventType.Press)
- elif tag is RELEASE:
- return super().get_events(_ti_core.EventType.Release)
+ return self.window.get_events(_ti_core.EventType.Any)
+ if tag is PRESS:
+ return self.window.get_events(_ti_core.EventType.Press)
+ if tag is RELEASE:
+ return self.window.get_events(_ti_core.EventType.Release)
raise Exception("unrecognized event tag")
def get_event(self, tag=None):
@@ -56,24 +64,36 @@ def get_event(self, tag=None):
"""
if tag is None:
- return super().get_event(_ti_core.EventType.Any)
- elif tag is PRESS:
- return super().get_event(_ti_core.EventType.Press)
- elif tag is RELEASE:
- return super().get_event(_ti_core.EventType.Release)
+ return self.window.get_event(_ti_core.EventType.Any)
+ if tag is PRESS:
+ return self.window.get_event(_ti_core.EventType.Press)
+ if tag is RELEASE:
+ return self.window.get_event(_ti_core.EventType.Release)
raise Exception("unrecognized event tag")
def is_pressed(self, *keys):
for k in keys:
- if super().is_pressed(k):
+ if self.window.is_pressed(k):
return True
return False
def get_canvas(self):
"""Returns a canvas handle. See :class`~taichi.ui.canvas.Canvas` """
- return Canvas(super().get_canvas())
+ return Canvas(self.window.get_canvas())
@property
def GUI(self):
"""Returns a IMGUI handle. See :class`~taichi.ui.ui.Gui` """
- return Gui(super().GUI())
+ return Gui(self.window.GUI())
+
+ def get_cursor_pos(self):
+ return self.window.get_cursor_pos()
+
+ def show(self):
+ return self.window.show()
+
+ def write_image(self, filename):
+ return self.window.write_image(filename)
+
+ def destroy(self):
+ return self.window.destroy()
diff --git a/python/ti.cpp b/python/ti.cpp
deleted file mode 100644
index f5ab3400cccf9..0000000000000
--- a/python/ti.cpp
+++ /dev/null
@@ -1,40 +0,0 @@
-#include
-#include
-#include
-#include
-#include "taichi/platform/windows/windows.h"
-#include
-#include
-#include
-
-void main(int argc, char **argv) {
- Py_SetProgramName(L"ti");
- Py_Initialize();
-
- std::vector argv_converted;
- std::vector argv_char;
- argv_converted.resize(argc);
- argv_char.resize(argc);
-
- for (int i = 0; i < argc; i++) {
- int buffer_len = 3 * std::strlen(argv[i]) + 2;
- // printf("len %d\n", buffer_len);
- // would rather be safe here... TODO: figure out the maximum converted
- // length
- argv_converted[i].resize(buffer_len);
- auto ret = mbstowcs(&argv_converted[i][0], argv[i], buffer_len);
- argv_char[i] = &argv_converted[i][0];
- }
- PySys_SetArgv(argc, &argv_char[0]);
- // TODO: implement release mode for this
- auto dir = getenv("TAICHI_REPO_DIR");
- if (dir == nullptr) {
- std::cout << "Please set TAICHI_REPO_DIR" << std::endl;
- exit(-1);
- }
- auto path = std::string(dir) + "/bin/ti";
- auto file = std::fopen(path.c_str(), "r");
- PyRun_SimpleFile(file, "ti");
- Py_Finalize();
- return;
-}
diff --git a/requirements.txt b/requirements.txt
deleted file mode 100644
index 93b7d71e9a97d..0000000000000
--- a/requirements.txt
+++ /dev/null
@@ -1,15 +0,0 @@
-GitPython
-astor
-autograd
-colorama
-coverage
-isort
-numpy
-pybind11
-pylint
-pytest
-pytest-rerunfailures
-pytest-xdist
-setuptools
-sourceinspect
-yapf==0.31.0
diff --git a/requirements_dev.txt b/requirements_dev.txt
index 2620b6b2fb3dd..718434b2d8d72 100644
--- a/requirements_dev.txt
+++ b/requirements_dev.txt
@@ -1,19 +1,16 @@
cmake
colorama
coverage
-numpy
Pillow
pybind11
GitPython
yapf==0.31.0
distro
-autograd
astor
sourceinspect
-pylint
-pytest
-pytest-xdist
-pytest-rerunfailures
-pytest-cov
-torch
isort
+pylint
+requests==2.26
+twine
+wheel
+astunparse
diff --git a/requirements_test.txt b/requirements_test.txt
new file mode 100644
index 0000000000000..84228242d1dc7
--- /dev/null
+++ b/requirements_test.txt
@@ -0,0 +1,8 @@
+pytest
+pytest-xdist
+pytest-rerunfailures
+pytest-cov
+numpy
+autograd
+requests==2.26
+matplotlib
diff --git a/scripts/generate_pylint_tags.py b/scripts/generate_pylint_tags.py
new file mode 100644
index 0000000000000..5413cc7ff3797
--- /dev/null
+++ b/scripts/generate_pylint_tags.py
@@ -0,0 +1,31 @@
+TAGS = {
+ 'C0121': True,
+ 'C0415': True,
+ 'W0611': True,
+ 'W0202': True,
+ 'W0621': True,
+ 'W0622': True,
+ 'W0401': True,
+ 'C0209': True,
+ 'W0404': True,
+ 'W0612': True,
+ 'E1101': True,
+ 'R0402': True,
+ 'R0201': True,
+ 'W0235': True,
+ 'R1705': True,
+ 'C0200': True,
+ 'R0205': True,
+ 'R1732': True,
+ 'W0101': True,
+ 'R1710': True,
+ 'R1703': True,
+ 'W0108': True,
+ 'W1309': True,
+ 'C0321': True,
+ 'C0325': True,
+}
+
+if __name__ == '__main__':
+ enabled = [kv[0] for kv in TAGS.items() if kv[1]]
+ print(','.join(enabled))
diff --git a/scripts/run-clang-tidy.py b/scripts/run_clang_tidy.py
similarity index 97%
rename from scripts/run-clang-tidy.py
rename to scripts/run_clang_tidy.py
index 8921b1859c8ac..229d91bfdf4d0 100644
--- a/scripts/run-clang-tidy.py
+++ b/scripts/run_clang_tidy.py
@@ -82,6 +82,7 @@ def get_tidy_invocation(f, clang_tidy_binary, checks, tmpdir, build_path,
config):
"""Gets a command line for clang-tidy."""
start = [clang_tidy_binary]
+ start.append('-warnings-as-errors=*')
if header_filter is not None:
start.append('-header-filter=' + header_filter)
if checks:
@@ -172,12 +173,12 @@ def run_tidy(args, tmpdir, build_path, queue, lock, failed_files):
output, err = proc.communicate()
if proc.returncode != 0:
failed_files.append(name)
- with lock:
- sys.stdout.write(' '.join(invocation) + '\n' +
- output.decode('utf-8'))
- if len(err) > 0:
- sys.stdout.flush()
- sys.stderr.write(err.decode('utf-8'))
+ with lock:
+ sys.stdout.write(' '.join(invocation) + '\n' +
+ output.decode('utf-8'))
+ if len(err) > 0:
+ sys.stdout.flush()
+ sys.stderr.write(err.decode('utf-8'))
queue.task_done()
@@ -324,6 +325,8 @@ def main():
task_queue.join()
if len(failed_files):
return_code = 1
+ else:
+ print("No errors detected, congratulations!")
except KeyboardInterrupt:
# This is a sad hack. Unfortunately subprocess goes
diff --git a/scripts/run_clang_tidy.sh b/scripts/run_clang_tidy.sh
index bc268d3e797d9..e5bd6813270ad 100644
--- a/scripts/run_clang_tidy.sh
+++ b/scripts/run_clang_tidy.sh
@@ -6,4 +6,6 @@ mkdir -p build_clang_tidy/
cd build_clang_tidy
cmake .. -DCMAKE_CXX_COMPILER=clang -DCMAKE_C_COMPILER=clang -DCMAKE_EXPORT_COMPILE_COMMANDS=ON
cd ..
-python3 scripts/run-clang-tidy.py $PWD/taichi -header-filter="$PWD/taichi/" -p build_clang_tidy -j16 -fix
+TAICHI_SRC=$PWD/taichi
+VAR=${1:-${TAICHI_SRC}}
+python3 scripts/run_clang_tidy.py $PWD/taichi -header-filter="$PWD/taichi/" -p build_clang_tidy -j16 -fix
diff --git a/setup.cfg b/setup.cfg
index bd804696ecc4f..e48864773ecf0 100644
--- a/setup.cfg
+++ b/setup.cfg
@@ -1,2 +1,3 @@
[metadata]
-description_file = README
+long_description = file: README.md
+long_description_content_type = text/markdown; charset=UTF-8
diff --git a/setup.py b/setup.py
index 274f795ce3f54..355e09679ef74 100644
--- a/setup.py
+++ b/setup.py
@@ -1,6 +1,6 @@
# Optional environment variables supported by setup.py:
-# DEBUG
-# build the C++ taichi_core extension with debug symbols.
+# {DEBUG, RELWITHDEBINFO, MINSIZEREL}
+# build the C++ taichi_core extension with various build types.
#
# TAICHI_CMAKE_ARGS
# extra cmake args for C++ taichi_core extension.
@@ -32,15 +32,25 @@
'Programming Language :: Python :: 3.7',
'Programming Language :: Python :: 3.8',
'Programming Language :: Python :: 3.9',
+ 'Programming Language :: Python :: 3.10',
]
+
+def get_version():
+ if os.getenv("RELEASE_VERSION"):
+ version = os.environ["RELEASE_VERSION"]
+ else:
+ version_file = os.path.join(os.path.dirname(__file__), 'version.txt')
+ with open(version_file, 'r') as f:
+ version = f.read().strip()
+ return version.lstrip("v")
+
+
project_name = os.getenv('PROJECT_NAME', 'taichi')
-TI_VERSION_MAJOR = 0
-TI_VERSION_MINOR = 8
-TI_VERSION_PATCH = 3
-version = f'{TI_VERSION_MAJOR}.{TI_VERSION_MINOR}.{TI_VERSION_PATCH}'
+version = get_version()
+TI_VERSION_MAJOR, TI_VERSION_MINOR, TI_VERSION_PATCH = version.split('.')
-data_files = glob.glob('python/lib/*')
+data_files = glob.glob('python/_lib/runtime/*')
print(data_files)
packages = find_packages('python')
print(packages)
@@ -65,13 +75,19 @@ def get_os_name():
return 'win'
elif name.lower().startswith('linux'):
return 'linux'
+ elif 'bsd' in name.lower():
+ return 'unix'
assert False, "Unknown platform name %s" % name
def remove_tmp(taichi_dir):
shutil.rmtree(os.path.join(taichi_dir, 'assets'), ignore_errors=True)
- shutil.rmtree(os.path.join(taichi_dir, 'examples'), ignore_errors=True)
- shutil.rmtree(os.path.join(taichi_dir, 'tests'), ignore_errors=True)
+
+
+def remove_files_with_extension(dir_name, extension):
+ for file in os.listdir(dir_name):
+ if file.endswith(extension):
+ os.remove(os.path.join(dir_name, file))
class CMakeExtension(Extension):
@@ -84,8 +100,6 @@ def run(self):
taichi_dir = os.path.join(package_dir, 'taichi')
remove_tmp(taichi_dir)
- shutil.copytree('tests/python', os.path.join(taichi_dir, 'tests'))
- shutil.copytree('examples', os.path.join(taichi_dir, 'examples'))
shutil.copytree('external/assets', os.path.join(taichi_dir, 'assets'))
egg_info.run(self)
@@ -110,7 +124,7 @@ def parse_cmake_args_from_env(self):
def run(self):
try:
- out = subprocess.check_output(['cmake', '--version'])
+ subprocess.check_call(['cmake', '--version'])
except OSError:
raise RuntimeError(
"CMake must be installed to build the following extensions: " +
@@ -132,8 +146,21 @@ def run(self):
f'-DTI_VERSION_PATCH={TI_VERSION_PATCH}',
]
- self.debug = os.getenv('DEBUG', '0') in ('1', 'ON')
- cfg = 'Debug' if self.debug else 'Release'
+ emscriptened = os.getenv('TI_EMSCRIPTENED', '0') in ('1', 'ON')
+ if emscriptened:
+ cmake_args += ['-DTI_EMSCRIPTENED=ON']
+
+ if shutil.which('ninja'):
+ cmake_args += ['-GNinja']
+
+ cfg = 'Release'
+ if (os.getenv('DEBUG', '0') in ('1', 'ON')):
+ cfg = 'Debug'
+ elif (os.getenv('RELWITHDEBINFO', '0') in ('1', 'ON')):
+ cfg = 'RelWithDebInfo'
+ elif (os.getenv('MINSIZEREL', '0') in ('1', 'ON')):
+ cfg = 'MinSizeRel'
+
build_args = ['--config', cfg]
cmake_args += ['-DCMAKE_BUILD_TYPE=' + cfg]
@@ -150,6 +177,7 @@ def run(self):
os.makedirs(self.build_temp, exist_ok=True)
print('-' * 10, 'Running CMake prepare', '-' * 40)
+ print(' '.join(['cmake', cmake_list_dir] + cmake_args))
subprocess.check_call(['cmake', cmake_list_dir] + cmake_args,
cwd=self.build_temp,
env=env)
@@ -164,36 +192,50 @@ def prepare_package(self):
# We need to make sure these additional files are ready for
# - develop mode: must exist in local python/taichi/lib/ folder
# - install mode: must exist in self.build_lib/taichi/lib
- taichi_lib_dir = 'taichi/lib'
- for target in (
- os.path.join(package_dir, taichi_lib_dir),
- os.path.join(self.build_lib, taichi_lib_dir),
- ):
- shutil.rmtree(target, ignore_errors=True)
- os.makedirs(target)
- if get_os_name() == 'linux':
- shutil.copy(os.path.join(self.build_temp, 'libtaichi_core.so'),
- os.path.join(target, 'taichi_core.so'))
- elif get_os_name() == 'osx':
- shutil.copy(
- os.path.join(self.build_temp, 'libtaichi_core.dylib'),
- os.path.join(target, 'taichi_core.so'))
- else:
- shutil.copy('runtimes/Release/taichi_core.dll',
- os.path.join(target, 'taichi_core.pyd'))
-
- if get_os_name() != 'osx':
- libdevice_path = 'external/cuda_libdevice/slim_libdevice.10.bc'
- print("copying libdevice:", libdevice_path)
- assert os.path.exists(libdevice_path)
- shutil.copy(libdevice_path,
- os.path.join(target, 'slim_libdevice.10.bc'))
-
- llvm_runtime_dir = 'taichi/runtime/llvm'
- for f in os.listdir(llvm_runtime_dir):
- if f.startswith('runtime_') and f.endswith('.bc'):
- print(f"Fetching runtime file {f} to {target} folder")
- shutil.copy(os.path.join(llvm_runtime_dir, f), target)
+ base_dir = package_dir if self.inplace else self.build_lib
+ taichi_lib_dir = os.path.join(base_dir, 'taichi', '_lib')
+
+ runtime_dir = os.path.join(taichi_lib_dir, "runtime")
+ core_dir = os.path.join(taichi_lib_dir, "core")
+ os.makedirs(runtime_dir, exist_ok=True)
+ os.makedirs(core_dir, exist_ok=True)
+
+ if (get_os_name() == 'linux' or get_os_name() == 'unix'
+ or get_os_name() == 'osx'):
+ remove_files_with_extension(core_dir, ".so")
+ else:
+ remove_files_with_extension(core_dir, ".pyd")
+ if get_os_name() == 'osx':
+ remove_files_with_extension(runtime_dir, ".dylib")
+ remove_files_with_extension(runtime_dir, ".bc")
+
+ if get_os_name() == 'linux' or get_os_name() == 'unix':
+ self.copy_file(os.path.join(self.build_temp, 'libtaichi_core.so'),
+ os.path.join(core_dir, 'taichi_core.so'))
+ elif get_os_name() == 'osx':
+ self.copy_file(
+ os.path.join(self.build_temp, 'libtaichi_core.dylib'),
+ os.path.join(core_dir, 'taichi_core.so'))
+ moltenvk_path = os.path.join(self.build_temp, 'libMoltenVK.dylib')
+ if os.path.exists(moltenvk_path):
+ self.copy_file(moltenvk_path,
+ os.path.join(runtime_dir, 'libMoltenVK.dylib'))
+ else:
+ self.copy_file('runtimes/taichi_core.dll',
+ os.path.join(core_dir, 'taichi_core.pyd'))
+
+ if get_os_name() != 'osx':
+ libdevice_path = 'external/cuda_libdevice/slim_libdevice.10.bc'
+ print("copying libdevice:", libdevice_path)
+ assert os.path.exists(libdevice_path)
+ self.copy_file(libdevice_path,
+ os.path.join(runtime_dir, 'slim_libdevice.10.bc'))
+
+ llvm_runtime_dir = 'taichi/runtime/llvm'
+ for f in os.listdir(llvm_runtime_dir):
+ if f.startswith('runtime_') and f.endswith('.bc'):
+ print(f"Fetching runtime file {f} to {taichi_lib_dir} folder")
+ self.copy_file(os.path.join(llvm_runtime_dir, f), runtime_dir)
class Clean(clean):
@@ -203,8 +245,8 @@ def run(self):
if os.path.exists(self.build_temp):
remove_tree(self.build_temp, dry_run=self.dry_run)
generated_folders = ('bin', 'dist', 'python/taichi/assets',
- 'python/taichi/lib', 'python/taichi/examples',
- 'python/taichi/tests', 'python/taichi.egg-info')
+ 'python/taichi/_lib/runtime',
+ 'python/taichi.egg-info')
for d in generated_folders:
if os.path.exists(d):
remove_tree(d, dry_run=self.dry_run)
@@ -212,7 +254,8 @@ def run(self):
'taichi/common/commit_hash.h', 'taichi/common/version.h'
]
generated_files += glob.glob('taichi/runtime/llvm/runtime_*.bc')
- generated_files += glob.glob('taichi/runtime/llvm/runtime_*.ll')
+ generated_files += glob.glob('python/taichi/_lib/core/*.so')
+ generated_files += glob.glob('python/taichi/_lib/core/*.pyd')
for f in generated_files:
if os.path.exists(f):
print(f'removing generated file {f}')
@@ -228,20 +271,18 @@ def run(self):
author='Taichi developers',
author_email='yuanmhu@gmail.com',
url='https://github.com/taichi-dev/taichi',
+ python_requires=">=3.6,<3.11",
install_requires=[
- 'numpy',
- 'pybind11>=2.5.0',
- 'sourceinspect>=0.0.4',
- 'colorama',
- 'astor',
+ 'numpy', 'sourceinspect>=0.0.4', 'colorama', 'astor',
+ 'astunparse;python_version<"3.9"'
],
- data_files=[('lib', data_files)],
+ data_files=[(os.path.join('_lib', 'runtime'), data_files)],
keywords=['graphics', 'simulation'],
license='MIT',
include_package_data=True,
entry_points={
'console_scripts': [
- 'ti=taichi.main:main',
+ 'ti=taichi._main:main',
],
},
classifiers=classifiers,
diff --git a/taichi/analysis/alias_analysis.cpp b/taichi/analysis/alias_analysis.cpp
index 24e3acfb8fa8d..5f5928997eea1 100644
--- a/taichi/analysis/alias_analysis.cpp
+++ b/taichi/analysis/alias_analysis.cpp
@@ -29,8 +29,17 @@ AliasResult alias_analysis(Stmt *var1, Stmt *var2) {
Stmt *origin1 = retrieve_local(var1);
Stmt *origin2 = retrieve_local(var2);
if (origin1 != nullptr && origin2 != nullptr) {
- if (origin1 == origin2)
+ if (origin1 == origin2) {
+ if (var1->is() && var2->is()) {
+ auto diff = value_diff_ptr_index(var1->cast()->offset,
+ var2->cast()->offset);
+ if (diff.is_diff_certain) {
+ return diff.diff_range == 0 ? AliasResult::same
+ : AliasResult::different;
+ }
+ }
return AliasResult::uncertain;
+ }
if (origin1->is() || origin2->is())
return AliasResult::different;
TI_ASSERT(origin1->is() &&
@@ -82,16 +91,36 @@ AliasResult alias_analysis(Stmt *var1, Stmt *var2) {
: AliasResult::uncertain;
}
+ TI_ASSERT(var1->width() == 1);
+ TI_ASSERT(var2->width() == 1);
+
if (var1->is() || var2->is()) {
if (!var1->is() || !var2->is())
return AliasResult::different;
- return AliasResult::uncertain;
+ auto ptr1 = var1->as();
+ auto ptr2 = var2->as();
+ if (ptr1->base_ptrs[0] != ptr2->base_ptrs[0]) {
+ auto base1 = ptr1->base_ptrs[0]->as();
+ auto base2 = ptr2->base_ptrs[0]->as();
+ if (base1->arg_id != base2->arg_id) {
+ return AliasResult::different;
+ }
+ }
+ TI_ASSERT(ptr1->indices.size() == ptr2->indices.size());
+ bool uncertain = false;
+ for (int i = 0; i < (int)ptr1->indices.size(); i++) {
+ auto diff = value_diff_ptr_index(ptr1->indices[i], ptr2->indices[i]);
+ if (!diff.is_diff_certain) {
+ uncertain = true;
+ } else if (diff.diff_range != 0) {
+ return AliasResult::different;
+ }
+ }
+ return uncertain ? AliasResult::uncertain : AliasResult::same;
}
// If both statements are GlobalPtrStmts or GetChStmts, we can check by
// SNode::id.
- TI_ASSERT(var1->width() == 1);
- TI_ASSERT(var2->width() == 1);
auto get_snode_id = [](Stmt *s) {
if (auto ptr = s->cast()) {
return ptr->snodes[0]->id;
diff --git a/taichi/analysis/bls_analyzer.cpp b/taichi/analysis/bls_analyzer.cpp
index 868df66abd849..e9e7e1f4f54d1 100644
--- a/taichi/analysis/bls_analyzer.cpp
+++ b/taichi/analysis/bls_analyzer.cpp
@@ -50,7 +50,7 @@ void BLSAnalyzer::record_access(Stmt *stmt, AccessFlag flag) {
for (int i = 0; i < num_indices; i++) {
auto diff = irpass::analysis::value_diff_loop_index(ptr->indices[i],
for_stmt_, i);
- if (diff.related_() && diff.coeff > 0) {
+ if (diff.related() && diff.coeff > 0) {
offsets[i].low = diff.low;
offsets[i].high = diff.high;
coeffs[i] = diff.coeff;
diff --git a/taichi/analysis/build_cfg.cpp b/taichi/analysis/build_cfg.cpp
index b4e2eda30b86f..1ae05aeb46d41 100644
--- a/taichi/analysis/build_cfg.cpp
+++ b/taichi/analysis/build_cfg.cpp
@@ -62,18 +62,18 @@ namespace lang {
class CFGBuilder : public IRVisitor {
public:
CFGBuilder()
- : current_block(nullptr),
- last_node_in_current_block(nullptr),
- current_stmt_id(-1),
- begin_location(-1),
- current_offload(nullptr),
- in_parallel_for(false) {
+ : current_block_(nullptr),
+ last_node_in_current_block_(nullptr),
+ current_stmt_id_(-1),
+ begin_location_(-1),
+ current_offload_(nullptr),
+ in_parallel_for_(false) {
allow_undefined_visitor = true;
invoke_default_visitor = true;
- graph = std::make_unique();
+ graph_ = std::make_unique();
// Make an empty start node.
- auto start_node = graph->push_back();
- prev_nodes.push_back(start_node);
+ auto start_node = graph_->push_back();
+ prev_nodes_.push_back(start_node);
}
void visit(Stmt *stmt) override {
@@ -93,18 +93,18 @@ class CFGBuilder : public IRVisitor {
* @return The node which is just created.
*/
CFGNode *new_node(int next_begin_location) {
- auto node = graph->push_back(
- current_block, begin_location, /*end_location=*/current_stmt_id,
- /*is_parallel_executed=*/in_parallel_for,
- /*prev_node_in_same_block=*/last_node_in_current_block);
- for (auto &prev_node : prev_nodes) {
+ auto node = graph_->push_back(
+ current_block_, begin_location_, /*end_location=*/current_stmt_id_,
+ /*is_parallel_executed=*/in_parallel_for_,
+ /*prev_node_in_same_block=*/last_node_in_current_block_);
+ for (auto &prev_node : prev_nodes_) {
// Now that the "(next node)" is created, we should insert edges
// "node... -> (next node)" here.
CFGNode::add_edge(prev_node, node);
}
- prev_nodes.clear();
- begin_location = next_begin_location;
- last_node_in_current_block = node;
+ prev_nodes_.clear();
+ begin_location_ = next_begin_location;
+ last_node_in_current_block_ = node;
return node;
}
@@ -125,7 +125,7 @@ class CFGBuilder : public IRVisitor {
*/
void visit(ContinueStmt *stmt) override {
// Don't put ContinueStmt in any CFGNodes.
- continues_in_current_loop.push_back(new_node(current_stmt_id + 1));
+ continues_in_current_loop_.push_back(new_node(current_stmt_id_ + 1));
}
/**
@@ -145,9 +145,9 @@ class CFGBuilder : public IRVisitor {
*/
void visit(WhileControlStmt *stmt) override {
// Don't put WhileControlStmt in any CFGNodes.
- auto node = new_node(current_stmt_id + 1);
- breaks_in_current_loop.push_back(node);
- prev_nodes.push_back(node);
+ auto node = new_node(current_stmt_id_ + 1);
+ breaks_in_current_loop_.push_back(node);
+ prev_nodes_.push_back(node);
}
/**
@@ -174,20 +174,20 @@ class CFGBuilder : public IRVisitor {
*/
void visit(FuncCallStmt *stmt) override {
auto node_before_func_call = new_node(-1);
- CFGFuncKey func_key = {stmt->func->func_key, in_parallel_for};
- if (node_func_begin.count(func_key) == 0) {
+ CFGFuncKey func_key = {stmt->func->func_key, in_parallel_for_};
+ if (node_func_begin_.count(func_key) == 0) {
// Generate CFG for the function.
TI_ASSERT(stmt->func->ir->is());
- auto func_begin_index = graph->size();
+ auto func_begin_index = graph_->size();
stmt->func->ir->accept(this);
- node_func_begin[func_key] = graph->nodes[func_begin_index].get();
- node_func_end[func_key] = graph->nodes.back().get();
+ node_func_begin_[func_key] = graph_->nodes[func_begin_index].get();
+ node_func_end_[func_key] = graph_->nodes.back().get();
}
- CFGNode::add_edge(node_before_func_call, node_func_begin[func_key]);
- prev_nodes.push_back(node_func_end[func_key]);
+ CFGNode::add_edge(node_before_func_call, node_func_begin_[func_key]);
+ prev_nodes_.push_back(node_func_end_[func_key]);
// Don't put FuncCallStmt in any CFGNodes.
- begin_location = current_stmt_id + 1;
+ begin_location_ = current_stmt_id_ + 1;
}
/**
@@ -219,27 +219,27 @@ class CFGBuilder : public IRVisitor {
auto before_if = new_node(-1);
CFGNode *true_branch_end = nullptr;
if (if_stmt->true_statements) {
- auto true_branch_begin = graph->size();
+ auto true_branch_begin = graph_->size();
if_stmt->true_statements->accept(this);
- CFGNode::add_edge(before_if, graph->nodes[true_branch_begin].get());
- true_branch_end = graph->back();
+ CFGNode::add_edge(before_if, graph_->nodes[true_branch_begin].get());
+ true_branch_end = graph_->back();
}
CFGNode *false_branch_end = nullptr;
if (if_stmt->false_statements) {
- auto false_branch_begin = graph->size();
+ auto false_branch_begin = graph_->size();
if_stmt->false_statements->accept(this);
- CFGNode::add_edge(before_if, graph->nodes[false_branch_begin].get());
- false_branch_end = graph->back();
+ CFGNode::add_edge(before_if, graph_->nodes[false_branch_begin].get());
+ false_branch_end = graph_->back();
}
- TI_ASSERT(prev_nodes.empty());
+ TI_ASSERT(prev_nodes_.empty());
if (if_stmt->true_statements)
- prev_nodes.push_back(true_branch_end);
+ prev_nodes_.push_back(true_branch_end);
if (if_stmt->false_statements)
- prev_nodes.push_back(false_branch_end);
+ prev_nodes_.push_back(false_branch_end);
if (!if_stmt->true_statements || !if_stmt->false_statements)
- prev_nodes.push_back(before_if);
+ prev_nodes_.push_back(before_if);
// Container statements don't belong to any CFGNodes.
- begin_location = current_stmt_id + 1;
+ begin_location_ = current_stmt_id_ + 1;
}
/**
@@ -262,34 +262,34 @@ class CFGBuilder : public IRVisitor {
* }
*/
void visit_loop(Block *body, CFGNode *before_loop, bool is_while_true) {
- int loop_stmt_id = current_stmt_id;
- auto backup_continues = std::move(continues_in_current_loop);
- auto backup_breaks = std::move(breaks_in_current_loop);
- continues_in_current_loop.clear();
- breaks_in_current_loop.clear();
+ int loop_stmt_id = current_stmt_id_;
+ auto backup_continues = std::move(continues_in_current_loop_);
+ auto backup_breaks = std::move(breaks_in_current_loop_);
+ continues_in_current_loop_.clear();
+ breaks_in_current_loop_.clear();
- auto loop_begin_index = graph->size();
+ auto loop_begin_index = graph_->size();
body->accept(this);
- auto loop_begin = graph->nodes[loop_begin_index].get();
+ auto loop_begin = graph_->nodes[loop_begin_index].get();
CFGNode::add_edge(before_loop, loop_begin);
- auto loop_end = graph->back();
+ auto loop_end = graph_->back();
CFGNode::add_edge(loop_end, loop_begin);
if (!is_while_true) {
- prev_nodes.push_back(before_loop);
- prev_nodes.push_back(loop_end);
+ prev_nodes_.push_back(before_loop);
+ prev_nodes_.push_back(loop_end);
}
- for (auto &node : continues_in_current_loop) {
+ for (auto &node : continues_in_current_loop_) {
CFGNode::add_edge(node, loop_begin);
- prev_nodes.push_back(node);
+ prev_nodes_.push_back(node);
}
- for (auto &node : breaks_in_current_loop) {
- prev_nodes.push_back(node);
+ for (auto &node : breaks_in_current_loop_) {
+ prev_nodes_.push_back(node);
}
// Container statements don't belong to any CFGNodes.
- begin_location = loop_stmt_id + 1;
- continues_in_current_loop = std::move(backup_continues);
- breaks_in_current_loop = std::move(backup_breaks);
+ begin_location_ = loop_stmt_id + 1;
+ continues_in_current_loop_ = std::move(backup_continues);
+ breaks_in_current_loop_ = std::move(backup_breaks);
}
void visit(WhileStmt *stmt) override {
@@ -297,19 +297,27 @@ class CFGBuilder : public IRVisitor {
}
void visit(RangeForStmt *stmt) override {
- auto old_in_parallel_for = in_parallel_for;
- if (!current_offload)
- in_parallel_for = true;
+ auto old_in_parallel_for = in_parallel_for_;
+ if (!current_offload_)
+ in_parallel_for_ = true;
visit_loop(stmt->body.get(), new_node(-1), false);
- in_parallel_for = old_in_parallel_for;
+ in_parallel_for_ = old_in_parallel_for;
}
void visit(StructForStmt *stmt) override {
- auto old_in_parallel_for = in_parallel_for;
- if (!current_offload)
- in_parallel_for = true;
+ auto old_in_parallel_for = in_parallel_for_;
+ if (!current_offload_)
+ in_parallel_for_ = true;
visit_loop(stmt->body.get(), new_node(-1), false);
- in_parallel_for = old_in_parallel_for;
+ in_parallel_for_ = old_in_parallel_for;
+ }
+
+ void visit(MeshForStmt *stmt) override {
+ auto old_in_parallel_for = in_parallel_for_;
+ if (!current_offload_)
+ in_parallel_for_ = true;
+ visit_loop(stmt->body.get(), new_node(-1), false);
+ in_parallel_for_ = old_in_parallel_for;
}
/**
@@ -320,6 +328,9 @@ class CFGBuilder : public IRVisitor {
* } -> node_tls_prologue;
* node_tls_prologue {
* ...
+ * } -> node_mesh_prologue;
+ * node_mesh_prologue:
+ * ...
* } -> node_bls_prologue;
* node_bls_prologue {
* ...
@@ -338,63 +349,74 @@ class CFGBuilder : public IRVisitor {
* }
*/
void visit(OffloadedStmt *stmt) override {
- current_offload = stmt;
+ current_offload_ = stmt;
if (stmt->tls_prologue) {
auto before_offload = new_node(-1);
- int offload_stmt_id = current_stmt_id;
- auto block_begin_index = graph->size();
+ int offload_stmt_id = current_stmt_id_;
+ auto block_begin_index = graph_->size();
stmt->tls_prologue->accept(this);
- prev_nodes.push_back(graph->back());
+ prev_nodes_.push_back(graph_->back());
+ // Container statements don't belong to any CFGNodes.
+ begin_location_ = offload_stmt_id + 1;
+ CFGNode::add_edge(before_offload, graph_->nodes[block_begin_index].get());
+ }
+ if (stmt->mesh_prologue) {
+ auto before_offload = new_node(-1);
+ int offload_stmt_id = current_stmt_id_;
+ auto block_begin_index = graph_->size();
+ stmt->mesh_prologue->accept(this);
+ prev_nodes_.push_back(graph_->back());
// Container statements don't belong to any CFGNodes.
- begin_location = offload_stmt_id + 1;
- CFGNode::add_edge(before_offload, graph->nodes[block_begin_index].get());
+ begin_location_ = offload_stmt_id + 1;
+ CFGNode::add_edge(before_offload, graph_->nodes[block_begin_index].get());
}
if (stmt->bls_prologue) {
auto before_offload = new_node(-1);
- int offload_stmt_id = current_stmt_id;
- auto block_begin_index = graph->size();
+ int offload_stmt_id = current_stmt_id_;
+ auto block_begin_index = graph_->size();
stmt->bls_prologue->accept(this);
- prev_nodes.push_back(graph->back());
+ prev_nodes_.push_back(graph_->back());
// Container statements don't belong to any CFGNodes.
- begin_location = offload_stmt_id + 1;
- CFGNode::add_edge(before_offload, graph->nodes[block_begin_index].get());
+ begin_location_ = offload_stmt_id + 1;
+ CFGNode::add_edge(before_offload, graph_->nodes[block_begin_index].get());
}
if (stmt->has_body()) {
auto before_offload = new_node(-1);
- int offload_stmt_id = current_stmt_id;
- auto block_begin_index = graph->size();
+ int offload_stmt_id = current_stmt_id_;
+ auto block_begin_index = graph_->size();
if (stmt->task_type == OffloadedStmt::TaskType::range_for ||
- stmt->task_type == OffloadedStmt::TaskType::struct_for) {
- in_parallel_for = true;
+ stmt->task_type == OffloadedStmt::TaskType::struct_for ||
+ stmt->task_type == OffloadedStmt::TaskType::mesh_for) {
+ in_parallel_for_ = true;
}
stmt->body->accept(this);
- in_parallel_for = false;
- prev_nodes.push_back(graph->back());
+ in_parallel_for_ = false;
+ prev_nodes_.push_back(graph_->back());
// Container statements don't belong to any CFGNodes.
- begin_location = offload_stmt_id + 1;
- CFGNode::add_edge(before_offload, graph->nodes[block_begin_index].get());
+ begin_location_ = offload_stmt_id + 1;
+ CFGNode::add_edge(before_offload, graph_->nodes[block_begin_index].get());
}
if (stmt->bls_epilogue) {
auto before_offload = new_node(-1);
- int offload_stmt_id = current_stmt_id;
- auto block_begin_index = graph->size();
+ int offload_stmt_id = current_stmt_id_;
+ auto block_begin_index = graph_->size();
stmt->bls_epilogue->accept(this);
- prev_nodes.push_back(graph->back());
+ prev_nodes_.push_back(graph_->back());
// Container statements don't belong to any CFGNodes.
- begin_location = offload_stmt_id + 1;
- CFGNode::add_edge(before_offload, graph->nodes[block_begin_index].get());
+ begin_location_ = offload_stmt_id + 1;
+ CFGNode::add_edge(before_offload, graph_->nodes[block_begin_index].get());
}
if (stmt->tls_epilogue) {
auto before_offload = new_node(-1);
- int offload_stmt_id = current_stmt_id;
- auto block_begin_index = graph->size();
+ int offload_stmt_id = current_stmt_id_;
+ auto block_begin_index = graph_->size();
stmt->tls_epilogue->accept(this);
- prev_nodes.push_back(graph->back());
+ prev_nodes_.push_back(graph_->back());
// Container statements don't belong to any CFGNodes.
- begin_location = offload_stmt_id + 1;
- CFGNode::add_edge(before_offload, graph->nodes[block_begin_index].get());
+ begin_location_ = offload_stmt_id + 1;
+ CFGNode::add_edge(before_offload, graph_->nodes[block_begin_index].get());
}
- current_offload = nullptr;
+ current_offload_ = nullptr;
}
/**
@@ -415,56 +437,56 @@ class CFGBuilder : public IRVisitor {
* graph->final_node = node_block_end;
*/
void visit(Block *block) override {
- auto backup_block = current_block;
- auto backup_last_node = last_node_in_current_block;
- auto backup_stmt_id = current_stmt_id;
+ auto backup_block = current_block_;
+ auto backup_last_node = last_node_in_current_block_;
+ auto backup_stmt_id = current_stmt_id_;
// |begin_location| must be -1 (indicating we are not building any CFGNode)
// when the |current_block| changes.
- TI_ASSERT(begin_location == -1);
- TI_ASSERT(prev_nodes.empty() || graph->size() == 1);
- current_block = block;
- last_node_in_current_block = nullptr;
- begin_location = 0;
+ TI_ASSERT(begin_location_ == -1);
+ TI_ASSERT(prev_nodes_.empty() || graph_->size() == 1);
+ current_block_ = block;
+ last_node_in_current_block_ = nullptr;
+ begin_location_ = 0;
for (int i = 0; i < (int)block->size(); i++) {
- current_stmt_id = i;
+ current_stmt_id_ = i;
block->statements[i]->accept(this);
}
- current_stmt_id = block->size();
+ current_stmt_id_ = block->size();
new_node(-1); // Each block has a deterministic last node.
- graph->final_node = (int)graph->size() - 1;
+ graph_->final_node = (int)graph_->size() - 1;
- current_block = backup_block;
- last_node_in_current_block = backup_last_node;
- current_stmt_id = backup_stmt_id;
+ current_block_ = backup_block;
+ last_node_in_current_block_ = backup_last_node;
+ current_stmt_id_ = backup_stmt_id;
}
static std::unique_ptr run(IRNode *root) {
CFGBuilder builder;
root->accept(&builder);
- if (!builder.graph->nodes[builder.graph->final_node]->empty()) {
+ if (!builder.graph_->nodes[builder.graph_->final_node]->empty()) {
// Make the final node empty (by adding an empty final node).
- builder.graph->push_back();
- CFGNode::add_edge(builder.graph->nodes[builder.graph->final_node].get(),
- builder.graph->back());
- builder.graph->final_node = (int)builder.graph->size() - 1;
+ builder.graph_->push_back();
+ CFGNode::add_edge(builder.graph_->nodes[builder.graph_->final_node].get(),
+ builder.graph_->back());
+ builder.graph_->final_node = (int)builder.graph_->size() - 1;
}
- return std::move(builder.graph);
+ return std::move(builder.graph_);
}
private:
- std::unique_ptr graph;
- Block *current_block;
- CFGNode *last_node_in_current_block;
- std::vector continues_in_current_loop;
- std::vector breaks_in_current_loop;
- int current_stmt_id;
- int begin_location;
- std::vector prev_nodes;
- OffloadedStmt *current_offload;
- bool in_parallel_for;
- std::unordered_map node_func_begin;
- std::unordered_map node_func_end;
+ std::unique_ptr graph_;
+ Block *current_block_;
+ CFGNode *last_node_in_current_block_;
+ std::vector continues_in_current_loop_;
+ std::vector breaks_in_current_loop_;
+ int current_stmt_id_;
+ int begin_location_;
+ std::vector prev_nodes_;
+ OffloadedStmt *current_offload_;
+ bool in_parallel_for_;
+ std::unordered_map node_func_begin_;
+ std::unordered_map node_func_end_;
};
namespace irpass::analysis {
diff --git a/taichi/analysis/clone.cpp b/taichi/analysis/clone.cpp
index dc2e1b230a7c1..e4c2476008da4 100644
--- a/taichi/analysis/clone.cpp
+++ b/taichi/analysis/clone.cpp
@@ -12,7 +12,7 @@ TLANG_NAMESPACE_BEGIN
class IRCloner : public IRVisitor {
private:
IRNode *other_node;
- std::unordered_map operand_map;
+ std::unordered_map operand_map_;
public:
enum Phase { register_operand_map, replace_operand } phase;
@@ -34,16 +34,16 @@ class IRCloner : public IRVisitor {
void generic_visit(Stmt *stmt) {
if (phase == register_operand_map)
- operand_map[stmt] = other_node->as();
+ operand_map_[stmt] = other_node->as();
else {
TI_ASSERT(phase == replace_operand);
auto other_stmt = other_node->as();
TI_ASSERT(stmt->num_operands() == other_stmt->num_operands());
for (int i = 0; i < stmt->num_operands(); i++) {
- if (operand_map.find(stmt->operand(i)) == operand_map.end())
+ if (operand_map_.find(stmt->operand(i)) == operand_map_.end())
other_stmt->set_operand(i, stmt->operand(i));
else
- other_stmt->set_operand(i, operand_map[stmt->operand(i)]);
+ other_stmt->set_operand(i, operand_map_[stmt->operand(i)]);
}
}
}
@@ -67,14 +67,6 @@ class IRCloner : public IRVisitor {
}
}
- void visit(FuncBodyStmt *stmt) override {
- generic_visit(stmt);
- auto other = other_node->as();
- other_node = other->body.get();
- stmt->body->accept(this);
- other_node = other;
- }
-
void visit(WhileStmt *stmt) override {
generic_visit(stmt);
auto other = other_node->as();
@@ -112,6 +104,7 @@ class IRCloner : public IRVisitor {
CLONE_BLOCK(tls_prologue)
CLONE_BLOCK(bls_prologue)
+ CLONE_BLOCK(mesh_prologue)
if (stmt->body) {
other_node = other->body.get();
diff --git a/taichi/analysis/count_statements.cpp b/taichi/analysis/count_statements.cpp
index 43796c3c501f7..78733878b0a3e 100644
--- a/taichi/analysis/count_statements.cpp
+++ b/taichi/analysis/count_statements.cpp
@@ -8,7 +8,7 @@ TLANG_NAMESPACE_BEGIN
class StmtCounter : public BasicStmtVisitor {
private:
StmtCounter() {
- counter = 0;
+ counter_ = 0;
allow_undefined_visitor = true;
invoke_default_visitor = true;
}
@@ -17,21 +17,21 @@ class StmtCounter : public BasicStmtVisitor {
public:
void preprocess_container_stmt(Stmt *stmt) override {
- counter++;
+ counter_++;
}
void visit(Stmt *stmt) override {
- counter++;
+ counter_++;
}
static int run(IRNode *root) {
StmtCounter stmt_counter;
root->accept(&stmt_counter);
- return stmt_counter.counter;
+ return stmt_counter.counter_;
}
private:
- int counter;
+ int counter_;
};
namespace irpass::analysis {
diff --git a/taichi/analysis/data_source_analysis.cpp b/taichi/analysis/data_source_analysis.cpp
index d1218f947d093..ec23e8be085ef 100644
--- a/taichi/analysis/data_source_analysis.cpp
+++ b/taichi/analysis/data_source_analysis.cpp
@@ -42,7 +42,7 @@ std::vector get_load_pointers(Stmt *load_stmt) {
Stmt *get_store_data(Stmt *store_stmt) {
// If store_stmt provides one data source, return the data.
- if (store_stmt->is()) {
+ if (store_stmt->is() && !store_stmt->ret_type->is()) {
// For convenience, return store_stmt instead of the const [0] it actually
// stores.
return store_stmt;
@@ -57,7 +57,7 @@ Stmt *get_store_data(Stmt *store_stmt) {
std::vector get_store_destination(Stmt *store_stmt) {
// If store_stmt provides some data sources, return the pointers of the data.
- if (store_stmt->is()) {
+ if (store_stmt->is() && !store_stmt->ret_type->is()) {
// The statement itself provides a data source (const [0]).
return std::vector(1, store_stmt);
} else if (auto local_store = store_stmt->cast()) {
@@ -67,7 +67,12 @@ std::vector get_store_destination(Stmt *store_stmt) {
} else if (auto atomic = store_stmt->cast()) {
return std::vector(1, atomic->dest);
} else if (auto external_func = store_stmt->cast()) {
- return external_func->output_stmts;
+ if (store_stmt->cast()->type ==
+ ExternalFuncCallStmt::BITCODE) {
+ return external_func->arg_stmts;
+ } else {
+ return external_func->output_stmts;
+ }
} else {
return std::vector();
}
diff --git a/taichi/analysis/gather_mesh_thread_local.cpp b/taichi/analysis/gather_mesh_thread_local.cpp
new file mode 100644
index 0000000000000..41bc86a8aa953
--- /dev/null
+++ b/taichi/analysis/gather_mesh_thread_local.cpp
@@ -0,0 +1,82 @@
+#include "taichi/ir/ir.h"
+#include "taichi/ir/snode.h"
+#include "taichi/ir/mesh.h"
+#include "taichi/ir/visitors.h"
+#include "taichi/ir/analysis.h"
+#include "taichi/ir/statements.h"
+
+TLANG_NAMESPACE_BEGIN
+
+using MeshElementTypeSet = std::unordered_set;
+
+class GatherMeshThreadLocal : public BasicStmtVisitor {
+ public:
+ using BasicStmtVisitor::visit;
+
+ GatherMeshThreadLocal(OffloadedStmt *offload_,
+ MeshElementTypeSet *owned_ptr_,
+ MeshElementTypeSet *total_ptr_,
+ bool optimize_mesh_reordered_mapping_) {
+ allow_undefined_visitor = true;
+ invoke_default_visitor = true;
+
+ this->offload = offload_;
+ this->owned_ptr = owned_ptr_;
+ this->total_ptr = total_ptr_;
+ this->optimize_mesh_reordered_mapping = optimize_mesh_reordered_mapping_;
+ }
+
+ static void run(OffloadedStmt *offload,
+ MeshElementTypeSet *owned_ptr,
+ MeshElementTypeSet *total_ptr,
+ const CompileConfig &config) {
+ TI_ASSERT(offload->task_type == OffloadedStmt::TaskType::mesh_for);
+ GatherMeshThreadLocal analyser(offload, owned_ptr, total_ptr,
+ config.optimize_mesh_reordered_mapping);
+ offload->accept(&analyser);
+ }
+
+ void visit(LoopIndexStmt *stmt) override {
+ if (stmt->is_mesh_index()) {
+ this->owned_ptr->insert(stmt->mesh_index_type());
+ }
+ }
+
+ void visit(MeshRelationAccessStmt *stmt) override {
+ if (mesh::element_order(stmt->from_type()) >
+ mesh::element_order(stmt->to_type)) {
+ this->total_ptr->insert(stmt->from_type());
+ } else {
+ this->owned_ptr->insert(stmt->from_type());
+ }
+ }
+
+ void visit(MeshIndexConversionStmt *stmt) override {
+ this->total_ptr->insert(stmt->idx_type);
+ if (optimize_mesh_reordered_mapping &&
+ stmt->conv_type == mesh::ConvType::l2r) {
+ this->owned_ptr->insert(stmt->idx_type);
+ }
+ }
+
+ OffloadedStmt *offload{nullptr};
+ MeshElementTypeSet *owned_ptr{nullptr};
+ MeshElementTypeSet *total_ptr{nullptr};
+ bool optimize_mesh_reordered_mapping{false};
+};
+
+namespace irpass::analysis {
+
+std::pair
+gather_mesh_thread_local(OffloadedStmt *offload, const CompileConfig &config) {
+ MeshElementTypeSet local_owned{};
+ MeshElementTypeSet local_total{};
+
+ GatherMeshThreadLocal::run(offload, &local_owned, &local_total, config);
+ return std::make_pair(local_owned, local_total);
+}
+
+} // namespace irpass::analysis
+
+TLANG_NAMESPACE_END
diff --git a/taichi/analysis/gather_meshfor_relation_types.cpp b/taichi/analysis/gather_meshfor_relation_types.cpp
new file mode 100644
index 0000000000000..21ce9ef60f328
--- /dev/null
+++ b/taichi/analysis/gather_meshfor_relation_types.cpp
@@ -0,0 +1,65 @@
+#include "taichi/ir/ir.h"
+#include "taichi/ir/snode.h"
+#include "taichi/ir/mesh.h"
+#include "taichi/ir/visitors.h"
+#include "taichi/ir/analysis.h"
+#include "taichi/ir/statements.h"
+
+TLANG_NAMESPACE_BEGIN
+
+namespace irpass::analysis {
+
+class GatherMeshforRelationTypes : public BasicStmtVisitor {
+ public:
+ using BasicStmtVisitor::visit;
+
+ GatherMeshforRelationTypes() {
+ allow_undefined_visitor = true;
+ invoke_default_visitor = true;
+ }
+
+ static void run(IRNode *root) {
+ GatherMeshforRelationTypes analyser;
+ root->accept(&analyser);
+ }
+
+ void visit(MeshForStmt *stmt) override {
+ TI_ASSERT(mesh_for == nullptr);
+ TI_ASSERT(stmt->major_to_types.size() == 0);
+ TI_ASSERT(stmt->minor_relation_types.size() == 0);
+ mesh_for = stmt;
+ stmt->body->accept(this);
+ mesh_for = nullptr;
+ }
+
+ void visit(MeshRelationAccessStmt *stmt) override {
+ if (auto from_stmt =
+ stmt->mesh_idx->cast()) { // major relation
+ TI_ASSERT(from_stmt->mesh_index_type() == mesh_for->major_from_type);
+ mesh_for->major_to_types.insert(stmt->to_type);
+ } else if (auto from_stmt =
+ stmt->mesh_idx
+ ->cast()) { // minor relation
+ TI_ASSERT(!from_stmt->is_size());
+ auto from_order = mesh::element_order(from_stmt->to_type);
+ auto to_order = mesh::element_order(stmt->to_type);
+ TI_ASSERT_INFO(from_order > to_order,
+ "Cannot access an indeterminate relation (E.g, Vert-Vert) "
+ "in a nested neighbor access");
+ mesh_for->minor_relation_types.insert(
+ mesh::relation_by_orders(from_order, to_order));
+ } else {
+ TI_NOT_IMPLEMENTED;
+ }
+ }
+
+ MeshForStmt *mesh_for{nullptr};
+};
+
+void gather_meshfor_relation_types(IRNode *node) {
+ GatherMeshforRelationTypes::run(node);
+}
+
+} // namespace irpass::analysis
+
+TLANG_NAMESPACE_END
diff --git a/taichi/analysis/gather_statements.cpp b/taichi/analysis/gather_statements.cpp
index 3264cfea72181..785afd0b767df 100644
--- a/taichi/analysis/gather_statements.cpp
+++ b/taichi/analysis/gather_statements.cpp
@@ -6,27 +6,27 @@ TLANG_NAMESPACE_BEGIN
class StmtSearcher : public BasicStmtVisitor {
private:
- std::function test;
- std::vector results;
+ std::function test_;
+ std::vector results_;
public:
using BasicStmtVisitor::visit;
- StmtSearcher(std::function test) : test(test) {
+ StmtSearcher(std::function test) : test_(test) {
allow_undefined_visitor = true;
invoke_default_visitor = true;
}
- void visit(Stmt *stmt) {
- if (test(stmt))
- results.push_back(stmt);
+ void visit(Stmt *stmt) override {
+ if (test_(stmt))
+ results_.push_back(stmt);
}
static std::vector run(IRNode *root,
const std::function &test) {
StmtSearcher searcher(test);
root->accept(&searcher);
- return searcher.results;
+ return searcher.results_;
}
};
diff --git a/taichi/analysis/gather_uniquely_accessed_pointers.cpp b/taichi/analysis/gather_uniquely_accessed_pointers.cpp
index 8febbe134ba36..69b4d767a2960 100644
--- a/taichi/analysis/gather_uniquely_accessed_pointers.cpp
+++ b/taichi/analysis/gather_uniquely_accessed_pointers.cpp
@@ -42,6 +42,10 @@ class LoopUniqueStmtSearcher : public BasicStmtVisitor {
loop_invariant_.insert(stmt);
}
+ void visit(ExternalTensorShapeAlongAxisStmt *stmt) override {
+ loop_invariant_.insert(stmt);
+ }
+
void visit(UnaryOpStmt *stmt) override {
if (loop_invariant_.count(stmt->operand) > 0) {
loop_invariant_.insert(stmt);
@@ -55,6 +59,21 @@ class LoopUniqueStmtSearcher : public BasicStmtVisitor {
}
}
+ void visit(DecorationStmt *stmt) override {
+ if (stmt->decoration.size() == 2 &&
+ stmt->decoration[0] ==
+ uint32_t(DecorationStmt::Decoration::kLoopUnique)) {
+ if (loop_unique_.find(stmt->operand) == loop_unique_.end()) {
+ // This decoration exists IFF we are looping over NDArray (or any other
+ // cases where the array index is linearized by the codegen) In that
+ // case the original loop dimensions have been reduced to 1D.
+ loop_unique_[stmt->operand] = stmt->decoration[1];
+ num_different_loop_indices = std::max(loop_unique_[stmt->operand] + 1,
+ num_different_loop_indices);
+ }
+ }
+ }
+
void visit(BinaryOpStmt *stmt) override {
if (loop_invariant_.count(stmt->lhs) > 0 &&
loop_invariant_.count(stmt->rhs) > 0) {
@@ -81,6 +100,10 @@ class LoopUniqueStmtSearcher : public BasicStmtVisitor {
}
}
+ bool is_partially_loop_unique(Stmt *stmt) const {
+ return loop_unique_.find(stmt) != loop_unique_.end();
+ }
+
bool is_ptr_indices_loop_unique(GlobalPtrStmt *stmt) const {
// Check if the address is loop-unique, i.e., stmt contains
// either a loop-unique index or all top-level loop indices.
@@ -108,6 +131,36 @@ class LoopUniqueStmtSearcher : public BasicStmtVisitor {
// b[i, i] is not loop-unique (because there's no j)
return current_num_different_loop_indices == num_different_loop_indices;
}
+
+ bool is_ptr_indices_loop_unique(ExternalPtrStmt *stmt) const {
+ // Check if the address is loop-unique, i.e., stmt contains
+ // either a loop-unique index or all top-level loop indices.
+ TI_ASSERT(num_different_loop_indices != -1);
+ std::vector loop_indices;
+ loop_indices.reserve(stmt->indices.size());
+ for (auto &index : stmt->indices) {
+ auto loop_unique_index = loop_unique_.find(index);
+ if (loop_unique_index != loop_unique_.end()) {
+ if (loop_unique_index->second == -1) {
+ // LoopUniqueStmt
+ return true;
+ } else {
+ // LoopIndexStmt
+ loop_indices.push_back(loop_unique_index->second);
+ }
+ }
+ }
+ std::sort(loop_indices.begin(), loop_indices.end());
+ auto current_num_different_loop_indices =
+ std::unique(loop_indices.begin(), loop_indices.end()) -
+ loop_indices.begin();
+
+ // for i, j in x:
+ // a[j, i] is loop-unique
+ // b[i, i] is not loop-unique (because there's no j)
+ // c[j, i, 1] is loop-unique
+ return current_num_different_loop_indices == num_different_loop_indices;
+ }
};
class UniquelyAccessedSNodeSearcher : public BasicStmtVisitor {
@@ -118,6 +171,10 @@ class UniquelyAccessedSNodeSearcher : public BasicStmtVisitor {
// one GlobalPtrStmt (or by definitely-same-address GlobalPtrStmts),
// and that GlobalPtrStmt's address is loop-unique.
std::unordered_map accessed_pointer_;
+ std::unordered_map rel_access_pointer_;
+
+ // Search any_arrs that are uniquely accessed. Maps: ArgID -> ExternalPtrStmt
+ std::unordered_map accessed_arr_pointer_;
public:
using BasicStmtVisitor::visit;
@@ -128,6 +185,33 @@ class UniquelyAccessedSNodeSearcher : public BasicStmtVisitor {
}
void visit(GlobalPtrStmt *stmt) override {
+ // mesh-for loop unique
+ if (stmt->indices.size() == 1 &&
+ stmt->indices[0]->is()) {
+ auto idx = stmt->indices[0]->as()->idx;
+ while (idx->is()) { // special case: l2g +
+ // g2r
+ idx = idx->as()->idx;
+ }
+ if (idx->is() &&
+ idx->as()->is_mesh_index()) { // from-end access
+ for (auto &snode : stmt->snodes.data) {
+ if (rel_access_pointer_.find(snode) ==
+ rel_access_pointer_.end()) { // not accessed by neibhours yet
+ accessed_pointer_[snode] = stmt;
+ } else { // accessed by neibhours, so it's not unique
+ accessed_pointer_[snode] = nullptr;
+ }
+ }
+ } else { // to-end access
+ for (auto &snode : stmt->snodes.data) {
+ rel_access_pointer_[snode] = stmt;
+ accessed_pointer_[snode] =
+ nullptr; // from-end access should not be unique
+ }
+ }
+ }
+ // Range-for / struct-for
for (auto &snode : stmt->snodes.data) {
auto accessed_ptr = accessed_pointer_.find(snode);
if (accessed_ptr == accessed_pointer_.end()) {
@@ -145,11 +229,70 @@ class UniquelyAccessedSNodeSearcher : public BasicStmtVisitor {
}
}
- static std::unordered_map run(IRNode *root) {
+ void visit(ExternalPtrStmt *stmt) override {
+ // A memory location of an ExternalPtrStmt depends on the indices
+ // If the accessed indices are loop unique,
+ // the accessed memory location is loop unique
+ for (auto base_ptr : stmt->base_ptrs.data) {
+ ArgLoadStmt *arg_load_stmt = base_ptr->as();
+ int arg_id = arg_load_stmt->arg_id;
+
+ auto accessed_ptr = accessed_arr_pointer_.find(arg_id);
+
+ bool stmt_loop_unique =
+ loop_unique_stmt_searcher_.is_ptr_indices_loop_unique(stmt);
+
+ if (!stmt_loop_unique) {
+ accessed_arr_pointer_[arg_id] = nullptr; // not loop-unique
+ } else {
+ if (accessed_ptr == accessed_arr_pointer_.end()) {
+ // First time using arr @ arg_id
+ accessed_arr_pointer_[arg_id] = stmt;
+ } else {
+ /**
+ * We know stmt->base_ptr and the previously recorded pointers
+ * are loop-unique. We need to figure out whether their loop-unique
+ * indicies are the same while ignoring the others.
+ * e.g. a[i, j, 1] and a[i, j, 2] are both uniquely accessed
+ * a[i, j, 1] and a[j, i, 2] are not uniquely accessed
+ * a[i, j + 1, 1] and a[i, j, 2] are not uniquely accessed
+ * This is a bit stricter than needed.
+ * e.g. a[i, j, i] and a[i, j, 0] are uniquely accessed
+ * However this is probably not common and improvements can be made
+ * in a future patch.
+ */
+ if (accessed_ptr->second) {
+ ExternalPtrStmt *other_ptr = accessed_ptr->second;
+ TI_ASSERT(stmt->indices.size() == other_ptr->indices.size());
+ for (int axis = 0; axis < stmt->indices.size(); axis++) {
+ Stmt *this_index = stmt->indices[axis];
+ Stmt *other_index = other_ptr->indices[axis];
+ // We only compare unique indices here.
+ // Since both pointers are loop-unique, all the unique indices
+ // need to be the same for both to be uniquely accessed
+ if (loop_unique_stmt_searcher_.is_partially_loop_unique(
+ this_index)) {
+ if (!irpass::analysis::same_value(this_index, other_index)) {
+ // Not equal -> not uniquely accessed
+ accessed_arr_pointer_[arg_id] = nullptr;
+ break;
+ }
+ }
+ }
+ }
+ }
+ }
+ }
+ }
+
+ static std::pair,
+ std::unordered_map>
+ run(IRNode *root) {
TI_ASSERT(root->is());
auto offload = root->as();
UniquelyAccessedSNodeSearcher searcher;
- if (offload->task_type == OffloadedTaskType::range_for) {
+ if (offload->task_type == OffloadedTaskType::range_for ||
+ offload->task_type == OffloadedTaskType::mesh_for) {
searcher.loop_unique_stmt_searcher_.num_different_loop_indices = 1;
} else if (offload->task_type == OffloadedTaskType::struct_for) {
searcher.loop_unique_stmt_searcher_.num_different_loop_indices =
@@ -160,7 +303,9 @@ class UniquelyAccessedSNodeSearcher : public BasicStmtVisitor {
}
root->accept(&searcher.loop_unique_stmt_searcher_);
root->accept(&searcher);
- return searcher.accessed_pointer_;
+
+ return std::make_pair(searcher.accessed_pointer_,
+ searcher.accessed_arr_pointer_);
}
};
@@ -180,13 +325,18 @@ class UniquelyAccessedBitStructGatherer : public BasicStmtVisitor {
void visit(OffloadedStmt *stmt) override {
if (stmt->task_type == OffloadedTaskType::range_for ||
+ stmt->task_type == OffloadedTaskType::mesh_for ||
stmt->task_type == OffloadedTaskType::struct_for) {
auto &loop_unique_bit_struct = result_[stmt];
auto loop_unique_ptr =
- irpass::analysis::gather_uniquely_accessed_pointers(stmt);
+ irpass::analysis::gather_uniquely_accessed_pointers(stmt).first;
for (auto &it : loop_unique_ptr) {
auto *snode = it.first;
auto *ptr1 = it.second;
+ if (ptr1 != nullptr && ptr1->indices.size() > 0 &&
+ ptr1->indices[0]->is()) {
+ continue;
+ }
if (snode->is_bit_level) {
// Find the nearest non-bit-level ancestor
while (snode->is_bit_level) {
@@ -229,7 +379,8 @@ const std::string GatherUniquelyAccessedBitStructsPass::id =
"GatherUniquelyAccessedBitStructsPass";
namespace irpass::analysis {
-std::unordered_map
+std::pair,
+ std::unordered_map>
gather_uniquely_accessed_pointers(IRNode *root) {
// TODO: What about SNodeOpStmts?
return UniquelyAccessedSNodeSearcher::run(root);
diff --git a/taichi/analysis/gather_used_atomics.cpp b/taichi/analysis/gather_used_atomics.cpp
index 1664be7688e6a..cb1d8d32150ba 100644
--- a/taichi/analysis/gather_used_atomics.cpp
+++ b/taichi/analysis/gather_used_atomics.cpp
@@ -8,7 +8,7 @@ TLANG_NAMESPACE_BEGIN
class UsedAtomicsSearcher : public BasicStmtVisitor {
private:
- std::unique_ptr> used_atomics;
+ std::unique_ptr> used_atomics_;
public:
using BasicStmtVisitor::visit;
@@ -16,13 +16,13 @@ class UsedAtomicsSearcher : public BasicStmtVisitor {
UsedAtomicsSearcher() {
allow_undefined_visitor = true;
invoke_default_visitor = true;
- used_atomics = std::make_unique>();
+ used_atomics_ = std::make_unique>();
}
void search_operands(Stmt *stmt) {
for (auto &op : stmt->get_operands()) {
if (op != nullptr && op->is()) {
- used_atomics->insert(op->as());
+ used_atomics_->insert(op->as());
}
}
}
@@ -38,7 +38,7 @@ class UsedAtomicsSearcher : public BasicStmtVisitor {
static std::unique_ptr> run(IRNode *root) {
UsedAtomicsSearcher searcher;
root->accept(&searcher);
- return std::move(searcher.used_atomics);
+ return std::move(searcher.used_atomics_);
}
};
diff --git a/taichi/analysis/has_store_or_atomic.cpp b/taichi/analysis/has_store_or_atomic.cpp
index 04d546673364c..391ee99fbba5a 100644
--- a/taichi/analysis/has_store_or_atomic.cpp
+++ b/taichi/analysis/has_store_or_atomic.cpp
@@ -8,14 +8,14 @@ TLANG_NAMESPACE_BEGIN
// Find if there is a store (or AtomicOpStmt).
class LocalStoreSearcher : public BasicStmtVisitor {
private:
- const std::vector &vars;
- bool result;
+ const std::vector &vars_;
+ bool result_;
public:
using BasicStmtVisitor::visit;
explicit LocalStoreSearcher(const std::vector &vars)
- : vars(vars), result(false) {
+ : vars_(vars), result_(false) {
for (auto var : vars) {
TI_ASSERT(var->is());
}
@@ -24,18 +24,18 @@ class LocalStoreSearcher : public BasicStmtVisitor {
}
void visit(LocalStoreStmt *stmt) override {
- for (auto var : vars) {
+ for (auto var : vars_) {
if (stmt->dest == var) {
- result = true;
+ result_ = true;
break;
}
}
}
void visit(AtomicOpStmt *stmt) override {
- for (auto var : vars) {
+ for (auto var : vars_) {
if (stmt->dest == var) {
- result = true;
+ result_ = true;
break;
}
}
@@ -44,7 +44,7 @@ class LocalStoreSearcher : public BasicStmtVisitor {
static bool run(IRNode *root, const std::vector &vars) {
LocalStoreSearcher searcher(vars);
root->accept(&searcher);
- return searcher.result;
+ return searcher.result_;
}
};
diff --git a/taichi/analysis/last_store_or_atomic.cpp b/taichi/analysis/last_store_or_atomic.cpp
index 744e4ef1c62d9..4c90a094a11b1 100644
--- a/taichi/analysis/last_store_or_atomic.cpp
+++ b/taichi/analysis/last_store_or_atomic.cpp
@@ -9,37 +9,37 @@ TLANG_NAMESPACE_BEGIN
// after the last store.
class LocalStoreForwarder : public BasicStmtVisitor {
private:
- Stmt *var;
- bool is_valid;
- Stmt *result;
+ Stmt *var_;
+ bool is_valid_;
+ Stmt *result_;
public:
using BasicStmtVisitor::visit;
explicit LocalStoreForwarder(Stmt *var)
- : var(var), is_valid(true), result(nullptr) {
+ : var_(var), is_valid_(true), result_(nullptr) {
TI_ASSERT(var->is());
allow_undefined_visitor = true;
invoke_default_visitor = true;
}
void visit(LocalStoreStmt *stmt) override {
- if (stmt->dest == var) {
- is_valid = true;
- result = stmt;
+ if (stmt->dest == var_) {
+ is_valid_ = true;
+ result_ = stmt;
}
}
void visit(AllocaStmt *stmt) override {
- if (stmt == var) {
- is_valid = true;
- result = stmt;
+ if (stmt == var_) {
+ is_valid_ = true;
+ result_ = stmt;
}
}
void visit(AtomicOpStmt *stmt) override {
- if (stmt->dest == var) {
- is_valid = false;
+ if (stmt->dest == var_) {
+ is_valid_ = false;
}
}
@@ -50,33 +50,33 @@ class LocalStoreForwarder : public BasicStmtVisitor {
std::pair true_branch(true, nullptr);
if (if_stmt->true_statements) {
// create a new LocalStoreForwarder instance
- true_branch = run(if_stmt->true_statements.get(), var);
+ true_branch = run(if_stmt->true_statements.get(), var_);
}
std::pair false_branch(true, nullptr);
if (if_stmt->false_statements) {
- false_branch = run(if_stmt->false_statements.get(), var);
+ false_branch = run(if_stmt->false_statements.get(), var_);
}
auto true_stmt = true_branch.second;
auto false_stmt = false_branch.second;
if (!true_branch.first || !false_branch.first) {
// at least one branch finally modifies the variable without storing
- is_valid = false;
+ is_valid_ = false;
} else if (true_stmt == nullptr && false_stmt == nullptr) {
// both branches don't modify the variable
return;
} else if (true_stmt == nullptr || false_stmt == nullptr) {
// only one branch modifies the variable
- is_valid = false;
+ is_valid_ = false;
} else {
TI_ASSERT(true_stmt->is());
TI_ASSERT(false_stmt->is());
if (true_stmt->as()->val !=
false_stmt->as()->val) {
// two branches finally store the variable differently
- is_valid = false;
+ is_valid_ = false;
} else {
- is_valid = true;
- result = true_stmt; // same as false_stmt
+ is_valid_ = true;
+ result_ = true_stmt; // same as false_stmt
}
}
}
@@ -85,33 +85,33 @@ class LocalStoreForwarder : public BasicStmtVisitor {
// the "last" store inside a loop to the local load statement.
// What we can do is just check if the loop doesn't modify the variable.
void visit(WhileStmt *stmt) override {
- if (irpass::analysis::has_store_or_atomic(stmt, {var})) {
- is_valid = false;
+ if (irpass::analysis::has_store_or_atomic(stmt, {var_})) {
+ is_valid_ = false;
}
}
void visit(RangeForStmt *stmt) override {
- if (irpass::analysis::has_store_or_atomic(stmt, {var})) {
- is_valid = false;
+ if (irpass::analysis::has_store_or_atomic(stmt, {var_})) {
+ is_valid_ = false;
}
}
void visit(StructForStmt *stmt) override {
- if (irpass::analysis::has_store_or_atomic(stmt, {var})) {
- is_valid = false;
+ if (irpass::analysis::has_store_or_atomic(stmt, {var_})) {
+ is_valid_ = false;
}
}
void visit(OffloadedStmt *stmt) override {
- if (irpass::analysis::has_store_or_atomic(stmt, {var})) {
- is_valid = false;
+ if (irpass::analysis::has_store_or_atomic(stmt, {var_})) {
+ is_valid_ = false;
}
}
static std::pair run(IRNode *root, Stmt *var) {
LocalStoreForwarder searcher(var);
root->accept(&searcher);
- return std::make_pair(searcher.is_valid, searcher.result);
+ return std::make_pair(searcher.is_valid_, searcher.result_);
}
};
diff --git a/taichi/analysis/mesh_bls_analyzer.cpp b/taichi/analysis/mesh_bls_analyzer.cpp
new file mode 100644
index 0000000000000..280636acf006e
--- /dev/null
+++ b/taichi/analysis/mesh_bls_analyzer.cpp
@@ -0,0 +1,118 @@
+#include "taichi/analysis/mesh_bls_analyzer.h"
+
+#include "taichi/system/profiler.h"
+#include "taichi/ir/analysis.h"
+
+namespace taichi {
+namespace lang {
+
+MeshBLSAnalyzer::MeshBLSAnalyzer(OffloadedStmt *for_stmt,
+ MeshBLSCaches *caches,
+ bool auto_mesh_local,
+ const CompileConfig &config)
+ : for_stmt_(for_stmt),
+ caches_(caches),
+ auto_mesh_local_(auto_mesh_local),
+ config_(config) {
+ TI_AUTO_PROF;
+ allow_undefined_visitor = true;
+ invoke_default_visitor = false;
+}
+
+void MeshBLSAnalyzer::record_access(Stmt *stmt, AccessFlag flag) {
+ if (!analysis_ok_) {
+ return;
+ }
+ if (!stmt->is())
+ return; // local alloca
+ auto ptr = stmt->as();
+ if (ptr->indices.size() != std::size_t(1) ||
+ !ptr->indices[0]->is())
+ return;
+ auto conv = ptr->indices[0]->as();
+ auto element_type = conv->idx_type;
+ auto conv_type = conv->conv_type;
+ auto idx = conv->idx;
+ if (conv_type == mesh::ConvType::g2r)
+ return;
+ for (int l = 0; l < stmt->width(); l++) {
+ auto snode = ptr->snodes[l];
+ if (!caches_->has(snode)) {
+ if (auto_mesh_local_ &&
+ (flag == AccessFlag::accumulate ||
+ (flag == AccessFlag::read && config_.arch == Arch::cuda)) &&
+ (!idx->is() ||
+ !idx->as()->is_mesh_index())) {
+ caches_->insert(snode);
+ } else {
+ continue;
+ }
+ }
+
+ if (!caches_->access(snode, element_type, conv_type, flag,
+ idx->as()->neighbor_idx)) {
+ analysis_ok_ = false;
+ break;
+ }
+ }
+}
+
+void MeshBLSAnalyzer::visit(GlobalLoadStmt *stmt) {
+ TI_ASSERT(stmt->width() == 1); // TODO: support vectorization
+ record_access(stmt->src, AccessFlag::read);
+}
+
+void MeshBLSAnalyzer::visit(GlobalStoreStmt *stmt) {
+ TI_ASSERT(stmt->width() == 1); // TODO: support vectorization
+ record_access(stmt->dest, AccessFlag::write);
+}
+
+void MeshBLSAnalyzer::visit(AtomicOpStmt *stmt) {
+ if (stmt->op_type == AtomicOpType::add) {
+ record_access(stmt->dest, AccessFlag::accumulate);
+ }
+}
+
+void MeshBLSAnalyzer::visit(Stmt *stmt) {
+ TI_ASSERT(!stmt->is_container_statement());
+}
+
+bool MeshBLSAnalyzer::run() {
+ const auto &block = for_stmt_->body;
+
+ for (int i = 0; i < (int)block->statements.size(); i++) {
+ block->statements[i]->accept(this);
+ }
+
+ return analysis_ok_;
+}
+
+namespace irpass {
+namespace analysis {
+
+std::unique_ptr initialize_mesh_local_attribute(
+ OffloadedStmt *offload,
+ bool auto_mesh_local,
+ const CompileConfig &config) {
+ TI_AUTO_PROF
+ TI_ASSERT(offload->task_type == OffloadedTaskType::mesh_for);
+ std::unique_ptr caches;
+ caches = std::make_unique();
+ for (auto snode : offload->mem_access_opt.get_snodes_with_flag(
+ SNodeAccessFlag::mesh_local)) {
+ caches->insert(snode);
+ }
+
+ MeshBLSAnalyzer bls_analyzer(offload, caches.get(), auto_mesh_local, config);
+ bool analysis_ok = bls_analyzer.run();
+ if (!analysis_ok) {
+ TI_ERROR("Mesh BLS analysis failed !");
+ }
+ return caches;
+}
+
+} // namespace analysis
+} // namespace irpass
+
+} // namespace lang
+} // namespace taichi
diff --git a/taichi/analysis/mesh_bls_analyzer.h b/taichi/analysis/mesh_bls_analyzer.h
new file mode 100644
index 0000000000000..f8db1c3d2b9ba
--- /dev/null
+++ b/taichi/analysis/mesh_bls_analyzer.h
@@ -0,0 +1,158 @@
+#pragma once
+
+#include "taichi/program/compile_config.h"
+#include "taichi/ir/visitors.h"
+#include "taichi/ir/statements.h"
+#include "taichi/ir/mesh.h"
+
+#include
+
+namespace taichi {
+namespace lang {
+
+class MeshBLSCache {
+ public:
+ using AccessFlag = taichi::lang::AccessFlag;
+ using Rec = std::map,
+ std::set>>;
+
+ SNode *snode{nullptr};
+ mesh::MeshElementType element_type;
+ mesh::ConvType conv_type;
+
+ bool initialized;
+ bool finalized;
+ bool loop_index;
+ int unique_accessed;
+ AccessFlag total_flags;
+
+ MeshBLSCache() = default;
+
+ MeshBLSCache(SNode *snode) : snode(snode) {
+ total_flags = AccessFlag(0);
+ initialized = false;
+ finalized = false;
+ loop_index = false;
+ unique_accessed = 0;
+ }
+
+ bool access(mesh::MeshElementType element_type,
+ mesh::ConvType conv_type,
+ AccessFlag flags,
+ Stmt *idx) {
+ if (!initialized) {
+ initialized = true;
+ this->conv_type = conv_type;
+ this->element_type = element_type;
+ } else {
+ if (this->conv_type != conv_type || this->element_type != element_type)
+ return false;
+ }
+ this->total_flags |= flags;
+ if (idx->is()) {
+ loop_index = true;
+ } else {
+ unique_accessed++;
+ }
+ return true;
+ }
+
+ void finalize(Rec &rec) {
+ TI_ASSERT(!finalized);
+ finalized = true;
+ if (initialized) {
+ const auto cache_type = std::make_pair(element_type, conv_type);
+ auto ptr = rec.find(cache_type);
+ if (ptr == rec.end()) {
+ ptr = rec.emplace(std::piecewise_construct,
+ std::forward_as_tuple(cache_type),
+ std::forward_as_tuple())
+ .first;
+ }
+ ptr->second.insert(std::make_pair(snode, total_flags));
+ }
+ }
+};
+
+class MeshBLSCaches {
+ public:
+ std::map caches;
+
+ using AccessFlag = MeshBLSCache::AccessFlag;
+ using Rec = MeshBLSCache::Rec;
+
+ void insert(SNode *snode) {
+ if (caches.find(snode) == caches.end()) {
+ caches.emplace(std::piecewise_construct, std::forward_as_tuple(snode),
+ std::forward_as_tuple(snode));
+ } else {
+ TI_ERROR("mesh::MeshBLSCaches for {} already exists.",
+ snode->node_type_name);
+ }
+ }
+
+ bool access(SNode *snode,
+ mesh::MeshElementType element_type,
+ mesh::ConvType conv_type,
+ AccessFlag flags,
+ Stmt *idx) {
+ if (caches.find(snode) == caches.end())
+ return false;
+ return caches.find(snode)->second.access(element_type, conv_type, flags,
+ idx);
+ }
+
+ Rec finalize() {
+ Rec rec;
+ for (auto &cache : caches) {
+ cache.second.finalize(rec);
+ }
+ return rec;
+ }
+
+ bool has(SNode *snode) {
+ return caches.find(snode) != caches.end();
+ }
+
+ MeshBLSCache &get(SNode *snode) {
+ TI_ASSERT(caches.find(snode) != caches.end());
+ return caches[snode];
+ }
+};
+
+// Figure out accessed SNodes, and their ranges in this for stmt
+class MeshBLSAnalyzer : public BasicStmtVisitor {
+ using BasicStmtVisitor::visit;
+
+ public:
+ MeshBLSAnalyzer(OffloadedStmt *for_stmt,
+ MeshBLSCaches *caches,
+ bool auto_mesh_local,
+ const CompileConfig &config);
+
+ void visit(GlobalPtrStmt *stmt) override {
+ }
+
+ // Do not eliminate global data access
+ void visit(GlobalLoadStmt *stmt) override;
+
+ void visit(GlobalStoreStmt *stmt) override;
+
+ void visit(AtomicOpStmt *stmt) override;
+
+ void visit(Stmt *stmt) override;
+
+ bool run();
+
+ private:
+ void record_access(Stmt *stmt, AccessFlag flag);
+
+ OffloadedStmt *for_stmt_{nullptr};
+ MeshBLSCaches *caches_{nullptr};
+ bool analysis_ok_{true};
+ bool auto_mesh_local_{false};
+ CompileConfig config_;
+};
+
+} // namespace lang
+} // namespace taichi
diff --git a/taichi/analysis/same_statements.cpp b/taichi/analysis/same_statements.cpp
index e53b4baa35a5a..4a8cd890ff1b8 100644
--- a/taichi/analysis/same_statements.cpp
+++ b/taichi/analysis/same_statements.cpp
@@ -12,9 +12,9 @@ TLANG_NAMESPACE_BEGIN
// Compare if two IRNodes are equivalent.
class IRNodeComparator : public IRVisitor {
private:
- IRNode *other_node;
+ IRNode *other_node_;
// map the id from this node to the other node
- std::unordered_map id_map;
+ std::unordered_map id_map_;
bool recursively_check_;
@@ -39,13 +39,13 @@ class IRNodeComparator : public IRVisitor {
const std::optional>
&possibly_modified_states,
IRBank *ir_bank)
- : other_node(other_node) {
+ : other_node_(other_node) {
allow_undefined_visitor = true;
invoke_default_visitor = true;
same = true;
if (id_map.has_value()) {
recursively_check_ = true;
- this->id_map = id_map.value();
+ this->id_map_ = id_map.value();
} else {
recursively_check_ = false;
}
@@ -66,9 +66,9 @@ class IRNodeComparator : public IRVisitor {
}
void map_id(int this_id, int other_id) {
- auto it = id_map.find(this_id);
- if (it == id_map.end()) {
- id_map[this_id] = other_id;
+ auto it = id_map_.find(this_id);
+ if (it == id_map_.end()) {
+ id_map_[this_id] = other_id;
} else if (it->second != other_id) {
same = false;
}
@@ -77,8 +77,8 @@ class IRNodeComparator : public IRVisitor {
void check_mapping(Stmt *this_stmt, Stmt *other_stmt) {
// get the corresponding id in the other node
// and check if it is other_stmt->id
- auto it = id_map.find(this_stmt->id);
- if (it != id_map.end()) {
+ auto it = id_map_.find(this_stmt->id);
+ if (it != id_map_.end()) {
if (it->second != other_stmt->id) {
same = false;
}
@@ -89,43 +89,43 @@ class IRNodeComparator : public IRVisitor {
if (this_stmt->id != other_stmt->id) {
same = false;
}
- id_map[this_stmt->id] = other_stmt->id;
+ id_map_[this_stmt->id] = other_stmt->id;
} else {
// recursively check them
- IRNode *backup_other_node = other_node;
- other_node = other_stmt;
+ IRNode *backup_other_node = other_node_;
+ other_node_ = other_stmt;
this_stmt->accept(this);
- other_node = backup_other_node;
+ other_node_ = backup_other_node;
}
}
void visit(Block *stmt_list) override {
- if (!other_node->is()) {
+ if (!other_node_->is()) {
same = false;
return;
}
- auto other = other_node->as();
+ auto other = other_node_->as();
if (stmt_list->size() != other->size()) {
same = false;
return;
}
for (int i = 0; i < (int)stmt_list->size(); i++) {
- other_node = other->statements[i].get();
+ other_node_ = other->statements[i].get();
stmt_list->statements[i]->accept(this);
if (!same)
break;
}
- other_node = other;
+ other_node_ = other;
}
void basic_check(Stmt *stmt) {
// type check
- if (typeid(*other_node) != typeid(*stmt)) {
+ if (typeid(*other_node_) != typeid(*stmt)) {
same = false;
return;
}
- auto other = other_node->as();
+ auto other = other_node_->as();
if (stmt == other) {
return;
}
@@ -240,78 +240,68 @@ class IRNodeComparator : public IRVisitor {
basic_check(stmt);
if (!same)
return;
- auto other = other_node->as();
+ auto other = other_node_->as();
if (stmt->true_statements) {
if (!other->true_statements) {
same = false;
return;
}
- other_node = other->true_statements.get();
+ other_node_ = other->true_statements.get();
stmt->true_statements->accept(this);
- other_node = other;
+ other_node_ = other;
}
if (stmt->false_statements && same) {
if (!other->false_statements) {
same = false;
return;
}
- other_node = other->false_statements.get();
+ other_node_ = other->false_statements.get();
stmt->false_statements->accept(this);
- other_node = other;
+ other_node_ = other;
}
}
- void visit(FuncBodyStmt *stmt) override {
- basic_check(stmt);
- if (!same)
- return;
- auto other = other_node->as();
- other_node = other->body.get();
- stmt->body->accept(this);
- other_node = other;
- }
-
void visit(WhileStmt *stmt) override {
basic_check(stmt);
if (!same)
return;
- auto other = other_node->as();
- other_node = other->body.get();
+ auto other = other_node_->as();
+ other_node_ = other->body.get();
stmt->body->accept(this);
- other_node = other;
+ other_node_ = other;
}
void visit(RangeForStmt *stmt) override {
basic_check(stmt);
if (!same)
return;
- auto other = other_node->as();
- other_node = other->body.get();
+ auto other = other_node_->as();
+ other_node_ = other->body.get();
stmt->body->accept(this);
- other_node = other;
+ other_node_ = other;
}
void visit(StructForStmt *stmt) override {
basic_check(stmt);
if (!same)
return;
- auto other = other_node->as();
- other_node = other->body.get();
+ auto other = other_node_->as();
+ other_node_ = other->body.get();
stmt->body->accept(this);
- other_node = other;
+ other_node_ = other;
}
void visit(OffloadedStmt *stmt) override {
basic_check(stmt);
if (!same)
return;
- auto other = other_node->as();
+ auto other = other_node_->as();
if (stmt->has_body()) {
TI_ASSERT(stmt->body);
TI_ASSERT(other->body);
- other_node = other->body.get();
+ other_node_ = other->body.get();
stmt->body->accept(this);
- other_node = other;
+ other_node_ = other;
}
}
diff --git a/taichi/analysis/value_diff.cpp b/taichi/analysis/value_diff.cpp
index b6d88e4b25eb7..bea01f39ef63a 100644
--- a/taichi/analysis/value_diff.cpp
+++ b/taichi/analysis/value_diff.cpp
@@ -9,18 +9,18 @@ namespace taichi {
namespace lang {
DiffRange operator+(const DiffRange &a, const DiffRange &b) {
- return DiffRange(a.related_() && b.related_(), a.coeff + b.coeff,
- a.low + b.low, a.high + b.high - 1);
+ return DiffRange(a.related() && b.related(), a.coeff + b.coeff, a.low + b.low,
+ a.high + b.high - 1);
}
DiffRange operator-(const DiffRange &a, const DiffRange &b) {
- return DiffRange(a.related_() && b.related_(), a.coeff - b.coeff,
+ return DiffRange(a.related() && b.related(), a.coeff - b.coeff,
a.low - b.high + 1, a.high - b.low);
}
DiffRange operator*(const DiffRange &a, const DiffRange &b) {
return DiffRange(
- a.related_() && b.related_() && a.coeff * b.coeff == 0,
+ a.related() && b.related() && a.coeff * b.coeff == 0,
fmax(a.low * b.coeff, a.coeff * b.low),
fmin(a.low * b.low,
fmin(a.low * (b.high - 1),
@@ -33,7 +33,7 @@ DiffRange operator*(const DiffRange &a, const DiffRange &b) {
DiffRange operator<<(const DiffRange &a, const DiffRange &b) {
return DiffRange(
- a.related_() && b.related_() && b.coeff == 0 && b.high - b.low == 1,
+ a.related() && b.related() && b.coeff == 0 && b.high - b.low == 1,
a.coeff << b.low, a.low << b.low, ((a.high - 1) << b.low) + 1);
}
@@ -113,7 +113,7 @@ class ValueDiffLoopIndex : public IRVisitor {
stmt->rhs->accept(this);
auto ret1 = results[stmt->lhs->instance_id];
auto ret2 = results[stmt->rhs->instance_id];
- if (ret1.related_() && ret2.related_()) {
+ if (ret1.related() && ret2.related()) {
if (stmt->op_type == BinaryOpType::add) {
results[stmt->instance_id] = ret1 + ret2;
} else if (stmt->op_type == BinaryOpType::sub) {
diff --git a/taichi/analysis/verify.cpp b/taichi/analysis/verify.cpp
index ed22468538304..cd76cb8e30316 100644
--- a/taichi/analysis/verify.cpp
+++ b/taichi/analysis/verify.cpp
@@ -12,51 +12,51 @@ TLANG_NAMESPACE_BEGIN
class IRVerifier : public BasicStmtVisitor {
private:
- Block *current_block;
- Stmt *current_container_stmt;
+ Block *current_block_;
+ Stmt *current_container_stmt_;
// each scope corresponds to an unordered_set
- std::vector> visible_stmts;
+ std::vector> visible_stmts_;
public:
using BasicStmtVisitor::visit;
explicit IRVerifier(IRNode *root)
- : current_block(nullptr), current_container_stmt(nullptr) {
+ : current_block_(nullptr), current_container_stmt_(nullptr) {
allow_undefined_visitor = true;
invoke_default_visitor = true;
if (!root->is())
- visible_stmts.emplace_back();
+ visible_stmts_.emplace_back();
if (root->is() && root->as()->is_container_statement()) {
- current_container_stmt = root->as();
+ current_container_stmt_ = root->as();
}
}
void basic_verify(Stmt *stmt) {
- TI_ASSERT_INFO(stmt->parent == current_block,
+ TI_ASSERT_INFO(stmt->parent == current_block_,
"stmt({})->parent({}) != current_block({})", stmt->id,
- fmt::ptr(stmt->parent), fmt::ptr(current_block));
+ fmt::ptr(stmt->parent), fmt::ptr(current_block_));
for (auto &op : stmt->get_operands()) {
if (op == nullptr)
continue;
bool found = false;
- for (int depth = (int)visible_stmts.size() - 1; depth >= 0; depth--) {
- if (visible_stmts[depth].find(op) != visible_stmts[depth].end()) {
+ for (int depth = (int)visible_stmts_.size() - 1; depth >= 0; depth--) {
+ if (visible_stmts_[depth].find(op) != visible_stmts_[depth].end()) {
found = true;
break;
}
}
TI_ASSERT_INFO(
found,
- "IR broken: stmt {} cannot have operand {}."
+ "IR broken: stmt {} {} cannot have operand {} {}."
" If you are using autodiff, please check"
" https://docs.taichi.graphics/lang/articles/advanced/"
"differentiable_programming#kernel-simplicity-rule"
" If it doesn't help, please report this bug by opening an issue at"
" https://github.com/taichi-dev/taichi to help us improve."
" Thanks in advance!",
- stmt->id, op->id);
+ stmt->type(), stmt->id, op->type(), op->id);
}
- visible_stmts.back().insert(stmt);
+ visible_stmts_.back().insert(stmt);
}
void preprocess_container_stmt(Stmt *stmt) override {
@@ -69,23 +69,25 @@ class IRVerifier : public BasicStmtVisitor {
void visit(Block *block) override {
TI_ASSERT_INFO(
- block->parent_stmt == current_container_stmt,
+ block->parent_stmt == current_container_stmt_,
"block({})->parent({}) != current_container_stmt({})", fmt::ptr(block),
block->parent_stmt ? block->parent_stmt->name() : "nullptr",
- current_container_stmt ? current_container_stmt->name() : "nullptr");
- auto backup_block = current_block;
- current_block = block;
- auto backup_container_stmt = current_container_stmt;
- visible_stmts.emplace_back();
+ current_container_stmt_ ? current_container_stmt_->name() : "nullptr");
+ auto backup_block = current_block_;
+ current_block_ = block;
+ auto backup_container_stmt = current_container_stmt_;
+ if (!block->parent_stmt || !block->parent_stmt->is())
+ visible_stmts_.emplace_back();
for (auto &stmt : block->statements) {
if (stmt->is_container_statement())
- current_container_stmt = stmt.get();
+ current_container_stmt_ = stmt.get();
stmt->accept(this);
if (stmt->is_container_statement())
- current_container_stmt = backup_container_stmt;
+ current_container_stmt_ = backup_container_stmt;
}
- current_block = backup_block;
- visible_stmts.pop_back();
+ current_block_ = backup_block;
+ if (!block->parent_stmt || !block->parent_stmt->is())
+ current_block_ = backup_block;
}
void visit(OffloadedStmt *stmt) override {
@@ -122,10 +124,13 @@ class IRVerifier : public BasicStmtVisitor {
if (stmt->loop->is()) {
TI_ASSERT(stmt->loop->as()->task_type ==
OffloadedStmt::TaskType::struct_for ||
+ stmt->loop->as()->task_type ==
+ OffloadedStmt::TaskType::mesh_for ||
stmt->loop->as()->task_type ==
OffloadedStmt::TaskType::range_for);
} else {
TI_ASSERT(stmt->loop->is() ||
+ stmt->loop->is() ||
stmt->loop->is());
}
}
diff --git a/taichi/program/aot_module_builder.cpp b/taichi/aot/module_builder.cpp
similarity index 59%
rename from taichi/program/aot_module_builder.cpp
rename to taichi/aot/module_builder.cpp
index 306b5a49997e9..b194d2ee384c6 100644
--- a/taichi/program/aot_module_builder.cpp
+++ b/taichi/aot/module_builder.cpp
@@ -1,4 +1,4 @@
-#include "taichi/program/aot_module_builder.h"
+#include "taichi/aot/module_builder.h"
#include "taichi/program/kernel.h"
namespace taichi {
@@ -12,12 +12,14 @@ void AotModuleBuilder::add(const std::string &identifier, Kernel *kernel) {
}
void AotModuleBuilder::add_field(const std::string &identifier,
+ const SNode *rep_snode,
bool is_scalar,
DataType dt,
std::vector shape,
int row_num,
int column_num) {
- add_per_backend_field(identifier, is_scalar, dt, shape, row_num, column_num);
+ add_field_per_backend(identifier, rep_snode, is_scalar, dt, shape, row_num,
+ column_num);
}
void AotModuleBuilder::add_kernel_template(const std::string &identifier,
@@ -29,5 +31,26 @@ void AotModuleBuilder::add_kernel_template(const std::string &identifier,
add_per_backend_tmpl(identifier, key, kernel);
}
+bool AotModuleBuilder::all_fields_are_dense_in_container(
+ const SNode *container) {
+ for (const auto &ch : container->ch) {
+ if (ch->type != SNodeType::place) {
+ return false;
+ }
+ }
+ const auto *parent = container->parent;
+ if (!parent) {
+ return false;
+ }
+ if (parent->type != SNodeType::root) {
+ return false;
+ }
+ return true;
+}
+
+void AotModuleBuilder::load(const std::string &output_dir) {
+ TI_ERROR("Aot loader not supported");
+}
+
} // namespace lang
} // namespace taichi
diff --git a/taichi/program/aot_module_builder.h b/taichi/aot/module_builder.h
similarity index 64%
rename from taichi/program/aot_module_builder.h
rename to taichi/aot/module_builder.h
index 15fb4013b8662..1b05af1875038 100644
--- a/taichi/program/aot_module_builder.h
+++ b/taichi/aot/module_builder.h
@@ -3,6 +3,11 @@
#include
#include
+#include "taichi/aot/module_data.h"
+#include "taichi/backends/device.h"
+#include "taichi/ir/snode.h"
+#include "taichi/aot/module_data.h"
+
namespace taichi {
namespace lang {
@@ -16,6 +21,7 @@ class AotModuleBuilder {
void add(const std::string &identifier, Kernel *kernel);
void add_field(const std::string &identifier,
+ const SNode *rep_snode,
bool is_scalar,
DataType dt,
std::vector shape,
@@ -26,6 +32,8 @@ class AotModuleBuilder {
const std::string &key,
Kernel *kernel);
+ virtual void load(const std::string &output_dir);
+
virtual void dump(const std::string &output_dir,
const std::string &filename) const = 0;
@@ -35,15 +43,27 @@ class AotModuleBuilder {
*/
virtual void add_per_backend(const std::string &identifier,
Kernel *kernel) = 0;
- virtual void add_per_backend_field(const std::string &identifier,
+ virtual void add_field_per_backend(const std::string &identifier,
+ const SNode *rep_snode,
bool is_scalar,
DataType dt,
std::vector shape,
int row_num,
int column_num) = 0;
+ virtual void add_ndarray_per_backend(const std::string &identifier,
+ bool is_scalar,
+ DataType dt,
+ std::vector shape,
+ int row_num,
+ int column_num) {
+ TI_NOT_IMPLEMENTED;
+ }
+
virtual void add_per_backend_tmpl(const std::string &identifier,
const std::string &key,
Kernel *kernel) = 0;
+
+ static bool all_fields_are_dense_in_container(const SNode *container);
};
} // namespace lang
diff --git a/taichi/aot/module_data.h b/taichi/aot/module_data.h
new file mode 100644
index 0000000000000..7c401e6148a26
--- /dev/null
+++ b/taichi/aot/module_data.h
@@ -0,0 +1,104 @@
+#pragma once
+
+#include
+#include
+
+#include "taichi/common/core.h"
+#include "taichi/common/serialization.h"
+
+namespace taichi {
+namespace lang {
+namespace aot {
+
+struct CompiledFieldData {
+ std::string field_name;
+ uint32_t dtype{0};
+ std::string dtype_name;
+ size_t mem_offset_in_parent{0};
+ std::vector shape;
+ bool is_scalar{false};
+ std::vector element_shape;
+
+ TI_IO_DEF(field_name,
+ dtype,
+ dtype_name,
+ mem_offset_in_parent,
+ shape,
+ is_scalar,
+ element_shape);
+};
+
+struct CompiledOffloadedTask {
+ std::string type;
+ std::string range_hint;
+ std::string name;
+ // Do we need to inline the source code?
+ std::string source_path;
+ int gpu_block_size{0};
+
+ TI_IO_DEF(type, range_hint, name, source_path, gpu_block_size);
+};
+
+struct ScalarArg {
+ std::string dtype_name;
+ // Unit: byte
+ size_t offset_in_args_buf{0};
+
+ TI_IO_DEF(dtype_name, offset_in_args_buf);
+};
+
+struct ArrayArg {
+ std::string dtype_name;
+ std::size_t field_dim{0};
+ // If |element_shape| is empty, it means this is a scalar
+ std::vector element_shape;
+ // Unit: byte
+ std::size_t shape_offset_in_args_buf{0};
+ // For Vulkan/OpenGL/Metal, this is the binding index
+ int bind_index{0};
+
+ TI_IO_DEF(dtype_name,
+ field_dim,
+ element_shape,
+ shape_offset_in_args_buf,
+ bind_index);
+};
+
+struct CompiledTaichiKernel {
+ std::vector tasks;
+ int args_count{0};
+ int rets_count{0};
+ size_t args_buffer_size{0};
+ size_t rets_buffer_size{0};
+
+ std::unordered_map scalar_args;
+ std::unordered_map arr_args;
+
+ TI_IO_DEF(tasks,
+ args_count,
+ rets_count,
+ args_buffer_size,
+ rets_buffer_size,
+ scalar_args,
+ arr_args);
+};
+
+struct ModuleData {
+ std::unordered_map kernels;
+ std::unordered_map kernel_tmpls;
+ std::vector fields;
+
+ size_t root_buffer_size;
+
+ void dump_json(std::string path) {
+ TextSerializer ts;
+ ts.serialize_to_json("aot_data", *this);
+ ts.write_to_file(path);
+ }
+
+ TI_IO_DEF(kernels, kernel_tmpls, fields, root_buffer_size);
+};
+
+} // namespace aot
+} // namespace lang
+} // namespace taichi
diff --git a/taichi/aot/module_loader.cpp b/taichi/aot/module_loader.cpp
new file mode 100644
index 0000000000000..9d168d3055da3
--- /dev/null
+++ b/taichi/aot/module_loader.cpp
@@ -0,0 +1,28 @@
+#include "taichi/aot/module_loader.h"
+
+#include "taichi/backends/vulkan/aot_module_loader_impl.h"
+#include "taichi/backends/metal/aot_module_loader_impl.h"
+
+namespace taichi {
+namespace lang {
+namespace aot {
+
+std::unique_ptr Module::load(const std::string &path,
+ Arch arch,
+ std::any mod_params) {
+ if (arch == Arch::vulkan) {
+#ifdef TI_WITH_VULKAN
+ return vulkan::make_aot_module(mod_params);
+#endif
+ } else if (arch == Arch::metal) {
+#ifdef TI_WITH_METAL
+ return metal::make_aot_module(mod_params);
+#endif
+ } else {
+ TI_NOT_IMPLEMENTED;
+ }
+}
+
+} // namespace aot
+} // namespace lang
+} // namespace taichi
diff --git a/taichi/aot/module_loader.h b/taichi/aot/module_loader.h
new file mode 100644
index 0000000000000..634f8462e92ec
--- /dev/null
+++ b/taichi/aot/module_loader.h
@@ -0,0 +1,127 @@
+#pragma once
+
+#include
+#include
+#include
+#include
+#include
+
+#include "taichi/aot/module_data.h"
+#include "taichi/backends/device.h"
+#include "taichi/ir/snode.h"
+#include "taichi/aot/module_data.h"
+
+namespace taichi {
+namespace lang {
+
+struct RuntimeContext;
+
+namespace aot {
+
+class TI_DLL_EXPORT Field {
+ public:
+ // Rule of 5 to make MSVC happy
+ Field() = default;
+ virtual ~Field() = default;
+ Field(const Field &) = delete;
+ Field &operator=(const Field &) = delete;
+ Field(Field &&) = default;
+ Field &operator=(Field &&) = default;
+};
+
+class TI_DLL_EXPORT Kernel {
+ public:
+ // Rule of 5 to make MSVC happy
+ Kernel() = default;
+ virtual ~Kernel() = default;
+ Kernel(const Kernel &) = delete;
+ Kernel &operator=(const Kernel &) = delete;
+ Kernel(Kernel &&) = default;
+ Kernel &operator=(Kernel &&) = default;
+
+ /**
+ * @brief Launches the kernel to the device
+ *
+ * This does not manage the device to host synchronization.
+ *
+ * @param ctx Host context
+ */
+ virtual void launch(RuntimeContext *ctx) = 0;
+};
+
+class TI_DLL_EXPORT Module {
+ public:
+ // Rule of 5 to make MSVC happy
+ Module() = default;
+ virtual ~Module() = default;
+ Module(const Module &) = delete;
+ Module &operator=(const Module &) = delete;
+ Module(Module &&) = default;
+ Module &operator=(Module &&) = default;
+
+ static std::unique_ptr load(const std::string &path,
+ Arch arch,
+ std::any mod_params);
+
+ // Module metadata
+ virtual Arch arch() const = 0;
+ virtual uint64_t version() const = 0;
+ virtual std::unique_ptr get_kernel(const std::string &name) = 0;
+ virtual std::unique_ptr get_field(const std::string &name) = 0;
+ virtual size_t get_root_size() const = 0;
+
+ protected:
+ virtual std::unique_ptr make_new_kernel(const std::string &name) = 0;
+
+ private:
+ std::unordered_map> loaded_kernels_;
+};
+
+// Only responsible for reporting device capabilities
+class TargetDevice : public Device {
+ public:
+ TargetDevice(Arch arch) {
+ // TODO: make this configurable
+ set_default_caps(arch);
+ }
+
+ void set_default_caps(Arch arch) {
+ if (arch == Arch::vulkan) {
+ set_cap(DeviceCapability::spirv_version, 0x10300);
+ }
+ }
+
+ DeviceAllocation allocate_memory(const AllocParams ¶ms) override {
+ TI_NOT_IMPLEMENTED;
+ }
+ void dealloc_memory(DeviceAllocation handle) override {
+ TI_NOT_IMPLEMENTED;
+ }
+ std::unique_ptr create_pipeline(
+ const PipelineSourceDesc &src,
+ std::string name = "Pipeline") override {
+ TI_NOT_IMPLEMENTED;
+ }
+ void *map_range(DevicePtr ptr, uint64_t size) override {
+ TI_NOT_IMPLEMENTED;
+ }
+ void *map(DeviceAllocation alloc) override {
+ TI_NOT_IMPLEMENTED;
+ }
+ void unmap(DevicePtr ptr) override {
+ TI_NOT_IMPLEMENTED;
+ }
+ void unmap(DeviceAllocation alloc) override {
+ TI_NOT_IMPLEMENTED;
+ }
+ void memcpy_internal(DevicePtr dst, DevicePtr src, uint64_t size) override {
+ TI_NOT_IMPLEMENTED;
+ }
+ Stream *get_compute_stream() override {
+ TI_NOT_IMPLEMENTED;
+ }
+};
+
+} // namespace aot
+} // namespace lang
+} // namespace taichi
diff --git a/taichi/program/arch.cpp b/taichi/backends/arch.cpp
similarity index 94%
rename from taichi/program/arch.cpp
rename to taichi/backends/arch.cpp
index 53848d8c58751..9de9c8e82f7e3 100644
--- a/taichi/program/arch.cpp
+++ b/taichi/backends/arch.cpp
@@ -1,6 +1,6 @@
-#include "taichi/program/arch.h"
+#include "taichi/backends/arch.h"
-TLANG_NAMESPACE_BEGIN
+namespace taichi {
std::string arch_name(Arch arch) {
switch (arch) {
@@ -79,4 +79,4 @@ int default_simd_width(Arch arch) {
}
}
-TLANG_NAMESPACE_END
+} // namespace taichi
diff --git a/taichi/program/arch.h b/taichi/backends/arch.h
similarity index 92%
rename from taichi/program/arch.h
rename to taichi/backends/arch.h
index b2627b5ff4d22..2d7cffde8950f 100644
--- a/taichi/program/arch.h
+++ b/taichi/backends/arch.h
@@ -4,7 +4,6 @@
#include "taichi/common/core.h"
namespace taichi {
-namespace lang {
enum class Arch : int {
#define PER_ARCH(x) x,
@@ -29,5 +28,4 @@ bool arch_use_host_memory(Arch arch);
int default_simd_width(Arch arch);
-} // namespace lang
} // namespace taichi
diff --git a/taichi/backends/cc/cc_kernel.h b/taichi/backends/cc/cc_kernel.h
index fc322d4375fdd..5ec18eebfeeac 100644
--- a/taichi/backends/cc/cc_kernel.h
+++ b/taichi/backends/cc/cc_kernel.h
@@ -24,7 +24,7 @@ class CCKernel {
}
void compile();
- void launch(Context *ctx);
+ void launch(RuntimeContext *ctx);
std::string get_object() {
return obj_path_;
}
diff --git a/taichi/backends/cc/cc_program.cpp b/taichi/backends/cc/cc_program.cpp
index 9c704c19c9cb7..49d9d1d6e866a 100644
--- a/taichi/backends/cc/cc_program.cpp
+++ b/taichi/backends/cc/cc_program.cpp
@@ -20,7 +20,7 @@ FunctionType CCProgramImpl::compile(Kernel *kernel, OffloadedStmt *) {
auto ker = codegen.compile();
auto ker_ptr = ker.get();
this->add_kernel(std::move(ker));
- return [ker_ptr](Context &ctx) { return ker_ptr->launch(&ctx); };
+ return [ker_ptr](RuntimeContext &ctx) { return ker_ptr->launch(&ctx); };
}
void CCProgramImpl::materialize_runtime(MemoryPool *memory_pool,
@@ -35,14 +35,13 @@ void CCProgramImpl::materialize_runtime(MemoryPool *memory_pool,
void CCProgramImpl::materialize_snode_tree(
SNodeTree *tree,
std::vector