taichi 和正常python 速度对比

例子1

import taichi as ti
import time
ti.init(arch=ti.vulkan)

@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

def is_prime_one(n: int):
    result = True
    for k in range(2, int(n ** 0.5) + 1):
        if n % k == 0:
            result = False
            break
    return result

def count_primes_one(n: int) -> int:
    count = 0
    for k in range(2, n):
        if is_prime_one(k):
            count += 1

    return count
start_time=time.time()
print(count_primes(1000000))
print(time.time()-start_time)
start_time=time.time()
print(count_primes_one(1000000))
print(time.time()-start_time)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43

例子二

import numpy as np
import taichi as ti
import time
ti.init(arch=ti.vulkan)


@ti.kernel
def taichi_forward_v0(
        out: ti.types.ndarray(ndim=2),
        w: ti.types.ndarray(ndim=2),
        k: ti.types.ndarray(ndim=2)):
    wr,wl=w.shape
    kr,kl=k.shape
    ti.simt.block.sync()
    for r in range(0,wr):
        for l in range(0,kl):

            for ll in range(0,kr):
                out[r,l]+=w[r,ll]*k[ll,l]



a = np.zeros([1333,1333])
b = np.zeros([1333,1333])
c =np.zeros([1333,1333])
start=time.time()
taichi_forward_v0(a, b, c)
print(time.time()-start)
start=time.time()
print(a,b@c)
print(time.time()-start)

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32

例子3

import numpy as np
import taichi as ti
import time

ti.init(arch=ti.vulkan)



@ti.kernel
def mult():
    a=ti.Vector([1,3,4]*100)
    b=ti.Vector([4,5,6]*100)
    c=a@b
    # print(c)
def mult1():
    a=np.array([1,3,4]*100)
    b=np.array([4,5,6]*100)
    c=a@b
    # print(c)
start=time.time()
for i in range(110000):
    mult()
print(time.time()-start)

start=time.time()
for i in range(110000):
    mult1()
print(time.time()-start)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28