[tvm4j] RPC Server

jvm/assembly/osx-x86_64-cpu/src/main/assembly/assembly.xml

@@ -6,7 +6,7 @@
   <includeBaseDirectory>false</includeBaseDirectory>
   <files>
     <file>
-      <source>../../../lib/libtvm_runtime.so</source>
+      <source>../../../lib/libtvm_runtime.dylib</source>
       <outputDirectory>lib/native</outputDirectory>
       <fileMode>0644</fileMode>
     </file>

jvm/core/pom.xml

@@ -20,18 +20,21 @@
       <id>osx-x86_64-cpu</id>
       <properties>
         <platform>osx-x86_64-cpu</platform>
+        <libtvm.so.filename>libtvm_runtime.dylib</libtvm.so.filename>
       </properties>
     </profile>
     <profile>
       <id>linux-x86_64-cpu</id>
       <properties>
         <platform>linux-x86_64-cpu</platform>
+        <libtvm.so.filename>libtvm_runtime.so</libtvm.so.filename>
       </properties>
     </profile>
     <profile>
       <id>linux-x86_64-gpu</id>
       <properties>
         <platform>linux-x86_64-gpu</platform>
+        <libtvm.so.filename>libtvm_runtime.so</libtvm.so.filename>
       </properties>
     </profile>
   </profiles>
@@ -88,7 +91,7 @@
           <threadCount>1</threadCount>
           <argLine>
             -Djava.library.path=${project.parent.basedir}/native/${platform}/target
-            -Dlibtvm.so.path=${project.parent.basedir}/../lib/libtvm_runtime.so
+            -Dlibtvm.so.path=${project.parent.basedir}/../lib/${libtvm.so.filename}
           </argLine>
         </configuration>
         <executions>

jvm/core/src/main/java/ml/dmlc/tvm/Base.java

@@ -80,7 +80,18 @@ public RefTVMValue() {
     if (tvmLibFilename == null || !new File(tvmLibFilename).isFile()
         || _LIB.nativeLibInit(tvmLibFilename) != 0) {
       try {
-        NativeLibraryLoader.extractResourceFileToTempDir("libtvm_runtime.so", new Action() {
+        String runtimeLibname;
+        String os = System.getProperty("os.name");
+        // ref: http://lopica.sourceforge.net/os.html
+        if (os.startsWith("Linux")) {
+          runtimeLibname = "libtvm_runtime.so";
+        } else if (os.startsWith("Mac")) {
+          runtimeLibname = "libtvm_runtime.dylib";
+        } else {
+          // TODO(yizhi) support windows later
+          throw new UnsatisfiedLinkError("Windows not supported currently");
+        }
+        NativeLibraryLoader.extractResourceFileToTempDir(runtimeLibname, new Action() {
           @Override public void invoke(File target) {
             System.err.println("Loading tvm runtime from " + target.getPath());
             checkCall(_LIB.nativeLibInit(target.getPath()));

jvm/core/src/main/java/ml/dmlc/tvm/Function.java

@@ -34,7 +34,7 @@ public class Function extends TVMValue {
    * @param name full function name.
    * @return TVM function.
    */
-  static Function getFunction(final String name) {
+  public static Function getFunction(final String name) {
     for (String fullName : listGlobalFuncNames()) {
       if (fullName.equals(name)) {
         return getGlobalFunc(fullName, true, false);

jvm/core/src/main/java/ml/dmlc/tvm/NativeLibraryLoader.java

@@ -30,7 +30,7 @@ class NativeLibraryLoader {
 
   static {
     try {
-      tempDir = File.createTempFile("tvm", "");
+      tempDir = File.createTempFile("tvm4j", "");
       if (!tempDir.delete() || !tempDir.mkdir()) {
         throw new IOException("Couldn't create directory " + tempDir.getAbsolutePath());
       }

jvm/core/src/main/java/ml/dmlc/tvm/TVMContext.java

@@ -17,12 +17,12 @@
 
 package ml.dmlc.tvm;
 
+import ml.dmlc.tvm.rpc.RPC;
+
 import java.util.HashMap;
 import java.util.Map;
 
 public class TVMContext {
-  private static final int RPC_SESS_MASK = 128;
-
   private static final Map<Integer, String> MASK2STR = new HashMap<Integer, String>();
   private static final Map<String, Integer> STR2MASK = new HashMap<String, Integer>();
 
@@ -169,9 +169,9 @@ public void sync() {
   }
 
   @Override public String toString() {
-    if (deviceType >= RPC_SESS_MASK) {
-      int tblId = deviceType / RPC_SESS_MASK - 1;
-      int devType = deviceType % RPC_SESS_MASK;
+    if (deviceType >= RPC.RPC_SESS_MASK) {
+      int tblId = deviceType / RPC.RPC_SESS_MASK - 1;
+      int devType = deviceType % RPC.RPC_SESS_MASK;
       return String.format("remote[%d]:%s(%d)", tblId, MASK2STR.get(devType), deviceId);
     }
     return String.format("%s(%d)", MASK2STR.get(deviceType), deviceId);

jvm/core/src/main/java/ml/dmlc/tvm/rpc/Client.java

@@ -0,0 +1,49 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package ml.dmlc.tvm.rpc;
+
+import ml.dmlc.tvm.Function;
+import ml.dmlc.tvm.TVMValue;
+
+public class Client {
+  /**
+   * Connect to RPC Server.
+   * @param url The url of the host.
+   * @param port The port to connect to.
+   * @param key Additional key to match server.
+   * @return The connected session.
+   */
+  public static RPCSession connect(String url, int port, String key) {
+    Function doConnect = RPC.getApi("_Connect");
+    if (doConnect == null) {
+      throw new RuntimeException("Please compile with USE_RPC=1");
+    }
+    TVMValue sess = doConnect.pushArg(url).pushArg(port).pushArg(key).invoke();
+    return new RPCSession(sess.asModule());
+  }
+
+  /**
+   * Connect to RPC Server.
+   * @param url The url of the host.
+   * @param port The port to connect to.
+   * @return The connected session.
+   */
+  public static RPCSession connect(String url, int port) {
+    return connect(url, port, "");
+  }
+}

jvm/core/src/main/java/ml/dmlc/tvm/rpc/RPC.java

@@ -0,0 +1,48 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package ml.dmlc.tvm.rpc;
+
+import ml.dmlc.tvm.Function;
+
+import java.util.HashMap;
+import java.util.Map;
+
+public class RPC {
+  public static final int RPC_MAGIC = 0xff271;
+  public static final int RPC_SESS_MASK = 128;
+
+  private static ThreadLocal<Map<String, Function>> apiFuncs
+      = new ThreadLocal<Map<String, Function>>() {
+          @Override
+          protected Map<String, Function> initialValue() {
+            return new HashMap<String, Function>();
+          }
+        };
+
+  static Function getApi(String name) {
+    Function func = apiFuncs.get().get(name);
+    if (func == null) {
+      func = Function.getFunction("contrib.rpc." + name);
+      if (func == null) {
+        return null;
+      }
+      apiFuncs.get().put(name, func);
+    }
+    return func;
+  }
+}