File size: 705 Bytes
c1af2fa |
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 |
#ifndef MetalContext_h
#define MetalContext_h
#include <atomic>
#include <ATen/Tensor.h>
namespace at::metal {
struct MetalInterface {
virtual ~MetalInterface() = default;
virtual bool is_metal_available() const = 0;
virtual at::Tensor& metal_copy_(at::Tensor& self, const at::Tensor& src)
const = 0;
};
extern std::atomic<const MetalInterface*> g_metal_impl_registry;
class MetalImplRegistrar {
public:
explicit MetalImplRegistrar(MetalInterface*);
};
at::Tensor& metal_copy_(at::Tensor& self, const at::Tensor& src);
} // namespace at::metal
namespace at::native {
bool is_metal_available();
} // namespace at::native
#endif /* MetalContext_h */
|