search

pytorch / pytorch

Tensors and Dynamic neural networks in Python with strong GPU acceleration
https://pytorch.org
Other
81.61k stars 21.89k forks source link

Compilation Fails with torch.sparse and "fullgraph=True" #130687

Open asglover opened 1 month ago

asglover commented 1 month ago

🐛 Describe the bug

When you have a function that uses torch.sparse, torch.compile can compile it, but not with "fullgraph=True"

Error logs

(e3nn_env) aglover@login36:~/e3nn_benchmarking> python minimal_sparse_compilation_error.py /global/u1/a/aglover/e3nn_benchmarking/minimal_sparse_compilation_error.py:21: UserWarning: Sparse CSR tensor support is in beta state. If you miss a functionality in the sparse tensor support, please submit a feature request to https://github.com/pytorch/pytorch/issues. (Triggered internally at /home/conda/feedstock_root/build_artifacts/libtorch_1718580740865/work/aten/src/ATen/SparseCsrTensorImpl.cpp:53.) z = torch.randn(size=(3,9), dtype=torch.float32, device=device).to_sparse_csr() 0.000072 I0713 17:21:46.222000 140632849619840 torch/_dynamo/logging.py:55] [0/0] Step 1: torchdynamo start tracing inner /global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/external_utils.py:34 I0713 17:21:46.233000 140632849619840 torch/_dynamo/convert_frame.py:538] [0/0] Restarting analysis due to _dynamo/symbolic_convert.py:142 in fail_and_restart_analysis I0713 17:21:46.234000 140632849619840 torch/_dynamo/logging.py:55] [0/0_1] Step 1: torchdynamo start tracing inner /global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/external_utils.py:34 I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822] WON'T CONVERT inner /global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/external_utils.py line 34 I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822] ========== TorchDynamo Stack Trace ========== I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822] Traceback (most recent call last): I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822] File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py", line 786, in _convert_frame I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822] result = inner_convert( I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822] File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py", line 400, in _convert_frame_assert I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822] return _compile( I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822] File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/contextlib.py", line 79, in inner I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822] return func(*args, kwds) I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822] File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py", line 676, in _compile I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822] guarded_code = compile_inner(code, one_graph, hooks, transform) I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822] File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/utils.py", line 262, in time_wrapper I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822] r = func(*args, *kwargs) I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822] File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py", line 535, in compile_inner I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822] out_code = transform_code_object(code, transform) I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822] File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/bytecode_transformation.py", line 1036, in transform_code_object I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822] transformations(instructions, code_options) I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822] File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py", line 165, in _fn I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822] return fn(args, kwargs) I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822] File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py", line 500, in transform I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822] tracer.run() I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822] File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 2149, in run I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822] super().run() I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822] File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 810, in run I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822] and self.step() I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822] File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 773, in step I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822] getattr(self, inst.opname)(inst) I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822] File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 484, in wrapper I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822] return handle_graph_break(self, inst, speculation.reason) I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822] File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 548, in handle_graph_break I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822] self.output.compile_subgraph(self, reason=reason) I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822] File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/output_graph.py", line 961, in compile_subgraph I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822] and all( I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822] File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/output_graph.py", line 962, in I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822] not isinstance( I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822] File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/variables/base.py", line 135, in instancecheck I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822] instance = instance.realize() I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822] File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/variables/lazy.py", line 58, in realize I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822] self._cache.realize(self.parents_tracker) I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822] File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/variables/lazy.py", line 24, in realize I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822] self.vt = VariableBuilder(tx, self.source)(self.value) I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822] File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/variables/builder.py", line 269, in call I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822] vt = self._wrap(value) I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822] File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/variables/builder.py", line 526, in _wrap I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822] keywords[k] = VariableBuilder( I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822] File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/variables/builder.py", line 269, in call I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822] vt = self._wrap(value) I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822] File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/variables/builder.py", line 402, in _wrap I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822] return type_dispatch(self, value) I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822] File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/variables/builder.py", line 1069, in wrap_tensor I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822] unimplemented( I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822] File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/exc.py", line 190, in unimplemented I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822] raise Unsupported(msg) I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822] torch._dynamo.exc.Unsupported: torch.compile does not support sparse Tensor with torch.sparse_csr layout I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822] I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822] from user code: I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822] File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/external_utils.py", line 36, in inner I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822] return fn(*args, kwargs) I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822] I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822] I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822] You can suppress this exception and fall back to eager by setting: I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822] import torch._dynamo I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822] torch._dynamo.config.suppress_errors = True I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822] I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822] I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822] ========== The above exception occurred while processing the following code ========== I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822] I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822] File "/global/u1/a/aglover/e3nn_benchmarking/minimal_sparse_compilation_error.py", line 38, in I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822] main() I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822] File "/global/u1/a/aglover/e3nn_benchmarking/minimal_sparse_compilation_error.py", line 29, in main I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822] print_performance(lambda: compiled_fullgraph_false(x, y), times=10, repeat=10) I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822] File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_inductor/utils.py", line 379, in print_performance I0713 17:21:46.240000 140632849619840 torch/_dynamo/convertframe.py:822] timings = torch.tensor([timed(fn, args, times, device) for in range(repeat)]) I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822] File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_inductor/utils.py", line 379, in I0713 17:21:46.240000 140632849619840 torch/_dynamo/convertframe.py:822] timings = torch.tensor([timed(fn, args, times, device) for in range(repeat)]) I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822] File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_inductor/utils.py", line 368, in timed I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822] result = model(example_inputs) I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822] File "/global/u1/a/aglover/e3nn_benchmarking/minimal_sparse_compilation_error.py", line 29, in I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822] print_performance(lambda: compiled_fullgraph_false(x, y), times=10, repeat=10) I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822] File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/external_utils.py", line 36, in inner I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822] return fn(args, kwargs) I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822] I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822] ========== I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822] Traceback (most recent call last): I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822] File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py", line 786, in _convert_frame I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822] result = inner_convert( I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822] File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py", line 400, in _convert_frame_assert I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822] return _compile( I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822] File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/contextlib.py", line 79, in inner I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822] return func(*args, kwds) I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822] File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py", line 676, in _compile I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822] guarded_code = compile_inner(code, one_graph, hooks, transform) I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822] File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/utils.py", line 262, in time_wrapper I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822] r = func(*args, *kwargs) I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822] File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py", line 535, in compile_inner I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822] out_code = transform_code_object(code, transform) I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822] File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/bytecode_transformation.py", line 1036, in transform_code_object I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822] transformations(instructions, code_options) I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822] File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py", line 165, in _fn I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822] return fn(args, kwargs) I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822] File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py", line 500, in transform I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822] tracer.run() I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822] File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 2149, in run I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822] super().run() I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822] File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 810, in run I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822] and self.step() I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822] File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 773, in step I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822] getattr(self, inst.opname)(inst) I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822] File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 484, in wrapper I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822] return handle_graph_break(self, inst, speculation.reason) I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822] File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 548, in handle_graph_break I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822] self.output.compile_subgraph(self, reason=reason) I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822] File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/output_graph.py", line 961, in compile_subgraph I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822] and all( I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822] File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/output_graph.py", line 962, in I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822] not isinstance( I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822] File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/variables/base.py", line 135, in instancecheck I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822] instance = instance.realize() I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822] File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/variables/lazy.py", line 58, in realize I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822] self._cache.realize(self.parents_tracker) I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822] File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/variables/lazy.py", line 24, in realize I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822] self.vt = VariableBuilder(tx, self.source)(self.value) I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822] File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/variables/builder.py", line 269, in call I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822] vt = self._wrap(value) I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822] File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/variables/builder.py", line 526, in _wrap I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822] keywords[k] = VariableBuilder( I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822] File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/variables/builder.py", line 269, in call I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822] vt = self._wrap(value) I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822] File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/variables/builder.py", line 402, in _wrap I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822] return type_dispatch(self, value) I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822] File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/variables/builder.py", line 1069, in wrap_tensor I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822] unimplemented( I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822] File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/exc.py", line 190, in unimplemented I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822] raise Unsupported(msg) I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822] torch._dynamo.exc.Unsupported: torch.compile does not support sparse Tensor with torch.sparse_csr layout I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822] I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822] from user code: I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822] File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/external_utils.py", line 36, in inner I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822] return fn(*args, kwargs) I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822] I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822] I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822] You can suppress this exception and fall back to eager by setting: I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822] import torch._dynamo I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822] torch._dynamo.config.suppress_errors = True I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822] I0713 17:21:46.243000 140632849619840 torch/_dynamo/logging.py:55] [1/0] Step 1: torchdynamo start tracing func_with_sparse /global/u1/a/aglover/e3nn_benchmarking/minimal_sparse_compilation_error.py:6 I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822] WON'T CONVERT func_with_sparse /global/u1/a/aglover/e3nn_benchmarking/minimal_sparse_compilation_error.py line 6 I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822] ========== TorchDynamo Stack Trace ========== I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822] Traceback (most recent call last): I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822] File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py", line 786, in _convert_frame I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822] result = inner_convert( I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822] File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py", line 400, in _convert_frame_assert I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822] return _compile( I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822] File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/contextlib.py", line 79, in inner I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822] return func(*args, *kwds) I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822] File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py", line 676, in _compile I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822] guarded_code = compile_inner(code, one_graph, hooks, transform) I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822] File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/utils.py", line 262, in time_wrapper I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822] r = func(args, kwargs) I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822] File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py", line 535, in compile_inner I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822] out_code = transform_code_object(code, transform) I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822] File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/bytecode_transformation.py", line 1036, in transform_code_object I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822] transformations(instructions, code_options) I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822] File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py", line 165, in _fn I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822] return fn(*args, kwargs) I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822] File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py", line 500, in transform I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822] tracer.run() I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822] File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 2149, in run I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822] super().run() I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822] File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 810, in run I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822] and self.step() I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822] File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 773, in step I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822] getattr(self, inst.opname)(inst) I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822] File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 242, in impl I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822] self.push(fn_var.call_function(self, self.popn(nargs), {})) I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822] File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/variables/builtin.py", line 535, in call_function I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822] args = [v.realize() for v in args] I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822] File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/variables/builtin.py", line 535, in I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822] args = [v.realize() for v in args] I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822] File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/variables/lazy.py", line 58, in realize I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822] self._cache.realize(self.parents_tracker) I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822] File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/variables/lazy.py", line 24, in realize I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822] self.vt = VariableBuilder(tx, self.source)(self.value) I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822] File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/variables/builder.py", line 269, in call I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822] vt = self._wrap(value) I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822] File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/variables/builder.py", line 402, in _wrap I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822] return type_dispatch(self, value) I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822] File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/variables/builder.py", line 1069, in wrap_tensor I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822] unimplemented( I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822] File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/exc.py", line 190, in unimplemented I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822] raise Unsupported(msg) I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822] torch._dynamo.exc.Unsupported: torch.compile does not support sparse Tensor with torch.sparse_csr layout I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822] I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822] from user code: I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822] File "/global/u1/a/aglover/e3nn_benchmarking/minimal_sparse_compilation_error.py", line 10, in func_with_sparse I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822] return (z @ torch.einsum("i,j->ij", x, y).reshape(-1)) I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822] I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822] I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822] You can suppress this exception and fall back to eager by setting: I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822] import torch._dynamo I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822] torch._dynamo.config.suppress_errors = True I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822] I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822] I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822] ========== The above exception occurred while processing the following code ========== I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822] I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822] File "/global/u1/a/aglover/e3nn_benchmarking/minimal_sparse_compilation_error.py", line 38, in I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822] main() I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822] File "/global/u1/a/aglover/e3nn_benchmarking/minimal_sparse_compilation_error.py", line 29, in main I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822] print_performance(lambda: compiled_fullgraph_false(x, y), times=10, repeat=10) I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822] File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_inductor/utils.py", line 379, in print_performance I0713 17:21:46.260000 140632849619840 torch/_dynamo/convertframe.py:822] timings = torch.tensor([timed(fn, args, times, device) for in range(repeat)]) I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822] File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_inductor/utils.py", line 379, in I0713 17:21:46.260000 140632849619840 torch/_dynamo/convertframe.py:822] timings = torch.tensor([timed(fn, args, times, device) for in range(repeat)]) I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822] File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_inductor/utils.py", line 368, in timed I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822] result = model(example_inputs) I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822] File "/global/u1/a/aglover/e3nn_benchmarking/minimal_sparse_compilation_error.py", line 29, in I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822] print_performance(lambda: compiled_fullgraph_false(x, y), times=10, repeat=10) I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822] File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/external_utils.py", line 36, in inner I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822] return fn(args, kwargs) I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822] File "/global/u1/a/aglover/e3nn_benchmarking/minimal_sparse_compilation_error.py", line 10, in func_with_sparse I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822] return (z @ torch.einsum("i,j->ij", x, y).reshape(-1)) I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822] I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822] ========== I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822] Traceback (most recent call last): I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822] File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py", line 786, in _convert_frame I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822] result = inner_convert( I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822] File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py", line 400, in _convert_frame_assert I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822] return _compile( I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822] File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/contextlib.py", line 79, in inner I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822] return func(*args, kwds) I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822] File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py", line 676, in _compile I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822] guarded_code = compile_inner(code, one_graph, hooks, transform) I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822] File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/utils.py", line 262, in time_wrapper I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822] r = func(*args, *kwargs) I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822] File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py", line 535, in compile_inner I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822] out_code = transform_code_object(code, transform) I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822] File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/bytecode_transformation.py", line 1036, in transform_code_object I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822] transformations(instructions, code_options) I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822] File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py", line 165, in _fn I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822] return fn(args, kwargs) I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822] File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py", line 500, in transform I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822] tracer.run() I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822] File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 2149, in run I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822] super().run() I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822] File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 810, in run I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822] and self.step() I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822] File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 773, in step I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822] getattr(self, inst.opname)(inst) I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822] File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 242, in impl I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822] self.push(fn_var.call_function(self, self.popn(nargs), {})) I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822] File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/variables/builtin.py", line 535, in call_function I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822] args = [v.realize() for v in args] I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822] File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/variables/builtin.py", line 535, in I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822] args = [v.realize() for v in args] I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822] File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/variables/lazy.py", line 58, in realize I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822] self._cache.realize(self.parents_tracker) I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822] File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/variables/lazy.py", line 24, in realize I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822] self.vt = VariableBuilder(tx, self.source)(self.value) I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822] File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/variables/builder.py", line 269, in call I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822] vt = self._wrap(value) I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822] File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/variables/builder.py", line 402, in _wrap I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822] return type_dispatch(self, value) I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822] File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/variables/builder.py", line 1069, in wrap_tensor I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822] unimplemented( I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822] File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/exc.py", line 190, in unimplemented I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822] raise Unsupported(msg) I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822] torch._dynamo.exc.Unsupported: torch.compile does not support sparse Tensor with torch.sparse_csr layout I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822] I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822] from user code: I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822] File "/global/u1/a/aglover/e3nn_benchmarking/minimal_sparse_compilation_error.py", line 10, in func_with_sparse I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822] return (z @ torch.einsum("i,j->ij", x, y).reshape(-1)) I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822] I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822] I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822] You can suppress this exception and fall back to eager by setting: I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822] import torch._dynamo I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822] torch._dynamo.config.suppress_errors = True I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822] 0.000150 I0713 17:21:46.286000 140632849619840 torch/_dynamo/logging.py:55] [0/0] Step 1: torchdynamo start tracing inner /global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/external_utils.py:34 Traceback (most recent call last): File "/global/u1/a/aglover/e3nn_benchmarking/minimal_sparse_compilation_error.py", line 38, in main() File "/global/u1/a/aglover/e3nn_benchmarking/minimal_sparse_compilation_error.py", line 35, in main print_performance(lambda: compiled_fullgraph_true(x, y), times=10, repeat=10) File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_inductor/utils.py", line 379, in printperformance timings = torch.tensor([timed(fn, args, times, device) for in range(repeat)]) File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/inductor/utils.py", line 379, in timings = torch.tensor([timed(fn, args, times, device) for in range(repeat)]) File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_inductor/utils.py", line 368, in timed result = model(example_inputs) File "/global/u1/a/aglover/e3nn_benchmarking/minimal_sparse_compilation_error.py", line 35, in print_performance(lambda: compiled_fullgraph_true(x, y), times=10, repeat=10) File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/eval_frame.py", line 451, in _fn return fn(args, kwargs) File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py", line 921, in catch_errors return callback(frame, cache_entry, hooks, frame_state, skip=1) File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py", line 400, in _convert_frame_assert return _compile( File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/contextlib.py", line 79, in inner return func(*args, *kwds) File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py", line 676, in _compile guarded_code = compile_inner(code, one_graph, hooks, transform) File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/utils.py", line 262, in time_wrapper r = func(args, kwargs) File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py", line 535, in compile_inner out_code = transform_code_object(code, transform) File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/bytecode_transformation.py", line 1036, in transform_code_object transformations(instructions, code_options) File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py", line 165, in _fn return fn(*args, **kwargs) File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py", line 500, in transform tracer.run() File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 2149, in run super().run() File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 810, in run and self.step() File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 773, in step getattr(self, inst.opname)(inst) File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 489, in wrapper return inner_fn(self, inst) File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 1238, in CALL_FUNCTION_EX isinstance(fn, GetAttrVariable) File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/variables/base.py", line 135, in instancecheck instance = instance.realize() File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/variables/lazy.py", line 58, in realize self._cache.realize(self.parents_tracker) File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/variables/lazy.py", line 24, in realize self.vt = VariableBuilder(tx, self.source)(self.value) File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/variables/builder.py", line 269, in call vt = self._wrap(value) File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/variables/builder.py", line 526, in _wrap keywords[k] = VariableBuilder( File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/variables/builder.py", line 269, in call vt = self._wrap(value) File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/variables/builder.py", line 402, in _wrap return type_dispatch(self, value) File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/variables/builder.py", line 1069, in wrap_tensor unimplemented( File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/exc.py", line 190, in unimplemented raise Unsupported(msg) torch._dynamo.exc.Unsupported: torch.compile does not support sparse Tensor with torch.sparse_csr layout

