initial commit

Author: duckki

Diff for: Readme.md

@@ -0,0 +1,82 @@
+# Fast Fourier Transform over Finite Field Performance in Mojo and Python
+
+[Fast Fourier Transform (FFT)](https://en.wikipedia.org/wiki/Fast_Fourier_transform) over [Finite Field](https://en.wikipedia.org/wiki/Finite_field) is used in Cryptography.
+The Cooley-Tukey FFT algorithm can be implemented in less than 20 lines of Python. With a slight modification, the domain can be changed to a finite field, instead of complex number.
+
+## Implementations
+
+### Python version
+
+The Python version is in [`python/fft-python.py`](python/fft-python.py). The algorithm looks as following:
+
+```python
+def fft_over_finite_field( P: list[FieldElement], w: FieldElement ) -> list[FieldElement]:
+    n = len(P)
+    if n == 1:
+        return P
+
+    w_square = w * w
+    Pe, Po = P[::2], P[1::2]
+    ye, yo = fft_over_finite_field(Pe, w_square), fft_over_finite_field(Po, w_square)
+    y = [FieldElement(0)] * n
+    w_power = FieldElement(1)
+    for j in range(n // 2):
+        u = ye[j]
+        v = (w_power * yo[j])
+        y[j] = u + v
+        y[j + n // 2] = u - v
+        w_power = w_power * w
+    return y
+```
+
+It depends on the `FieldElement` class, which is defined in the [`python/field.py`](python/field.py). The characteristic prime of the field is hardcoded as `3 * 2**30 + 1` for simplicity.
+
+
+### Mojo version
+
+The Mojo version was implemented using the DynamicVector.
+
+```mojo
+fn fft_over_finite_field( P: DynamicVector[FieldElement], w: FieldElement ) \
+                        -> DynamicVector[FieldElement]:
+    let n = len(P)
+    if n == 1:
+        return P
+
+    let w_square = w * w
+    let Pe = half_vector( P, 0 )
+    let Po = half_vector( P, 1 )
+    let ye = fft_over_finite_field( Pe, w_square )
+    let yo = fft_over_finite_field( Po, w_square )
+    var y = make_vector( n )
+    var w_power = FieldElement(1)
+    for j in range(n // 2):
+        let u = ye[j]
+        let v = (w_power * yo[j])
+        y[j] = u + v
+        y[j + n // 2] = u - v
+        w_power = w_power * w
+    return y
+```
+
+I factored out `make_vector` and `half_vector` functions since Mojo does not yet support list comprehension. The `FieldElement` struct is defined in the [`mojo/field.mojo`](mojo/field.mojo)
+
+
+### Python `galoios` module
+
+I also tried the Python `galois` module that was available as a pip package. It implements the finite field math and uses Numpy as backend. Thus, one could use Numpy's FFT implementation over finite field elements.
+
+
+## Performance
+
+I generated a random array of size 1024 * 256 and ran FFT in each version. It was measured on my M1 MacBook Pro laptop.
+
+| Implementation | Runtime (secs) |
+| --- | --- |
+| Python standalone | 4.3244 |
+| Python `galois` module | 9.5888 |
+| Mojo | 0.0573 |
+
+Mojo's speedup was 75x over the standalone Python implementation.
+
+The `galois` module turned out to be much slower even if it used Numpy's FFT implementation. I guess Numpy couldn't use the optimized C kernel, since it has to interact with `galois`'s `GF` class.

Diff for: mojo/fft-mojo.mojo

@@ -0,0 +1,96 @@
+from random import random_ui64
+from field import FieldElement
+import benchmark
+
+#=============================================================================
+# Helper functions
+
+fn make_vector( size: Int ) -> DynamicVector[FieldElement]:
+    let zero = FieldElement(0)
+    var vec = DynamicVector[FieldElement]()
+    vec.resize( size, zero )
+    return vec
+
+fn generate_random_values( size: Int ) raises -> DynamicVector[FieldElement]:
+    var values = make_vector( size )
+    for i in range(size):
+        let r: Int = int(random_ui64(0, FieldElement.characteristic()-1))
+        values[i] = FieldElement(r)
+    return values
+
+fn to_str( vec: DynamicVector[FieldElement] ) -> String:
+    var s: String = "["
+    let l = len(vec)
+    if l < 6:
+        for i in range(l):
+            if i > 0:
+                s += ", "
+            s += str(vec[i])
+    else:
+        for i in range(3):
+            if i > 0:
+                s += ", "
+            s += str(vec[i])
+        s += " ... "
+        for i in range(3):
+            if i > 0:
+                s += ", "
+            s += str(vec[l-3-1+i])
+    return s + "]"
+
+
+#=============================================================================
+# Implementation of Cooley-Tukey FFT over Finite Field
+
+fn half_vector( vec: DynamicVector[FieldElement], start: Int ) -> DynamicVector[FieldElement]:
+    var result = make_vector( len(vec) // 2 )
+    debug_assert( start + (len(result)-1) * 2 < len(vec), "out of bounds" )
+    for i in range(len(result)):
+        result[i] = vec[start + i*2]
+    return result
+
+fn fft_over_finite_field( P: DynamicVector[FieldElement], w: FieldElement ) \
+                        -> DynamicVector[FieldElement]:
+    let n = len(P)
+    if n == 1:
+        return P
+
+    let w_square = w * w
+    let Pe = half_vector( P, 0 )
+    let Po = half_vector( P, 1 )
+    let ye = fft_over_finite_field( Pe, w_square )
+    let yo = fft_over_finite_field( Po, w_square )
+    var y = make_vector( n )
+    var w_power = FieldElement(1)
+    for j in range(n // 2):
+        let u = ye[j]
+        let v = (w_power * yo[j])
+        y[j] = u + v
+        y[j + n // 2] = u - v
+        w_power = w_power * w
+    return y
+
+
+#=============================================================================
+# Main program
+
+let size: Int = 1024 * 256
+let group_element_power: Int = 2 ** 30 - size
+let g = FieldElement.primitive_element() ** (3 * group_element_power)
+
+fn main() raises:
+    print( "size:", size )
+
+    # print( "generating random values..." )
+    let values = generate_random_values( size )
+    # print( "values:", to_str(values) )
+
+    # let p = fft_over_finite_field( values, g )
+    # print( "p:", to_str(p) )
+
+    @parameter
+    fn bench():
+        _ = fft_over_finite_field( values, g )
+
+    let report = benchmark.run[bench]( 1, 3, 4, 10 )
+    print( "mean runtime:", report.mean("s") )

Diff for: mojo/field.mojo

@@ -0,0 +1,80 @@
+# Ported from starkware101 tutorial's Python version[1].
+# [1]: https://github.com/starkware-industries/stark101
+
+@value
+struct FieldElement(CollectionElement, Stringable):
+
+    # ------------------------------------------------------------------------
+    # Static methods
+
+    @staticmethod
+    fn characteristic() -> Int:
+        return 3 * 2**30 + 1
+
+    @staticmethod
+    fn primitive_element() -> Self:
+        return Self(5)
+
+
+    # ------------------------------------------------------------------------
+    # Fields and basic methods
+
+    var val: Int
+
+    fn __init__(inout self, val: Int):
+        self.val = val % Self.characteristic()
+
+    fn __str__(self) -> String:
+        return String(self.val)
+
+
+    # ------------------------------------------------------------------------
+    # Arithmetic operators
+
+    fn __eq__(self, other: Self) -> Bool:
+        return self.val == other.val
+
+    fn __neg__(self) -> Self:
+        return Self(-self.val)
+
+    fn __add__(self, other: Self) -> Self:
+        return Self(self.val + other.val)
+
+    fn __sub__(self, other: Self) -> Self:
+        return Self(self.val - other.val)
+
+    fn __mul__(self, other: Self) -> Self:
+        return Self(self.val * other.val)
+
+    fn __imul__(inout self, other: Self):
+        self.val = (self.val * other.val) % Self.characteristic()
+
+    fn inverse(self) -> Self:
+        var t: Int = 0
+        var new_t: Int = 1
+        var r = FieldElement.characteristic()
+        var new_r = self.val
+        while new_r != 0:
+            let quotient = r // new_r
+            t, new_t = new_t, (t - (quotient * new_t))
+            r, new_r = new_r, (r - (quotient * new_r))
+        debug_assert( r == 1, "inverse() failed" )
+        return Self(t)
+
+    fn __truediv__(self, other: Self) -> Self:
+        return self * other.inverse()
+
+    fn __pow__(self, _n: Int) -> Self:
+        debug_assert( _n >= 0, "unexpected negative argument" )
+        var cur_pow = self
+        var res = FieldElement(1)
+        var n = _n
+        while n > 0:
+            if n % 2 != 0:
+                res *= cur_pow
+            n = n // 2
+            cur_pow *= cur_pow
+        return res
+
+    fn __pow__(self, n: Self) -> Self:
+        return self.__pow__(n.val)

Diff for: python/fft-python-galois.py

@@ -0,0 +1,20 @@
+import galois
+import numpy as np
+from utility import measure_runtime
+
+# The prime number for the field: 3 * 2^30 + 1 (= 3221225473)
+p = 3 * 2**30 + 1
+
+GF = galois.GF(p, repr="power")
+# print(GF.properties)
+
+size = 1024 * 256
+print( "size:", size )
+
+# print( "generating random values..." )
+values = GF.Random( size )
+# print( "values:", values )
+
+p, runtime = measure_runtime( np.fft.fft, values )
+# print( "p:", p )
+print( "runtime:", runtime )

Diff for: python/fft-python.py

@@ -0,0 +1,43 @@
+from field import FieldElement
+from utility import measure_runtime
+import galois
+
+#=============================================================================
+# Implementation of Cooley-Tukey FFT over Finite Field
+
+def fft_over_finite_field( P: list[FieldElement], w: FieldElement ) -> list[FieldElement]:
+    n = len(P)
+    if n == 1:
+        return P
+
+    w_square = w * w
+    Pe, Po = P[::2], P[1::2]
+    ye, yo = fft_over_finite_field(Pe, w_square), fft_over_finite_field(Po, w_square)
+    y = [FieldElement(0)] * n
+    w_power = FieldElement(1)
+    for j in range(n // 2):
+        u = ye[j]
+        v = (w_power * yo[j])
+        y[j] = u + v
+        y[j + n // 2] = u - v
+        w_power = w_power * w
+    return y
+
+
+#=============================================================================
+# Main program
+
+size = 1024 * 256
+group_element_power = 2 ** 30 - size
+g = FieldElement.generator() ** (3 * group_element_power)
+
+print( "size:", size )
+
+# print( "generating random values..." )
+GF = galois.GF(FieldElement.k_modulus)
+values = [FieldElement(int(x)) for x in GF.Random(size)]
+# print( "values:", values[0:3], "...", values[-3:] )
+
+p, runtime = measure_runtime( fft_over_finite_field, values, g )
+# print( "p:", p[0:3], "...", p[-3:] )
+print( "runtime:", runtime )