Skip to content

Instantly share code, notes, and snippets.

@diogobaeder
Created September 16, 2023 04:08
Show Gist options
  • Save diogobaeder/147554f534ef0076f0b464de5c299d7c to your computer and use it in GitHub Desktop.
Save diogobaeder/147554f534ef0076f0b464de5c299d7c to your computer and use it in GitHub Desktop.
NumPy parallelism
#!/usr/bin/env python3
import sys
import timeit
from concurrent.futures import ThreadPoolExecutor, wait
import numpy as np
N_THREADS = int(sys.argv[1])
ITERATIONS = 100
def work(arr: np.array) -> int:
arr = arr * 2
result = arr.sum()
return result
def main() -> None:
arr = np.arange(1_000_000, dtype=np.int64)
with ThreadPoolExecutor(max_workers=N_THREADS) as executor:
futures = [
executor.submit(work, arr.copy())
for _ in range(N_THREADS)
]
wait(futures)
if __name__ == '__main__':
elapsed = timeit.timeit('main()', number=ITERATIONS, globals=globals())
print(
'Took', elapsed, 'seconds to run with',
N_THREADS, 'threads and', ITERATIONS, 'iterations.')
@diogobaeder
Copy link
Author

Now, there's a new problem, I changed arr = arr * 2 for arr *= 2 which is fine with Python objects, but not with NumPy.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment