Constraint based upper bound memory planning. (#264)

Summary:
Pull Request resolved: #264

Before this PR, we treat inputs that passed in for tracing as the upper bound tensor itself. We use their shape as the uppper bound. However, 1. this approach doesn't respect user provided constraints. 2.  this approach cannot handle certain data-dependent op such as tensor.item() because we don't know the exact value at export time.

This PR adds ConstraintBasedSymShapeEval pass, which uses the rangeAnalysis infra in symbolic shape. This change has several implications. User must provide a concrete integer as upper bound for

1. dynamic input tensors.

2. outputs of data dependent operations (e.g. nonzero, .item()).

Otherwise, upper bound memory planning cannot proceed because by default the uppper bound is infinity.

Reviewed By: angelayi

Differential Revision: D49158290

fbshipit-source-id: bab88241e1a14a07903cab5b6e20e290edf2acc5

Author: ydwu4

exir/capture/TARGETS

@@ -37,5 +37,6 @@ python_library(
         "//executorch/exir:pass_manager",
         "//executorch/exir:tracer",
         "//executorch/exir/passes:lib",
+        "//executorch/exir/passes:sym_shape_eval_pass",
     ],
 )

exir/capture/_config.py

@@ -10,6 +10,7 @@
 from executorch.exir.dynamic_shape import DynamicMemoryPlanningMode
 from executorch.exir.pass_manager import PassType
 from executorch.exir.passes import MemoryPlanningPass, ToOutVarPass
+from executorch.exir.passes.sym_shape_eval_pass import HintBasedSymShapeEvalPass
 from executorch.exir.tracer import ExirDynamoConfig
 from torch.fx._compatibility import compatibility
 
@@ -63,3 +64,4 @@ class ExecutorchBackendConfig:
     # If provided, the minimum alignment of delegate data in the program. Must
     # be a power of 2. If not provided, uses the value in the schema file.
     delegate_alignment: Optional[int] = None
+    sym_shape_eval_pass: PassType = HintBasedSymShapeEvalPass()

exir/emit/test/test_emit.py

@@ -21,6 +21,7 @@
 from executorch.exir.emit import emit_program  # noqa
 from executorch.exir.error import InternalError
 from executorch.exir.passes.const_prop_pass import ConstPropPass
