How Taichi Accelerates Prime Counting by 200× Over Pure Python

Previously I mentioned using the numba library to accelerate Python; here I introduce another acceleration library, taichi.

First, we write a function to count the number of primes within a given range:

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

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

Running the normal (non‑accelerated) version on my computer to count primes between 0 and 10 million takes:

2 minutes 10 seconds.

Using Taichi for acceleration:

import taichi as ti
ti.init()

@ti.func
def is_prime1(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_prime1(n:int) -> int:
    count = 0
    for k in range(2, n):
        if is_prime1(k):
            count += 1
    return count

print(count_prime1(10000000))

Time taken:

1.09 seconds, about a 200× speedup.

Now using numba for acceleration:

from numba import jit
@jit(nopython=True)
def is_prime2(n:int):
    result = True
    for k in range(2, int(n**0.5)+1):
        if n % k == 0:
            result = False
            break
    return result

@jit(nopython=True)
def count_prime2(n:int) -> int:
    count = 0
    for k in range(2, n):
        if is_prime2(k):
            count += 1
    return count

print(count_prime2(10000000))

Time taken:

5 seconds. Thus Taichi is about five times faster than Numba for this problem – impressive!

Reference

https://tech.sina.com.cn/csj/2022-09-09/doc-imqqsmrn8434391.shtml