chunk the mandelbrot loop

Author: ev-br

e2e/mandelbrot/mandelbrot.png

@@ -60,10 +60,11 @@ def sq2(a):
 
 
 @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]
+def step(n, c, Z, N, horizon, chunksize):
+    for _ in range(chunksize):
+        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
 
 
@@ -75,8 +76,11 @@ 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 = 20
+    n_chunks = maxiter // chunksize
+
+    for n in range(n_chunks):
+        Z, N = step(n, c, Z, N, horizon, chunksize)
 
     N = np.where(N == maxiter-1, 0, N)     # N[N == maxiter-1] = 0
     return Z, N