lightgbmlib.i 12.1 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
                                           int start_iteration,
                                           int num_iteration,
40
                                           int feature_importance_type,
41
42
                                           int64_t buffer_len,
                                           int64_t* out_len) {
43
    char* dst = new char[buffer_len];
44
    int result = LGBM_BoosterSaveModelToString(handle, start_iteration, num_iteration, feature_importance_type, buffer_len, out_len, dst);
45
46
47
48
49
    // Reallocate to use larger length
    if (*out_len > buffer_len) {
      delete [] dst;
      int64_t realloc_len = *out_len;
      dst = new char[realloc_len];
50
      result = LGBM_BoosterSaveModelToString(handle, start_iteration, num_iteration, feature_importance_type, realloc_len, out_len, dst);
51
52
53
54
55
56
57
    }
    if (result != 0) {
      return nullptr;
    }
    return dst;
  }

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

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

93
    int ret = LGBM_BoosterPredictForMatSingleRow(handle, data0, data_type, ncol, is_row_major, predict_type, start_iteration,
94
                                                 num_iteration, parameter, out_len, out_result);
95

96
    jenv->ReleasePrimitiveArrayCritical(data, data0, JNI_ABORT);
97

98
99
    return ret;
  }
100

101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
  /*! \brief Even faster variant of `LGBM_BoosterPredictForMatSingle`.
   *
   * Uses `LGBM_BoosterPredictForMatSingleRowFast` which is faster
   * than `LGBM_BoosterPredictForMatSingleRow` and the trick of
   * `LGBM_BoosterPredictForMatSingle` to capture the Java data array
   * using `GetPrimitiveArrayCritical`, which can yield faster access
   * to the array if the JVM passes the actual address to the C++ side
   * instead of performing a copy.
   */
  int LGBM_BoosterPredictForMatSingleRowFastCriticalSWIG(JNIEnv *jenv,
                                                         jdoubleArray data,
                                                         FastConfigHandle handle,
                                                         int64_t* out_len,
                                                         double* out_result) {
    double* data0 = (double*)jenv->GetPrimitiveArrayCritical(data, 0);

117
    int ret = LGBM_BoosterPredictForMatSingleRowFast(handle, data0, out_len, out_result);
118
119
120
121
122
123

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

    return ret;
  }

124
  int LGBM_BoosterPredictForCSRSingle(JNIEnv *jenv,
125
126
127
128
129
130
131
132
133
                                      jintArray indices,
                                      jdoubleArray values,
                                      int numNonZeros,
                                      BoosterHandle handle,
                                      int indptr_type,
                                      int data_type,
                                      int64_t nelem,
                                      int64_t num_col,
                                      int predict_type,
134
                                      int start_iteration,
135
136
137
138
                                      int num_iteration,
                                      const char* parameter,
                                      int64_t* out_len,
                                      double* out_result) {
139
140
141
142
143
    // 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
144

145
146
147
    jboolean isCopy;
    int* indices0 = (int*)jenv->GetPrimitiveArrayCritical(indices, &isCopy);
    double* values0 = (double*)jenv->GetPrimitiveArrayCritical(values, &isCopy);
148

149
    int32_t ind[2] = { 0, numNonZeros };
150

151
    int ret = LGBM_BoosterPredictForCSRSingleRow(handle, ind, indptr_type, indices0, values0, data_type, 2,
152
                                                 nelem, num_col, predict_type, start_iteration, num_iteration, parameter, out_len, out_result);
153

154
    jenv->ReleasePrimitiveArrayCritical(values, values0, JNI_ABORT);
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
    jenv->ReleasePrimitiveArrayCritical(indices, indices0, JNI_ABORT);

    return ret;
  }

  /*! \brief Even faster variant of `LGBM_BoosterPredictForCSRSingle`.
   *
   * Uses `LGBM_BoosterPredictForCSRSingleRowFast` which is faster
   * than `LGBM_BoosterPredictForMatSingleRow` and the trick of
   * `LGBM_BoosterPredictForCSRSingle` to capture the Java data array
   * using `GetPrimitiveArrayCritical`, which can yield faster access
   * to the array if the JVM passes the actual address to the C++ side
   * instead of performing a copy.
   */
  int LGBM_BoosterPredictForCSRSingleRowFastCriticalSWIG(JNIEnv *jenv,
                                                         jintArray indices,
                                                         jdoubleArray values,
                                                         int numNonZeros,
                                                         FastConfigHandle handle,
                                                         int indptr_type,
                                                         int64_t nelem,
                                                         int64_t* out_len,
                                                         double* out_result) {
    // 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

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

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

    int ret = LGBM_BoosterPredictForCSRSingleRowFast(handle, ind, indptr_type, indices0, values0, 2,
191
                                                     nelem, out_len, out_result);
192
193

    jenv->ReleasePrimitiveArrayCritical(values, values0, JNI_ABORT);
194
    jenv->ReleasePrimitiveArrayCritical(indices, indices0, JNI_ABORT);
195

196
197
    return ret;
  }
