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