Jason Ansel
8531d247ba
[fx] Optimize TracerBase.create_arg and Graph._gen_python_code ( #148292 )
...
Before: 19502951 function calls (18702776 primitive calls) in 8.533 seconds
After: 16402551 function calls (15602452 primitive calls) in 7.701 seconds
Pull Request resolved: https://github.com/pytorch/pytorch/pull/148292
Approved by: https://github.com/oulgen
ghstack dependencies: #148243 , #148260 , #148261 , #148303 , #148288
2025-03-04 02:42:23 +00:00
Jason Ansel
5eb0337cfd
[fx] Optimizations for node name generation ( #148288 )
...
Before:

After:

Pull Request resolved: https://github.com/pytorch/pytorch/pull/148288
Approved by: https://github.com/oulgen
ghstack dependencies: #148243 , #148260 , #148261 , #148303
2025-03-04 02:42:23 +00:00
Jason Ansel
edaff88f69
[fx] Move map_aggregate to C++ ( #148243 )
...
Microbenchmarking `fx.symbolic_trace(lambda x: functools.reduce(operator.add, [x, *range(100000)]))`, before:
```
30603618 function calls (29403419 primitive calls) in 13.744 seconds
```
after:
```
25203549 function calls (24403352 primitive calls) in 12.090 seconds
```
Pull Request resolved: https://github.com/pytorch/pytorch/pull/148243
Approved by: https://github.com/oulgen
2025-03-02 22:42:31 +00:00
Aaron Orenstein
db4ce78d46
PEP585: More UP006 fixes ( #146392 )
...
This should be the final PR before we can enable RUFF UP006.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/146392
Approved by: https://github.com/justinchuby , https://github.com/albanD , https://github.com/Skylion007
2025-02-20 06:18:13 +00:00
Aaron Orenstein
1f8ff94d4f
PEP585: Add noqa to necessary tests ( #146391 )
...
Pull Request resolved: https://github.com/pytorch/pytorch/pull/146391
Approved by: https://github.com/justinchuby , https://github.com/Skylion007
2025-02-12 15:29:50 +00:00
Aaron Orenstein
35c8c31f11
Fix for failure in D68425364 ( #145304 )
...
Summary: Back out change from #145166 which causes an internal model to fail.
Differential Revision: D68459095
Pull Request resolved: https://github.com/pytorch/pytorch/pull/145304
Approved by: https://github.com/izaitsevfb
2025-01-22 23:33:02 +00:00
Aaron Orenstein
0b2a3687b9
PEP585 update - torch/fx ( #145166 )
...
See #145101 for details.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/145166
Approved by: https://github.com/bobrenjc93
2025-01-20 18:11:54 +00:00
Aaron Gokaslan
cb66146f2b
[BE]: Update literal typing for torch/fx/graph nodelist ( #144650 )
...
Mentioned in discussion for #144631
Pull Request resolved: https://github.com/pytorch/pytorch/pull/144650
Approved by: https://github.com/jansel
2025-01-12 21:02:13 +00:00
Jason Ansel
fd382f1269
Micro-optimization in Graph.nodes.__iter__ ( #144631 )
...
This generates slightly better code (removing a generator frame) and
drops a redundant assert.
```py
>>> import timeit
>>> def a():
... yield from range(3)
...
>>> def b():
... return range(3)
...
>>> timeit.timeit(lambda: [*a()])
0.2714634328149259
>>> timeit.timeit(lambda: [*b()])
0.12076826114207506
>>>
```
Pull Request resolved: https://github.com/pytorch/pytorch/pull/144631
Approved by: https://github.com/oulgen , https://github.com/Skylion007
2025-01-12 17:46:46 +00:00
Marvin Kim
b1b0afb8e8
[BE] Add type annotation to eliminate_dead_code ( #142251 )
...
Test Plan: CI
Reviewed By: evanleed
D-ifferential Revision: D66887283
Pull Request resolved: https://github.com/pytorch/pytorch/pull/142251
Approved by: https://github.com/ezyang , https://github.com/Skylion007
2024-12-10 17:09:21 +00:00
PyTorch MergeBot
75530885ba
Revert "[BE] Add type annotation to eliminate_dead_code ( #142251 )"
...
This reverts commit 3d04de6b2f .
Reverted https://github.com/pytorch/pytorch/pull/142251 on behalf of https://github.com/jeanschmidt due to checking if reverting will fix 'FAILED [5.0221s] test_dataloader.py::TestIndividualWorkerQueue::test_ind_worker_queue' on windows ([comment](https://github.com/pytorch/pytorch/pull/142251#issuecomment-2531706362 ))
2024-12-10 13:57:00 +00:00
Marvin Kim
3d04de6b2f
[BE] Add type annotation to eliminate_dead_code ( #142251 )
...
Test Plan: CI
Reviewed By: evanleed
Differential Revision: D66887283
Pull Request resolved: https://github.com/pytorch/pytorch/pull/142251
Approved by: https://github.com/ezyang , https://github.com/Skylion007
2024-12-10 09:27:29 +00:00
Avik Chaudhuri
74eb92ed6e
fix deep copy of empty graph ( #141660 )
...
Differential Revision: D66532131
Pull Request resolved: https://github.com/pytorch/pytorch/pull/141660
Approved by: https://github.com/ezyang
2024-12-02 22:03:13 +00:00
Sherlock Huang
071d48c56e
Add output_node util function to fx.Graph ( #139770 )
...
Summary: A util function for access output node for FX graph
Test Plan: OSS CI
Differential Revision: D65486457
Pull Request resolved: https://github.com/pytorch/pytorch/pull/139770
Approved by: https://github.com/ezyang , https://github.com/Chillee
2024-11-07 18:54:59 +00:00
Matthew Francis-Landau
a7f49de485
Fixes issue with enums in a tuple for dynamo ( #133123 )
...
Currently when tuples values are encountered in dynamo, they are encoded using `repr(arg)`. This causes an issue if one of the values inside of the tuple will not be properly encoded. In this case, if an enum is contained inside of a tuple, it will cause invalid python code to be generated
Pull Request resolved: https://github.com/pytorch/pytorch/pull/133123
Approved by: https://github.com/jansel
2024-10-21 23:45:11 +00:00
Xuehai Pan
abbd71d29d
[BE][Easy] enable PYFMT for torch.fx ( #138443 )
...
Reproduce command:
```bash
ghstack checkout https://github.com/pytorch/pytorch/pull/138443
git checkout HEAD~1 torch/
lintrunner -a --take "PYFMT" --all-files
```
Pull Request resolved: https://github.com/pytorch/pytorch/pull/138443
Approved by: https://github.com/ezyang
2024-10-21 19:15:49 +00:00
Tom Ritchford
c0582fd0f8
Remove unused Python variables in torch/[b-z]* ( #136963 )
...
Pull Request resolved: https://github.com/pytorch/pytorch/pull/136963
Approved by: https://github.com/ezyang
2024-10-19 16:45:22 +00:00
Jason Ansel
1f15c0c7a5
[fx] Replace _snake_case with a regexp ( #135822 )
...
~2x speedup on this function, though saves <0.5s overall
Pull Request resolved: https://github.com/pytorch/pytorch/pull/135822
Approved by: https://github.com/oulgen
ghstack dependencies: #135787 , #135788 , #135820 , #135821
2024-09-13 00:18:41 +00:00
Jason Ansel
a72124add9
[fx] Minor optimization in create_arg ( #135821 )
...
Pull Request resolved: https://github.com/pytorch/pytorch/pull/135821
Approved by: https://github.com/oulgen
ghstack dependencies: #135787 , #135788 , #135820
2024-09-13 00:18:41 +00:00
Jason Ansel
86335e9135
[reland 3/3][fx] Bypass custom __setattr__ in Node.__init__ ( #135735 )
...
Relands #135079 whcih was reverted by #135562
I broke this up into three parts to test internally.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/135735
Approved by: https://github.com/oulgen
2024-09-12 05:50:39 +00:00
Ivan Zaitsev
440f8f57af
Revert "[fx] Bypass custom __setattr__ in Node.__init__ ( #135079 )" ( #135562 )
...
This reverts commit 66da3b3b2a .
#135079 breaks internal tests and needs to be reverted. Revert with mergebot doesn't work as this PR is technically part of the stack, but, according to @jansel, it should be possible to revert it individually.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/135562
Approved by: https://github.com/jansel , https://github.com/seemethere
2024-09-10 18:07:11 +00:00
Riley Dulin
24482e5c68
[torch][fx] Set maximum warning count during fx.Graph.lint ( #135069 )
...
Summary:
resnet152 spent about 15 minutes writing warning messages in _unlift
during `to_executorch` because they're all written to unbuffered stderr
by the `warnings` module.
These warnings are almost always about get_attr nodes referencing a
non-existent name:
```lang=py
warnings.warn(f'Node {node} target {node.target} {atom} of {seen_qualname} does '
'not reference an nn.Module, nn.Parameter, or buffer, which is '
'what \'get_attr\' Nodes typically target'
)
```
I'm not aware of a way to configure the warnings module to write this out
at most once, so I'm just going to disable the lint for now.
Test Plan:
Re-ran resnet152 with Executorch and the XNNPackBackend, it is much faster now
Differential Revision: D62156090
Pull Request resolved: https://github.com/pytorch/pytorch/pull/135069
Approved by: https://github.com/yushangdi
2024-09-06 16:41:59 +00:00
Jason Ansel
66da3b3b2a
[fx] Bypass custom __setattr__ in Node.__init__ ( #135079 )
...
Before:

After:

Pull Request resolved: https://github.com/pytorch/pytorch/pull/135079
Approved by: https://github.com/oulgen
ghstack dependencies: #135070 , #135076 , #135082 , #135084
2024-09-06 06:11:46 +00:00
Aaron Orenstein
ed86ac2f25
[BE] typing for decorators - fx/_compatibility ( #134054 )
...
Summary: See #131429
Test Plan: unit tests pass
Differential Revision: D61493706
Pull Request resolved: https://github.com/pytorch/pytorch/pull/134054
Approved by: https://github.com/oulgen
2024-08-26 04:00:27 +00:00
Jiashen Cao
fa8c34301a
[ts-migration]: Quantized ops to standard ops pass. ( #133026 )
...
#### Description
Transform quantized operation properly. Add de/quantization before and after the quantized operation.
#### Test Plan
`pytest test/export/test_converter.py -s -k test_ts2ep_convert_quantized_model`
Pull Request resolved: https://github.com/pytorch/pytorch/pull/133026
Approved by: https://github.com/angelayi
2024-08-08 23:10:17 +00:00
PyTorch MergeBot
6f99e97f0a
Revert "[ts-migration]: Support quantized operation transformation ( #131915 )"
...
This reverts commit 0e8541766f .
Reverted https://github.com/pytorch/pytorch/pull/131915 on behalf of https://github.com/ezyang due to test broken on windows 0e8541766f ([comment](https://github.com/pytorch/pytorch/pull/131915#issuecomment-2275974907 ))
2024-08-08 14:30:35 +00:00
Jiashen Cao
0e8541766f
[ts-migration]: Support quantized operation transformation ( #131915 )
...
#### Description
Transform quantized operation properly. Add de/quantization before and after the quantized operation.
#### Test Plan
`pytest test/export/test_converter.py -s -k test_ts2ep_convert_quantized_model`
Pull Request resolved: https://github.com/pytorch/pytorch/pull/131915
Approved by: https://github.com/angelayi
2024-08-08 06:34:53 +00:00
David Berard
5973aec671
[fx] python_code(verbose=True): show size/strides for all tensors ( #132192 )
...
python_code(verbose=True) (or print_readable()) generates a string with the code representing the fx graph, with extra annotations indicating the size or stride of the tensor. Currently, it'll only shows sizes/strides for FakeTensors provided in metadata. For subclass tensors like NestedTensor, the outer class (provided in the node metadata) will be a non-FakeTensor and the inner tensors will be fake. This PR expands the conditional to show sizes/strides for all tensors, not just FakeTensors.
Testing: I ran this test script (below), ran it with `TORCH_LOGS=+dynamo` and found in the logs the graph shown below - we see that the input nested tensor has sizes and strides associated with it. Also, I stacked a diff on top of this one that forces the readable graph to be generated whenever PT2 is in use in tests, which should hopefully find any issues; https://github.com/pytorch/pytorch/pull/132195 shows no significant failures except for preexisting failures.
test script:
```python
import torch
def fn(x):
return x.cos()
nt = torch.nested.nested_tensor_from_jagged(
torch.randn(10, 10),
torch.tensor([0, 1, 3, 6, 10]),
)
torch.compile(fn)(nt)
```
logs excerpt:
```
[0/0] [__graph_code] TRACED GRAPH
[0/0] [__graph_code] ===== __compiled_fn_1 =====
[0/0] [__graph_code] /data/users/dberard/pytorch/torch/fx/_lazy_graph_module.py class GraphModule(torch.nn.M
[0/0] [__graph_code] def forward(self, L_x_: "f32[4, zf1, 10][10*zf1, 10, 1]cpu", zf1: "Sym(zf1)"):
[0/0] [__graph_code] l_x_ = L_x_
[0/0] [__graph_code]
[0/0] [__graph_code] # File: /data/users/dberard/scripts/nt_print_graph.py:4 in fn, code: return x.c
[0/0] [__graph_code] cos: "f32[4, zf1, 10][10*zf1, 10, 1]cpu" = l_x_.cos(); l_x_ = None
[0/0] [__graph_code] return (cos,)
```
Pull Request resolved: https://github.com/pytorch/pytorch/pull/132192
Approved by: https://github.com/Chillee
2024-08-03 02:54:32 +00:00
YangQun1
589aef4bb0
Fix py codegen to delete values that don't have any users ( #131028 )
...
Fixes #131025
Pull Request resolved: https://github.com/pytorch/pytorch/pull/131028
Approved by: https://github.com/ezyang
2024-08-01 03:18:37 +00:00
PyTorch MergeBot
945bf78894
Revert "[BE] typing for decorators - fx/_compatibility ( #131568 )"
...
This reverts commit 193f62fde9 .
Reverted https://github.com/pytorch/pytorch/pull/131568 on behalf of https://github.com/clee2000 due to same as https://github.com/pytorch/pytorch/pull/131572#issuecomment-2254328359 but I clicked the wrong link by accident. This is where it actually starts ([comment](https://github.com/pytorch/pytorch/pull/131568#issuecomment-2254330781 ))
2024-07-28 03:43:39 +00:00
Aaron Orenstein
193f62fde9
[BE] typing for decorators - fx/_compatibility ( #131568 )
...
See #131429
Pull Request resolved: https://github.com/pytorch/pytorch/pull/131568
Approved by: https://github.com/justinchuby , https://github.com/oulgen , https://github.com/zou3519
2024-07-25 22:24:19 +00:00
PyTorch MergeBot
c3679bed35
Revert "Fix py codegen to delete values that don't have any users ( #131028 )"
...
This reverts commit 91aba7baac .
Reverted https://github.com/pytorch/pytorch/pull/131028 on behalf of https://github.com/clee2000 due to broke inductor/test_triton_kernels inductor/test_triton_kernels.py::KernelTests::test_triton_kernel_functionalize [GH job link](https://github.com/pytorch/pytorch/actions/runs/10094659640/job/27915271250 ) [HUD commit link](91aba7baac ) ([comment](https://github.com/pytorch/pytorch/pull/131028#issuecomment-2251058374 ))
2024-07-25 17:42:18 +00:00
YangQun1
91aba7baac
Fix py codegen to delete values that don't have any users ( #131028 )
...
Fixes #131025
Pull Request resolved: https://github.com/pytorch/pytorch/pull/131028
Approved by: https://github.com/ezyang
2024-07-25 13:04:23 +00:00
PyTorch MergeBot
8ffd109a00
Revert "Fix py codegen to delete values that don't have any users ( #131028 )"
...
This reverts commit 466c167b71 .
Reverted https://github.com/pytorch/pytorch/pull/131028 on behalf of https://github.com/atalman due to breaks CI ([comment](https://github.com/pytorch/pytorch/pull/131028#issuecomment-2247771530 ))
2024-07-24 12:21:43 +00:00
YangQun1
466c167b71
Fix py codegen to delete values that don't have any users ( #131028 )
...
Fixes #131025
Pull Request resolved: https://github.com/pytorch/pytorch/pull/131028
Approved by: https://github.com/ezyang
2024-07-24 01:03:56 +00:00
Aaron Orenstein
5a0068cc69
[BE] mypy: disallow untyped decorators ( #131428 )
...
Untyped decorators strip the types from their decorated function so even if the underlying function is fully typed then callers to it don't get any benefit from type annotations.
Step 1 - Enable the error and override in all the offending files.
#131429
Pull Request resolved: https://github.com/pytorch/pytorch/pull/131428
Approved by: https://github.com/justinchuby , https://github.com/oulgen
2024-07-23 21:50:55 +00:00
Shangdi Yu
ea4b80e6d6
[FX][export] strict DCE pass, check schema for node impurity ( #130552 )
...
Fixes the failure in `test/export/test_export_training_ir_to_run_decomp.py ` caused by dead code elimination removing node with side effects.
For background, in export, we may want to export higher-level IRs that are not functional, so we need to check for side effects more carefully.
A call_function node is impure if it has at least one mutable argument.
Fixed the tests below:
test_to_module_with_mutated_buffer_multiple_update_sub_later
test_export_input_mutation_static_shape
test_buffer_util
Another attempt modifying the original DCE pass is made in PR #130395 , but it breaks some other tests, so here we add a flag and use it for export only.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/130552
Approved by: https://github.com/pianpwk
2024-07-12 15:43:27 +00:00
Chen Lai
721a798886
add bits16 to graph dtype_abbrs ( #130339 )
...
As title, patch the dtype in torch.fx.graph
Pull Request resolved: https://github.com/pytorch/pytorch/pull/130339
Approved by: https://github.com/angelayi
2024-07-09 19:58:51 +00:00
angelayi
e9c6e8369c
Torchbind call method + effects support ( #128397 )
...
Adds effect token support to torchbind method calls by allowing `with_effects` to take in `torch.ops._higher_order_ops.call_torchbind` as an input.
Here is the print from `TORCH_LOGS="aot" python test/export/test_torchbind.py -k test_compile_obj_torchbind_op`:
```python
def forward(self, arg0_1: "f32[0]", arg1_1: "f32[2]", arg2_1):
# File: /data/users/angelayi/pytorch2/test/export/test_torchbind.py:1266 in f, code: torch.ops._TorchScriptTesting.queue_push(tq, x.cos())
cos: "f32[2]" = torch.ops.aten.cos.default(arg1_1)
with_effects = torch._higher_order_ops.effects.with_effects(arg0_1, torch.ops._TorchScriptTesting.queue_push.default, arg2_1, cos); arg0_1 = cos = None
getitem: "f32[0]" = with_effects[0]; with_effects = None
# File: /data/users/angelayi/pytorch2/test/export/test_torchbind.py:1267 in f, code: torch.ops._TorchScriptTesting.queue_push(tq, x.cos() + 1)
cos_1: "f32[2]" = torch.ops.aten.cos.default(arg1_1)
add: "f32[2]" = torch.ops.aten.add.Tensor(cos_1, 1); cos_1 = None
with_effects_1 = torch._higher_order_ops.effects.with_effects(getitem, torch.ops._TorchScriptTesting.queue_push.default, arg2_1, add); getitem = add = None
getitem_2: "f32[0]" = with_effects_1[0]; with_effects_1 = None
# File: /data/users/angelayi/pytorch2/test/export/test_torchbind.py:1268 in f, code: torch.ops._TorchScriptTesting.queue_pop(tq)
with_effects_2 = torch._higher_order_ops.effects.with_effects(getitem_2, torch.ops._TorchScriptTesting.queue_pop.default, arg2_1); getitem_2 = None
getitem_4: "f32[0]" = with_effects_2[0]; with_effects_2 = None
# File: /data/users/angelayi/pytorch2/test/export/test_torchbind.py:1269 in f, code: torch.ops._TorchScriptTesting.queue_push(tq, x.sin())
sin: "f32[2]" = torch.ops.aten.sin.default(arg1_1); arg1_1 = None
with_effects_3 = torch._higher_order_ops.effects.with_effects(getitem_4, torch.ops._TorchScriptTesting.queue_push.default, arg2_1, sin); getitem_4 = sin = None
getitem_6: "f32[0]" = with_effects_3[0]; with_effects_3 = None
# File: /data/users/angelayi/pytorch2/test/export/test_torchbind.py:1270 in f, code: return tq.pop(), tq.pop() + tq.size(), tq
with_effects_4 = torch._higher_order_ops.effects.with_effects(getitem_6, torch.ops._higher_order_ops.call_torchbind, arg2_1, 'pop'); getitem_6 = None
getitem_8: "f32[0]" = with_effects_4[0]
getitem_9: "f32[2]" = with_effects_4[1]; with_effects_4 = None
with_effects_5 = torch._higher_order_ops.effects.with_effects(getitem_8, torch.ops._higher_order_ops.call_torchbind, arg2_1, 'pop'); getitem_8 = None
getitem_10: "f32[0]" = with_effects_5[0]
getitem_11: "f32[2]" = with_effects_5[1]; with_effects_5 = None
with_effects_6 = torch._higher_order_ops.effects.with_effects(getitem_10, torch.ops._higher_order_ops.call_torchbind, arg2_1, 'size'); getitem_10 = arg2_1 = None
getitem_12: "f32[0]" = with_effects_6[0]; with_effects_6 = None
add_1: "f32[2]" = torch.ops.aten.add.Tensor(getitem_11, 0); getitem_11 = None
return (getitem_12, getitem_9, add_1)
```
In order to support this, this PR makes the following changes:
* Adds `FakeScriptObject` to `CustomObjArgument`, which will be put on the `meta["val"]` of nodes representing torchbind objects.
* Adds pickle/deepcopy support to FunctionSchema.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/128397
Approved by: https://github.com/ydwu4 , https://github.com/zou3519
2024-06-14 21:28:17 +00:00
chilli
c486e2ab64
Add coloring to fx graph print out ( #128476 )
...
Note: Won't land immediately, at least I'll need to add a color option to the field. But curious if any tests fail.
Old:
<img width="1294" alt="image" src="https://github.com/pytorch/pytorch/assets/6355099/c3a750ed-5e54-4621-b2e4-be5481be15b6 ">
New:
<img width="1303" alt="image" src="https://github.com/pytorch/pytorch/assets/6355099/3a1f1adc-6f3a-413e-8b87-ee53da9bf4ed ">
Pull Request resolved: https://github.com/pytorch/pytorch/pull/128476
Approved by: https://github.com/ezyang
2024-06-13 23:39:04 +00:00
Oguz Ulgen
5b5d269d34
Speed up fx graph iteration by implementing it in C++ ( #128288 )
...
Before this change
```
python benchmarks/dynamo/microbenchmarks/fx_microbenchmarks.py
iterating over 100000000 FX nodes took 19.5s (5132266 nodes/s)
```
After this change
```
python benchmarks/dynamo/microbenchmarks/fx_microbenchmarks.py
iterating over 100000000 FX nodes took 3.4s (29114001 nodes/s)
```
5.7x improvement
Differential Revision: [D58343997](https://our.internmc.facebook.com/intern/diff/D58343997 )
Pull Request resolved: https://github.com/pytorch/pytorch/pull/128288
Approved by: https://github.com/jansel , https://github.com/albanD
2024-06-11 05:48:31 +00:00
Aaron Orenstein
038b927590
Flip default value for mypy disallow_untyped_defs [7/11] ( #127844 )
...
See #127836 for details.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/127844
Approved by: https://github.com/oulgen
ghstack dependencies: #127842 , #127843
2024-06-08 18:49:45 +00:00
Xuehai Pan
8b08b0f340
[BE] enable ruff rule Q from flake8-quotes ( #127713 )
...
Enable [ruff rule `Q`](https://docs.astral.sh/ruff/rules/#flake8-quotes-q ) from flake8-quotes. Fixes:
- [avoidable-escaped-quote (Q003)](https://docs.astral.sh/ruff/rules/avoidable-escaped-quote/#avoidable-escaped-quote-q003 )
- [unnecessary-escaped-quote (Q004)](https://docs.astral.sh/ruff/rules/unnecessary-escaped-quote/#unnecessary-escaped-quote-q004 )
Pull Request resolved: https://github.com/pytorch/pytorch/pull/127713
Approved by: https://github.com/ezyang
2024-06-02 23:25:26 +00:00
Sheng Fu
bbeb0906c4
Register creak_node_hook ( #126671 )
...
Differential Revision: D57469157
Pull Request resolved: https://github.com/pytorch/pytorch/pull/126671
Approved by: https://github.com/angelayi
2024-05-24 23:32:15 +00:00
angelayi
8be4c1bc2f
[export] Add metadata for nodes insert_deferred_runtime_asserts ( #125414 )
...
Fixes [internal error](https://fb.workplace.com/groups/1075192433118967/permalink/1416709435633930/ ).
The issue is that the asserting nodes added in the `insert_deferred_runtime_assertion` pass do not contain metadata that the ExportedProgram requires the graph to have. One solution to fix this is to retrace the entire module, or another solution is to manually add back this metadata.
This diff implements the latter solution (manually add back the metadata) through hooking into fx.graph's `create_node` function, and adding export-specific metadata for every node that is created. The reason I did this is so that the `insert_deferred_runtime_assertion` does not have to know about what metadata export wants.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/125414
Approved by: https://github.com/zhxchen17 , https://github.com/BoyuanFeng
2024-05-07 23:15:21 +00:00
Edward Z. Yang
ecd62746e3
Also pull size/stride info from example_value ( #125505 )
...
Signed-off-by: Edward Z. Yang <ezyang@meta.com>
Pull Request resolved: https://github.com/pytorch/pytorch/pull/125505
Approved by: https://github.com/jansel
2024-05-05 22:27:46 +00:00
Simon Fan
43a7ab2a21
[compiled autograd] introduce verbose logs, add autograd node info to graph ( #124954 )
...
- sets it as a fake stack trace as we don't have a generic comment feature
- when verbose is disabled, still adds a contextmanager and flag checks. the alternative is to use MACROS, but that wouldn't be usable with TORCH_LOGS
Pull Request resolved: https://github.com/pytorch/pytorch/pull/124954
Approved by: https://github.com/jansel
2024-04-27 01:10:37 +00:00
Sherlock Huang
c2f687f32c
Option to include stride and device annotation in gm.print_readable() ( #123690 )
...
Summary:
Sample output for gm.print_readable(include_stride=True, include_device=True)
```
getitem_21: "i32[1200][1]cuda:0" = auto_functionalized_4[1]
copy_2: "f32[2, 60][60, 1]cuda:1" = ....
```
Test Plan: CI
Differential Revision: D55949129
Pull Request resolved: https://github.com/pytorch/pytorch/pull/123690
Approved by: https://github.com/Chillee
2024-04-11 06:53:10 +00:00
Oguz Ulgen
03b13851d9
[FX] Add side table to FX Graph for O(1) op/target query ( #121565 )
...
Pull Request resolved: https://github.com/pytorch/pytorch/pull/121565
Approved by: https://github.com/jansel
2024-04-07 18:51:05 +00:00
Oguz Ulgen
7c5e29ae71
Back out "Support triton.language.dtype with torch.compile ( #121690 )" ( #122108 )
...
Summary: Some hard to deal with package import/export related problems. Lets revert and start with clean slate.
Test Plan: CI
Differential Revision: D55024877
Pull Request resolved: https://github.com/pytorch/pytorch/pull/122108
Approved by: https://github.com/ezyang
2024-03-18 20:50:28 +00:00