xref: /aosp_15_r20/external/pytorch/torch/csrc/jit/passes/metal_rewrite.h (revision da0073e96a02ea20f0ac840b70461e3646d07c45)
1 #pragma once
2 #include <torch/csrc/jit/api/module.h>
3 #include <torch/csrc/jit/ir/ir.h>
4 #include <string>
5 #include <vector>
6 
7 namespace torch::jit {
8 TORCH_API void metalInsertPrePackedOps(std::shared_ptr<Graph>& graph);
9 TORCH_API void metalInsertPrePackedOps(script::Module& module);
10 TORCH_API void metalFusePrePackedConvWithClamp(script::Module& module);
11 TORCH_API void metalFoldPrePackingOps(script::Module& module);
12 TORCH_API script::Module metalOptimizeForMobile(
13     const script::Module& module,
14     const std::vector<std::string>& preserved_methods);
15 } // namespace torch::jit
16