mirror of
https://github.com/pytorch/pytorch.git
synced 2025-10-20 21:14:14 +08:00
70d7638b0d
Remove some clang-tidy suppression in torch/csrc/jit by applying fixes or refactoring. Pull Request resolved: https://github.com/pytorch/pytorch/pull/152271 Approved by: https://github.com/Skylion007, https://github.com/malfet Co-authored-by: Aaron Gokaslan <aaronGokaslan@gmail.com>
115 lines
3.7 KiB
C++
115 lines
3.7 KiB
C++
#pragma once
|
|
|
|
#include <ATen/core/builtin_function.h>
|
|
#include <ATen/core/stack.h>
|
|
#include <torch/csrc/jit/backends/backend_interface.h>
|
|
#include <torch/custom_class.h>
|
|
|
|
namespace torch::jit {
|
|
namespace {
|
|
inline c10::FunctionSchema getIsAvailableSchema() {
|
|
c10::Argument self("self", c10::AnyType::get());
|
|
c10::Argument available("available", c10::BoolType::get());
|
|
c10::FunctionSchema preprocessor_schema(
|
|
"is_available",
|
|
/*overload_name=*/"",
|
|
/*arguments=*/{self},
|
|
/*returns=*/{available});
|
|
return preprocessor_schema;
|
|
}
|
|
|
|
constexpr static auto kBackendsNamespace = "__backends__";
|
|
|
|
inline c10::FunctionSchema getCompileSchema() {
|
|
c10::Argument self("self", c10::AnyType::get());
|
|
c10::Argument mod("processed", c10::AnyType::get());
|
|
auto any_dict_ty =
|
|
c10::DictType::create(c10::StringType::get(), c10::AnyType::get());
|
|
c10::Argument method_compile_spec("method_compile_spec", any_dict_ty);
|
|
c10::Argument handles("handles", any_dict_ty);
|
|
|
|
c10::FunctionSchema compile_schema(
|
|
"compile",
|
|
/*overload_name=*/"",
|
|
/*arguments=*/{self, mod, method_compile_spec},
|
|
/*returns=*/{handles});
|
|
return compile_schema;
|
|
}
|
|
|
|
inline c10::FunctionSchema getExecuteSchema() {
|
|
auto any_list_ty = c10::ListType::create(c10::AnyType::get());
|
|
c10::Argument self("self", c10::AnyType::get());
|
|
c10::Argument handle("handle", c10::AnyType::get());
|
|
c10::Argument input("input", any_list_ty);
|
|
c10::Argument output("output", any_list_ty);
|
|
return c10::FunctionSchema(
|
|
"execute",
|
|
/*overload_name=*/"",
|
|
/*arguments=*/{self, handle, input},
|
|
/*returns=*/{output});
|
|
}
|
|
|
|
template <typename TBackendInterface>
|
|
std::function<void(Stack&)> getIsAvailableFunc() {
|
|
return [](Stack& stack) {
|
|
auto self = pop(stack).toCustomClass<TBackendInterface>();
|
|
auto ret = self->is_available();
|
|
push(stack, ret);
|
|
};
|
|
}
|
|
|
|
template <typename TBackendInterface>
|
|
std::function<void(Stack&)> getCompileFunc() {
|
|
return [](Stack& stack) {
|
|
auto method_compile_spec = pop(stack).toGenericDict();
|
|
auto processed = pop(stack);
|
|
auto self = pop(stack).toCustomClass<TBackendInterface>();
|
|
auto ret = self->compile(processed, method_compile_spec);
|
|
push(stack, ret);
|
|
};
|
|
}
|
|
|
|
template <typename TBackendInterface>
|
|
std::function<void(Stack&)> getExecuteFunc() {
|
|
return [](Stack& stack) {
|
|
auto args = pop(stack);
|
|
auto handle = pop(stack);
|
|
auto self = pop(stack);
|
|
auto backend = self.toCustomClass<TBackendInterface>();
|
|
auto res = backend->execute(handle, args.toList());
|
|
push(stack, res);
|
|
};
|
|
}
|
|
} // namespace
|
|
|
|
// Static registration API for backends.
|
|
template <class TBackendInterface>
|
|
class backend {
|
|
static_assert(
|
|
std::is_base_of_v<PyTorchBackendInterface, TBackendInterface>,
|
|
"torch::jit::backend<T> requires T to inherit from PyTorchBackendInterface");
|
|
std::string backend_name_;
|
|
|
|
public:
|
|
// Registers a new backend with /p name, and the given /p preprocess
|
|
// function.
|
|
backend(const std::string& name) : backend_name_(name) {
|
|
static auto cls = torch::class_<TBackendInterface>(kBackendsNamespace, name)
|
|
.def(torch::init<>())
|
|
._def_unboxed(
|
|
"is_available",
|
|
getIsAvailableFunc<TBackendInterface>(),
|
|
getIsAvailableSchema())
|
|
._def_unboxed(
|
|
"compile",
|
|
getCompileFunc<TBackendInterface>(),
|
|
getCompileSchema())
|
|
._def_unboxed(
|
|
"execute",
|
|
getExecuteFunc<TBackendInterface>(),
|
|
getExecuteSchema());
|
|
}
|
|
};
|
|
|
|
} // namespace torch::jit
|