christoph/wamr
1
Fork 0
Code Releases Activity

wasi-nn: Improve TPU support (#2447)

Browse Source 
1. Allow TPU and GPU support at the same time.
2. Add Dockerfile to run example with [Coral USB](https://coral.ai/products/accelerator/).
This commit is contained in:
tonibofarull
2023-08-14 14:03:56 +02:00
committed by GitHub
parent 923e8d6471
commit b45d014112
10 changed files with 130 additions and 37 deletions
Download Patch File Download Diff File Expand all files Collapse all files

59
core/iwasm/libraries/wasi-nn/src/wasi_nn_tensorflowlite.cpp
View File

@ -16,11 +16,11 @@
#include <tensorflow/lite/optional_debug_tools.h>
#include <tensorflow/lite/error_reporter.h>
#if defined(WASI_NN_ENABLE_GPU)
#if WASM_ENABLE_WASI_NN_GPU != 0
#include <tensorflow/lite/delegates/gpu/delegate.h>
#endif
#if defined(WASI_NN_ENABLE_EXTERNAL_DELEGATE)
#if WASM_ENABLE_WASI_NN_EXTERNAL_DELEGATE != 0
#include <tensorflow/lite/delegates/external/external_delegate.h>
#endif
@ -130,8 +130,8 @@ tensorflowlite_load(void *tflite_ctx, graph_builder_array *builder,
return invalid_argument;
}
if (target != cpu && target != gpu) {
NN_ERR_PRINTF("Only CPU and GPU target is supported.");
if (target != cpu && target != gpu && target != tpu) {
NN_ERR_PRINTF("Only CPU, GPU and TPU target is supported.");
return invalid_argument;
}
@ -195,7 +195,7 @@ tensorflowlite_init_execution_context(void *tflite_ctx, graph g,
switch (tfl_ctx->models[g].target) {
case gpu:
{
#if defined(WASI_NN_ENABLE_GPU)
#if WASM_ENABLE_WASI_NN_GPU != 0
NN_WARN_PRINTF("GPU enabled.");
// https://www.tensorflow.org/lite/performance/gpu
TfLiteGpuDelegateOptionsV2 options =
@ -216,10 +216,19 @@ tensorflowlite_init_execution_context(void *tflite_ctx, graph g,
NN_ERR_PRINTF("Error when enabling GPU delegate.");
use_default = true;
}
#elif defined(WASI_NN_ENABLE_EXTERNAL_DELEGATE)
#else
NN_WARN_PRINTF("GPU not enabled.");
use_default = true;
#endif
break;
}
case tpu:
{
#if WASM_ENABLE_WASI_NN_EXTERNAL_DELEGATE != 0
NN_WARN_PRINTF("external delegation enabled.");
TfLiteExternalDelegateOptions options =
TfLiteExternalDelegateOptionsDefault(WASI_NN_EXT_DELEGATE_PATH);
TfLiteExternalDelegateOptionsDefault(
WASM_WASI_NN_EXTERNAL_DELEGATE_PATH);
tfl_ctx->delegate = TfLiteExternalDelegateCreate(&options);
if (tfl_ctx->delegate == NULL) {
NN_ERR_PRINTF("Error when generating External delegate.");
@ -233,7 +242,7 @@ tensorflowlite_init_execution_context(void *tflite_ctx, graph g,
use_default = true;
}
#else
NN_WARN_PRINTF("GPU not enabled.");
NN_WARN_PRINTF("External delegate not enabled.");
use_default = true;
#endif
break;
@ -438,19 +447,35 @@ tensorflowlite_destroy(void *tflite_ctx)
*/
TFLiteContext *tfl_ctx = (TFLiteContext *)tflite_ctx;
if (tfl_ctx->delegate != NULL) {
#if defined(WASI_NN_ENABLE_GPU)
TfLiteGpuDelegateV2Delete(tfl_ctx->delegate);
#elif defined(WASI_NN_ENABLE_EXTERNAL_DELEGATE)
TfLiteExternalDelegateDelete(tfl_ctx->delegate);
#endif
}
NN_DBG_PRINTF("Freeing memory.");
for (int i = 0; i < MAX_GRAPHS_PER_INST; ++i) {
tfl_ctx->models[i].model.reset();
if (tfl_ctx->models[i].model_pointer)
if (tfl_ctx->models[i].model_pointer) {
if (tfl_ctx->delegate) {
switch (tfl_ctx->models[i].target) {
case gpu:
{
#if WASM_ENABLE_WASI_NN_GPU != 0
TfLiteGpuDelegateV2Delete(tfl_ctx->delegate);
#else
NN_ERR_PRINTF("GPU delegate delete but not enabled.");
#endif
break;
}
case tpu:
{
#if WASM_ENABLE_WASI_NN_EXTERNAL_DELEGATE != 0
TfLiteExternalDelegateDelete(tfl_ctx->delegate);
#else
NN_ERR_PRINTF(
"External delegate delete but not enabled.");
#endif
break;
}
}
}
wasm_runtime_free(tfl_ctx->models[i].model_pointer);
}
tfl_ctx->models[i].model_pointer = NULL;
}
for (int i = 0; i < MAX_GRAPH_EXEC_CONTEXTS_PER_INST; ++i) {