MojoでProject Euler 31（3）

この問題はSIMDが使えます。1コアでいくつかのデータに対して同時に処理を行えます。いくつ同時に処理を行えるかは、次のように分かります。

alias nelts = simdwidthof[DType.int32]()

手元では8でした。
まず、領域を確保しておいて、

    var p = DTypePointer[DType.int32].alloc(N+1)

8つ一組のデータを用意します。

    var ones = SIMD[DType.int32, nelts](1)

こうすると、1が8つになります。

    p.store[width=nelts](0, ones)

とすると、0から8つに1が格納されます。

    var e = self.a.load[width=nelts](0)

とすると、0から8つの値が取り出せます。また、8つのデータ同士で演算ができます。

$N = 10^8$ で前のコードより4倍くらい速かったです。全てが並列化の効果ではないですが、たぶん8つ一斉にDの剰余を取るところが支配的だと思います。

from math import max
import sys


#################### constatns ####################

# number of elements
alias nelts = simdwidthof[DType.int32]()
alias D: SIMD[DType.int32, 1] = 10**9+7


#################### Polynomial ####################

struct Polynomial:
    var a: DTypePointer[DType.int32]
    var L: Int
    
    fn __init__(inout self, a: DTypePointer[DType.int32], L: Int):
        self.a = a
        self.L = L
    
    fn __del__(owned self):
        self.a.free()
        pass
    
    fn __str__(self) -> String:
        var s = str(self.a[0])
        for i in range(1, self.L):
            s += " " + str(self.a[i])
        return s
    
    fn value(self, i: Int) -> SIMD[DType.int32, 1]:
        return self.a.load(i)
    
    fn divide(inout self, e: Int):
        fn get_width(w: Int) -> Int:
            var width = nelts
            while w < width:
                width >>= 1
            return width
        
        var width = get_width(e)
        if e >= nelts:
            for i in range(e, self.L-width, width):
                var e1 = self.a.load[width=nelts](i)
                var e2 = self.a.load[width=nelts](i-e)
                self.a.store[width=nelts](i, (e1+e2)%D)
        elif e >= 4:
            for i in range(e, self.L-width, width):
                var e1 = self.a.load[width=4](i)
                var e2 = self.a.load[width=4](i-e)
                self.a.store[width=4](i, (e1+e2)%D)
        elif e >= 2:
            for i in range(e, self.L-width, width):
                var e1 = self.a.load[width=2](i)
                var e2 = self.a.load[width=2](i-e)
                self.a.store[width=2](i, (e1+e2)%D)
        
        var first = max(e, e+(self.L-e-width-1)//width*width+width)
        for j in range(first, self.L):
            var e1 = self.a.load(j)
            var e2 = self.a.load(j-e)
            self.a.store(j, (e1+e2)%D)


#################### process ####################

fn collect_coins(N: Int) -> List[Int]:
    var ns = List[Int](1, 2, 5);
    var ps = List[Int]()
    for e in range(19):
        for n in ns:
            var p = n[] * 10**e
            if p > N:
                break
            ps.append(p)
    return ps

# 1/(1-x)
fn initialize(N: Int) -> Polynomial:
    var L = (N//nelts+1)*nelts
    var ones = SIMD[DType.int32, nelts](1)
    var a = SIMD[DType.int32, nelts](1)
    var p = DTypePointer[DType.int32].alloc(L)
    for i in range(0, N+1, nelts):
        p.store[width=nelts](i, a)
    return Polynomial(p, L)

fn f(N: Int) -> SIMD[DType.int32, 1]:
    var p = initialize(N)
    var ps = collect_coins(N)
    for i in range(1, ps.size):
        p.divide(ps[i])
    
    return p.value(N)

fn main() raises:
    var args = sys.argv()
    var N = atol(args[1])
    print(f(N))