Taichi 计算加速

[wizardyhnr]

无意间发现 Taichi 计算加速还挺快的。官网例子如下。不知道更复杂的案例效果怎么样。

#!/usr/bin/env python
#prime_taichi.py
import taichi as ti
ti.init(arch=ti.gpu)
@ti.func
def is_prime(n: int):
    result = True
    for k in range(2, int(n ** 0.5) + 1):
        if n % k == 0:
            result = False
            break
    return result

@ti.kernel
def count_primes(n: int) -> int:
    count = 0
    for k in range(2, n):
        if is_prime(k):
            count += 1

    return count

#!/usr/bin/env python
#test.py
"""Count the prime numbers in the range [1, n]
"""
from prime_taichi import count_primes
import time

start: int = time.time()
print(count_primes(10000000))
print(time.time()-start)

#Taichi 版本，笔记本 2070S

[Taichi] version 1.5.0, llvm 15.0.4, commit 7b885c28, linux, python 3.10.6
[Taichi] Starting on arch=cuda
664579
0.20201754570007324

#纯 Python 版本

664579
83.8845808506012

[dayeye2006199]

这个和 numba 之类的是否有相似之处？

[wizardyhnr]

都是 JIT 加速，性能好像比 numba 好不少。
file:///dev/shm/performance-comparisons-for-python-libraries-in-parallel-computing-and-physical-simulation-1.pdf
https://docs.taichi-lang.org/blog/taichi-compared-to-cub-cupy-numba

[huoshen]

Taichi 是能加速, 因为运行的时候不是 python 了, 是他编译的静态语言, 但限制挺多的, 有些不支持递归, 有的不支持嵌套, 有的并行运行, 会导致额外的非原子内存操作导致数据不一致. 想额外吐槽一下 Taichi, 这学期上的一个 CG 课, 教授指定用的就是 Taichi, 官网上文档都不全, 然后各种 Machine dependent bug 搞得我心态都快炸了.

[t133]

taichi 跟 jax 哪个快有没有都用过的

[oldshensheep]

怎么不试试神奇的 GraalVM ，用 GraalVM 直接运行 Python 文件：
664579
8.049000024795532
用 CPython 运行：
664579
115.89525604248047
GraalVM 不能使用 GPU 加速，实验在甲骨文的 ARM 服务器上运行，快了十几倍
优点就是不用改代码，不过不能使用 GPU