run preprocess pte file

Author: lucylq

examples/models/flamingo/export_preprocess_lib.py

@@ -8,6 +8,7 @@
 
 import torch
 from executorch.exir import EdgeCompileConfig, ExecutorchBackendConfig, to_edge
+from executorch.exir.passes.sym_shape_eval_pass import ConstraintBasedSymShapeEvalPass
 from executorch.exir.program._program import ExecutorchProgramManager
 
 from executorch.extension.llm.custom_ops import preprocess_custom_ops  # noqa
@@ -76,5 +77,9 @@ def lower_to_executorch_preprocess(
         exported_program, compile_config=EdgeCompileConfig(_check_ir_validity=False)
     )
 
-    et_program = edge_program.to_executorch(ExecutorchBackendConfig())
+    et_program = edge_program.to_executorch(
+        ExecutorchBackendConfig(
+            sym_shape_eval_pass=ConstraintBasedSymShapeEvalPass(),
+        )
+    )
     return et_program

examples/models/flamingo/run_preprocess.py

@@ -1,17 +1,18 @@
+import torch
+from executorch.extension.pybindings.portable_lib import _load_for_executorch
+from executorch.extension.pybindings import portable_lib
+from executorch.extension.llm.custom_ops import sdpa_with_kv_cache
+model = _load_for_executorch("preprocess.pte")
 
+image = torch.ones(3, 800, 600)
+target_size = torch.tensor([448, 336])
+canvas_size = torch.tensor([448, 448])
 
-class NativePreprocessRunner():
-  """
-  Runs preprocess via ExecuTorch with provided pte file.
-  """
-  def __init__(self, pte_file):
-    super().__init__()
-    self.model = _load_for_executorch(pte_file)
-  
-  def forward(
-    self,
-    image: torch.Tensor,
-    target_size: torch.Tensor,
-    canvas_size: torch.Tensor,
-  ):
-    return self.model.forward((image, target_size, canvas_size))
+res = model.forward((image, target_size, canvas_size))
+print(res[0].shape)
+
+image2 = torch.ones(3, 200, 200)
+target_size2 = torch.tensor([224, 224])
+canvas_size2 = torch.tensor([224, 224])
+res2 = model.forward((image2, target_size2, canvas_size2))
+print(res2[0].shape)

extension/llm/custom_ops/preprocess_custom_ops.py

@@ -15,6 +15,10 @@
 
 # Register and define tile_crop and out variant.
 preprocess_op_lib.define("tile_crop(Tensor input, int tile_size) -> Tensor")
+from torch.fx.experimental.symbolic_shapes import _constrain_range_for_size, ShapeEnv
+
+# Keep this in sync with model config.
+MAX_NUM_TILES = 4
 
 
 @impl(preprocess_op_lib, "tile_crop", dispatch_key="CompositeExplicitAutograd")
@@ -56,6 +60,11 @@ def tile_crop_out_impl(
 # Register meta kernel to prevent export tracing into the tile_crop impl.
 @torch.library.register_fake("preprocess::tile_crop")
 def tile_crop(output: torch.Tensor, tile_size: int) -> torch.Tensor:
-    # Returned tensor is of size [n, 3, 224, 224], where n is the number of tiles.
-    # We should export with n = max_num_tiles. Set 50 for now.
-    return torch.empty([50, output.size(0), 224, 224])
+    # Returned tensor is of size [n, 3, 224, 224], where n = number of tiles.
+    # Use an unbacked symint to create an upper-bounded dynamic shape output.
+    # Otherwise, output is set to a static shape, and we can only output
+    # tensors of shape [MAX_NUM_TILES, 3, 224, 224].
+    shape_env = ShapeEnv()
+    s0 = shape_env.create_unbacked_symint()
+    _constrain_range_for_size(s0, 0, MAX_NUM_TILES)
+    return torch.empty([s0, output.size(0), tile_size, tile_size])