[Scala] add CustomOp support #4118

Ldpe2G · 2016-12-06T09:13:37Z

follows: https://github.com/dmlc/mxnet/blob/master/python/mxnet/operator.py
@Javelinjs Could you help to review the code?
Or anyone else help to review this code will be appreciated.
I am not sure whether I am implementating it right, for me this is the reasonable way I can come up with.

yzhliu

I simply reviewed the code conventions. This is a great work. I'll look into and check the correctness and design latter tomorrow. My instant feel is that the C codes are a little bit complex, and may contain risk of memory leak.

yzhliu · 2016-12-07T17:12:37Z

scala-package/core/src/main/scala/ml/dmlc/mxnet/Operator.scala

+   *            array of aux shapes calculated from in_shape,
+   *            in the same order as declared in listAuxiliaryStates().
+   */
+  def inferShape(inShape: Array[Array[Int]]):


seems you shall use Array[Shape] instead.

yzhliu · 2016-12-07T17:13:40Z

scala-package/core/src/main/scala/ml/dmlc/mxnet/Operator.scala

+    } catch {
+      case ex: Throwable => {
+        success = false
+        new Throwable().printStackTrace()


try ... catch ..., set success, return 0/1
hmm... this is somewhat wired

or change to return boolean?

ex.printStackTrace()

yzhliu · 2016-12-07T17:26:44Z

scala-package/native/src/main/native/ml_dmlc_mxnet_native_c_api.cc

+    // set CustomOpProp.kwargs
+    std::string opPropKey(opType);
+    if (globalOpPropMap.find(opPropKey) == globalOpPropMap.end()) {
+      std::cout << "CustomOpProp: " << opPropKey << " not found" << std::endl;


can we use dmlc::log ?

yzhliu · 2016-12-07T17:31:41Z

scala-package/native/src/main/native/ml_dmlc_mxnet_native_c_api.cc

+          auto delEntry = [](void *state) {
+            std::string key((char *)state);
+            bool success = true;
+            // if (globalOpMap.find(key) == globalOpMap.end()) {


debug codes?

actually I use global variable map to store the customop and customopprop object, and want to
remove them from the map and release them in the del fuction. But this delEntry function seems to be called during the execution, and cause seom problem , so I comment the code.

I'm not quite sure what happens here. ping @piiswrong

yzhliu · 2016-12-07T17:32:29Z

scala-package/native/src/main/native/ml_dmlc_mxnet_native_c_api.cc

+
+  CustomOpPropCreator creator = static_cast<bool(*)(const char*, const int, const char**, const char**, CustomOpPropInfo*)>(creatorLambda);
+  int ret = MXCustomOpRegister(regName, creator);
+  return ret;


simply return MXCustomOpRegister(regName, creator);

yzhliu · 2016-12-07T17:33:21Z

scala-package/native/src/main/native/ml_dmlc_mxnet_native_c_api.cc

+          if (obj != NULL) {
+            jobjectArray jauxs = (jobjectArray)obj;
+            int len = globalEnv->GetArrayLength(jauxs);
+            (*auxs) = new char *[len+1];


new but never delete ?

the auxs pointer was used to store the result, and passed in by the caller

Yes, but you shall not do in this way. We need to find a way to delete them finally, e.g., use an object like mxnet::common::ThreadLocalStore to manage it. Or, I'm not sure, is there a way to convert java String to char *, but let jvm handle the reference and automatically do GC?

ok. I take a look.

I think the memory release should be handle by the caller side? Because the caller call this function and pass a pointer for example auxs, then after the function invocation the auxs contains the results the caller needed, then if you handle the memory release inside the function, how do you know when you release the memory, the caller side won't need that anymore?

I figure out a way to handle the memory release 64a65cc
, what's your opinion?

yzhliu · 2016-12-08T18:04:21Z

scala-package/native/src/main/native/ml_dmlc_mxnet_native_c_api.cc

+JNIEXPORT jint JNICALL Java_ml_dmlc_mxnet_LibInfo_mxCustomOpRegister
+  (JNIEnv *env, jobject obj, jstring jregName, jobject jopProp) {
+
+  globalEnv = env;


pls refer to KVStoreServerControllerFunc. I dont think what here does is correct.

I will take a look.

why I use global variable here is because, if you want to use a local variable inside a c++ lambda,
this variable needs to be captured, then you can not convert this lambda to function pointer.

oh, I think might understand what you mean.

yzhliu · 2016-12-08T18:29:59Z

scala-package/native/src/main/native/ml_dmlc_mxnet_native_c_api.cc

+          if (obj != NULL) {
+            jobjectArray jauxs = (jobjectArray)obj;
+            int len = globalEnv->GetArrayLength(jauxs);
+            (*auxs) = new char *[len+1];


Yes, but you shall not do in this way. We need to find a way to delete them finally, e.g., use an object like mxnet::common::ThreadLocalStore to manage it. Or, I'm not sure, is there a way to convert java String to char *, but let jvm handle the reference and automatically do GC?

yzhliu · 2016-12-08T18:53:39Z

scala-package/native/src/main/native/ml_dmlc_mxnet_native_c_api.cc

+
+JNIEnv *globalEnv;
+std::unordered_map<std::string, jobject> globalOpPropMap;
+std::unordered_map<std::string, jobject> globalOpMap;


I wonder whether we have a better way to deal with the registry here, can we make the jobject, i.e., make the user defined class a closure, like what python code does ?

what I actually tried to implement is that because when those callback function was called, it was called in C++ side, then inside the callback function it needs to find the specific java object then
it can call the java object function, so I store the object in a map during initialization, then retrieve them with their name.

yzhliu · 2016-12-08T19:14:34Z

scala-package/native/src/main/native/ml_dmlc_mxnet_native_c_api.cc

+          auto delEntry = [](void *state) {
+            std::string key((char *)state);
+            bool success = true;
+            // if (globalOpMap.find(key) == globalOpMap.end()) {


I'm not quite sure what happens here. ping @piiswrong

Ldpe2G · 2016-12-12T03:14:12Z

@Javelinjs I have tested customop with multithread, and the results on mnist dataset seems fine.

yzhliu

Pls also fix the code conventions, according to http://107.23.36.180:8080/job/mxnet/759/console

yzhliu · 2016-12-12T15:26:03Z

scala-package/core/src/main/scala/ml/dmlc/mxnet/Operator.scala

+    } catch {
+      case ex: Throwable => {
+        success = false
+        new Throwable().printStackTrace()


ex.printStackTrace()

yzhliu · 2016-12-12T15:48:44Z

scala-package/native/src/main/native/ml_dmlc_mxnet_native_c_api.cc

@@ -1711,3 +1715,505 @@ JNIEXPORT jint JNICALL Java_ml_dmlc_mxnet_LibInfo_mxRtcFree
  int ret = MXRtcFree(handle);
  return ret;
 }
+
+std::unordered_map<std::string, jobject> globalOpPropMap;


better to give comments to indicate what those global variables for.

yzhliu · 2016-12-12T15:49:37Z

scala-package/native/src/main/native/ml_dmlc_mxnet_native_c_api.cc

+    // set CustomOpProp.kwargs
+    std::string opPropKey(opType);
+    if (globalOpPropMap.find(opPropKey) == globalOpPropMap.end()) {
+      LOG(FATAL) << "CustomOpProp: " << opPropKey << " not found";


LOG(FATAL) will die immediately. Is it important to return a state in this function? if so, you may change the log level to WARN

yzhliu · 2016-12-12T16:36:04Z

scala-package/native/src/main/native/ml_dmlc_mxnet_native_c_api.cc

+          jint *rdepsArr = env->GetIntArrayElements(jrdeps, NULL);
+
+          (*numDeps) = env->GetArrayLength(jrdeps);
+          (* rdeps) = new int[(* numDeps)];


you can remove these parentheses.

yzhliu · 2016-12-12T17:00:54Z

scala-package/native/src/main/native/ml_dmlc_mxnet_native_c_api.cc

+                                                       tagsArr,
+                                                       reqsArr,
+                                                       *(const_cast<bool*>(&isTrain)));
+                if (result == false) success = false;


success = env->call...

same for the others

yzhliu · 2016-12-12T17:08:01Z

scala-package/native/src/main/native/ml_dmlc_mxnet_native_c_api.cc

+            (*auxs) = new char *[1];
+            (*auxs)[0] = NULL;
+          }
+          dmlc::ThreadLocalStore<Ref<char>>::Get()->setData(*auxs);


OK, let's remove this. So we have no chance to get rid of memory leak for such data transfer from jvm side?

yes, I think so.

Ldpe2G · 2016-12-14T01:03:46Z

@Javelinjs I think I have done all the changes and made this pr satified the requirements. Is there anything else I missed?

yzhliu · 2016-12-15T16:50:39Z

Please also take care once this feature is merged to nnvm branch.

Ldpe2G · 2016-12-16T01:02:48Z

@Javelinjs ok

yzhliu suggested changes Dec 7, 2016

View reviewed changes

yzhliu suggested changes Dec 8, 2016

View reviewed changes

yzhliu suggested changes Dec 12, 2016

View reviewed changes

Ldpe2G added 11 commits December 15, 2016 15:56

[Scala] add CustomOp support

1d6a263

fix code style

1ea58d3

fix some code style

c92382b

fix

8a4346c

change globalEnv to local env

a814af4

fix potential memory leak

c69523c

fix some problems

9fb9394

fix thread safty problem

8c4a3a6

add comment, fix code style

bb202a7

fix code style

a93b3ef

remove unuse headfile

43002ae

yzhliu merged commit 87c6404 into apache:master Dec 15, 2016

[Scala] add CustomOp support #4118

[Scala] add CustomOp support #4118

Conversation

Ldpe2G commented Dec 6, 2016

yzhliu left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

yzhliu Dec 7, 2016 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Ldpe2G Dec 8, 2016 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Ldpe2G Dec 10, 2016 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Ldpe2G commented Dec 12, 2016

yzhliu left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Ldpe2G commented Dec 14, 2016

yzhliu commented Dec 15, 2016

Ldpe2G commented Dec 16, 2016

yzhliu Dec 7, 2016 •

edited

Loading

Ldpe2G Dec 8, 2016 •

edited

Loading

Ldpe2G Dec 10, 2016 •

edited

Loading