[RFC]: Functionalize assertions (#103757)

The idea here is to create do a graph mutation to:
* Create an initial dependency token at the beginning of the program.
* Replace non-functional version of assertion statements to functional version.
* The functional version of assertion statement will:
  * Accept a dependency token from output of previous functional assertion statement (or the initial dependency token if there isn't any).
  * Generate a dependency token as the output of assertion statement.
  * Augment the output to include the dependency token generated by last assertion statement.

The goal here is to:
* Form an explicit dependency chain and avoid potential reordering during other passes of compiling.
* Make the assertions a part of overall execution graph will affect the final output (or it could potentially be DCEed).

**NOTE:**
* Currently only cover `contrain_range` and WIP to support other assertions. Send out this PR to collect feedback first.
* Here it only focus on implementation itself. Will integrate it with current export in future PR.

Pull Request resolved: #103757
Approved by: https://github.com/avikchaudhuri

Author: XuanQi

test/export/test_passes.py

@@ -3,16 +3,20 @@
 
 import torch
 from torch.testing._internal.common_utils import run_tests, TestCase
+from torch.testing import FileCheck
 from torch._dynamo.eval_frame import is_dynamo_supported
 from torch._export import export, dynamic_dim
-from torch._export.constraints import constrain_as_value
+from torch._export.constraints import constrain_as_value, constrain_as_size
 from torch._export.passes import (
     ReplaceViewOpsWithViewCopyOpsPass,
 )
 from torch._export.passes.replace_view_ops_with_view_copy_ops_pass import (
     is_view_op,
     get_view_copy_of_view_op,
 )
+from torch._export.passes.functionalize_side_effectful_ops_pass import (
+    _FunctionalizeSideEffectfulOpsPass,
+)
 from functorch.experimental.control_flow import cond
 
 
@@ -335,6 +339,81 @@ def forward(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
         real_result = m(x, y)
         self.assertTrue(torch._dynamo.utils.same(real_result, dynamo_result))
 
+    def test_functionalize_inline_contraints(self) -> None:
+        def f(x):
+            a = x.item()
+            constrain_as_size(a, 4, 7)
+            return torch.empty((a, 4))
+
+        ep = torch._export.export(f, (torch.tensor([7]),))
+        gm = ep.graph_module
+        FileCheck().check_count(
+            "torch.ops.aten.sym_constrain_range.default",
+            1,
+            exactly=True,
+        ).run(gm.code)
+
+        gm = ep.transform(_FunctionalizeSideEffectfulOpsPass()).graph_module
+
+        with self.assertRaisesRegex(
+            RuntimeError,
+            r"_local_scalar_dense_default is outside of inline constraint \[4, 7\]",
+        ) as cm:
+            gm(torch.tensor([20]))
+
+        inp = torch.tensor([5])
+        res, dep_token = gm(inp)
+        self.assertEqual(res.shape, torch.Size([5, 4]))
+        self.assertEqual(dep_token.shape, torch.Size([]))
+
+        FileCheck().check_count(
+            "torch.ops.aten._functional_sym_constrain_range", 1, exactly=True
+        ).run(gm.code)
+        FileCheck().check_count(
+            "torch.ops.aten.sym_constrain_range.default", 0, exactly=True
+        ).run(gm.code)
+
+        dep_token_node = next(n for n in gm.graph.nodes if n.name == "dep_token3")
+        constrain_node = next(
+            n
+            for n in gm.graph.nodes
+            if n.target == torch.ops.aten._functional_sym_constrain_range
+        )
+        self.assertEqual(constrain_node.kwargs["dep_token"], dep_token_node)
+
+    def test_functionalize_input_constraints(self) -> None:
+        def f(x):
+            return x * 2
+
+        inp = torch.zeros(4, 8)
+        ep = torch._export.export(
+            f,
+            (inp,),
+            constraints=[
+                dynamic_dim(inp, 0) < 10,
+                dynamic_dim(inp, 0) >= 3,
+            ],
+        )
+        FileCheck().check_count(
+            "torch.ops.aten._assert_async.msg", 3, exactly=True
+        ).run(ep.graph_module.code)
+
+        gm = ep.transform(_FunctionalizeSideEffectfulOpsPass()).graph_module
+        with self.assertRaisesRegex(
+            RuntimeError,
+            r"Input arg0_1.shape\[0\] is outside of specified dynamic range \[3, 9\]",
+        ):
+            gm(torch.ones(11, 8))
+
+        inp = torch.ones(6, 8)
+        self.assertEqual(gm(inp)[0], f(inp))
+        FileCheck().check_count(
+            "torch.ops.aten._functional_assert_async.msg", 3, exactly=True
+        ).run(gm.code)
+        FileCheck().check_count(
+            "torch.ops.aten._assert_async.msg", 0, exactly=True
+        ).run(gm.code)
+
 
 if __name__ == '__main__':
     run_tests()

torch/_export/pass_base.py

@@ -1,4 +1,5 @@
 import operator
+import traceback
 import typing
 from contextlib import nullcontext
 from typing import Any, Callable, Dict, List, Optional, Tuple, Union
@@ -38,6 +39,11 @@ class ExportPassBase(PassBase):
     transformations.
     """
 
+    @staticmethod
+    def _create_dummy_node_metadata():
+        return NodeMetadata({"stack_trace": traceback.format_exc(-1)})
+
+
     class ExportTracer(PythonKeyTracer):
         """
         Tracer used to create nodes during the retracing part of the ExportPassBase

torch/_export/passes/add_runtime_assertions_for_constraints_pass.py

@@ -13,7 +13,6 @@
 import torch.fx
 from torch.fx.experimental.symbolic_shapes import SymInt
 from torch._export.pass_base import ExportPassBase, ProxyValue, PassResult
-from torch._export.pass_infra.node_metadata import NodeMetadata
 from torch._subclasses.fake_tensor import FakeTensor
 
 
@@ -188,11 +187,6 @@ def _insert_equality_assert_inplace(
                 assert_msg
             )
 
-    def _create_dummy_node_metadata(self):
-        return NodeMetadata({
-            "stack_trace": traceback.format_exc(-1)
-        })
-
     def _insert_assert_async_inplace(self, graph, operator, args, assert_msg):
         """
         Inserts assert_async call_function nodes in the graph. This function is

torch/_export/passes/functionalize_side_effectful_ops_pass.py

@@ -0,0 +1,94 @@
+import copy
+from typing import Dict, Optional, Tuple, List
+
+import torch
+from torch._export.pass_base import ExportPassBase, PassResult, Argument
+from torch._export.pass_infra.node_metadata import NodeMetadata
+from torch._export.pass_infra.proxy_value import ProxyValue
+from torch._ops import OpOverload
+
+aten = torch.ops.aten
+
+_NON_FUNCTIONAL_TO_FUNCTIONAL_SIDE_EFFECTFUL_FUNCS: Dict[OpOverload, OpOverload] = {
+    aten.sym_constrain_range.default: aten._functional_sym_constrain_range,
+    aten._assert_async.msg: aten._functional_assert_async.msg,
+}
+
+
+class _FunctionalizeSideEffectfulOpsPass(ExportPassBase):
+    """
+    Functionalize ops with side effect in graph module by replacing the op with
+    functional version of it. A new dependency token (`dep_token`) will be
+    created and propagated through functional ops to output.
+    For example:
+    ```
+    def f(x):
+        sym_constrain_range(x.shape[0], min=1, max=3)
+        return x.add(3)
+    ```
+    Will be transformed to:
+    ```
+    def f(x):
+        dep_token0 = _make_dep_token()
+        dep_token1 = _functional_sym_constrain_range(
+            x.shape[0], min=1, max=3, dep_token=dep_token0
+        )
+
+        return x.add(3), dep_token1
+    ```
+    """
+
+    def __init__(self) -> None:
+        super().__init__()
+        self._dep_token: Optional[ProxyValue] = None
+        self._next_dep_token_index: Optional[int] = None
+
+    def call(self, graph_module: torch.fx.GraphModule) -> PassResult:
+        # Early return if no non-functional assertions.
+        if not any(
+            n.target in _NON_FUNCTIONAL_TO_FUNCTIONAL_SIDE_EFFECTFUL_FUNCS
+            for n in graph_module.graph.nodes
+        ):
+            return PassResult(graph_module=graph_module, modified=False)
+
+        gm = copy.deepcopy(graph_module)
+        self._dep_token = None
+        self._next_dep_token_index = None
+        return super().call(gm)
+
+    def call_operator(
+        self,
+        op: OpOverload,
+        args: Tuple[Argument, ...],
+        kwargs: Dict[str, Argument],
+        meta: NodeMetadata,
+    ) -> ProxyValue:
+        if op not in _NON_FUNCTIONAL_TO_FUNCTIONAL_SIDE_EFFECTFUL_FUNCS:
+            return super().call_operator(op, args, kwargs, meta)
+
+        if self._dep_token is None:
+            self._dep_token = super().call_operator(
+                aten._make_dep_token,
+                args=(),
+                kwargs={},
+                meta=self._create_dummy_node_metadata(),
+            )
+            self._dep_token.node.name = "dep_token0"
+            self._next_dep_token_index = 1
+
+        self._dep_token = super().call_operator(
+            _NON_FUNCTIONAL_TO_FUNCTIONAL_SIDE_EFFECTFUL_FUNCS[op],
+            args=args,
+            kwargs={**kwargs, "dep_token": self._dep_token},
+            meta=meta,
+        )
+        assert self._next_dep_token_index is not None
+        self._dep_token.node.name = f"dep_token{self._next_dep_token_index}"
+        self._next_dep_token_index += 1
+
+        return self._dep_token
+
+    def output(self, results: List[Argument], meta: NodeMetadata) -> ProxyValue:
+        assert self._dep_token is not None
+
+        return super().output(results=(*results, self._dep_token), meta=meta)

torch/_meta_registrations.py

@@ -376,11 +376,34 @@ def assert_async_meta(val, assert_msg):
     return
 
 
+@register_meta(aten._make_dep_token.default)
+def make_dep_token(
+    *,
+    dtype=None,
+    layout=None,
+    device=None,
+    pin_memory=None,
+    memory_format=None,
+):
+    return torch.empty([], device="meta")
+
+
 @register_meta(aten.sym_constrain_range.default)
 def sym_constrain_range(size, min, max):
     constrain_range(size, min=min, max=max)
 
 
+@register_meta(aten._functional_sym_constrain_range.default)
+def functional_sym_constrain_range(size, min, max, dep_token):
+    aten.sym_constrain_range(size, min=min, max=max)
+    return dep_token
+
+
+@register_meta(aten._functional_assert_async.msg)
+def functional_assert_async_meta(val, assert_msg, dep_token):
+    return dep_token
+
+
 # From aten/src/ATen/native/LinearAlgebraUtils.h
 def squareCheckInputs(self: Tensor, f_name: str):
     assert (