from user code: File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/external_utils.py", line 36, in inner return fn(*args, **kwargs)

You can suppress this exception and fall back to eager by setting: import torch._dynamo torch._dynamo.config.suppress_errors = True

I0713 17:21:46.396000 140632849619840 torch/_dynamo/utils.py:320] TorchDynamo compilation metrics: I0713 17:21:46.396000 140632849619840 torch/_dynamo/utils.py:320] Function, Runtimes (s) I0713 17:21:46.396000 140632849619840 torch/_dynamo/utils.py:320] _compile..compile_inner, 0.0000

Minified repro

import torch from torch._inductor.utils import print_performance from functools import partial import logging

def func_with_sparse(x : torch.tensor , y : torch.tensor , z : torch.tensor): """ This peforms a contraction with a CSR matrix """ return (z @ torch.einsum("i,j->ij", x, y).reshape(-1))

def main(): torch._dynamo.config.verbose=True torch._logging.set_logs(dynamo = logging.INFO)

device = "cuda"

### Tensors, 2 dense, 1 sparse
x = torch.rand( 3,  dtype=torch.float32, device=device)
y = torch.rand( 3,  dtype=torch.float32, device=device)
z = torch.randn(size=(3,9), dtype=torch.float32, device=device).to_sparse_csr()

### Works with sparse arg 
func_with_sparse_arg = partial(func_with_sparse, z=z)
print_performance(lambda: func_with_sparse_arg(x, y), times=10, repeat=10)

