Skip to content

Commit

Permalink
chunk the mandelbrot loop
Browse files Browse the repository at this point in the history
  • Loading branch information
ev-br committed Jul 17, 2023
1 parent 781ad00 commit 12be7fb
Show file tree
Hide file tree
Showing 2 changed files with 13 additions and 7 deletions.
Binary file modified e2e/mandelbrot/mandelbrot.png
Loading
Sorry, something went wrong. Reload?
Sorry, we cannot display this file.
Sorry, this file is invalid so it cannot be displayed.
20 changes: 13 additions & 7 deletions e2e/mandelbrot/mandelbrot.py
Original file line number Diff line number Diff line change
Expand Up @@ -59,11 +59,13 @@ def sq2(a):
return z


@torch.compile
def step(n, c, Z, N, horizon):
I = abs2(Z) < horizon**2
N = np.where(I, n, N) # N[I] = n
Z = np.where(I[..., None], sq2(Z) + c, Z) # Z[I] = Z[I]**2 + C[I]
@torch.compile(dynamic=True)
def step(n0, c, Z, N, horizon, chunksize):
for j in range(chunksize):
n = n0 + j
I = abs2(Z) < horizon**2
N = np.where(I, n, N) # N[I] = n
Z = np.where(I[..., None], sq2(Z) + c, Z) # Z[I] = Z[I]**2 + C[I]
return Z, N


Expand All @@ -75,8 +77,12 @@ def mandelbrot_c(xmin, xmax, ymin, ymax, xn, yn, horizon=2**10, maxiter=5):
N = np.zeros(c.shape[:-1], dtype='int')
Z = np.zeros_like(c, dtype='float32')

for n in range(maxiter):
Z, N = step(n, c, Z, N, horizon)
chunksize=10
n_chunks = maxiter // chunksize

for i_chunk in range(n_chunks):
n0 = i_chunk*chunksize
Z, N = step(n0, c, Z, N, horizon, chunksize)

N = np.where(N == maxiter-1, 0, N) # N[N == maxiter-1] = 0
return Z, N
Expand Down

0 comments on commit 12be7fb

Please sign in to comment.