ffi.rst 6.98 KB
Newer Older
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
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
.. currentmodule:: dgl

DGL Foreign Function Interface (FFI)
====================================

We all like Python because it is easy to manipulate. We all like C because it
is fast, reliable and typed. To have the merits of both ends, DGL is mostly in
python, for quick prototyping, while lowers the performance-critical part to C.
Thus, DGL developers frequently face the scenario to write a C routine and has
it exposed to python, via a mechanism called *Foreign Function Interface (FFI)*.

There are many FFI solutions out there. In DGL, we want to keep it simple,
intuitive and efficient for critical use cases. That's why when we came across the
FFI solution in the TVM project, we immediately fell for it. It exploits the idea of
functional programming so that it exposes only a dozens of C APIs and new APIs
can be built upon it.

We decided to borrow the idea (shamelessly). For example, to define a C
API that is exposed to python is only a few lines of codes:

.. code:: c++

   // file: calculator.cc (put it in dgl/src folder)
   #include <dgl/runtime/packed_func.h>
   #include <dgl/runtime/registry.h>

   DGL_REGISTER_GLOBAL("calculator.MyAdd")
   .set_body([] (DGLArgs args, DGLRetValue* rv) {
       int a = args[0];
       int b = args[1];
       *rv = a * b;
     });

Compile and build the library. On the python side, create a
``calculator.py`` file under ``dgl/python/dgl/``

.. code:: python

   # file: calculator.py
   from ._ffi.function import _init_api

   def add(a, b):
     # MyAdd has been registered via `_ini_api` call below
     return MyAdd(a, b)

   _init_api("dgl.calculator")

The trick is that the FFI system first masks the type information of the
function arguments, so all the C function calls can go through one C API
(``DGLFuncCall``). The type information is retrieved in the function body by
static conversion, and we will do runtime type check to make sure that the type
conversion is correct. The overhead of such back-and-forth is negligible as
long as the function call is not too light (the above example is actually a bad
one). TVM's `PackedFunc
document <https://docs.tvm.ai/dev/runtime.html#packedfunc>`_ has more details.

Defining new types
------------------

``DGLArgs`` and ``DGLRetValue`` only support a limited number of types:

* Numerical values: int, float, double, ...
* string
* Function (in the form of PackedFunc)
* NDArray

Though limited, the above type system is very powerful because it supports
function as a first-class citizen. For example, if you want to return multiple
values, you can return a PackedFunc which returns each value given an integer
index. However, in many cases, new types are still desired to ease the
development process:

* The argument/return value is a composition of collections (e.g. dictionary of
  dictionary of list).
* Sometimes we just want to have a notion of "structure" (e.g. given an apple,
  get its color by ``apple.color``).

To achieve this, we introduce the Object type system. For example, to define a
new type ``Calculator``:

.. code:: c++

   // file: calculator.cc
   #include <dgl/packed_func_ext.h>
   using namespace runtime;
   class CalculatorObject : public Object {
    public:
     std::string brand;
     int price;
     
     void VisitAttrs(AttrVisitor *v) final {
       v->Visit("brand", &brand);
       v->Visit("price", &price);
     }

     static constexpr const char* _type_key = "Calculator";
     DGL_DECLARE_OBJECT_TYPE_INFO(CalculatorObject, Object);
   };

   // This is to define a reference class (the wrapper of an object shared pointer).
   // A minimal implementation is as follows, but you could define extra methods.
   class Calculator : public ObjectRef {
    public:
     const CalculatorObject* operator->() const {
       return static_cast<const CalculatorObject*>(obj_.get());
     }
     using ContainerType = CalculatorObject;
   };

   DGL_REGISTER_GLOBAL("calculator.CreateCaculator")
   .set_body([] (DGLArgs args, DGLRetValue* rv) {
     std::string brand = args[0];
     int price = args[1];
     auto o = std::make_shared<CalculatorObject>();
     o->brand = brand;
     o->price = price;
     *rv = o;
   }

On the python side:

.. code:: python

   # file: calculator.py
   from dgl._ffi.object import register_object, ObjectBase
   from ._ffi.function import _init_api

   @register_object
   class Calculator(ObjectBase):
     @staticmethod
     def create(brand, price):
       # invoke a C API, the return value is of `Calculator` type
       return CreateCalculator(brand, price)

   _init_api("dgl.calculator")

We can then simply create ``Calculator`` object by:

.. code:: python

   calc = Calculator.create("casio", 100)

What is nice about this object is that, it defines a visitor pattern that is
essentially a reflection mechanism to get its internal attributes. For example,
you can print the calculator's brand and by simply accessing its attributes.

.. code:: python

   print(calc.brand)
   print(calc.price)

The reflection is indeed a little bit slow due to the string key lookup. To
speed it up, you could define an attribute access API:

.. code:: c++

   // file: calculator.cc
   DGL_REGISTER_GLOBAL("calculator.CaculatorGetBrand")
   .set_body([] (DGLArgs args, DGLRetValue* rv) {
     Calculator calc = args[0];
     *rv = calc->brand;
   }

Containers
----------

Containers are also objects. For example, the C API below accepts a list of
integers and return their sum:

.. code:: c++

   // in file: calculator.cc
   #include <dgl/runtime/container.h>
   using namespace runtime;
   DGL_REGISTER_GLOBAL("calculator.Sum")
   .set_body([] (DGLArgs args, DGLRetValue* rv) {
     // All the DGL supported values are represented as a ValueObject, which
     //   contains a data field.
     List<Value> values = args[0];
     int sum = 0;
     for (int i = 0; i < values.size(); ++i) {
       sum += static_cast<int>(values[i]->data);
     }
   }

Invoking this API is simple -- just pass a python list of integers. DGL FFI will
automatically convert python list/tuple/dictionary to the corresponding object
type.

.. code:: python

   # in file: calculator.py
   from ._ffi.function import _init_api

   Sum([0, 1, 2, 3, 4, 5])

   _init_api("dgl.calculator")

The elements in the containers can be any objects, which allows the containers
to be composed. Below is an API that accepts a list of calculators and print
out their price:

.. code:: c++

   // in file: calculator.cc
   #include <iostream>
   #include <dgl/runtime/container.h>
   using namespace runtime;
   DGL_REGISTER_GLOBAL("calculator.PrintCalculators")
   .set_body([] (DGLArgs args, DGLRetValue* rv) {
     List<Calculator> calcs = args[0];
     for (int i = 0; i < calcs.size(); ++i) {
       std::cout << calcs[i]->price << std::endl;
     }
   }

Please note that containers are NOT meant for passing a large collection of
items from/to C APIs. It will be quite slow in these cases. It is recommended
to benchmark first. As an alternative, use NDArray for a large collection of
Zihao Ye's avatar
Zihao Ye committed
220
221
numerical values and use ``dgl.batch`` to batch a lot of ``DGLGraph``'s into 
a single ``DGLGraph``.