### Suceeds compilation (fullgraph=False) with sparse arg
compiled_fullgraph_false = torch.compile(func_with_sparse_arg, fullgraph=False, backend='eager') 
print_performance(lambda: compiled_fullgraph_false(x, y), times=10, repeat=10)

torch._dynamo.reset()

### Fails compilation (fullgraph=True) with sparse arg
compiled_fullgraph_true = torch.compile(func_with_sparse_arg, fullgraph=True, backend='eager') 
print_performance(lambda: compiled_fullgraph_true(x, y), times=10, repeat=10)

if name == "main": main()

Versions

(e3nn_env) aglover@login36:~/e3nn_benchmarking> python3 collect_env.py Collecting environment information... PyTorch version: 2.3.1.post300 Is debug build: False CUDA used to build PyTorch: 12.0 ROCM used to build PyTorch: N/A

OS: SUSE Linux Enterprise Server 15 SP4 (x86_64) GCC version: (SUSE Linux) 12.3.0 Clang version: Could not collect CMake version: version 3.20.4 Libc version: glibc-2.31

Python version: 3.10.14 | packaged by conda-forge | (main, Mar 20 2024, 12:45:18) [GCC 12.3.0] (64-bit runtime) Python platform: Linux-5.14.21-150400.24.81_12.0.87-cray_shasta_c-x86_64-with-glibc2.31 Is CUDA available: True CUDA runtime version: 12.5.82 CUDA_MODULE_LOADING set to: LAZY GPU models and configuration: GPU 0: NVIDIA A100-PCIE-40GB Nvidia driver version: 525.105.17 cuDNN version: Could not collect HIP runtime version: N/A MIOpen runtime version: N/A Is XNNPACK available: True

