lightgbmlib.i 11 KB
Newer Older
1
2
3
4
/*!
 * Copyright (c) 2018 Microsoft Corporation. All rights reserved.
 * Licensed under the MIT License. See LICENSE file in the project root for license information.
 */
5
6
7
/* lightgbmlib.i */
%module lightgbmlib
%ignore LGBM_BoosterSaveModelToString;
8
%ignore LGBM_BoosterGetEvalNames;
9
%ignore LGBM_BoosterGetFeatureNames;
10
11
12
13
%{
/* Includes the header in the wrapper code */
#include "../include/LightGBM/export.h"
#include "../include/LightGBM/utils/log.h"
14
#include "../include/LightGBM/utils/common.h"
15
16
17
#include "../include/LightGBM/c_api.h"
%}

18
19
20
%include "various.i"
%include "carrays.i"
%include "cpointer.i"
21
%include "stdint.i"
22
23
24
25
26
27

/* Note: instead of using array_functions for string array we apply a typemap instead.
   Future char** parameter names should be added to the typemap.
*/
%apply char **STRING_ARRAY { char **feature_names, char **out_strs }

28
29
30
31
/* header files */
%include "../include/LightGBM/export.h"
%include "../include/LightGBM/c_api.h"

32
33
%typemap(in, numinputs = 0) JNIEnv *jenv %{
  $1 = jenv;
34
35
%}

