WIP: e2e: compile Rougier's mandelbrot implementation #165

ev-br · 2023-07-16T20:52:48Z

This is an exercise to gauge the workarounds needed to compile an iterative algorithm with complex numbers and fancy indexing.

e2e/mandelbrot/mandelbrot.py

ev-br · 2023-07-17T12:54:07Z

The last commit measures elapsed time versus the original NumPy version.

Generally, there is a bit of speedup:

numpy:    elapsed= 0.9471771717071533
compiled: elapsed= 0.30139660835266113   speedup =  3.142627174486552

However, there is a bit of an issue with max_iter.

If we compile the whole mandelbrot_c simulation, with the for n in range(max_iter) loop, dynamo unrolls the loop and it may even fail compiling with an OOM
If we only compile the inner step (which is what this PR does), we hit

[2023-07-17 15:46:41,923] torch._dynamo.convert_frame: [WARNING] torch._dynamo hit config.cache_size_limit (64)

and the peformance drops to worse the original NumPy version.

Also this may serve as a nice demo for transformation of a numpy program needed to make it compile.

lezcano · 2023-07-17T13:22:01Z

[2023-07-17 15:46:41,923] torch._dynamo.convert_frame: [WARNING] torch._dynamo hit config.cache_size_limit (64)

This is coming from some input being dynamic, and dynamo not being able to trace it properly. Are you in the last nightly? I think in the last nightly this shouldn't happen. Otherwise, can you post the graph it traces after the first and after the second iteration?

ev-br · 2023-07-17T14:02:43Z

This does not happen if I put max_iter to be less than 64 :-).

And yes, it compiles max_iter versions of the inner loop :-).

lezcano · 2023-07-17T14:03:40Z

Sure, but the point is that it should be able to trace max_iter symbolically and generate a generic kernels with max_iter as an input.

ev-br · 2023-07-17T14:04:40Z

Let me retry with the latest nightly (am running a few days stale one)

ev-br · 2023-07-17T14:14:04Z

Note that loop variable is not traced:

@torch.compile
def step(....):
    ....

for n in range(max_iter):
    step(....)

and the trace is

$ TORCH_LOGS="graph_code" python mandelbrot.py 

