Taichi 是否可以结合Ray或者Dask等分布式计算的框架?
尝试了一下在Ray里套皮貌似寄了~
import ray
import taichi as ti
# ti.init(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(l: int, r: int) -> int:
count = 0
for k in range(l, r):
if is_prime(k):
count += 1
return count
@ray.remote
def count_wrapper(l:int, r:int) -> int:
return count_primes(l, r)
ray.init()
futures = []
for i in range(2):
worker = count_wrapper.remote(100*i, 100*i+100)
futures.append(worker)
res = ray.get(futures)
print(res)