Skip to content

Commit fcefd10

Browse files
tugsbayasgalanfacebook-github-bot
authored andcommitted
Move legacy range constraint calculator to executorch to unblock pytorch CI (#2925)
Summary: Pull Request resolved: #2925 To unblock: D55730289 Reviewed By: angelayi Differential Revision: D55892116 fbshipit-source-id: 657a3f580871ed24c3095deccc8111d0f3002fa2
1 parent 3708e74 commit fcefd10

File tree

1 file changed

+33
-1
lines changed

1 file changed

+33
-1
lines changed

exir/program/_program.py

Lines changed: 33 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -48,7 +48,6 @@
4848
unsafe_remove_auto_functionalized_pass,
4949
)
5050
from torch.export.exported_program import (
51-
_get_updated_range_constraints,
5251
ConstantArgument,
5352
ExportGraphSignature,
5453
InputKind,
@@ -64,6 +63,39 @@
6463
Val = Any
6564

6665

66+
def _get_updated_range_constraints(gm):
67+
def get_shape_env(gm):
68+
vals = [
69+
node.meta["val"]
70+
for node in gm.graph.nodes
71+
if node.meta.get("val", None) is not None
72+
]
73+
from torch._guards import detect_fake_mode # type: ignore[21]
74+
75+
fake_mode = detect_fake_mode(vals)
76+
if fake_mode is not None:
77+
return fake_mode.shape_env
78+
for v in vals:
79+
if isinstance(v, torch.SymInt):
80+
return v.node.shape_env
81+
82+
shape_env = get_shape_env(gm)
83+
if shape_env is None:
84+
return {}
85+
range_constraints = {
86+
k: v
87+
for k, v in shape_env.var_to_range.items()
88+
if k not in shape_env.replacements
89+
}
90+
# Only when we have an unbacked symint, and it's used as constructor inputs,
91+
# runtime_var_to_range will make a difference compated to var_to_range.
92+
# e.g. [2, oo) -> [0, oo)
93+
for k, v in shape_env.var_to_range.items():
94+
if k not in shape_env.replacements:
95+
range_constraints[k] = v
96+
return range_constraints
97+
98+
6799
def _get_updated_graph_signature(
68100
old_signature: ExportGraphSignature,
69101
new_gm: torch.fx.GraphModule,

0 commit comments

Comments
 (0)