[2023-07-17 17:11:26,956] torch._dynamo.output_graph.__graph_code: [DEBUG] TRACED GRAPH
[2023-07-17 17:11:26,956] torch._dynamo.output_graph.__graph_code: [DEBUG]  ===== __compiled_fn_0 =====
[2023-07-17 17:11:26,956] torch._dynamo.output_graph.__graph_code: [DEBUG]  <eval_with_key>.0 class GraphModule(torch.nn.Module):
[2023-07-17 17:11:26,956] torch._dynamo.output_graph.__graph_code: [DEBUG]     def forward(self, L_c_ : torch.Tensor, L_Z_ : torch.Tensor, L_N_ : torch.Tensor):
[2023-07-17 17:11:26,956] torch._dynamo.output_graph.__graph_code: [DEBUG]         l_c_ = L_c_
[2023-07-17 17:11:26,956] torch._dynamo.output_graph.__graph_code: [DEBUG]         l_z_ = L_Z_
[2023-07-17 17:11:26,956] torch._dynamo.output_graph.__graph_code: [DEBUG]         l_n_ = L_N_
[2023-07-17 17:11:26,956] torch._dynamo.output_graph.__graph_code: [DEBUG]         
[2023-07-17 17:11:26,956] torch._dynamo.output_graph.__graph_code: [DEBUG]         # File: mandelbrot.py:51, code: return a[..., 0]**2 + a[..., 1]**2
[2023-07-17 17:11:26,956] torch._dynamo.output_graph.__graph_code: [DEBUG]         getitem = l_z_[(Ellipsis, 0)]
[2023-07-17 17:11:26,956] torch._dynamo.output_graph.__graph_code: [DEBUG]         pow_1 = getitem ** 2;  getitem = None
[2023-07-17 17:11:26,956] torch._dynamo.output_graph.__graph_code: [DEBUG]         getitem_1 = l_z_[(Ellipsis, 1)]
[2023-07-17 17:11:26,956] torch._dynamo.output_graph.__graph_code: [DEBUG]         pow_2 = getitem_1 ** 2;  getitem_1 = None
[2023-07-17 17:11:26,956] torch._dynamo.output_graph.__graph_code: [DEBUG]         add = pow_1 + pow_2;  pow_1 = pow_2 = None
[2023-07-17 17:11:26,956] torch._dynamo.output_graph.__graph_code: [DEBUG]         
[2023-07-17 17:11:26,956] torch._dynamo.output_graph.__graph_code: [DEBUG]         # File: mandelbrot.py:64, code: I = abs2(Z) < horizon**2
[2023-07-17 17:11:26,956] torch._dynamo.output_graph.__graph_code: [DEBUG]         lt = add < 1048576;  add = None
[2023-07-17 17:11:26,956] torch._dynamo.output_graph.__graph_code: [DEBUG]         
[2023-07-17 17:11:26,956] torch._dynamo.output_graph.__graph_code: [DEBUG]         # File: mandelbrot.py:65, code: N = np.where(I, n, N)                         # N[I] = n
[2023-07-17 17:11:26,956] torch._dynamo.output_graph.__graph_code: [DEBUG]         wrapped_where = torch__dynamo_utils_wrapped_where(lt, 0, l_n_);  l_n_ = None
[2023-07-17 17:11:26,956] torch._dynamo.output_graph.__graph_code: [DEBUG]         
[2023-07-17 17:11:26,956] torch._dynamo.output_graph.__graph_code: [DEBUG]         # File: mandelbrot.py:66, code: Z = np.where(I[..., None], sq2(Z) + c, Z)     # Z[I] = Z[I]**2 + C[I]
[2023-07-17 17:11:26,956] torch._dynamo.output_graph.__graph_code: [DEBUG]         getitem_2 = lt[(Ellipsis, None)];  lt = None
[2023-07-17 17:11:26,956] torch._dynamo.output_graph.__graph_code: [DEBUG]         
[2023-07-17 17:11:26,956] torch._dynamo.output_graph.__graph_code: [DEBUG]         # File: mandelbrot.py:56, code: z = np.empty_like(a)
[2023-07-17 17:11:26,956] torch._dynamo.output_graph.__graph_code: [DEBUG]         wrapped_empty_like = torch__dynamo_utils_wrapped_empty_like(l_z_)
[2023-07-17 17:11:26,956] torch._dynamo.output_graph.__graph_code: [DEBUG]         
[2023-07-17 17:11:26,956] torch._dynamo.output_graph.__graph_code: [DEBUG]         # File: mandelbrot.py:57, code: z[..., 0] = a[..., 0]**2 - a[..., 1]**2
[2023-07-17 17:11:26,956] torch._dynamo.output_graph.__graph_code: [DEBUG]         getitem_3 = l_z_[(Ellipsis, 0)]
[2023-07-17 17:11:26,956] torch._dynamo.output_graph.__graph_code: [DEBUG]         pow_3 = getitem_3 ** 2;  getitem_3 = None
[2023-07-17 17:11:26,956] torch._dynamo.output_graph.__graph_code: [DEBUG]         getitem_4 = l_z_[(Ellipsis, 1)]
[2023-07-17 17:11:26,956] torch._dynamo.output_graph.__graph_code: [DEBUG]         pow_4 = getitem_4 ** 2;  getitem_4 = None
[2023-07-17 17:11:26,956] torch._dynamo.output_graph.__graph_code: [DEBUG]         sub = pow_3 - pow_4;  pow_3 = pow_4 = None
[2023-07-17 17:11:26,956] torch._dynamo.output_graph.__graph_code: [DEBUG]         wrapped___setitem__ = torch__dynamo_utils_wrapped___setitem__(wrapped_empty_like, (Ellipsis, 0), sub);  sub = None
[2023-07-17 17:11:26,956] torch._dynamo.output_graph.__graph_code: [DEBUG]         
[2023-07-17 17:11:26,956] torch._dynamo.output_graph.__graph_code: [DEBUG]         # File: mandelbrot.py:58, code: z[..., 1] = 2 * a[..., 0] * a[..., 1]
[2023-07-17 17:11:26,956] torch._dynamo.output_graph.__graph_code: [DEBUG]         getitem_5 = l_z_[(Ellipsis, 0)]
[2023-07-17 17:11:26,956] torch._dynamo.output_graph.__graph_code: [DEBUG]         mul = 2 * getitem_5;  getitem_5 = None
[2023-07-17 17:11:26,956] torch._dynamo.output_graph.__graph_code: [DEBUG]         getitem_6 = l_z_[(Ellipsis, 1)]
[2023-07-17 17:11:26,956] torch._dynamo.output_graph.__graph_code: [DEBUG]         mul_1 = mul * getitem_6;  mul = getitem_6 = None
[2023-07-17 17:11:26,956] torch._dynamo.output_graph.__graph_code: [DEBUG]         wrapped___setitem___1 = torch__dynamo_utils_wrapped___setitem___1(wrapped_empty_like, (Ellipsis, 1), mul_1);  mul_1 = None
[2023-07-17 17:11:26,956] torch._dynamo.output_graph.__graph_code: [DEBUG]         
[2023-07-17 17:11:26,956] torch._dynamo.output_graph.__graph_code: [DEBUG]         # File: mandelbrot.py:66, code: Z = np.where(I[..., None], sq2(Z) + c, Z)     # Z[I] = Z[I]**2 + C[I]
[2023-07-17 17:11:26,956] torch._dynamo.output_graph.__graph_code: [DEBUG]         add_1 = wrapped_empty_like + l_c_;  wrapped_empty_like = l_c_ = None
[2023-07-17 17:11:26,956] torch._dynamo.output_graph.__graph_code: [DEBUG]         wrapped_where_1 = torch__dynamo_utils_wrapped_where_1(getitem_2, add_1, l_z_);  getitem_2 = add_1 = l_z_ = None
[2023-07-17 17:11:26,956] torch._dynamo.output_graph.__graph_code: [DEBUG]         return (wrapped_where_1, wrapped_where)
[2023-07-17 17:11:26,956] torch._dynamo.output_graph.__graph_code: [DEBUG]         
[2023-07-17 17:11:26,956] torch._dynamo.output_graph.__graph_code: [DEBUG] 
[2023-07-17 17:11:32,564] torch._dynamo.output_graph.__graph_code: [DEBUG] TRACED GRAPH
[2023-07-17 17:11:32,564] torch._dynamo.output_graph.__graph_code: [DEBUG]  ===== __compiled_fn_1 =====
[2023-07-17 17:11:32,564] torch._dynamo.output_graph.__graph_code: [DEBUG]  <eval_with_key>.115 class GraphModule(torch.nn.Module):
[2023-07-17 17:11:32,564] torch._dynamo.output_graph.__graph_code: [DEBUG]     def forward(self, L_c_ : torch.Tensor, L_Z_ : torch.Tensor, L_N_ : torch.Tensor):
[2023-07-17 17:11:32,564] torch._dynamo.output_graph.__graph_code: [DEBUG]         l_c_ = L_c_
[2023-07-17 17:11:32,564] torch._dynamo.output_graph.__graph_code: [DEBUG]         l_z_ = L_Z_
[2023-07-17 17:11:32,564] torch._dynamo.output_graph.__graph_code: [DEBUG]         l_n_ = L_N_
[2023-07-17 17:11:32,564] torch._dynamo.output_graph.__graph_code: [DEBUG]         
[2023-07-17 17:11:32,564] torch._dynamo.output_graph.__graph_code: [DEBUG]         # File: mandelbrot.py:51, code: return a[..., 0]**2 + a[..., 1]**2
[2023-07-17 17:11:32,564] torch._dynamo.output_graph.__graph_code: [DEBUG]         getitem = l_z_[(Ellipsis, 0)]
[2023-07-17 17:11:32,564] torch._dynamo.output_graph.__graph_code: [DEBUG]         pow_1 = getitem ** 2;  getitem = None
[2023-07-17 17:11:32,564] torch._dynamo.output_graph.__graph_code: [DEBUG]         getitem_1 = l_z_[(Ellipsis, 1)]
[2023-07-17 17:11:32,564] torch._dynamo.output_graph.__graph_code: [DEBUG]         pow_2 = getitem_1 ** 2;  getitem_1 = None
[2023-07-17 17:11:32,564] torch._dynamo.output_graph.__graph_code: [DEBUG]         add = pow_1 + pow_2;  pow_1 = pow_2 = None
[2023-07-17 17:11:32,564] torch._dynamo.output_graph.__graph_code: [DEBUG]         
[2023-07-17 17:11:32,564] torch._dynamo.output_graph.__graph_code: [DEBUG]         # File: mandelbrot.py:64, code: I = abs2(Z) < horizon**2
[2023-07-17 17:11:32,564] torch._dynamo.output_graph.__graph_code: [DEBUG]         lt = add < 1048576;  add = None
[2023-07-17 17:11:32,564] torch._dynamo.output_graph.__graph_code: [DEBUG]         
[2023-07-17 17:11:32,564] torch._dynamo.output_graph.__graph_code: [DEBUG]         # File: mandelbrot.py:65, code: N = np.where(I, n, N)                         # N[I] = n
[2023-07-17 17:11:32,564] torch._dynamo.output_graph.__graph_code: [DEBUG]         wrapped_where = torch__dynamo_utils_wrapped_where(lt, 1, l_n_);  l_n_ = None
[2023-07-17 17:11:32,564] torch._dynamo.output_graph.__graph_code: [DEBUG]         
[2023-07-17 17:11:32,564] torch._dynamo.output_graph.__graph_code: [DEBUG]         # File: mandelbrot.py:66, code: Z = np.where(I[..., None], sq2(Z) + c, Z)     # Z[I] = Z[I]**2 + C[I]
[2023-07-17 17:11:32,564] torch._dynamo.output_graph.__graph_code: [DEBUG]         getitem_2 = lt[(Ellipsis, None)];  lt = None
[2023-07-17 17:11:32,564] torch._dynamo.output_graph.__graph_code: [DEBUG]         
[2023-07-17 17:11:32,564] torch._dynamo.output_graph.__graph_code: [DEBUG]         # File: mandelbrot.py:56, code: z = np.empty_like(a)
[2023-07-17 17:11:32,564] torch._dynamo.output_graph.__graph_code: [DEBUG]         wrapped_empty_like = torch__dynamo_utils_wrapped_empty_like(l_z_)
[2023-07-17 17:11:32,564] torch._dynamo.output_graph.__graph_code: [DEBUG]         
[2023-07-17 17:11:32,564] torch._dynamo.output_graph.__graph_code: [DEBUG]         # File: mandelbrot.py:57, code: z[..., 0] = a[..., 0]**2 - a[..., 1]**2
[2023-07-17 17:11:32,564] torch._dynamo.output_graph.__graph_code: [DEBUG]         getitem_3 = l_z_[(Ellipsis, 0)]
[2023-07-17 17:11:32,564] torch._dynamo.output_graph.__graph_code: [DEBUG]         pow_3 = getitem_3 ** 2;  getitem_3 = None
[2023-07-17 17:11:32,564] torch._dynamo.output_graph.__graph_code: [DEBUG]         getitem_4 = l_z_[(Ellipsis, 1)]
[2023-07-17 17:11:32,564] torch._dynamo.output_graph.__graph_code: [DEBUG]         pow_4 = getitem_4 ** 2;  getitem_4 = None
[2023-07-17 17:11:32,564] torch._dynamo.output_graph.__graph_code: [DEBUG]         sub = pow_3 - pow_4;  pow_3 = pow_4 = None
[2023-07-17 17:11:32,564] torch._dynamo.output_graph.__graph_code: [DEBUG]         wrapped___setitem__ = torch__dynamo_utils_wrapped___setitem__(wrapped_empty_like, (Ellipsis, 0), sub);  sub = None
[2023-07-17 17:11:32,564] torch._dynamo.output_graph.__graph_code: [DEBUG]         
[2023-07-17 17:11:32,564] torch._dynamo.output_graph.__graph_code: [DEBUG]         # File: mandelbrot.py:58, code: z[..., 1] = 2 * a[..., 0] * a[..., 1]
[2023-07-17 17:11:32,564] torch._dynamo.output_graph.__graph_code: [DEBUG]         getitem_5 = l_z_[(Ellipsis, 0)]
[2023-07-17 17:11:32,564] torch._dynamo.output_graph.__graph_code: [DEBUG]         mul = 2 * getitem_5;  getitem_5 = None
[2023-07-17 17:11:32,564] torch._dynamo.output_graph.__graph_code: [DEBUG]         getitem_6 = l_z_[(Ellipsis, 1)]
[2023-07-17 17:11:32,564] torch._dynamo.output_graph.__graph_code: [DEBUG]         mul_1 = mul * getitem_6;  mul = getitem_6 = None
[2023-07-17 17:11:32,564] torch._dynamo.output_graph.__graph_code: [DEBUG]         wrapped___setitem___1 = torch__dynamo_utils_wrapped___setitem___1(wrapped_empty_like, (Ellipsis, 1), mul_1);  mul_1 = None
[2023-07-17 17:11:32,564] torch._dynamo.output_graph.__graph_code: [DEBUG]         
[2023-07-17 17:11:32,564] torch._dynamo.output_graph.__graph_code: [DEBUG]         # File: mandelbrot.py:66, code: Z = np.where(I[..., None], sq2(Z) + c, Z)     # Z[I] = Z[I]**2 + C[I]
[2023-07-17 17:11:32,564] torch._dynamo.output_graph.__graph_code: [DEBUG]         add_1 = wrapped_empty_like + l_c_;  wrapped_empty_like = l_c_ = None
[2023-07-17 17:11:32,564] torch._dynamo.output_graph.__graph_code: [DEBUG]         wrapped_where_1 = torch__dynamo_utils_wrapped_where_1(getitem_2, add_1, l_z_);  getitem_2 = add_1 = l_z_ = None
[2023-07-17 17:11:32,564] torch._dynamo.output_graph.__graph_code: [DEBUG]         return (wrapped_where_1, wrapped_where)