36
37
%inline %{
  char * LGBM_BoosterSaveModelToStringSWIG(BoosterHandle handle,
38
39
40
41
                                           int start_iteration,
                                           int num_iteration,
                                           int64_t buffer_len,
                                           int64_t* out_len) {
42
    char* dst = new char[buffer_len];
43
    int result = LGBM_BoosterSaveModelToString(handle, start_iteration, num_iteration, buffer_len, out_len, dst);
44
45
46
47
48
49
50
51
52
53
54
55
56
    // Reallocate to use larger length
    if (*out_len > buffer_len) {
      delete [] dst;
      int64_t realloc_len = *out_len;
      dst = new char[realloc_len];
      result = LGBM_BoosterSaveModelToString(handle, start_iteration, num_iteration, realloc_len, out_len, dst);
    }
    if (result != 0) {
      return nullptr;
    }
    return dst;
  }

57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
  char * LGBM_BoosterDumpModelSWIG(BoosterHandle handle,
                                   int start_iteration,
                                   int num_iteration,
                                   int64_t buffer_len,
                                   int64_t* out_len) {
    char* dst = new char[buffer_len];
    int result = LGBM_BoosterDumpModel(handle, start_iteration, num_iteration, buffer_len, out_len, dst);
    // Reallocate to use larger length
    if (*out_len > buffer_len) {
      delete [] dst;
      int64_t realloc_len = *out_len;
      dst = new char[realloc_len];
      result = LGBM_BoosterDumpModel(handle, start_iteration, num_iteration, realloc_len, out_len, dst);
    }
    if (result != 0) {
      return nullptr;
    }
    return dst;
  }

77
  int LGBM_BoosterPredictForMatSingle(JNIEnv *jenv,
78
79
80
81
82
83
84
85
86
87
88
                                      jdoubleArray data,
                                      BoosterHandle handle,
                                      int data_type,
                                      int ncol,
                                      int is_row_major,
                                      int predict_type,
                                      int num_iteration,
                                      const char* parameter,
                                      int64_t* out_len,
                                      double* out_result) {
    double* data0 = (double*)jenv->GetPrimitiveArrayCritical(data, 0);
89

90
    int ret = LGBM_BoosterPredictForMatSingleRow(handle, data0, data_type, ncol, is_row_major, predict_type,
91
                                                 num_iteration, parameter, out_len, out_result);
92

93
    jenv->ReleasePrimitiveArrayCritical(data, data0, JNI_ABORT);
94

95
96
    return ret;
  }
97

98
  int LGBM_BoosterPredictForCSRSingle(JNIEnv *jenv,
99
100
101
102
103
104
105
106
107
108
109
110
111
                                      jintArray indices,
                                      jdoubleArray values,
                                      int numNonZeros,
                                      BoosterHandle handle,
                                      int indptr_type,
                                      int data_type,
                                      int64_t nelem,
                                      int64_t num_col,
                                      int predict_type,
                                      int num_iteration,
                                      const char* parameter,
                                      int64_t* out_len,
                                      double* out_result) {
112
113
114
115
116
    // Alternatives
    // - GetIntArrayElements: performs copy
    // - GetDirectBufferAddress: fails on wrapped array
    // Some words of warning for GetPrimitiveArrayCritical
    // https://stackoverflow.com/questions/23258357/whats-the-trade-off-between-using-getprimitivearraycritical-and-getprimitivety
117

118
119
120
    jboolean isCopy;
    int* indices0 = (int*)jenv->GetPrimitiveArrayCritical(indices, &isCopy);
    double* values0 = (double*)jenv->GetPrimitiveArrayCritical(values, &isCopy);
121

122
    int32_t ind[2] = { 0, numNonZeros };
123

124
    int ret = LGBM_BoosterPredictForCSRSingleRow(handle, ind, indptr_type, indices0, values0, data_type, 2,
125
126
                                                 nelem, num_col, predict_type, num_iteration, parameter, out_len, out_result);

127
128
    jenv->ReleasePrimitiveArrayCritical(values, values0, JNI_ABORT);
    jenv->ReleasePrimitiveArrayCritical(indices, indices0, JNI_ABORT);
129

130
131
    return ret;
  }
132

133
  #include <functional>
134
135
  #include <vector>

136
137
138
139
140
141
142
  struct CSRDirect {
          jintArray indices;
          jdoubleArray values;
          int* indices0;
          double* values0;
          int size;
  };
143

144
145
146
147
148
149
150
  int LGBM_DatasetCreateFromCSRSpark(JNIEnv *jenv,
                                     jobjectArray arrayOfSparseVector,
                                     int num_rows,
                                     int64_t num_col,
                                     const char* parameters,
                                     const DatasetHandle reference,
                                     DatasetHandle* out) {
151
152
153
154
155
156
157
158
159
160
161
    jclass sparseVectorClass = jenv->FindClass("org/apache/spark/ml/linalg/SparseVector");
    jmethodID sparseVectorIndices = jenv->GetMethodID(sparseVectorClass, "indices", "()[I");
    jmethodID sparseVectorValues = jenv->GetMethodID(sparseVectorClass, "values", "()[D");

    std::vector<CSRDirect> jniCache;
    jniCache.reserve(num_rows);

    // this needs to be done ahead of time as row_func is invoked from multiple threads
    // these threads would have to be registered with the JVM and also unregistered.
    // It is not clear if that can be achieved with OpenMP
    for (int i = 0; i < num_rows; i++) {
162
      // get the row
163
164
      jobject objSparseVec = jenv->GetObjectArrayElement(arrayOfSparseVector, i);

165
166
      // get the size, indices and values
      auto indices = (jintArray)jenv->CallObjectMethod(objSparseVec, sparseVectorIndices);
167
168
169
      if (jenv->ExceptionCheck()) {
        return -1;
      }
170
      auto values = (jdoubleArray)jenv->CallObjectMethod(objSparseVec, sparseVectorValues);
171
172
173
      if (jenv->ExceptionCheck()) {
        return -1;
      }
174
      int size = jenv->GetArrayLength(indices);
175

176
177
178
179
180
181
      // Note: when testing on larger data (e.g. 288k rows per partition and 36mio rows total)
      // using GetPrimitiveArrayCritical resulted in a dead-lock
      // lock arrays
      // int* indices0 = (int*)jenv->GetPrimitiveArrayCritical(indices, 0);
      // double* values0 = (double*)jenv->GetPrimitiveArrayCritical(values, 0);
      // in test-usecase an alternative to GetPrimitiveArrayCritical as it performs copies
182
      int* indices0 = (int *)jenv->GetIntArrayElements(indices, 0);
183
      double* values0 = jenv->GetDoubleArrayElements(values, 0);
184

185
      jniCache.push_back({indices, values, indices0, values0, size});
186
187
188
189
    }

    // type is important here as we want a std::function, rather than a lambda
    std::function<void(int idx, std::vector<std::pair<int, double>>& ret)> row_func = [&](int row_num, std::vector<std::pair<int, double>>& ret) {
190
      auto& jc = jniCache[row_num];
191
192
193
      ret.clear();  // reset size, but not free()
      ret.reserve(jc.size);  // make sure we have enough allocated

194
195
196
197
      // copy data
      int* indices0p = jc.indices0;
      double* values0p = jc.values0;
      int* indices0e = indices0p + jc.size;
198

199
200
      for (; indices0p != indices0e; ++indices0p, ++values0p)
        ret.emplace_back(*indices0p, *values0p);
201
202
203
204
205
    };

    int ret = LGBM_DatasetCreateFromCSRFunc(&row_func, num_rows, num_col, parameters, reference, out);

    for (auto& jc : jniCache) {
206
207
208
      // jenv->ReleasePrimitiveArrayCritical(jc.values, jc.values0, JNI_ABORT);
      // jenv->ReleasePrimitiveArrayCritical(jc.indices, jc.indices0, JNI_ABORT);
      jenv->ReleaseDoubleArrayElements(jc.values, jc.values0, JNI_ABORT);
209
      jenv->ReleaseIntArrayElements(jc.indices, (jint *)jc.indices0, JNI_ABORT);
210
211
212
    }

    return ret;
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
  }
%}

