xref: /aosp_15_r20/external/pytorch/torch/csrc/jit/tensorexpr/external_functions_core.h (revision da0073e96a02ea20f0ac840b70461e3646d07c45)
1 #pragma once
2 
3 #include <ATen/ATen.h>
4 #include <ATen/Parallel.h>
5 #include <torch/csrc/Export.h>
6 #include <cstdint>
7 
8 namespace torch::jit::tensorexpr {
9 
10 #ifdef C10_MOBILE
11 extern "C" {
12 #endif
13 void DispatchParallel(
14     int8_t* func,
15     int64_t start,
16     int64_t stop,
17     int8_t* packed_data) noexcept;
18 
19 TORCH_API void nnc_aten_free(size_t bufs_num, void** ptrs) noexcept;
20 
21 #ifdef C10_MOBILE
22 } // extern "C"
23 #endif
24 
25 } // namespace torch::jit::tensorexpr
26