1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48
|
#include <torch/csrc/lazy/backend/backend_interface.h>
#include <torch/csrc/lazy/core/internal_ops/ltc_ops.h>
#include <utility>
namespace torch::lazy {
namespace {
std::atomic<const BackendImplInterface*> backend_impl_registry;
} // namespace
bool hasBackend() {
return !!backend_impl_registry.load();
}
const BackendImplInterface* getBackend() {
auto* interface = backend_impl_registry.load();
TORCH_CHECK(interface, "Lazy tensor backend not registered.");
return interface;
}
BackendRegistrar::BackendRegistrar(
const BackendImplInterface* backend_impl_interface) {
backend_impl_registry.store(backend_impl_interface);
}
// Get IrBuilder from backend. Use TorchScriptIrBuilder by default
const IrBuilder* getIrBuilder() {
static const IrBuilder* builder = getBackend()->GetIrBuilder();
return builder;
}
std::unique_ptr<LoweringContext> LoweringContext::Create(
const std::string& name,
BackendDevice device,
c10::ArrayRef<const Node*> post_order,
Util::EmissionMap emit_status) {
return getBackend()->CreateLoweringContext(
name, std::move(device), post_order, std::move(emit_status));
}
std::unique_ptr<LoweringContext> LoweringContext::Create(
const std::string& name,
BackendDevice device) {
return getBackend()->CreateLoweringContext(name, std::move(device));
}
} // namespace torch::lazy
|