Add manual device placement for graphs

Author: lantiga

examples/js/mobilenet.js

@@ -49,7 +49,7 @@ async function run(filenames) {
   const buffer = fs.readFileSync(graph_filename, {'flag': 'r'});
 
   console.log("Setting graph");
-  redis.call('AI.SET', 'GRAPH', 'mobilenet', 'TF', buffer);
+  redis.call('AI.SET', 'GRAPH', 'mobilenet', 'TF', 'GPU', buffer);
 
   const image_height = 224;
   const image_width = 224;

examples/models/load_model.sh

@@ -1,7 +1,7 @@
 REDIS_CLI=../../deps/redis/src/redis-cli
 
 echo "SET GRAPH"
-$REDIS_CLI -x AI.SET GRAPH foo TF < graph.pb
+$REDIS_CLI -x AI.SET GRAPH foo TF GPU < graph.pb
 
 echo "SET TENSORS"
 $REDIS_CLI AI.SET TENSOR a FLOAT 1 2 VALUES 2 3

examples/models/load_yolo.sh

@@ -13,7 +13,7 @@ IMAGE_WIDTH=224
 IMAGE_HEIGHT=224
 
 echo "SET GRAPH"
-$REDIS_CLI -x AI.SET GRAPH $GRAPH_KEY TF < $GRAPH_FILE
+$REDIS_CLI -x AI.SET GRAPH $GRAPH_KEY TF GPU < $GRAPH_FILE
 
 # TODO: cast tensor, change shape of tensor (NHWC, NCHW)
 #       instead of casting, we could specify the type of data provided in the blob

src/backends/tensorflow.c

@@ -136,7 +136,8 @@ TF_Tensor* RAI_TFTensorFromTensor(RAI_Tensor* t){
 }
 
 
-RAI_Graph *RAI_GraphCreateTF(const char *prefix, RAI_Backend backend,
+RAI_Graph *RAI_GraphCreateTF(const char *prefix,
+                             RAI_Backend backend, RAI_Device device,
                              const char *graphdef, size_t graphlen) {
   TF_Graph* graph = TF_NewGraph();
 
@@ -160,9 +161,31 @@ RAI_Graph *RAI_GraphCreateTF(const char *prefix, RAI_Backend backend,
   TF_DeleteBuffer(buffer);
   TF_DeleteStatus(status);
 
-  TF_Status *sessionStatus = TF_NewStatus();
+  TF_Status *optionsStatus = TF_NewStatus();
 
   TF_SessionOptions *sessionOptions = TF_NewSessionOptions();
+
+  // For setting config options in session from the C API see:
+  // https://github.com/tensorflow/tensorflow/issues/13853
+  // import tensorflow as tf
+  // config = tf.ConfigProto()
+  // config.intra_op_parallelism_threads = 1
+  // serialized = config.SerializeToString()
+  // result = list(map(hex, serialized))
+
+  // TODO: complain if device is GPU and GPU not available?
+  if (device == RAI_DEVICE_CPU) {
+    uint8_t config[9] = {0x0a, 0x07, 0x0a, 0x03, 0x47, 0x50, 0x55, 0x10, 0x00};
+    TF_SetConfig(sessionOptions, (void *)config, 9, status);
+  }
+
+  if (TF_GetCode(optionsStatus) != TF_OK) {
+    // TODO: free memory
+    return NULL;
+  }
+  TF_DeleteStatus(optionsStatus);
+
+  TF_Status *sessionStatus = TF_NewStatus();
   TF_Session *session = TF_NewSession(graph, sessionOptions, sessionStatus);
 
   if (TF_GetCode(sessionStatus) != TF_OK) {

src/backends/tensorflow.h

@@ -11,7 +11,8 @@ RAI_Tensor* RAI_TensorCreateFromTFTensor(TF_Tensor *tensor);
 
 TF_Tensor* RAI_TFTensorFromTensor(RAI_Tensor* t);
 
-RAI_Graph *RAI_GraphCreateTF(const char *prefix, RAI_Backend backend,
+RAI_Graph *RAI_GraphCreateTF(const char *prefix,
+                             RAI_Backend backend, RAI_Device device,
                              const char *graphdef, size_t graphlen);
 
 void RAI_GraphFreeTF(RAI_Graph* graph);

src/config.h

@@ -11,6 +11,12 @@ typedef enum {
   RAI_BACKEND_ONNXRUNTIME,
 } RAI_Backend;
 
+typedef enum {
+  RAI_DEVICE_CPU = 0,
+  // TODO: multi GPU
+  RAI_DEVICE_GPU,
+} RAI_Device;
+
 //#define RAI_COPY_RUN_INPUT
 #define RAI_COPY_RUN_OUTPUT
 

src/graph.c

@@ -37,10 +37,11 @@ int RAI_GraphInit(RedisModuleCtx* ctx) {
   return RedisAI_GraphType != NULL;
 }
 
-RAI_Graph *RAI_GraphCreate(const char *prefix, RAI_Backend backend,
+RAI_Graph *RAI_GraphCreate(const char *prefix,
+                           RAI_Backend backend, RAI_Device device,
                            const char *graphdef, size_t graphlen) {
   if (backend == RAI_BACKEND_TENSORFLOW) {
-    return RAI_GraphCreateTF(prefix, backend, graphdef, graphlen);
+    return RAI_GraphCreateTF(prefix, backend, device, graphdef, graphlen);
   }
 
   return NULL;

src/graph.h

@@ -16,7 +16,7 @@
 extern RedisModuleType *RedisAI_GraphType;
 
 int RAI_GraphInit(RedisModuleCtx* ctx);
-RAI_Graph* RAI_GraphCreate(const char* prefix, RAI_Backend backend, const char* graphdef, size_t graphlen);
+RAI_Graph* RAI_GraphCreate(const char* prefix, RAI_Backend backend, RAI_Device device, const char* graphdef, size_t graphlen);
 void RAI_GraphFree(RAI_Graph* graph);
 RAI_GraphRunCtx* RAI_RunCtxCreate(RAI_Graph* graph);
 int RAI_RunCtxAddInput(RAI_GraphRunCtx* gctx, const char* inputName, RAI_Tensor* inputTensor);

src/redisai.c

@@ -392,7 +392,7 @@ int RedisAI_Get_Tensor_RedisCommand(RedisModuleCtx *ctx, RedisModuleString **arg
 int RedisAI_Set_Graph_RedisCommand(RedisModuleCtx *ctx, RedisModuleString **argv, int argc) {
   RedisModule_AutoMemory(ctx);
 
-  if ((argc != 4) && (argc != 5)) return RedisModule_WrongArity(ctx);
+  if ((argc != 5) && (argc != 6)) return RedisModule_WrongArity(ctx);
 
   const char* bckstr;
   int backend;
@@ -410,16 +410,29 @@ int RedisAI_Set_Graph_RedisCommand(RedisModuleCtx *ctx, RedisModuleString **argv
     return RedisModule_ReplyWithError(ctx, "ERR unsupported backend");
   }
 
+  const char* devicestr;
+  int device;
+  devicestr = RedisModule_StringPtrLen(argv[3], NULL);
+  if (strcasecmp(devicestr, "CPU") == 0) {
+    device = RAI_DEVICE_CPU;
+  }
+  else if (strcasecmp(devicestr, "GPU") == 0) {
+    device = RAI_DEVICE_GPU;
+  }
+  else {
+    return RedisModule_ReplyWithError(ctx, "ERR unsupported device");
+  }
+
   RAI_Graph *graph = NULL;
 
   size_t graphlen;
-  const char *graphdef = RedisModule_StringPtrLen(argv[3], &graphlen);
+  const char *graphdef = RedisModule_StringPtrLen(argv[4], &graphlen);
   const char *prefix = "";
-  if (argc == 5) {
-    const char *prefix = RedisModule_StringPtrLen(argv[4], NULL);
+  if (argc == 6) {
+    const char *prefix = RedisModule_StringPtrLen(argv[5], NULL);
   }
 
-  graph = RAI_GraphCreate(prefix, backend, graphdef, graphlen);
+  graph = RAI_GraphCreate(prefix, backend, device, graphdef, graphlen);
 
   if(graph == NULL){
     return RedisModule_ReplyWithError(ctx, "ERR failed creating the graph");

src/redisai.h

@@ -13,6 +13,7 @@ typedef struct RAI_Graph RAI_Graph;
 typedef struct RAI_GraphRunCtx RAI_GraphRunCtx;
 
 typedef enum RAI_Backend RAI_Backend;
+typedef enum RAI_Device RAI_Device;
 
 RAI_Tensor* MODULE_API_FUNC(RedisAI_TensorCreate)(const char* dataTypeStr, long long* dims, int ndims);
 size_t MODULE_API_FUNC(RedisAI_TensorLength)(RAI_Tensor* t);
@@ -30,7 +31,7 @@ long long MODULE_API_FUNC(RedisAI_TensorDim)(RAI_Tensor* t, int dim);
 size_t MODULE_API_FUNC(RedisAI_TensorByteSize)(RAI_Tensor* t);
 char* MODULE_API_FUNC(RedisAI_TensorData)(RAI_Tensor* t);
 
-RAI_Graph* MODULE_API_FUNC(RedisAI_GraphCreate)(const char* prefix, RAI_Backend backend, const char* graphdef, size_t graphlen);
+RAI_Graph* MODULE_API_FUNC(RedisAI_GraphCreate)(const char* prefix, RAI_Backend backend, RAI_Device device, const char* graphdef, size_t graphlen);
 void MODULE_API_FUNC(RedisAI_GraphFree)(RAI_Graph* graph);
 RAI_GraphRunCtx* MODULE_API_FUNC(RedisAI_RunCtxCreate)(RAI_Graph* graph);
 int MODULE_API_FUNC(RedisAI_RunCtxAddInput)(RAI_GraphRunCtx* gctx, const char* inputName, RAI_Tensor* inputTensor);