Register preprocess in 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:
https://github.com/pytorch/pytorch/pull/135933
Pull Request resolved: https://github.com/pytorch/executorch/pull/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 L
lucylq committed
8a8e876d9db900f6fb9d5334de24bd7af0a04fa8
Parent: f7954f6
Committed by Facebook GitHub Bot <facebook-github-bot@users.noreply.github.com>
on 9/16/2024, 10:45:00 PM