Taichi 计算加速

查看 17|回复 0
作者: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
您需要登录后才可以回帖 登录 | 立即注册

返回顶部