CPU: Architecture: x86_64 CPU op-mode(s): 32-bit, 64-bit Address sizes: 48 bits physical, 48 bits virtual Byte Order: Little Endian CPU(s): 256 On-line CPU(s) list: 0-255 Vendor ID: AuthenticAMD Model name: AMD EPYC 7713 64-Core Processor CPU family: 25 Model: 1 Thread(s) per core: 2 Core(s) per socket: 64 Socket(s): 2 Stepping: 1 Frequency boost: enabled CPU max MHz: 2000.0000 CPU min MHz: 1500.0000 BogoMIPS: 3992.60 Flags: fpu vme de pse tsc msr pae mce cx8 apic sep mtrr pge mca cmov pat pse36 clflush mmx fxsr sse sse2 ht syscall nx mmxext fxsr_opt pdpe1gb rdtscp lm constant_tsc rep_good nopl nonstop_tsc cpuid extd_apicid aperfmperf rapl pni pclmulqdq monitor ssse3 fma cx16 pcid sse4_1 sse4_2 x2apic movbe popcnt aes xsave avx f16c rdrand lahf_lm cmp_legacy svm extapic cr8_legacy abm sse4a misalignsse 3dnowprefetch osvw ibs skinit wdt tce topoext perfctr_core perfctr_nb bpext perfctr_llc mwaitx cpb cat_l3 cdp_l3 invpcid_single hw_pstate ssbd mba ibrs ibpb stibp vmmcall fsgsbase bmi1 avx2 smep bmi2 erms invpcid cqm rdt_a rdseed adx smap clflushopt clwb sha_ni xsaveopt xsavec xgetbv1 xsaves cqm_llc cqm_occup_llc cqm_mbm_total cqm_mbm_local clzero irperf xsaveerptr rdpru wbnoinvd amd_ppin arat npt lbrv svm_lock nrip_save tsc_scale vmcb_clean flushbyasid decodeassists pausefilter pfthreshold v_vmsave_vmload vgif v_spec_ctrl umip pku ospke vaes vpclmulqdq rdpid overflow_recov succor smca fsrm Virtualization: AMD-V L1d cache: 4 MiB (128 instances) L1i cache: 4 MiB (128 instances) L2 cache: 64 MiB (128 instances) L3 cache: 512 MiB (16 instances) NUMA node(s): 2 NUMA node0 CPU(s): 0-63,128-191 NUMA node1 CPU(s): 64-127,192-255 Vulnerability Gather data sampling: Not affected Vulnerability Itlb multihit: Not affected Vulnerability L1tf: Not affected Vulnerability Mds: Not affected Vulnerability Meltdown: Not affected Vulnerability Mmio stale data: Not affected Vulnerability Retbleed: Not affected Vulnerability Spec rstack overflow: Mitigation; safe RET Vulnerability Spec store bypass: Mitigation; Speculative Store Bypass disabled via prctl and seccomp Vulnerability Spectre v1: Mitigation; usercopy/swapgs barriers and __user pointer sanitization Vulnerability Spectre v2: Mitigation; Retpolines, IBPB conditional, IBRS_FW, STIBP always-on, RSB filling, PBRSB-eIBRS Not affected Vulnerability Srbds: Not affected Vulnerability Tsx async abort: Not affected

