Parallelizing scrypt key-derivation function

Question

To review my use of multiprocessing, I don't think it is at all necessary to understand the algorithm, but it's the scrypt key-derivation function.

This uses hashlib.pbkdf2_hmac which was added in Python 3.4. To run it using earlier versions, you can paste in pbkdf2_hmac from here.

I've split the below to multiple code blocks for readability, but it's a single file.

---

Mostly uninteresting stuff, including algorithm internals. Feel free to review if you like, but I have tests that ensure correctness.

(My actual code inlines the first three short functions for speed.)

#!/usr/bin/env python

# Originally based on https://github.com/ricmoo/pyscrypt
# but no line from that survives here and mistakes are likely my own

from hashlib import pbkdf2_hmac
import multiprocessing as mp
import struct
import time

try:
    xrange(1)
except:
    xrange = range

def array_overwrite(source, s_start, dest, d_start, length):
    dest[d_start:d_start + length] = source[s_start:s_start + length]

def blockxor(source, s_start, dest, d_start, length):
    for i in xrange(length):
        dest[d_start + i] ^= source[s_start + i]

def R(X, destination, a1, a2, b):
    a = (X[a1] + X[a2]) & 0xffffffff
    X[destination] ^= ((a << b) | (a >> (32 - b)))

def salsa20_8(B):
    x = B[:16]

    for i in xrange(4):
        R(x, 4, 0,12, 7);R(x, 8, 4, 0, 9);R(x,12, 8, 4,13);R(x, 0,12, 8,18)
        R(x, 9, 5, 1, 7);R(x,13, 9, 5, 9);R(x, 1,13, 9,13);R(x, 5, 1,13,18)
        R(x,14,10, 6, 7);R(x, 2,14,10, 9);R(x, 6, 2,14,13);R(x,10, 6, 2,18)
        R(x, 3,15,11, 7);R(x, 7, 3,15, 9);R(x,11, 7, 3,13);R(x,15,11, 7,18)
        R(x, 1, 0, 3, 7);R(x, 2, 1, 0, 9);R(x, 3, 2, 1,13);R(x, 0, 3, 2,18)
        R(x, 6, 5, 4, 7);R(x, 7, 6, 5, 9);R(x, 4, 7, 6,13);R(x, 5, 4, 7,18)
        R(x,11,10, 9, 7);R(x, 8,11,10, 9);R(x, 9, 8,11,13);R(x,10, 9, 8,18)
        R(x,12,15,14, 7);R(x,13,12,15, 9);R(x,14,13,12,13);R(x,15,14,13,18)

    for i in xrange(16):
        B[i] = (x[i] + B[i]) & 0xffffffff

def blockmix_salsa8(BY, Yi, r):
    start = (2 * r - 1) * 16
    X = BY[start:start+16]

    for i in xrange(2 * r):
        blockxor(BY, i * 16, X, 0, 16)
        salsa20_8(X)
        array_overwrite(X, 0, BY, Yi + (i * 16), 16)

    for i in xrange(r):
        array_overwrite(BY, Yi + (i * 2) * 16, BY, i * 16, 16)
        array_overwrite(BY, Yi + (i*2 + 1) * 16, BY, (i + r) * 16, 16)

def smix(B, Bi, r, N, V, X):
    array_overwrite(B, Bi, X, 0, 32 * r)

    for i in xrange(N):
        array_overwrite(X, 0, V, i * (32 * r), 32 * r)
        blockmix_salsa8(X, 32 * r, r)

    for i in xrange(N):
        j = X[(2 * r - 1) * 16] & (N - 1)
        blockxor(V, j * (32 * r), X, 0, 32 * r)
        blockmix_salsa8(X, 32 * r, r)

    array_overwrite(X, 0, B, Bi, 32 * r)

---

Parallelization. This is the part I'm interested in getting reviewed.

def smix_mp(args):
    B, r, N = args
    B  = list(B)
    XY = [0] * (64 * r)
    V  = [0] * (32 * r * N)
    smix(B, 0, r, N, V, XY)
    return B


scrypt_pool = mp.Pool()
def scrypt_mp(password, salt, N, r, p, olen=64, parallel=True):
    B  = pbkdf2_hmac('sha256', password, salt, 1, p * 128 * r)
    B  = struct.unpack('<%dI' % (len(B) // 4), B)

    if parallel:
        work = scrypt_pool.imap(
            smix_mp,
            [(B[i*32*r:(i+1)*32*r], r, N) for i in xrange(p)]
        )
    else:
        work = map(
            smix_mp,
            [(B[i*32*r:(i+1)*32*r], r, N) for i in xrange(p)]
        )

    B = []
    for i in work:
        B += i

    B = struct.pack('<%dI' % len(B), *B)
    return pbkdf2_hmac('sha256', password, B, 1, olen)

---

Test code.

password, salt, N, r, p = b'pass', b'salt', 2**8, 8, 2
reps = 5

t = time.time()
for i in xrange(reps):
    scrypt_mp(password, salt, N, r, p, parallel=False)
print(time.time()- t)

t = time.time()
for i in xrange(reps):
    scrypt_mp(password, salt, N, r, p, parallel=True)
print(time.time()- t)

---

On my computer with python3 the test code outputs:

5.988052845001221
4.6062071323394775

So only a ~25% speedup.

With pypy (and N=2**10 because it's much faster):

5.42319893837
5.31443500519

Practically no benefit!

Profiling shows that with these parameters more than 90% of the time is spent in salsa20_8, so I would expect it to parallelize better. With larger values of N it does approach 40% on CPython, which is reasonable, but I'm wondering if I could improve the setup costs of multiprocessing somehow. Perhaps using another way to pass the data to/from workers than the two-way list conversions currently there?

Answer 1

I was able to shave some milliseconds off by moving the struct.pack call inside the map calls, which let me get rid of the explicit loop:

def smix_mp(args):
    B, r, N = args
    B  = list(B)
    XY = [0] * (64 * r)
    V  = [0] * (32 * r * N)
    smix(B, 0, r, N, V, XY)
    return struct.pack('<%dI' % len(B), *B)


scrypt_pool = mp.Pool()
def scrypt_mp(password, salt, N, r, p, olen=64, parallel=True):
    B  = pbkdf2_hmac('sha256', password, salt, 1, p * 128 * r)
    B  = struct.unpack('<%dI' % (len(B) // 4), B)

    if parallel:
        work = scrypt_pool.imap(
            smix_mp,
            [(B[i*32*r:(i+1)*32*r], r, N) for i in xrange(p)]
        )
    else:
        work = map(
            smix_mp,
            [(B[i*32*r:(i+1)*32*r], r, N) for i in xrange(p)]
        )

    B = b''.join(work)
    return pbkdf2_hmac('sha256', password, B, 1, olen)

I also found this PyPy issue that leaves me less confident there's any performance still on the table.