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
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
|
import torch._C
import contextlib
import ctypes
import sys
import types
import os.path
import torch.jit
# Query `hasattr` only once.
_SET_GLOBAL_FLAGS = hasattr(sys, 'getdlopenflags') and hasattr(sys, 'setdlopenflags')
@contextlib.contextmanager
def dl_open_guard():
"""
Context manager to set the RTLD_GLOBAL dynamic linker flag while we open a
shared library to load custom operators.
"""
if _SET_GLOBAL_FLAGS:
old_flags = sys.getdlopenflags()
sys.setdlopenflags(old_flags | ctypes.RTLD_GLOBAL)
yield
if _SET_GLOBAL_FLAGS:
sys.setdlopenflags(old_flags)
# _OpNamespace is a subclass of ModuleType because the torch script
# allows attribute lookups on modules only. Since we want torch.ops.foo.bar()
# to work from script, we need to ensure ops and foo are modules
class _OpNamespace(types.ModuleType):
"""
An op namespace to dynamically bind Operators into Python.
Say a user has created a custom Operator called "my_namespace::my_op". To
call this op, the user will write torch.ops.my_namespace.my_op(...).
At startup, this operation will not yet be bound into Python. Instead, the
following sequence of magic tricks will occur:
1. `torch.ops.my_namespace` will invoke the `__getattr__` magic method
on the `torch.ops` object, which will create a new `_OpNamespace`
object called `my_namespace` and set it as an attribute on the `ops`
object.
2. `torch.ops.my_namespace.my_op` will then invoke `__getattr__` on
the `my_namespace` object, which will retrieve the operation via
`torch.get_operation`, a function bound from C++, and then in a similar
fashion bind this new object onto the `my_namespace` object.
3. `torch.ops.my_namespace.my_op(...)` then calls this new operation
and subsequent accesses will incur no further lookup (the namespace and
operation will already exist).
"""
def __init__(self, name):
super(_OpNamespace, self).__init__('torch.ops.' + name)
self.name = name
def __getattr__(self, op_name):
# Get the op `my_namespace::my_op` if available. This will also check
# for overloads and raise an exception if there are more than one.
qualified_op_name = '{}::{}'.format(self.name, op_name)
op = torch._C._jit_get_operation(qualified_op_name)
# let the script frontend know that op is identical to the builtin op
# with qualified_op_name
torch.jit._register_builtin(op, qualified_op_name)
setattr(self, op_name, op)
return op
class _Ops(types.ModuleType):
def __init__(self):
super(_Ops, self).__init__('torch.ops')
self.loaded_libraries = set()
def __getattr__(self, name):
# Here we are creating `torch.ops.my_namespace`
namespace = _OpNamespace(name)
setattr(self, name, namespace)
return namespace
def load_library(self, path):
"""
Loads a shared library from the given path into the current process.
The library being loaded may run global initialization code to register
custom operators with the PyTorch JIT runtime. This allows dynamically
loading custom operators. For this, you should compile your operator
and the static registration code into a shared library object, and then
call ``torch.ops.load_library('path/to/libcustom.so')`` to load the
shared object.
After the library is loaded, it is added to the
``torch.ops.loaded_libraries`` attribute, a set that may be inspected
for the paths of all libraries loaded using this function.
Arguments:
path (str): A path to a shared library to load.
"""
path = os.path.realpath(path)
with dl_open_guard():
# Import the shared library into the process, thus running its
# static (global) initialization code in order to register custom
# operators with the JIT.
ctypes.CDLL(path)
self.loaded_libraries.add(path)
# The ops "namespace"
ops = _Ops()
|
