xref: /aosp_15_r20/external/pytorch/aten/src/ATen/metal/Context.cpp (revision da0073e96a02ea20f0ac840b70461e3646d07c45) 
1 #include <atomic>
2 
3 #include <ATen/Tensor.h>
4 #include <ATen/metal/Context.h>
5 
6 namespace at::metal {
7 
8 std::atomic<const MetalInterface*> g_metal_impl_registry;
9 
MetalImplRegistrar(MetalInterface * impl)10 MetalImplRegistrar::MetalImplRegistrar(MetalInterface* impl) {
11   g_metal_impl_registry.store(impl);
12 }
13 
metal_copy_(at::Tensor & self,const at::Tensor & src)14 at::Tensor& metal_copy_(at::Tensor& self, const at::Tensor& src) {
15   auto p = at::metal::g_metal_impl_registry.load();
16   if (p) {
17     return p->metal_copy_(self, src);
18   }
19   AT_ERROR("Metal backend was not linked to the build");
20 }
21 } // namespace at::metal
22 
23 namespace at::native {
is_metal_available()24 bool is_metal_available() {
25   auto p = at::metal::g_metal_impl_registry.load();
26   return p ? p->is_metal_available() : false;
27 }
28 
29 } // namespace at::native
30