%pointer_functions(int, intp)
%pointer_functions(long, longp)
%pointer_functions(double, doublep)
%pointer_functions(float, floatp)
%pointer_functions(int64_t, int64_tp)
%pointer_functions(int32_t, int32_tp)

%pointer_cast(int64_t *, long *, int64_t_to_long_ptr)
%pointer_cast(int64_t *, double *, int64_t_to_double_ptr)
%pointer_cast(int32_t *, int *, int32_t_to_int_ptr)
%pointer_cast(long *, int64_t *, long_to_int64_t_ptr)
%pointer_cast(double *, int64_t *, double_to_int64_t_ptr)
%pointer_cast(int *, int32_t *, int_to_int32_t_ptr)
229
230

%pointer_cast(double *, void *, double_to_voidp_ptr)
231
%pointer_cast(float *, void *, float_to_voidp_ptr)
232
233
234
%pointer_cast(int *, void *, int_to_voidp_ptr)
%pointer_cast(int32_t *, void *, int32_t_to_voidp_ptr)
%pointer_cast(int64_t *, void *, int64_t_to_voidp_ptr)
235
236

/* Custom pointer manipulation template */
237
%define %pointer_manipulation(TYPE, NAME)
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
%{
  static TYPE *new_##NAME() { %}
  %{  TYPE* NAME = new TYPE; return NAME; %}
  %{}

  static void delete_##NAME(TYPE *self) { %}
  %{  if (self) delete self; %}
  %{}
  %}

TYPE *new_##NAME();
void  delete_##NAME(TYPE *self);

%enddef

253
%define %pointer_dereference(TYPE, NAME)
254
255
256
257
258
259
260
261
262
263
264
%{
  static TYPE NAME ##_value(TYPE *self) {
    TYPE NAME = *self;
    return NAME;
  }
%}

TYPE NAME##_value(TYPE *self);

%enddef

265
%define %pointer_handle(TYPE, NAME)
266
267
268
269
270
271
272
273
274
275
%{
  static TYPE* NAME ##_handle() { %}
  %{ TYPE* NAME = new TYPE; *NAME = (TYPE)operator new(sizeof(int*)); return NAME; %}
  %{}
%}

TYPE *NAME##_handle();

%enddef

276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
%define %long_array_functions(TYPE,NAME)
%{
  static TYPE *new_##NAME(int64_t nelements) { %}
  %{  return new TYPE[nelements](); %}
  %{}

  static void delete_##NAME(TYPE *ary) { %}
  %{  delete [] ary; %}
  %{}

  static TYPE NAME##_getitem(TYPE *ary, int64_t index) {
    return ary[index];
  }
  static void NAME##_setitem(TYPE *ary, int64_t index, TYPE value) {
    ary[index] = value;
  }
  %}

TYPE *new_##NAME(int64_t nelements);
void delete_##NAME(TYPE *ary);
TYPE NAME##_getitem(TYPE *ary, int64_t index);
void NAME##_setitem(TYPE *ary, int64_t index, TYPE value);

%enddef

%long_array_functions(double, doubleArray)
%long_array_functions(float, floatArray)
%long_array_functions(int, intArray)
%long_array_functions(long, longArray)

306
307
308
309
310
311
312
%pointer_manipulation(void*, voidpp)

/* Allow dereferencing of void** to void* */
%pointer_dereference(void*, voidpp)

/* Allow retrieving handle to void** */
%pointer_handle(void*, voidpp)
313
314

%include "StringArray_API_extensions.i"