Taichi 计算加速

无意间发现 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