Make DimConstraints create actionable message #100103

ydwu4 · 2023-04-26T16:28:55Z

This pr makes summary of dimension constraints actionable. Before the pr, it will print:

torch.fx.experimental.symbolic_shapes: [WARNING] Summary of dimension constraints:
The following dimensions have been specialized and CANNOT be dynamic.
NOTE: Specializations will happen by default with `assume_static_by_default=True`.
        L['c'].size()[1] == 3
        L['a'].size()[2] == 3
        L['a'].size()[1] == 3
        L['b'].size()[2] == 2
        L['b'].size()[1] == 2
        L['c'].size()[2] == 3

The following dimensions CAN be dynamic.
You can use the following code to specify the constraints they must satisfy:
'''
constraints=[
        dynamic_dim(L['c'], 0) == dynamic_dim(L['a'], 0),
        2 <= dynamic_dim(L['b'], 0),
        2 <= dynamic_dim(L['a'], 0),
]
'''

Users need to initialize the L environment manually and copy the constraints over. After the pr, we have:

[2023-04-26 05:43:12,849] torch._dynamo.eval_frame: [WARNING] Summary of dimension constraints:
The following dimensions have been specialized and CANNOT be dynamic.
NOTE: Specializations will happen by default with `assume_static_by_default=True`.
'''
def specializations(a, b, c):
    return (a.size()[2] == 3 and
    c.size()[1] == 3 and
    a.size()[1] == 3 and
    c.size()[2] == 3 and
    b.size()[2] == 2 and
    b.size()[1] == 2)
   
'''

The following dimensions CAN be dynamic.
You can use the following code to specify the constraints they must satisfy:
'''
def specify_constraints(a, b, c):
    return [
        2 <= dynamic_dim(b, 0),
        dynamic_dim(c, 0) == dynamic_dim(a, 0),
        2 <= dynamic_dim(a, 0),
    ]
'''

, where dynamic_constraints has the same input signature as users code. This allow users to copy-paste and run the code to generate the constraints before exporting as shown below:

def specify_constraints(a, b, c):
    return [
        2 <= dynamic_dim(b, 0),
        dynamic_dim(c, 0) == dynamic_dim(a, 0),
        2 <= dynamic_dim(a, 0),
    ]
torch._dynamo.export(my_dyn_fn, x, y, z, constraints=specify_constriants(x, y, z))

Implementation-wise, this pr also

changes shape_env.produce_guards to produce_guards_and_constraints,
adds contraints_export_fn hooks,
The purpose is to surface the DimConstraints to dynamo.export, where we could reliably get the original function's signature.