198

199
  #include <functional>
200
201
  #include <vector>

202
203
204
205
206
207
208
  struct CSRDirect {
          jintArray indices;
          jdoubleArray values;
          int* indices0;
          double* values0;
          int size;
  };
209

210
211
212
213
214
215
216
  int LGBM_DatasetCreateFromCSRSpark(JNIEnv *jenv,
                                     jobjectArray arrayOfSparseVector,
                                     int num_rows,
                                     int64_t num_col,
                                     const char* parameters,
                                     const DatasetHandle reference,
                                     DatasetHandle* out) {
217
218
219
220
221
222
223
224
225
226
227
    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++) {
228
      // get the row
229
230
      jobject objSparseVec = jenv->GetObjectArrayElement(arrayOfSparseVector, i);

231
232
      // get the size, indices and values
      auto indices = (jintArray)jenv->CallObjectMethod(objSparseVec, sparseVectorIndices);
233
234
235
      if (jenv->ExceptionCheck()) {
        return -1;
      }
236
      auto values = (jdoubleArray)jenv->CallObjectMethod(objSparseVec, sparseVectorValues);
237
238
239
      if (jenv->ExceptionCheck()) {
        return -1;
      }
240
      int size = jenv->GetArrayLength(indices);
241

242
243
244
245
246
247
      // 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
248
      int* indices0 = (int *)jenv->GetIntArrayElements(indices, 0);
249
      double* values0 = jenv->GetDoubleArrayElements(values, 0);
250

251
      jniCache.push_back({indices, values, indices0, values0, size});
252
253
254
255
    }

    // 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) {
256
      auto& jc = jniCache[row_num];
257
258
259
      ret.clear();  // reset size, but not free()
      ret.reserve(jc.size);  // make sure we have enough allocated

260
261
262
263
      // copy data
      int* indices0p = jc.indices0;
      double* values0p = jc.values0;
      int* indices0e = indices0p + jc.size;
264

265
266
      for (; indices0p != indices0e; ++indices0p, ++values0p)
        ret.emplace_back(*indices0p, *values0p);
267
268
269
270
271
    };

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

    for (auto& jc : jniCache) {
272
273
274
      // jenv->ReleasePrimitiveArrayCritical(jc.values, jc.values0, JNI_ABORT);
      // jenv->ReleasePrimitiveArrayCritical(jc.indices, jc.indices0, JNI_ABORT);
      jenv->ReleaseDoubleArrayElements(jc.values, jc.values0, JNI_ABORT);
275
      jenv->ReleaseIntArrayElements(jc.indices, (jint *)jc.indices0, JNI_ABORT);
276
277
278
    }

    return ret;
279
280
281
  }
%}

282

283
%include "pointer_manipulation.i"
284
%include "StringArray_API_extensions.i"
285
%include "ChunkedArray_API_extensions.i"