xref: /aosp_15_r20/external/executorch/extension/llm/custom_ops/op_tile_crop_aot.cpp (revision 523fa7a60841cd1ecfb9cc4201f1ca8b03ed023a)
1 /*
2  * Copyright (c) Meta Platforms, Inc. and affiliates.
3  * All rights reserved.
4  *
5  * This source code is licensed under the BSD-style license found in the
6  * LICENSE file in the root directory of this source tree.
7  */
8 
9 #include <executorch/extension/aten_util/make_aten_functor_from_et_functor.h>
10 #include <executorch/extension/kernel_util/make_boxed_from_unboxed_functor.h>
11 #include <executorch/extension/llm/custom_ops/op_tile_crop.h>
12 
13 #include <torch/library.h>
14 
15 namespace torch {
16 namespace executor {
17 
18 namespace native {
19 
20 Tensor&
tile_crop_out_no_context(const Tensor & input,int64_t tile_size,Tensor & out)21 tile_crop_out_no_context(const Tensor& input, int64_t tile_size, Tensor& out) {
22   exec_aten::RuntimeContext context{};
23   return tile_crop_out_impl(context, input, tile_size, out);
24 }
25 
tile_crop_aten(const at::Tensor & input,int64_t tile_size)26 at::Tensor tile_crop_aten(const at::Tensor& input, int64_t tile_size) {
27   // max_num_tiles = 4, num_channels = 3.
28   auto output = at::empty({4, 3, tile_size, tile_size});
29 
30   WRAP_TO_ATEN(torch::executor::native::tile_crop_out_no_context, 2)
31   (input, tile_size, output);
32   return output;
33 }
34 
35 } // namespace native
36 } // namespace executor
37 } // namespace torch
38 
TORCH_LIBRARY(preprocess,m)39 TORCH_LIBRARY(preprocess, m) {
40   m.def("tile_crop(Tensor input, int tile_size) -> Tensor");
41   m.def(
42       "tile_crop.out(Tensor input, int tile_size, *, Tensor(a!) out) -> Tensor(a!)");
43 }
44 
TORCH_LIBRARY_IMPL(preprocess,CompositeExplicitAutograd,m)45 TORCH_LIBRARY_IMPL(preprocess, CompositeExplicitAutograd, m) {
46   m.impl("tile_crop", torch::executor::native::tile_crop_aten);
47   m.impl(
48       "tile_crop.out",
49       WRAP_TO_ATEN(torch::executor::native::tile_crop_out_no_context, 2));
50 }
51