^C to stop the output with compiled_fn_2, compiled_fn_3 etc

this is with the latest nightly:

$ python -c'import torch; print(torch.__version__)'
2.1.0.dev20230717+cpu

lezcano · 2023-07-17T14:47:05Z

Can you run the same experiment manually changing the calls to np.where to torch.where and passing PyTorch tensors, see if the same thing happens?

ev-br · 2023-07-17T14:49:50Z

Meanwhile, the last commit trivially chunks the loop to compile a chunk of 20 steps:

numpy:    elapsed= 3.1579246520996094
compiled: elapsed= 0.6164162158966064   speedup =  5.123039548053192

ev-br · 2023-07-17T14:53:09Z

Can you run the same experiment manually changing the calls to np.where to torch.where and passing PyTorch tensors, see if the same thing happens?

Yes, same:

[2023-07-17 17:51:45,295] torch._dynamo.output_graph.__graph_code: [DEBUG] TRACED GRAPH
[2023-07-17 17:51:45,295] torch._dynamo.output_graph.__graph_code: [DEBUG]  ===== __compiled_fn_0 =====
[2023-07-17 17:51:45,295] torch._dynamo.output_graph.__graph_code: [DEBUG]  <eval_with_key>.0 class GraphModule(torch.nn.Module):
[2023-07-17 17:51:45,295] torch._dynamo.output_graph.__graph_code: [DEBUG]     def forward(self, L_c_ : torch.Tensor, L_Z_ : torch.Tensor, L_N_ : torch.Tensor):
[2023-07-17 17:51:45,295] torch._dynamo.output_graph.__graph_code: [DEBUG]         l_c_ = L_c_
[2023-07-17 17:51:45,295] torch._dynamo.output_graph.__graph_code: [DEBUG]         l_z_ = L_Z_
[2023-07-17 17:51:45,295] torch._dynamo.output_graph.__graph_code: [DEBUG]         l_n_ = L_N_
[2023-07-17 17:51:45,295] torch._dynamo.output_graph.__graph_code: [DEBUG]         
[2023-07-17 17:51:45,295] torch._dynamo.output_graph.__graph_code: [DEBUG]         # File: mandelbrot.py:65, code: N = torch.as_tensor(N)
[2023-07-17 17:51:45,295] torch._dynamo.output_graph.__graph_code: [DEBUG]         as_tensor = torch.as_tensor(l_n_);  l_n_ = None
[2023-07-17 17:51:45,295] torch._dynamo.output_graph.__graph_code: [DEBUG]         
[2023-07-17 17:51:45,295] torch._dynamo.output_graph.__graph_code: [DEBUG]         # File: mandelbrot.py:66, code: Z = torch.as_tensor(Z)
[2023-07-17 17:51:45,295] torch._dynamo.output_graph.__graph_code: [DEBUG]         as_tensor_1 = torch.as_tensor(l_z_);  l_z_ = None
[2023-07-17 17:51:45,295] torch._dynamo.output_graph.__graph_code: [DEBUG]         
[2023-07-17 17:51:45,295] torch._dynamo.output_graph.__graph_code: [DEBUG]         # File: mandelbrot.py:51, code: return a[..., 0]**2 + a[..., 1]**2
[2023-07-17 17:51:45,295] torch._dynamo.output_graph.__graph_code: [DEBUG]         getitem = as_tensor_1[(Ellipsis, 0)]
[2023-07-17 17:51:45,295] torch._dynamo.output_graph.__graph_code: [DEBUG]         pow_1 = getitem ** 2;  getitem = None
[2023-07-17 17:51:45,295] torch._dynamo.output_graph.__graph_code: [DEBUG]         getitem_1 = as_tensor_1[(Ellipsis, 1)]
[2023-07-17 17:51:45,295] torch._dynamo.output_graph.__graph_code: [DEBUG]         pow_2 = getitem_1 ** 2;  getitem_1 = None
[2023-07-17 17:51:45,295] torch._dynamo.output_graph.__graph_code: [DEBUG]         add = pow_1 + pow_2;  pow_1 = pow_2 = None
[2023-07-17 17:51:45,295] torch._dynamo.output_graph.__graph_code: [DEBUG]         
[2023-07-17 17:51:45,295] torch._dynamo.output_graph.__graph_code: [DEBUG]         # File: mandelbrot.py:68, code: I = abs2(Z) < horizon**2
[2023-07-17 17:51:45,295] torch._dynamo.output_graph.__graph_code: [DEBUG]         lt = add < 1048576;  add = None
[2023-07-17 17:51:45,295] torch._dynamo.output_graph.__graph_code: [DEBUG]         
[2023-07-17 17:51:45,295] torch._dynamo.output_graph.__graph_code: [DEBUG]         # File: mandelbrot.py:69, code: N = torch.where(I, n, N)                         # N[I] = n
[2023-07-17 17:51:45,295] torch._dynamo.output_graph.__graph_code: [DEBUG]         where = torch.where(lt, 0, as_tensor);  as_tensor = None
[2023-07-17 17:51:45,295] torch._dynamo.output_graph.__graph_code: [DEBUG]         
[2023-07-17 17:51:45,295] torch._dynamo.output_graph.__graph_code: [DEBUG]         # File: mandelbrot.py:70, code: Z = torch.where(I[..., None], sq2(Z) + c, Z)     # Z[I] = Z[I]**2 + C[I]
[2023-07-17 17:51:45,295] torch._dynamo.output_graph.__graph_code: [DEBUG]         getitem_2 = lt[(Ellipsis, None)];  lt = None
[2023-07-17 17:51:45,295] torch._dynamo.output_graph.__graph_code: [DEBUG]         
[2023-07-17 17:51:45,295] torch._dynamo.output_graph.__graph_code: [DEBUG]         # File: mandelbrot.py:56, code: z = torch.empty_like(a)
[2023-07-17 17:51:45,295] torch._dynamo.output_graph.__graph_code: [DEBUG]         empty_like = torch.empty_like(as_tensor_1)
[2023-07-17 17:51:45,295] torch._dynamo.output_graph.__graph_code: [DEBUG]         
[2023-07-17 17:51:45,295] torch._dynamo.output_graph.__graph_code: [DEBUG]         # File: mandelbrot.py:57, code: z[..., 0] = a[..., 0]**2 - a[..., 1]**2
[2023-07-17 17:51:45,295] torch._dynamo.output_graph.__graph_code: [DEBUG]         getitem_3 = as_tensor_1[(Ellipsis, 0)]
[2023-07-17 17:51:45,295] torch._dynamo.output_graph.__graph_code: [DEBUG]         pow_3 = getitem_3 ** 2;  getitem_3 = None
[2023-07-17 17:51:45,295] torch._dynamo.output_graph.__graph_code: [DEBUG]         getitem_4 = as_tensor_1[(Ellipsis, 1)]
[2023-07-17 17:51:45,295] torch._dynamo.output_graph.__graph_code: [DEBUG]         pow_4 = getitem_4 ** 2;  getitem_4 = None
[2023-07-17 17:51:45,295] torch._dynamo.output_graph.__graph_code: [DEBUG]         sub = pow_3 - pow_4;  pow_3 = pow_4 = None
[2023-07-17 17:51:45,295] torch._dynamo.output_graph.__graph_code: [DEBUG]         empty_like[(Ellipsis, 0)] = sub;  setitem = empty_like;  sub = None
[2023-07-17 17:51:45,295] torch._dynamo.output_graph.__graph_code: [DEBUG]         
[2023-07-17 17:51:45,295] torch._dynamo.output_graph.__graph_code: [DEBUG]         # File: mandelbrot.py:58, code: z[..., 1] = 2 * a[..., 0] * a[..., 1]
[2023-07-17 17:51:45,295] torch._dynamo.output_graph.__graph_code: [DEBUG]         getitem_5 = as_tensor_1[(Ellipsis, 0)]
[2023-07-17 17:51:45,295] torch._dynamo.output_graph.__graph_code: [DEBUG]         mul = 2 * getitem_5;  getitem_5 = None
[2023-07-17 17:51:45,295] torch._dynamo.output_graph.__graph_code: [DEBUG]         getitem_6 = as_tensor_1[(Ellipsis, 1)]
[2023-07-17 17:51:45,295] torch._dynamo.output_graph.__graph_code: [DEBUG]         mul_1 = mul * getitem_6;  mul = getitem_6 = None
[2023-07-17 17:51:45,295] torch._dynamo.output_graph.__graph_code: [DEBUG]         empty_like[(Ellipsis, 1)] = mul_1;  setitem_1 = empty_like;  mul_1 = None
[2023-07-17 17:51:45,295] torch._dynamo.output_graph.__graph_code: [DEBUG]         
[2023-07-17 17:51:45,295] torch._dynamo.output_graph.__graph_code: [DEBUG]         # File: mandelbrot.py:70, code: Z = torch.where(I[..., None], sq2(Z) + c, Z)     # Z[I] = Z[I]**2 + C[I]
[2023-07-17 17:51:45,295] torch._dynamo.output_graph.__graph_code: [DEBUG]         add_1 = empty_like + l_c_;  empty_like = l_c_ = None
[2023-07-17 17:51:45,295] torch._dynamo.output_graph.__graph_code: [DEBUG]         where_1 = torch.where(getitem_2, add_1, as_tensor_1);  getitem_2 = add_1 = as_tensor_1 = None
[2023-07-17 17:51:45,295] torch._dynamo.output_graph.__graph_code: [DEBUG]         return (where_1, where)
[2023-07-17 17:51:45,295] torch._dynamo.output_graph.__graph_code: [DEBUG]         
[2023-07-17 17:51:45,295] torch._dynamo.output_graph.__graph_code: [DEBUG] 
[2023-07-17 17:51:51,077] torch._dynamo.output_graph.__graph_code: [DEBUG] TRACED GRAPH
[2023-07-17 17:51:51,077] torch._dynamo.output_graph.__graph_code: [DEBUG]  ===== __compiled_fn_1 =====
[2023-07-17 17:51:51,077] torch._dynamo.output_graph.__graph_code: [DEBUG]  <eval_with_key>.115 class GraphModule(torch.nn.Module):
[2023-07-17 17:51:51,077] torch._dynamo.output_graph.__graph_code: [DEBUG]     def forward(self, L_c_ : torch.Tensor, L_Z_ : torch.Tensor, L_N_ : torch.Tensor):
[2023-07-17 17:51:51,077] torch._dynamo.output_graph.__graph_code: [DEBUG]         l_c_ = L_c_
[2023-07-17 17:51:51,077] torch._dynamo.output_graph.__graph_code: [DEBUG]         l_z_ = L_Z_
[2023-07-17 17:51:51,077] torch._dynamo.output_graph.__graph_code: [DEBUG]         l_n_ = L_N_
[2023-07-17 17:51:51,077] torch._dynamo.output_graph.__graph_code: [DEBUG]         
[2023-07-17 17:51:51,077] torch._dynamo.output_graph.__graph_code: [DEBUG]         # File: mandelbrot.py:65, code: N = torch.as_tensor(N)
[2023-07-17 17:51:51,077] torch._dynamo.output_graph.__graph_code: [DEBUG]         as_tensor = torch.as_tensor(l_n_);  l_n_ = None
[2023-07-17 17:51:51,077] torch._dynamo.output_graph.__graph_code: [DEBUG]         
[2023-07-17 17:51:51,077] torch._dynamo.output_graph.__graph_code: [DEBUG]         # File: mandelbrot.py:66, code: Z = torch.as_tensor(Z)
[2023-07-17 17:51:51,077] torch._dynamo.output_graph.__graph_code: [DEBUG]         as_tensor_1 = torch.as_tensor(l_z_);  l_z_ = None
[2023-07-17 17:51:51,077] torch._dynamo.output_graph.__graph_code: [DEBUG]         
[2023-07-17 17:51:51,077] torch._dynamo.output_graph.__graph_code: [DEBUG]         # File: mandelbrot.py:51, code: return a[..., 0]**2 + a[..., 1]**2
[2023-07-17 17:51:51,077] torch._dynamo.output_graph.__graph_code: [DEBUG]         getitem = as_tensor_1[(Ellipsis, 0)]
[2023-07-17 17:51:51,077] torch._dynamo.output_graph.__graph_code: [DEBUG]         pow_1 = getitem ** 2;  getitem = None
[2023-07-17 17:51:51,077] torch._dynamo.output_graph.__graph_code: [DEBUG]         getitem_1 = as_tensor_1[(Ellipsis, 1)]
[2023-07-17 17:51:51,077] torch._dynamo.output_graph.__graph_code: [DEBUG]         pow_2 = getitem_1 ** 2;  getitem_1 = None
[2023-07-17 17:51:51,077] torch._dynamo.output_graph.__graph_code: [DEBUG]         add = pow_1 + pow_2;  pow_1 = pow_2 = None
[2023-07-17 17:51:51,077] torch._dynamo.output_graph.__graph_code: [DEBUG]         
[2023-07-17 17:51:51,077] torch._dynamo.output_graph.__graph_code: [DEBUG]         # File: mandelbrot.py:68, code: I = abs2(Z) < horizon**2
[2023-07-17 17:51:51,077] torch._dynamo.output_graph.__graph_code: [DEBUG]         lt = add < 1048576;  add = None
[2023-07-17 17:51:51,077] torch._dynamo.output_graph.__graph_code: [DEBUG]         
[2023-07-17 17:51:51,077] torch._dynamo.output_graph.__graph_code: [DEBUG]         # File: mandelbrot.py:69, code: N = torch.where(I, n, N)                         # N[I] = n
[2023-07-17 17:51:51,077] torch._dynamo.output_graph.__graph_code: [DEBUG]         where = torch.where(lt, 1, as_tensor);  as_tensor = None
[2023-07-17 17:51:51,077] torch._dynamo.output_graph.__graph_code: [DEBUG]         
[2023-07-17 17:51:51,077] torch._dynamo.output_graph.__graph_code: [DEBUG]         # File: mandelbrot.py:70, code: Z = torch.where(I[..., None], sq2(Z) + c, Z)     # Z[I] = Z[I]**2 + C[I]
[2023-07-17 17:51:51,077] torch._dynamo.output_graph.__graph_code: [DEBUG]         getitem_2 = lt[(Ellipsis, None)];  lt = None
[2023-07-17 17:51:51,077] torch._dynamo.output_graph.__graph_code: [DEBUG]         
[2023-07-17 17:51:51,077] torch._dynamo.output_graph.__graph_code: [DEBUG]         # File: mandelbrot.py:56, code: z = torch.empty_like(a)
[2023-07-17 17:51:51,077] torch._dynamo.output_graph.__graph_code: [DEBUG]         empty_like = torch.empty_like(as_tensor_1)
[2023-07-17 17:51:51,077] torch._dynamo.output_graph.__graph_code: [DEBUG]         
[2023-07-17 17:51:51,077] torch._dynamo.output_graph.__graph_code: [DEBUG]         # File: mandelbrot.py:57, code: z[..., 0] = a[..., 0]**2 - a[..., 1]**2
[2023-07-17 17:51:51,077] torch._dynamo.output_graph.__graph_code: [DEBUG]         getitem_3 = as_tensor_1[(Ellipsis, 0)]
[2023-07-17 17:51:51,077] torch._dynamo.output_graph.__graph_code: [DEBUG]         pow_3 = getitem_3 ** 2;  getitem_3 = None
[2023-07-17 17:51:51,077] torch._dynamo.output_graph.__graph_code: [DEBUG]         getitem_4 = as_tensor_1[(Ellipsis, 1)]
[2023-07-17 17:51:51,077] torch._dynamo.output_graph.__graph_code: [DEBUG]         pow_4 = getitem_4 ** 2;  getitem_4 = None
[2023-07-17 17:51:51,077] torch._dynamo.output_graph.__graph_code: [DEBUG]         sub = pow_3 - pow_4;  pow_3 = pow_4 = None
[2023-07-17 17:51:51,077] torch._dynamo.output_graph.__graph_code: [DEBUG]         empty_like[(Ellipsis, 0)] = sub;  setitem = empty_like;  sub = None
[2023-07-17 17:51:51,077] torch._dynamo.output_graph.__graph_code: [DEBUG]         
[2023-07-17 17:51:51,077] torch._dynamo.output_graph.__graph_code: [DEBUG]         # File: mandelbrot.py:58, code: z[..., 1] = 2 * a[..., 0] * a[..., 1]
[2023-07-17 17:51:51,077] torch._dynamo.output_graph.__graph_code: [DEBUG]         getitem_5 = as_tensor_1[(Ellipsis, 0)]
[2023-07-17 17:51:51,077] torch._dynamo.output_graph.__graph_code: [DEBUG]         mul = 2 * getitem_5;  getitem_5 = None
[2023-07-17 17:51:51,077] torch._dynamo.output_graph.__graph_code: [DEBUG]         getitem_6 = as_tensor_1[(Ellipsis, 1)]
[2023-07-17 17:51:51,077] torch._dynamo.output_graph.__graph_code: [DEBUG]         mul_1 = mul * getitem_6;  mul = getitem_6 = None
[2023-07-17 17:51:51,077] torch._dynamo.output_graph.__graph_code: [DEBUG]         empty_like[(Ellipsis, 1)] = mul_1;  setitem_1 = empty_like;  mul_1 = None
[2023-07-17 17:51:51,077] torch._dynamo.output_graph.__graph_code: [DEBUG]         
[2023-07-17 17:51:51,077] torch._dynamo.output_graph.__graph_code: [DEBUG]         # File: mandelbrot.py:70, code: Z = torch.where(I[..., None], sq2(Z) + c, Z)     # Z[I] = Z[I]**2 + C[I]
[2023-07-17 17:51:51,077] torch._dynamo.output_graph.__graph_code: [DEBUG]         add_1 = empty_like + l_c_;  empty_like = l_c_ = None
[2023-07-17 17:51:51,077] torch._dynamo.output_graph.__graph_code: [DEBUG]         where_1 = torch.where(getitem_2, add_1, as_tensor_1);  getitem_2 = add_1 = as_tensor_1 = None
[2023-07-17 17:51:51,077] torch._dynamo.output_graph.__graph_code: [DEBUG]         return (where_1, where)
[2023-07-17 17:51:51,077] torch._dynamo.output_graph.__graph_code: [DEBUG]         
[2023-07-17 17:51:51,077] torch._dynamo.output_graph.__graph_code: [DEBUG] 
[2023-07-17 17:51:51,280] torch._dynamo.output_graph.__graph_code: [DEBUG] TRACED GRAPH
[2023-07-17 17:51:51,280] torch._dynamo.output_graph.__graph_code: [DEBUG]  ===== __compiled_fn_2 =====
[2023-07-17 17:51:51,280] torch._dynamo.output_graph.__graph_code: [DEBUG]  <eval_with_key>.121 class GraphModule(torch.nn.Module):
[2023-07-17 17:51:51,280] torch._dynamo.output_graph.__graph_code: [DEBUG]     def forward(self, L_c_ : torch.Tensor, L_Z_ : torch.Tensor, L_N_ : torch.Tensor):
[2023-07-17 17:51:51,280] torch._dynamo.output_graph.__graph_code: [DEBUG]         l_c_ = L_c_
[2023-07-17 17:51:51,280] torch._dynamo.output_graph.__graph_code: [DEBUG]         l_z_ = L_Z_
[2023-07-17 17:51:51,280] torch._dynamo.output_graph.__graph_code: [DEBUG]         l_n_ = L_N_
[2023-07-17 17:51:51,280] torch._dynamo.output_graph.__graph_code: [DEBUG]         
[2023-07-17 17:51:51,280] torch._dynamo.output_graph.__graph_code: [DEBUG]         # File: mandelbrot.py:65, code: N = torch.as_tensor(N)
[2023-07-17 17:51:51,280] torch._dynamo.output_graph.__graph_code: [DEBUG]         as_tensor = torch.as_tensor(l_n_);  l_n_ = None
[2023-07-17 17:51:51,280] torch._dynamo.output_graph.__graph_code: [DEBUG]         
[2023-07-17 17:51:51,280] torch._dynamo.output_graph.__graph_code: [DEBUG]         # File: mandelbrot.py:66, code: Z = torch.as_tensor(Z)
[2023-07-17 17:51:51,280] torch._dynamo.output_graph.__graph_code: [DEBUG]         as_tensor_1 = torch.as_tensor(l_z_);  l_z_ = None
[2023-07-17 17:51:51,280] torch._dynamo.output_graph.__graph_code: [DEBUG]         
[2023-07-17 17:51:51,280] torch._dynamo.output_graph.__graph_code: [DEBUG]         # File: mandelbrot.py:51, code: return a[..., 0]**2 + a[..., 1]**2
[2023-07-17 17:51:51,280] torch._dynamo.output_graph.__graph_code: [DEBUG]         getitem = as_tensor_1[(Ellipsis, 0)]
[2023-07-17 17:51:51,280] torch._dynamo.output_graph.__graph_code: [DEBUG]         pow_1 = getitem ** 2;  getitem = None
[2023-07-17 17:51:51,280] torch._dynamo.output_graph.__graph_code: [DEBUG]         getitem_1 = as_tensor_1[(Ellipsis, 1)]
[2023-07-17 17:51:51,280] torch._dynamo.output_graph.__graph_code: [DEBUG]         pow_2 = getitem_1 ** 2;  getitem_1 = None
[2023-07-17 17:51:51,280] torch._dynamo.output_graph.__graph_code: [DEBUG]         add = pow_1 + pow_2;  pow_1 = pow_2 = None
[2023-07-17 17:51:51,280] torch._dynamo.output_graph.__graph_code: [DEBUG]         
[2023-07-17 17:51:51,280] torch._dynamo.output_graph.__graph_code: [DEBUG]         # File: mandelbrot.py:68, code: I = abs2(Z) < horizon**2
[2023-07-17 17:51:51,280] torch._dynamo.output_graph.__graph_code: [DEBUG]         lt = add < 1048576;  add = None
[2023-07-17 17:51:51,280] torch._dynamo.output_graph.__graph_code: [DEBUG]         
[2023-07-17 17:51:51,280] torch._dynamo.output_graph.__graph_code: [DEBUG]         # File: mandelbrot.py:69, code: N = torch.where(I, n, N)                         # N[I] = n
[2023-07-17 17:51:51,280] torch._dynamo.output_graph.__graph_code: [DEBUG]         where = torch.where(lt, 2, as_tensor);  as_tensor = None
[2023-07-17 17:51:51,280] torch._dynamo.output_graph.__graph_code: [DEBUG]         
[2023-07-17 17:51:51,280] torch._dynamo.output_graph.__graph_code: [DEBUG]         # File: mandelbrot.py:70, code: Z = torch.where(I[..., None], sq2(Z) + c, Z)     # Z[I] = Z[I]**2 + C[I]
[2023-07-17 17:51:51,280] torch._dynamo.output_graph.__graph_code: [DEBUG]         getitem_2 = lt[(Ellipsis, None)];  lt = None
[2023-07-17 17:51:51,280] torch._dynamo.output_graph.__graph_code: [DEBUG]         
[2023-07-17 17:51:51,280] torch._dynamo.output_graph.__graph_code: [DEBUG]         # File: mandelbrot.py:56, code: z = torch.empty_like(a)
[2023-07-17 17:51:51,280] torch._dynamo.output_graph.__graph_code: [DEBUG]         empty_like = torch.empty_like(as_tensor_1)
[2023-07-17 17:51:51,280] torch._dynamo.output_graph.__graph_code: [DEBUG]         
[2023-07-17 17:51:51,280] torch._dynamo.output_graph.__graph_code: [DEBUG]         # File: mandelbrot.py:57, code: z[..., 0] = a[..., 0]**2 - a[..., 1]**2
[2023-07-17 17:51:51,280] torch._dynamo.output_graph.__graph_code: [DEBUG]         getitem_3 = as_tensor_1[(Ellipsis, 0)]
[2023-07-17 17:51:51,280] torch._dynamo.output_graph.__graph_code: [DEBUG]         pow_3 = getitem_3 ** 2;  getitem_3 = None
[2023-07-17 17:51:51,280] torch._dynamo.output_graph.__graph_code: [DEBUG]         getitem_4 = as_tensor_1[(Ellipsis, 1)]
[2023-07-17 17:51:51,280] torch._dynamo.output_graph.__graph_code: [DEBUG]         pow_4 = getitem_4 ** 2;  getitem_4 = None
[2023-07-17 17:51:51,280] torch._dynamo.output_graph.__graph_code: [DEBUG]         sub = pow_3 - pow_4;  pow_3 = pow_4 = None
[2023-07-17 17:51:51,280] torch._dynamo.output_graph.__graph_code: [DEBUG]         empty_like[(Ellipsis, 0)] = sub;  setitem = empty_like;  sub = None
[2023-07-17 17:51:51,280] torch._dynamo.output_graph.__graph_code: [DEBUG]         
[2023-07-17 17:51:51,280] torch._dynamo.output_graph.__graph_code: [DEBUG]         # File: mandelbrot.py:58, code: z[..., 1] = 2 * a[..., 0] * a[..., 1]
[2023-07-17 17:51:51,280] torch._dynamo.output_graph.__graph_code: [DEBUG]         getitem_5 = as_tensor_1[(Ellipsis, 0)]
[2023-07-17 17:51:51,280] torch._dynamo.output_graph.__graph_code: [DEBUG]         mul = 2 * getitem_5;  getitem_5 = None
[2023-07-17 17:51:51,280] torch._dynamo.output_graph.__graph_code: [DEBUG]         getitem_6 = as_tensor_1[(Ellipsis, 1)]
[2023-07-17 17:51:51,280] torch._dynamo.output_graph.__graph_code: [DEBUG]         mul_1 = mul * getitem_6;  mul = getitem_6 = None
[2023-07-17 17:51:51,280] torch._dynamo.output_graph.__graph_code: [DEBUG]         empty_like[(Ellipsis, 1)] = mul_1;  setitem_1 = empty_like;  mul_1 = None
[2023-07-17 17:51:51,280] torch._dynamo.output_graph.__graph_code: [DEBUG]         
[2023-07-17 17:51:51,280] torch._dynamo.output_graph.__graph_code: [DEBUG]         # File: mandelbrot.py:70, code: Z = torch.where(I[..., None], sq2(Z) + c, Z)     # Z[I] = Z[I]**2 + C[I]
[2023-07-17 17:51:51,280] torch._dynamo.output_graph.__graph_code: [DEBUG]         add_1 = empty_like + l_c_;  empty_like = l_c_ = None
[2023-07-17 17:51:51,280] torch._dynamo.output_graph.__graph_code: [DEBUG]         where_1 = torch.where(getitem_2, add_1, as_tensor_1);  getitem_2 = add_1 = as_tensor_1 = None
[2023-07-17 17:51:51,280] torch._dynamo.output_graph.__graph_code: [DEBUG]         return (where_1, where)

