Use tensorflow-plugins/libmetal_plugin.dylib in tensorflow C API

Hello

I am trying to use Mac M1 GPU in tensorflow C API.

My code calls tensorflow C API via Rust bindings.

   match Library::load("libmetal_plugin.dylib") {
        Ok(lib) => println!("Loaded plugin successfully. {:?}", lib.op_list()),
        Err(e) => println!("WARNING: Unable to load plugin. {}", e),
    };

    // Load the saved model exported in python
    let mut graph = Graph::new();
    let bundle = SavedModelBundle::load(&SessionOptions::new(),
        &["serve"],
        &mut graph,
        export_dir).expect("Unable to load model from disk");

    println!("{:?}", bundle.session.device_list().unwrap() );

tensorflow::Library::load internally calls tensorflow C API TF_LoadLibrary to load the plugin. And it does successfully load the plugin.

Then I tried to enumerate the devices by device_list(). Only CPU is there.

[Device { name: "/job:localhost/replica:0/task:0/device:CPU:0", device_type: "CPU", memory_bytes: 268435456, incarnation: 13748960769752595872 }]

GPU does work in tensorflow when I use python to train. Is there some extra work needed to make it work in tensorflow C API?

Tensorflow C library version 2.9.0

Found the solution, I should use TF_LoadPluggableDeviceLibrary API, but it crashes after loading the plugin

2022-08-17 08:16:29.397578: I tensorflow/cc/saved_model/reader.cc:43] Reading SavedModel
2022-08-17 08:16:29.401477: I tensorflow/cc/saved_model/reader.cc:81] Reading meta graph with tags { serve }
2022-08-17 08:16:29.401500: I tensorflow/cc/saved_model/reader.cc:122] Reading SavedModel debug info (if present) 
Metal device set to: Apple M1

systemMemory: 16.00 GB
maxCacheSize: 5.33 GB

2022-08-17 08:16:29.411250: I tensorflow/core/common_runtime/pluggable_device/pluggable_device_factory.cc:306] Could not identify NUMA node of platform GPU ID 0, defaulting to 0. Your kernel may not have been built with NUMA support.
2022-08-17 08:16:29.411398: I tensorflow/core/common_runtime/pluggable_device/pluggable_device_factory.cc:272] Created TensorFlow device (/job:localhost/replica:0/task:0/device:GPU:0 with 0 MB memory) -> physical PluggableDevice (device: 0, name: METAL, pci bus id: <undefined>)
2022-08-17 08:16:29.422914: I tensorflow/compiler/mlir/mlir_graph_optimization_pass.cc:354] MLIR V1 optimization pass is not enabled
2022-08-17 08:16:29.427004: I tensorflow/cc/saved_model/loader.cc:227] Restoring SavedModel bundle.
2022-08-17 08:16:29.429435: W tensorflow/core/platform/profile_utils/cpu_utils.cc:128] Failed to get CPU frequency: 0 Hz
2022-08-17 08:16:29.440099: I tensorflow/core/grappler/optimizers/custom_graph_optimizer_registry.cc:113] Plugin optimizer for device_type GPU is enabled.
zsh: segmentation fault 
Use tensorflow-plugins/libmetal_plugin.dylib in tensorflow C API
 
 
Q