mirror of
https://github.com/pytorch/pytorch.git
synced 2025-10-21 05:34:18 +08:00
1a73255102
Fixes #ISSUE_NUMBER Pull Request resolved: https://github.com/pytorch/pytorch/pull/138976 Approved by: https://github.com/Skylion007
40 lines
1.1 KiB
C++
40 lines
1.1 KiB
C++
#pragma once
|
|
|
|
#include <torch/csrc/jit/api/module.h>
|
|
|
|
#include <ATen/core/jit_type.h>
|
|
|
|
#include <functional>
|
|
|
|
namespace torch::jit {
|
|
|
|
using DebugHandleType = int64_t;
|
|
|
|
using NodeToDebugHandle = std::unordered_map<Node*, DebugHandleType>;
|
|
|
|
using BackendDebugHandleGenerator =
|
|
std::function<NodeToDebugHandle(const std::shared_ptr<Graph>&)>;
|
|
|
|
namespace detail {
|
|
|
|
using BackendPreprocessFunction = std::function<c10::IValue(
|
|
const Module&,
|
|
const c10::Dict<IValue, IValue>&,
|
|
const BackendDebugHandleGenerator& generate_debug_handles)>;
|
|
|
|
TORCH_API void registerBackendPreprocessFunction(
|
|
const std::string& name,
|
|
const BackendPreprocessFunction& preprocess);
|
|
|
|
bool hasBackendPreprocessFunction(const std::string& name);
|
|
|
|
BackendPreprocessFunction getBackendPreprocessFunction(const std::string& name);
|
|
|
|
TORCH_API Module codegen_backend_module(
|
|
const std::string& backend_name,
|
|
const Module& orig_module,
|
|
const c10::Dict<IValue, IValue>& method_compile_spec,
|
|
const c10::DictTypePtr& any_dict_ty);
|
|
} // namespace detail
|
|
} // namespace torch::jit
|