Versions of relevant libraries: [pip3] numpy==1.26.4 [pip3] torch==2.3.1.post300 [conda] libmagma 2.7.2 h173bb3b_2 conda-forge [conda] libmagma_sparse 2.7.2 h173bb3b_3 conda-forge [conda] libtorch 2.3.1 cuda120_h2b0da52_300 conda-forge [conda] mkl 2023.2.0 h84fe81f_50496 conda-forge [conda] numpy 1.26.4 py310hb13e2d6_0 conda-forge [conda] pytorch 2.3.1 cuda120_py310h2c91c31_300 conda-forge

cc @alexsamardzic @nikitaved @pearu @cpuhrsch @amjames @bhosmer @jcaip @ezyang @anijain2305 @chauhang @penguinwu

asglover commented 1 month ago

I think this is probably something that your team is aware of, but I was asked to file a ticket. If there are any additional experiments or things you would like me to check, please let me know.

anijain2305 commented 1 month ago

cc @pearu We have very limited support for torch.compile + sparse tensors. I don't think anybody is working on it at the moment.

ezyang commented 1 month ago

Actually, @aartbik has been working on it. But don't expect speedups even if we can capture it https://github.com/pytorch/pytorch/pull/117907

asglover commented 1 month ago

Hi @ezyang Thanks for the info!

Hi @aartbik thank you for doing this work on the sparse compilers, it is much appreciated.

If you happen to be looking for real world applications / test cases for your compiler, I can share our application. We are going to solve it by writing a custom CUDA kernel, but the kernel could absolutely be generated by a sufficiently powerful sparse compiler.

Again, thank you for your work on this infrastructure!

aartbik commented 1 month ago

@asglover I am very interested in getting more real world applications / test cases for our compiler, which is ramping up (see https://github.com/MPACT-ORG/mpact-compiler). Would you mind contacting me by email to make that work?

asglover commented 1 month ago

I'll write up a little summary of our application and the ways a compiler could help. Right off the bat, I would like to introduce @mitkotak who is looking at this problem specifically from a compiler perspective, and is very knowledgeable about compiler specifics. I'm focusing on the characteristics of the sparsity and trying to work out and implement by hand what I think should be performant. I'll email you when I get the time to put something together.

aartbik commented 1 month ago

Thank you @asglover and @mitkotak. Looking forward to working together on sparse!