+from executorch.exir.passes.sym_shape_eval_pass import ConstraintBasedSymShapeEvalPass
 from executorch.exir.print_program import pretty_print, print_program  # noqa
 from executorch.exir.schema import (
     Bool,
@@ -46,6 +47,8 @@
 from functorch.experimental import control_flow
 from torch import nn
 
+from torch.export import dynamic_dim
+
 
 class TestEmit(unittest.TestCase):
     @classmethod
@@ -1050,6 +1053,41 @@ def make_program(
             merged_program2.execution_plan[2], merged_program.execution_plan[2]
         )
 
+    def test_upper_bound_memory_planning_respect_input_constraints(self) -> None:
+        def func(k: torch.Tensor) -> torch.Tensor:
+            k = torch.cat((k, torch.ones(1, 4)))
+            return k
+
+        k = torch.rand(2, 4)
+        constraints = [
+            dynamic_dim(k, 0) <= 3,
+        ]
+        captured = exir.capture(
+            func,
+            (k,),
+            exir.CaptureConfig(pt2_mode=True, enable_aot=True),
+            constraints=constraints,  # enable_aot=False works
+        )
+        edge = captured.to_edge()
+        from executorch.exir.passes import MemoryPlanningPass
+
+        config = exir.ExecutorchBackendConfig(
+            sym_shape_eval_pass=ConstraintBasedSymShapeEvalPass(),
+            memory_planning_pass=MemoryPlanningPass(
+                memory_planning_algo="greedy",
+                # allow_lifetime_and_storage_overlap: bool = False,
+                alloc_graph_input=True,
+                alloc_graph_output=False,
+            ),
+        )
+
+        exe_prog = edge.to_executorch(config)
+        program = exe_prog.program
+        exir.print_program.pretty_print(exe_prog.program.execution_plan)
+        execution_plan = program.execution_plan[0]
+        self.check_tensor_buffer_loc(0, execution_plan.values, 0, 1, 0)
+        self.check_tensor_buffer_loc(1, execution_plan.values, 0, 1, 48)
+
     def test_emit_prims(self) -> None:
         class Simple(torch.nn.Module):
             def __init__(self) -> None:

exir/passes/sym_shape_eval_pass.py

@@ -10,7 +10,7 @@
 import torch.utils._pytree as pytree
 from executorch.exir.dialects._ops import ops as exir_ops
 from executorch.exir.pass_base import PassBase, PassResult
-from executorch.exir.sym_util import eval_expr, eval_shape
+from executorch.exir.sym_util import eval_expr, eval_shape, eval_upper_bound
 from executorch.exir.tensor import TensorSpec
 from sympy import Integer
 from torch.fx import GraphModule
@@ -96,3 +96,38 @@ def get_val(arg):
                         spec.shape = concrete_shape
                         spec.stride = concrete_spec
         return PassResult(graph_module, True)
+
+
+class ConstraintBasedSymShapeEvalPass(PassBase):
+    """
+    If we enable dynamic shape tracing, a tensor's shape may become a symbolic
+    formula. We should convert those symbolic formula to concrete value for
+    static/upperbound tensors so we can properly do memory planning for them.
+
+    Not inherit from ExportPass since we simply need a way to iterate thru
+    every node's output. PassBase is easier for that purpose.
+    """
+
+    def call(self, graph_module: GraphModule):
+        for subgm in graph_module.modules():
+            if not isinstance(subgm, GraphModule):
+                continue
+            for node in subgm.graph.nodes:
+                for spec in pytree.tree_flatten(node.meta.get("spec", []))[0]:
+                    # Node for function like aten.sym_size does not have spec
+                    if isinstance(spec, TensorSpec):
+                        concrete_shape = [eval_upper_bound(s) for s in spec.shape]
+                        concrete_stride = [eval_upper_bound(s) for s in spec.stride]
+                        if any(not isinstance(s, int) for s in concrete_shape) or any(
+                            not isinstance(s, int) for s in concrete_stride
+                        ):
+                            raise RuntimeError(
+                                f"Cannot evalute the shape upper bound of a dynamic-shaped tensor to a concrete bounded integer. Got tensor spec: {spec}."
+                                f"The upper bound shape we get {concrete_shape}, the upper bound stride we get {concrete_stride}"
+                                "This tensor could either be from 1. a data-dependent operation such as nonzero. Or 2. an input, whose don't have a constraint for the upper bound."
+                                "Please use export's constrain_as_size() or constrain_as_value() apis and set a concrete upper bound to resolve this."
+                            )
+
+                        spec.shape = concrete_shape  # pyre-ignore[8]: Attribute `stride` declared in class `TensorSpec` has type `Tuple[int]` but is used as type `List[Optional[int]]`
+                        spec.stride = concrete_stride  # pyre-ignore[8]: Attribute `stride` declared in class `TensorSpec` has type `Tuple[int]` but is used as type `List[Optional[int]]`
+        return PassResult(graph_module, True)

exir/program/_program.py

@@ -19,7 +19,6 @@
 from executorch.exir.passes import (
     aten_to_edge_passes,
     EdgeToBackendOpsPass,
-    HintBasedSymShapeEvalPass,
     OpReplacePass,
 )
 from executorch.exir.passes.remove_assert_async_pass import RemoveAssertAsyncPass
@@ -309,7 +308,7 @@ def edge_to_executorch_passes(config: ExecutorchBackendConfig) -> List[PassType]
         SpecPropPass(),
         EdgeToBackendOpsPass(),
         RemoveAssertAsyncPass(),
-        HintBasedSymShapeEvalPass(),
+        config.sym_shape_eval_pass,
         config.to_out_var_pass,
         config.memory_planning_pass,
     ]

exir/sym_util.py

@@ -4,11 +4,12 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
-from typing import Optional, Set, Union
+from typing import List, Optional, Set, Union
 
 import sympy
 
 import torch
+from torch.utils._sympy.value_ranges import bound_sympy, ValueRanges
 
 
 def eval_expr(symint: Union[int, torch.SymInt]) -> Optional[int]:
@@ -28,6 +29,32 @@ def eval_expr(symint: Union[int, torch.SymInt]) -> Optional[int]:
     return int(output)
 
 
+def eval_upper_bound(maybe_symint: Union[int, torch.SymInt]) -> Optional[int]:
+    """
+    Evaluate a symint to its uppper bound value. Returns None if symint's symoblic expr's
+    upper bound can not be evaluated to valid integer according to the constraints in shape_env.
+    """
+    if isinstance(maybe_symint, int):
+        return maybe_symint
+    node = maybe_symint.node
+    shape_env = node.shape_env
+    expr = node.expr
+    var_range: ValueRanges = bound_sympy(expr, shape_env.var_to_range)
+    upper_bound = var_range.upper
+    if isinstance(upper_bound, sympy.Integer):
+        concrete_upper = int(var_range.upper)  # pyre-ignore
+        assert isinstance(
+            concrete_upper, int
+        ), f"Expect upper bound to be a concrete int but got {concrete_upper}"
+        return concrete_upper
+    elif isinstance(upper_bound, sympy.oo):
+        return None
+    else:
+        raise RuntimeError(
+            f"Expect upper bound to be sympy.Integer or sympy.oo. but got {upper_bound}"
+        )
+
+
 def eval_shape(shape):
     """
     Shape maybe immutable so we return a new shape. Return None for
@@ -39,6 +66,13 @@ def eval_shape(shape):
     return new_shape
 
 
+def eval_shape_upper_bound(shape) -> List[int]:
+    new_shape = []
+    for _, s in enumerate(shape):
+        new_shape.append(eval_upper_bound(s))
+    return new_shape
+
+
 def collect_free_symbols(shape) -> Set[sympy.Symbol]:
     symset = set()
     for sz in shape: