Support runtime assertion for inline constraints #100763

ydwu4 · 2023-05-05T22:28:44Z

This pr does the following:

previously, inline constraints is not properly set for tensor output data-dependent ops such as a.nonzero because of its return value is not symint. This pr just uses all the unbacked symbols i.e.those start with "i"/"f" in create_unbacked_sym* functions. Note that these symbols are guaranteed to be a super set of inline user constraints.
add inline assertions support by checking.

Currently, it only deal with tensor, SymInt, SymFloat, SymBool output data-dependent ops and ignore the rest. It's good enough for now as we only have a limited number of data-dependent ops (.item and .nonzero are explicitly tested).

The examples for graph that is added assertions is shown below:

class ExportGraphModule(torch.nn.Module):
    def forward(self, x):
        arg0: i64[s0], = fx_pytree.tree_flatten_spec(([x], {}), self._in_spec)
        nonzero_default: i64[i0, 1] = torch.ops.aten.nonzero.default(arg0);  arg0 = None
        return pytree.tree_unflatten([nonzero_default], self._out_spec)
        
class GraphModule(torch.nn.Module):
    def forward(self, x):
        arg0: i64[s0], = fx_pytree.tree_flatten_spec(([x], {}), self._in_spec)
        sym_size: Sym(s0) = torch.ops.aten.sym_size(arg0, 0)
        nonzero_default: i64[i1, 1] = torch.ops.aten.nonzero.default(arg0);  arg0 = None
        sym_size_1: Sym(i1) = torch.ops.aten.sym_size(nonzero_default, 0)
        ge: Sym(i1 >= 3) = sym_size_1 >= 3
        scalar_tensor_default: f32[] = torch.ops.aten.scalar_tensor.default(ge);  ge = None
        _assert_async_msg = torch.ops.aten._assert_async.msg(scalar_tensor_default, 'nonzero_default.shape[0] is outside of inline constraint [3, 5].');  scalar_tensor_default = None
        le: Sym(i1 <= 5) = sym_size_1 <= 5;  sym_size_1 = None
        scalar_tensor_default_1: f32[] = torch.ops.aten.scalar_tensor.default(le);  le = None
        _assert_async_msg_1 = torch.ops.aten._assert_async.msg(scalar_tensor_default_1, 'nonzero_default.shape[0] is outside of inline constraint [3, 5].');  scalar_tensor_default_1 = None
        return pytree.tree_unflatten([nonzero_default], self._out_spec)

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

pytorch-bot · 2023-05-05T22:28:50Z

🔗 Helpful Links

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

📄 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 17be12d:
💚 Looks good so far! There are no failures yet. 💚

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

tugsbayasgalan · 2023-05-08T17:49:59Z

test/dynamo/test_export.py

+        capture_scalar_outputs=True,
+    )
+    def test_export_preserve_constraints_as_metadata_tensor(self):
+        from torch._export.constraints import constrain_as_value


Can you move it at the top of file? We recently ran into circular import issue, so i wanted to make sure it doesn't happen again?

tugsbayasgalan · 2023-05-08T17:51:04Z

test/export/test_passes.py

+        self.assertEqual(num_assert, 3)
+        self.assertEqual(num_scalar_tensor, 3)
+
+        with self.assertRaisesRegex(RuntimeError, r"^_local_scalar_dense_default.*\[2, 5\].$"):


This error message is bit hard to understand from user's perspective, how about we just dump the stacktrace from user code instead?

tugsbayasgalan · 2023-05-08T17:52:13Z

test/export/test_passes.py

+        new_inp = torch.tensor([1, 1, 1, 1])
+        self.assertEqual(mod(new_inp), new_gm(new_inp))
+
+    # FIXME: support control flow operators for the pass


should be fixed after #100836

tugsbayasgalan · 2023-05-08T17:53:08Z

torch/_export/passes/add_runtime_assertions_for_constraints_pass.py

+        if upper < math.inf:
+            self._inser_assert_async(operator.le, proxy, upper, assert_msg)
+
+    def _inser_assert_async(self, operator, l, r, assert_msg):


nit: _insert_assert_async

tugsbayasgalan · 2023-05-08T17:53:52Z

torch/_export/passes/add_runtime_assertions_for_constraints_pass.py

+        if "val" in meta:
+            val = meta["val"]
+
+            def add_assertions(val):


Can you leave a comment on explaining the high level approach? Especially the part why we need to accumulate messages.

tugsbayasgalan · 2023-05-08T17:54:55Z

torch/_export/passes/add_runtime_assertions_for_constraints_pass.py

+                        cbs, msgs = add_assertions(sym)
+                        for cb, msg in zip(cbs, msgs):
+                            def sym_size_cb(proxy, assert_msg, dim):
+                                dim_proxy = super(AddRuntimeAssertionsForConstraintsPass, self).call_operator(


Why can't it be just super().call_operator()?

I'm not sure.. Seems related to some weird name resolution issue.

ydwu4 · 2023-05-09T02:01:28Z

@pytorchbot merge

pytorchmergebot · 2023-05-09T02:03:28Z

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 May 5, 2023

ydwu4 requested review from tugsbayasgalan and guangy10 and removed request for tugsbayasgalan May 5, 2023 22:28

github-actions bot added ciflow/inductor module: dynamo labels May 5, 2023

ydwu4 requested review from angelayi and tugsbayasgalan May 5, 2023 22:28

tugsbayasgalan approved these changes May 8, 2023

View reviewed changes

ydwu4 force-pushed the runtime_assertions branch from a3967ba to f21baac Compare May 8, 2023 19:08

Support runtime assertion for inline constraints

17be12d

ydwu4 force-pushed the runtime_assertions branch from f21baac to 17be12d Compare May 8, 2023 23:38

pytorch-bot bot added the ciflow/trunk Trigger trunk jobs on your pull request label May 9, 2023

pytorchmergebot added the merging label May 9, 2023

pytorchmergebot added Merged and removed merging labels May 9, 2023

pytorchmergebot closed this in 26cd958 May 9, 2023

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Support runtime assertion for inline constraints #100763

Support runtime assertion for inline constraints #100763

ydwu4 commented May 5, 2023 •

edited by pytorch-bot bot

Loading

pytorch-bot bot commented May 5, 2023 •

edited

Loading

tugsbayasgalan May 8, 2023

tugsbayasgalan May 8, 2023

ydwu4 May 8, 2023

tugsbayasgalan May 8, 2023

tugsbayasgalan May 8, 2023

tugsbayasgalan May 8, 2023

tugsbayasgalan May 8, 2023

ydwu4 May 8, 2023

ydwu4 commented May 9, 2023

pytorchmergebot commented May 9, 2023

Support runtime assertion for inline constraints #100763

Support runtime assertion for inline constraints #100763

Conversation

ydwu4 commented May 5, 2023 • edited by pytorch-bot bot Loading

pytorch-bot bot commented May 5, 2023 • edited Loading

🔗 Helpful Links

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

✅ No Failures

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 commented May 9, 2023

pytorchmergebot commented May 9, 2023

Merge started

ydwu4 commented May 5, 2023 •

edited by pytorch-bot bot

Loading

pytorch-bot bot commented May 5, 2023 •

edited

Loading