The alternative to the above is to get the function signature before creating SHAPE_ENV guard (https://github.com/pytorch/pytorch/blob/main/torch/_dynamo/output_graph.py#L227) and pass it to DimConstraints, but I couldn't recover the signature before creating SHAPE_ENV because the frame's f_globals/locals don't contain the original function.

cc @soumith @voznesenskym @penguinwu @anijain2305 @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @Xia-Weiwen @wenzhe-nrv @jiayisunx @desertfire

pytorch-bot · 2023-04-26T16:28:58Z

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/100103

📄 Preview Python docs built from this PR
📄 Preview C++ docs built from this PR
❓ Need help or want to give feedback on the CI? Visit the bot commands wiki or our office hours

Note: Links to docs will display an error until the docs builds have been completed.

✅ No Failures

As of commit 9e0d86a:
💚 Looks good so far! There are no failures yet. 💚

This comment was automatically generated by Dr. CI and updates every 15 minutes.

avikchaudhuri

I'll wait for the tests to finish before approving.

avikchaudhuri · 2023-04-26T22:24:17Z

test/test_dynamic_shapes.py

+        action_code = dim_constraints.prettify_results(inspect.signature(dummy_f))
+        static_code, dynamic_code = re.findall(r"```(.*?)```", action_code, re.DOTALL)
+        expected_static = '''
+def static_constraints(x0, x1, x2, x3, x4, x5, x6, x7, x8, x9, x11, x12):


This function doesn't do anything. One idea would be to join the constraints with and (still on separate lines). Then perhaps the user can call it to check specialized inputs.

While you're at it, call this specializations and call the other specify_constraints, since the API calls the result constraints.

yeah, sounds good to me. For specializations call, it seems to possibly contain assignment operator "=": https://github.com/pytorch/pytorch/blob/main/torch/fx/experimental/symbolic_shapes.py#L1638. Should we just turn it into "==" and in the _static_results?

avikchaudhuri · 2023-04-26T22:24:43Z

test/test_proxy_tensor.py

-        gm.shape_env.produce_guards(fx_placeholder_vals(gm), names, _simplified=True, constraint_inputs=None)
-    )
+    guards = gm.shape_env.produce_guards_and_constraints(fx_placeholder_vals(gm),
+                                                         names, _simplified=True, constraint_inputs=None)[0]


funny formatting

Yeah, I can change it to other format. lintrunner -a somehow didn't work.

avikchaudhuri · 2023-04-26T22:27:55Z

torch/fx/experimental/symbolic_shapes.py

@@ -1994,11 +2008,13 @@ def create_symbol(
            self.var_to_sources[r].append(source)
        return r

-    # Generates a list of guards strings which, when evaluated in a context that
+    # Generates a tuple of 1. a list of guards strings which, when evaluated in a context that


Maybe break 1. and 2. into separate lines, always find it hard to read otherwise.

guangy10 · 2023-04-27T01:35:27Z

torch/_dynamo/guards.py

-        guards = output_graph.shape_env.produce_guards(
+        (
+            guards,
+            shape_constraints,


These are only input shape constraints, not include inline shape constraints due to unbacked symint right?

Yep i think we should name it dim_constraints to be consistent with original naming.

guangy10 · 2023-04-27T01:36:33Z

torch/fx/experimental/symbolic_shapes.py

        if self._static_results:
+            sorted_static_results = sorted(self._static_results)


Is it sorted by tensor first then by dim?

It's just sorted lexicographically to have deterministic results.

torch/fx/experimental/symbolic_shapes.py

tugsbayasgalan · 2023-04-27T15:09:28Z

torch/_dynamo/eval_frame.py

@@ -800,6 +805,7 @@ def result_capturing_wrapper(*graph_inputs):
            hooks=Hooks(
                guard_export_fn=guard_export_print,
                guard_fail_fn=None,
+                constraint_export_fn=constraint_export,


I feel like we don't need a hook here. In here (https://github.com/pytorch/pytorch/pull/100103/files#diff-c663df72d7a648204ff256e604fe3c9b393e7e2d58204e3bc818d90a99712911L768), you already have access to fake tensors where you can just access the shape_env directly.

tugsbayasgalan · 2023-04-27T15:10:40Z

torch/_dynamo/guards.py

-        guards = output_graph.shape_env.produce_guards(
+        (
+            guards,
+            shape_constraints,


Yep i think we should name it dim_constraints to be consistent with original naming.

tugsbayasgalan · 2023-04-27T15:11:16Z

torch/_dynamo/hooks.py

@@ -7,3 +7,4 @@
 class Hooks:
    guard_export_fn: Optional[Callable[[Set["Guard"]], None]] = None
    guard_fail_fn: Optional[Callable[[Tuple["GuardFail"]], None]] = None
+    constraint_export_fn: Optional[Callable[["DimConstraints"], None]] = None


doesn't seem like we need it if we just access the fake tensor and get shape env from there.

So we just add dim_constraints as part of ShapeEnv, right? That's sounds better to me.

tugsbayasgalan · 2023-04-27T15:13:24Z

test/test_dynamic_shapes.py

+    x4.size()[0] == 8
+    x5.size()[1] = 22
+    x7.size()[3] == 96
+    x8.size()[1] == 22


I am bit confused why it needs to be wrapped in function. We don't expect user to execute this function right? Maybe it is better if we just print it as it is

The naming of function signature may be different from user program:

args = [t1, t2, t3] def f(x0, x1, x2): # computations on x0, x1, x2 pass torchdynamo.export(f, *args)

In this case, directly returned constraints may look like this:

constraints = [ 2 <= dynamic_dim(x1, 1), dynamic_dim(x2, 1) == dynamic_dim(x1, 1), dynamic_dim(x1, 1) == dynamic_dim(x0, 1), ]

It's not directly copy-pastable: users need to find the mapping between t1, t2, t3 and x0, x1, x2. If we export a function with the same signature, we could directly call specify_constriants like below:

def specify_constraints(x0, x1, x2): return [ 2 <= dynamic_dim(x1, 1), dynamic_dim(x2, 1) == dynamic_dim(x1, 1), dynamic_dim(x1, 1) == dynamic_dim(x0, 1), ] constraints = specify_constraitns(*args)

We could use the same calling convention as original f. Does it make sense?

I guess I am curious why the static constraints like x7.size()[3] == 96 are wrapped inside function that doesn't return anything. For dynamic constraints, it totally makes sense :)

Ohh, sorry I thought you were referring to static ones. For the static constraints, we'll call it specializations and make it return a boolean in this revision. It seems to be prettier than the L[""] and more consistent with specify_constraints.

tugsbayasgalan · 2023-04-27T17:40:51Z

torch/_dynamo/eval_frame.py

@@ -827,6 +827,24 @@ def result_capturing_wrapper(*graph_inputs):
    assert out_guards is not None, "Failed to produce guards during tracing"
    assert fake_mode is not None

+    if tracing_mode == "symbolic":


I think it will always be a fake tensor received to the export backend? So probably no need for check here

ok, let's just use the already instantiated fake_mode then.

tugsbayasgalan · 2023-04-27T17:43:18Z

torch/_dynamo/eval_frame.py

+                if "example_value" in node.meta:
+                    return node.meta[
+                        "example_value"
+                    ].fake_mode.shape_env.dim_constraints


Why not just do it here? (

pytorch/torch/_dynamo/eval_frame.py

Line 789 in ae0eb23

fake_mode = _guards.detect_fake_mode(inner_example_inputs)

)

Thanks for pointing out! That makes the code even simpler.

guangy10 · 2023-04-27T19:12:50Z

test/test_dynamic_shapes.py

+def specializations(x0, x1, x2, x3, x4, x5, x6, x7, x8, x9, x11, x12):
+    return (x0.size()[0] == 8 and
+    x1.size()[2] == 96 and
+    x11.size()[1] == 1 and
+    x12.size()[2] == 3 and
+    x2.size()[0] == 8 and
+    x3.size()[0] == 8 and
+    x4.size()[0] == 8 and
+    x5.size()[1] == 22 and
+    x7.size()[3] == 96 and
+    x8.size()[1] == 22)


This is to inform users that those dims are specialized. Users don't need to paste and use it anywhere explicitly right?

@tugsbayasgalan These are also turned into runtime check correct? If it's not stored in metadata, how do you find this info

ydwu4 · 2023-04-27T21:20:27Z

@pytorchbot merge

pytorchmergebot · 2023-04-27T21:24:13Z

Merge started

Your change will be merged once all checks pass (ETA 0-4 Hours).

Learn more about merging in the wiki.

Questions? Feedback? Please reach out to the PyTorch DevX Team

Advanced Debugging

Check the merge workflow status
here

pytorch-bot bot added the release notes: fx release notes category label Apr 26, 2023

github-actions bot added ciflow/inductor module: dynamo labels Apr 26, 2023

ydwu4 force-pushed the actionable_constraints branch 3 times, most recently from 55688d2 to 2765afe Compare April 26, 2023 21:55

avikchaudhuri reviewed Apr 26, 2023

View reviewed changes

guangy10 reviewed Apr 27, 2023

View reviewed changes

torch/fx/experimental/symbolic_shapes.py Show resolved Hide resolved

guangy10 approved these changes Apr 27, 2023

View reviewed changes

tugsbayasgalan reviewed Apr 27, 2023

View reviewed changes

ydwu4 force-pushed the actionable_constraints branch from 2765afe to 1cf6e21 Compare April 27, 2023 17:11

ydwu4 requested review from avikchaudhuri, tugsbayasgalan and guangy10 April 27, 2023 17:11

ydwu4 added the ciflow/trunk Trigger trunk jobs on your pull request label Apr 27, 2023

ydwu4 requested a review from gmagogsfm April 27, 2023 17:17

tugsbayasgalan reviewed Apr 27, 2023

View reviewed changes

Make DimConstraints create actionable message

9e0d86a

ydwu4 force-pushed the actionable_constraints branch from 1cf6e21 to 9e0d86a Compare April 27, 2023 18:34

ydwu4 requested a review from tugsbayasgalan April 27, 2023 18:35

tugsbayasgalan approved these changes Apr 27, 2023

View reviewed changes

guangy10 reviewed Apr 27, 2023

View reviewed changes

pytorchmergebot added the merging label Apr 27, 2023

pytorchmergebot added the Merged label Apr 27, 2023

pytorchmergebot closed this in 61917a0 Apr 27, 2023

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Make DimConstraints create actionable message #100103

Make DimConstraints create actionable message #100103

ydwu4 commented Apr 26, 2023 •

edited

pytorch-bot bot commented Apr 26, 2023 •

edited

avikchaudhuri left a comment

avikchaudhuri Apr 26, 2023

ydwu4 Apr 26, 2023

avikchaudhuri Apr 26, 2023

ydwu4 Apr 26, 2023

avikchaudhuri Apr 26, 2023

guangy10 Apr 27, 2023

tugsbayasgalan Apr 27, 2023

ydwu4 Apr 27, 2023

guangy10 Apr 27, 2023

ydwu4 Apr 27, 2023

tugsbayasgalan Apr 27, 2023

tugsbayasgalan Apr 27, 2023

tugsbayasgalan Apr 27, 2023

ydwu4 Apr 27, 2023

tugsbayasgalan Apr 27, 2023

ydwu4 Apr 27, 2023 •

edited

tugsbayasgalan Apr 27, 2023

ydwu4 Apr 27, 2023

tugsbayasgalan Apr 27, 2023

ydwu4 Apr 27, 2023 •

edited

tugsbayasgalan Apr 27, 2023

ydwu4 Apr 27, 2023

guangy10 Apr 27, 2023

ydwu4 commented Apr 27, 2023

pytorchmergebot commented Apr 27, 2023

		if self._static_results:
		sorted_static_results = sorted(self._static_results)

Make DimConstraints create actionable message #100103

Make DimConstraints create actionable message #100103

Conversation

ydwu4 commented Apr 26, 2023 • edited

pytorch-bot bot commented Apr 26, 2023 • edited

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/100103

✅ No Failures

avikchaudhuri left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

ydwu4 Apr 27, 2023 • edited

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

ydwu4 Apr 27, 2023 • edited

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

ydwu4 commented Apr 27, 2023

pytorchmergebot commented Apr 27, 2023

Merge started

ydwu4 commented Apr 26, 2023 •

edited

pytorch-bot bot commented Apr 26, 2023 •

edited

ydwu4 Apr 27, 2023 •

edited

ydwu4 Apr 27, 2023 •

edited