Register preprocess in pytorch (pytorch#5350)

Summary:
- Following https://pytorch.org/executorch/stable/kernel-library-custom-aten-kernel.html, use WRAP_TO_ATEN to register preprocess in pytorch
- Create a separate `op_tile_crop_aot.py` that registers the C++ aot library into Python. Inside export_preprocess, use `op_tile_crop_aot.py` instead of `preprocess_custom_ops.py`, which is the pure python lib. Otherwise, we end up loading the C++ library when the python one already exists.

Note, include these PyTorch changes for AOTI export:
pytorch/pytorch#135933

Pull Request resolved: pytorch#5350

Test Plan:
```
>>> import torch
>>> from executorch.extension.llm.custom_ops import sdpa_with_kv_cache  # noqa # usort: skip
>>> x = torch._export.aot_load("/home/lfq/local/executorch/aoti_preprocess.so", "cpu")
>>> img = torch.ones([3, 600, 800])
>>> canvas_size = torch.tensor([448, 448])
>>> target_size = torch.tensor([336, 448])
>>> res = x(img, target_size, canvas_size)
>>> res[0].shape
torch.Size([4, 3, 224, 224])
>>> res[1]
tensor([2, 2])
>>>
```

Reviewed By: larryliu0820

Differential Revision: D62651605

Pulled By: lucylq

fbshipit-source-id: bdf5b46033ebbd73d10307ab58219743a73fd6fd

Author: lucylq

examples/models/flamingo/preprocess/export_preprocess_lib.py

@@ -11,7 +11,7 @@
 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
+from executorch.extension.llm.custom_ops import op_tile_crop_aot  # noqa
 
 from torch.export import Dim, ExportedProgram
 from torchtune.models.clip.inference._transform import _CLIPImageTransform

extension/llm/custom_ops/CMakeLists.txt

@@ -75,6 +75,7 @@ if(EXECUTORCH_BUILD_KERNELS_CUSTOM_AOT)
   add_library(
     custom_ops_aot_lib SHARED
     ${_custom_ops__srcs} ${CMAKE_CURRENT_SOURCE_DIR}/op_sdpa_aot.cpp
+    ${CMAKE_CURRENT_SOURCE_DIR}/op_tile_crop_aot.cpp
     ${CMAKE_CURRENT_SOURCE_DIR}/op_tile_crop.cpp
   )
   target_include_directories(

extension/llm/custom_ops/op_tile_crop_aot.cpp

@@ -0,0 +1,50 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#include <executorch/extension/aten_util/make_aten_functor_from_et_functor.h>
+#include <executorch/extension/kernel_util/make_boxed_from_unboxed_functor.h>
+#include <executorch/extension/llm/custom_ops/op_tile_crop.h>
+
+#include <torch/library.h>
+
+namespace torch {
+namespace executor {
+
+namespace native {
+
+Tensor&
+tile_crop_out_no_context(const Tensor& input, int64_t tile_size, Tensor& out) {
+  exec_aten::RuntimeContext context{};
+  return tile_crop_out_impl(context, input, tile_size, out);
+}
+
+at::Tensor tile_crop_aten(const at::Tensor& input, int64_t tile_size) {
+  // max_num_tiles = 4, num_channels = 3.
+  auto output = at::empty({4, 3, tile_size, tile_size});
+
+  WRAP_TO_ATEN(torch::executor::native::tile_crop_out_no_context, 2)
+  (input, tile_size, output);
+  return output;
+}
+
+} // namespace native
+} // namespace executor
+} // namespace torch
+
+TORCH_LIBRARY(preprocess, m) {
+  m.def("tile_crop(Tensor input, int tile_size) -> Tensor");
+  m.def(
+      "tile_crop.out(Tensor input, int tile_size, *, Tensor(a!) out) -> Tensor(a!)");
+}
+
+TORCH_LIBRARY_IMPL(preprocess, CompositeExplicitAutograd, m) {
+  m.impl("tile_crop", torch::executor::native::tile_crop_aten);
+  m.impl(
+      "tile_crop.out",
+      WRAP_TO_ATEN(torch::executor::native::tile_crop_out_no_context, 2));
+}

extension/llm/custom_ops/op_tile_crop_aot.py

@@ -0,0 +1,38 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+import logging
+from pathlib import Path
+
+import torch
+
+try:
+    tile_crop = torch.ops.preprocess.tile_crop.default
+    assert tile_crop is not None
+except:
+    libs = list(Path(__file__).parent.resolve().glob("libcustom_ops_aot_lib.*"))
+    assert len(libs) == 1, f"Expected 1 library but got {len(libs)}"
+    logging.info(f"Loading custom ops library: {libs[0]}")
+    torch.ops.load_library(libs[0])
+    tile_crop = torch.ops.preprocess.tile_crop.default
+    assert tile_crop is not None
+
+preprocess_ops_lib = torch.library.Library("preprocess", "IMPL")
+
+MAX_NUM_TILES = 4
+
+
+# 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 = 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].
+    ctx = torch._custom_ops.get_ctx()
+    s0 = ctx.create_unbacked_symint()
+    torch._constrain_as_size(s0, 0, MAX_NUM_TILES)
+    return torch.empty([s0, output.size(0), tile_size, tile_size])

extension/llm/custom_ops/targets.bzl

@@ -35,7 +35,10 @@ def define_common_targets():
             name = "custom_ops_aot_lib" + mkl_dep,
             srcs = [
                 "op_sdpa_aot.cpp",
+                "op_tile_crop_aot.cpp",
+                "op_tile_crop.cpp",
             ],
+            headers = ["op_tile_crop.h"],
             visibility = [
                 "//executorch/...",
                 "@EXECUTORCH_CLIENTS",