^C

lezcano · 2023-07-21T20:16:42Z

we probably want to pass cpu as the device, at least in the repo.

e2e/kmeans/kmeans.py

ev-br commented Jul 16, 2023

View reviewed changes

e2e/mandelbrot/mandelbrot.py Outdated Show resolved Hide resolved

ev-br requested a review from lezcano July 16, 2023 20:56

ev-br force-pushed the mandelbrot_compile branch from ed73c4d to dc37920 Compare July 17, 2023 12:17

ev-br mentioned this pull request Jul 17, 2023

e2e : add an inner op of the kmeans algorithm #164

Merged

ev-br added 4 commits July 17, 2023 15:57

e2e/mandelbrot: replace fancy indexing with np.where

34e7c80

e2e: compile the Rougier's mandelbrot implementation

127d25f

e2e: measure speedup

d727eff

DOC: briefly update the readme

781ad00

ev-br force-pushed the mandelbrot_compile branch from 72e8933 to 781ad00 Compare July 17, 2023 13:39

chunk the mandelbrot loop

12be7fb

ev-br force-pushed the mandelbrot_compile branch from bb3e125 to 12be7fb Compare July 17, 2023 17:24

e2e: kmeans: remove a workaround

58df095

ev-br commented Jul 22, 2023

View reviewed changes

e2e/kmeans/kmeans.py Outdated Show resolved Hide resolved

Update e2e/kmeans/kmeans.py

db1710a

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

WIP: e2e: compile Rougier's mandelbrot implementation #165

WIP: e2e: compile Rougier's mandelbrot implementation #165

ev-br commented Jul 16, 2023

ev-br commented Jul 17, 2023 •

edited

Loading

lezcano commented Jul 17, 2023

ev-br commented Jul 17, 2023 •

edited

Loading

lezcano commented Jul 17, 2023

ev-br commented Jul 17, 2023

ev-br commented Jul 17, 2023 •

edited

Loading

lezcano commented Jul 17, 2023

ev-br commented Jul 17, 2023

ev-br commented Jul 17, 2023

lezcano commented Jul 21, 2023

WIP: e2e: compile Rougier's mandelbrot implementation #165

Are you sure you want to change the base?

WIP: e2e: compile Rougier's mandelbrot implementation #165

Conversation

ev-br commented Jul 16, 2023

ev-br commented Jul 17, 2023 • edited Loading

lezcano commented Jul 17, 2023

ev-br commented Jul 17, 2023 • edited Loading

lezcano commented Jul 17, 2023

ev-br commented Jul 17, 2023

ev-br commented Jul 17, 2023 • edited Loading

lezcano commented Jul 17, 2023

ev-br commented Jul 17, 2023

ev-br commented Jul 17, 2023

lezcano commented Jul 21, 2023

ev-br commented Jul 17, 2023 •

edited

Loading

ev-br commented Jul 17, 2023 •

edited

Loading

ev-br commented Jul 17, 2023 •

edited

Loading