Open ev-br opened 1 year ago
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
.
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[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.
[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?
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 :-).
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.
Let me retry with the latest nightly (am running a few days stale one)
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
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?
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
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
we probably want to pass cpu
as the device, at least in the repo.
This is an exercise to gauge the workarounds needed to compile an iterative algorithm with complex numbers and fancy indexing.