Rename SymShapeEval pass to HintBasedSymShapeEval pass and add warning for it. (#377)

Summary:
Pull Request resolved: #377

As titled. The original SymShapeEvalPass doesn't respect user constraints and is based on hints of the symbolics. This diff changes its name to HintBasedSymShapeEval pass to make it clearer.

Reviewed By: angelayi

Differential Revision: D49335965

fbshipit-source-id: 8dbacc666276c46653dbaa2b8d23c81bee70c507

Author: ydwu4

exir/passes/__init__.py

@@ -51,7 +51,7 @@
 from executorch.exir.passes.replace_sym_size_op_pass import ReplaceSymSizeOpPass
 from executorch.exir.passes.scalar_to_tensor_pass import ScalarToTensorPass
 from executorch.exir.passes.spec_prop_pass import SpecPropPass
-from executorch.exir.passes.sym_shape_eval_pass import SymShapeEvalPass
+from executorch.exir.passes.sym_shape_eval_pass import HintBasedSymShapeEvalPass
 from executorch.exir.passes.sym_to_tensor_pass import SymToTensorPass
 from torch import fx
 from torch._subclasses import FakeTensor
@@ -65,7 +65,7 @@
     "OpReplacePass",
     "EdgeToBackendOpsPass",
     "MemoryFormatOpsPass",
-    "SymShapeEvalPass",
+    "HintBasedSymShapeEvalPass",
 ]
 
 Argument = Optional[
@@ -510,5 +510,5 @@ def propagate_dynamic_shape(
     """
     return [
         SpecPropPass(),
-        SymShapeEvalPass(),
+        HintBasedSymShapeEvalPass(),
     ]

exir/passes/sym_shape_eval_pass.py

@@ -32,12 +32,23 @@ def nonzero(args, kwargs) -> List[Optional[int]]:
     return [eval_expr(args[0].shape[0]), len(args[0].shape)]
 
 
-class SymShapeEvalPass(PassBase):
+class HintBasedSymShapeEvalPass(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.
 
+    HintBasedSymShapeEvalPass evalutes the symbolic expression of shapes based
+    on its hint, which is a concrete integer that backs the sym expression. The original
+    hint comes from the sizes of the inputs that user uses for tracing and hints of
+    symbolic expressions are propagated via meta tensor computation.
+    For example, when export f(x), we use x = torch.ones(3, 4) as an exmaple input to f and
+    suppose we constrain both dimensions of x as dynamic. We'll have two symbols s0, s1 created
+    and they are backed up with hints 3 and 4 respectively. If there is a y = x[0] operation in f,
+    the shape of y is inferred to be s1, which is backed up with hint 4.
+
+    Warning: if you're using torch.export with constrain API, this method doesn't respect the input constraints.
+
     Not inherit from ExportPass since we simply need a way to iterate thru
     every node's output. PassBase is easier for that purpose.
     """

exir/program/_program.py

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

exir/tests/test_dynamic_shape_propagation.py

@@ -7,7 +7,7 @@
 from unittest import TestCase
 
 from executorch import exir
-from executorch.exir.passes import DebugPass, SpecPropPass, SymShapeEvalPass
+from executorch.exir.passes import DebugPass, HintBasedSymShapeEvalPass, SpecPropPass
 from executorch.exir.tests.models import Repeat
 
 
@@ -23,7 +23,7 @@ def test_repeat(self):
             exir.CaptureConfig(enable_dynamic_shape=True),
         ).to_edge(exir.EdgeCompileConfig(_check_ir_validity=False))
 
-        new_prog = prog.transform(SpecPropPass(), SymShapeEvalPass())
+        new_prog = prog.transform(SpecPropPass(), HintBasedSymShapeEvalPass())
 
         gm = new_prog.exported_program.graph_module
 

exir/tests/test_passes.py

@@ -24,11 +24,11 @@
 from executorch.exir.passes import (
     dead_code_elimination_pass,
     DebugPass,
+    HintBasedSymShapeEvalPass,
     MemoryPlanningPass,
     propagate_dynamic_shape,
     RemoveNoopPass,
     ReplaceSymSizeOpPass,
-    SymShapeEvalPass,
     ToOutVarPass,
 )
 from executorch.exir.passes.const_prop_pass import ConstPropPass
@@ -586,7 +586,7 @@ def test_alloc_node_spec(self) -> None:
         ).to_edge(exir.EdgeCompileConfig(_check_ir_validity=False))
         passes = [
             SpecPropPass(),
-            SymShapeEvalPass(),
+            HintBasedSymShapeEvalPass(),
             ToOutVarPass(),
             MemoryPlanningPass("greedy"),
         ]